modifying default index API to accept None (#1308)

pytorch · May 17, 2021 · ccfba91 · ccfba91
1 parent 3858bae
commit ccfba91
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 3 deletions.
diff --git a/test/test_vocab.py b/test/test_vocab.py
@@ -71,6 +71,10 @@ def test_default_index(self):
         v.set_default_index(0)
         self.assertEqual(v['not in vocab'], 0)
 
+        v.set_default_index(None)
+        with self.assertRaises(RuntimeError):
+            v['not in vocab']
+
     def test_default_index_jit(self):
         token_to_freq = {'<unk>': 2, 'a': 2, 'b': 2}
         sorted_by_freq_tuples = sorted(token_to_freq.items(), key=lambda x: x[1], reverse=True)
@@ -80,6 +84,10 @@ def test_default_index_jit(self):
         v_jit = torch.jit.script(v)
         self.assertEqual(v_jit['not in vocab'], 0)
 
+        v_jit.set_default_index(None)
+        with self.assertRaises(RuntimeError):
+            v_jit['not in vocab']
+
     def test_vocab_insert_token(self):
         c = OrderedDict({'<unk>': 2, 'a': 2})
 

diff --git a/torchtext/csrc/vocab.cpp b/torchtext/csrc/vocab.cpp
@@ -46,7 +46,7 @@ int64_t Vocab::__getitem__(const c10::string_view &token) const {
   return default_index_.value();
 }
 
-void Vocab::set_default_index(int64_t index) { default_index_ = index; }
+void Vocab::set_default_index(c10::optional<int64_t> index) { default_index_ = index; }
 
 c10::optional<int64_t> Vocab::get_default_index() const {
   return default_index_;

diff --git a/torchtext/csrc/vocab.h b/torchtext/csrc/vocab.h
@@ -27,7 +27,7 @@ struct Vocab : torch::CustomClassHolder {
   int64_t __len__() const;
   int64_t __getitem__(const c10::string_view &token) const;
   bool __contains__(const c10::string_view &token) const;
-  void set_default_index(int64_t index);
+  void set_default_index(c10::optional<int64_t> index);
   c10::optional<int64_t> get_default_index() const;
   void reassign_token(const std::string &token, const int64_t &index);
   void insert_token(const std::string &token, const int64_t &index);

diff --git a/torchtext/vocab.py b/torchtext/vocab.py
@@ -84,7 +84,7 @@ def __getitem__(self, token: str) -> int:
         return self.vocab[token]
 
     @torch.jit.export
-    def set_default_index(self, index: int) -> None:
+    def set_default_index(self, index: Optional[int]) -> None:
         r"""
         Args:
             index: Value of default index. This index will be returned when OOV token is queried.