Adding dunder methods to cpp Vectors (#852)

* Adding dunder methods to cpp Vectors

* Fix length

* Fixed pytest errors

* Fixed stylecheck

Author: Nayef211

test/experimental/test_vectors.py

@@ -16,6 +16,11 @@
 
 
 class TestVectors(TorchtextTestCase):
+    def tearDown(self):
+        super().tearDown()
+        torch._C._jit_clear_class_registry()
+        torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()
+
     def test_empty_vectors(self):
         tokens = []
         vectors = []

torchtext/csrc/vectors.cpp

@@ -41,25 +41,28 @@ struct Vectors : torch::CustomClassHolder {
     }
   }
 
-  torch::Tensor GetItem(const std::string &token) const {
+  torch::Tensor __getitem__(const std::string &token) const {
     if (stovectors_.find(token) != stovectors_.end()) {
       return stovectors_.at(token);
     }
     return unk_tensor_;
   }
 
-  void AddItem(const std::string &token, const torch::Tensor &vector) {
+  void __setitem__(const std::string &token, const torch::Tensor &vector) {
     stovectors_[token] = vector;
   }
+
+  int64_t __len__() { return stovectors_.size(); }
 };
 
 // Registers our custom class with torch.
 static auto vectors =
     torch::class_<Vectors>("torchtext", "Vectors")
         .def(torch::init<std::vector<std::string>, std::vector<torch::Tensor>,
                          torch::Tensor>())
-        .def("GetItem", &Vectors::GetItem)
-        .def("AddItem", &Vectors::AddItem)
+        .def("__getitem__", &Vectors::__getitem__)
+        .def("__setitem__", &Vectors::__setitem__)
+        .def("__len__", &Vectors::__len__)
         .def_pickle(
             // __getstate__
             [](const c10::intrusive_ptr<Vectors> &self)

torchtext/experimental/vectors.py

@@ -256,7 +256,7 @@ def __getitem__(self, token: str) -> Tensor:
         Returns:
             vector (Tensor): a tensor (the vector) corresponding to the associated token.
         """
-        return self.vectors.GetItem(token)
+        return self.vectors[token]
 
     @torch.jit.export
     def __setitem__(self, token: str, vector: Tensor):
@@ -271,7 +271,16 @@ def __setitem__(self, token: str, vector: Tensor):
         if vector.dtype != torch.float:
             raise TypeError("`vector` should be of data type `torch.float` but it's of type " + vector.dtype)
 
-        self.vectors.AddItem(token, vector.float())
+        self.vectors[token] = vector.float()
+
+    @torch.jit.export
+    def __len__(self):
+        r"""Get length of vectors object.
+
+        Returns:
+            length (int): the length of the vectors.
+        """
+        return len(self.vectors)
 
 
 CHECKSUMS_GLOVE = {