Merge pull request #20 from floscha/embeddings-contains-method

Contains method for word vectors
PetrochukM · Apr 12, 2018 · 7e0840f · 7e0840f
2 parents b0aa81d + f1b4d49
commit 7e0840f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/tests/word_to_vector/test_fast_text.py b/tests/word_to_vector/test_fast_text.py
@@ -22,6 +22,9 @@ def test_fasttext_simple(mock_urlretrieve):
     assert 'the' in vectors.stoi
     assert len(vectors) == 1
 
+    # Test implementation of __contains()__
+    assert 'the' in vectors
+
     # Test with the unknown characters
     assert len(vectors['漢字']) == 300
 

diff --git a/torchnlp/word_to_vector/pretrained_word_vectors.py b/torchnlp/word_to_vector/pretrained_word_vectors.py
@@ -71,6 +71,9 @@ def __init__(self,
         self.name = name
         self.cache(name, cache, url=url)
 
+    def __contains__(self, token):
+        return token in self.stoi
+
     def __getitem__(self, token):
         if token in self.stoi:
             return self.vectors[self.stoi[token]]