Expose ngram_indices and subword_indices

Author: NianhengWu

Cargo.toml

@@ -21,7 +21,7 @@ features = ["extension-module"]
 [dependencies]
 itertools = "0.8"
 failure = "0.1"
-finalfusion = "0.8"
+finalfusion = "0.8.2"
 libc = "0.2"
 ndarray = "0.12"
 numpy = "0.6"

src/vocab.rs

@@ -1,7 +1,7 @@
 use std::cell::RefCell;
 use std::rc::Rc;
 
-use finalfusion::chunks::vocab::WordIndex;
+use finalfusion::chunks::vocab::{NGramIndices, SubwordIndices, VocabWrap, WordIndex};
 use finalfusion::prelude::*;
 use pyo3::class::sequence::PySequenceProtocol;
 use pyo3::exceptions;
@@ -34,6 +34,28 @@ impl PyVocab {
             }
         })
     }
+
+    fn ngram_indices(&self, word: &str) -> PyResult<Option<Vec<(String, usize)>>> {
+        let embeds = self.embeddings.borrow();
+        match embeds.vocab() {
+            VocabWrap::FastTextSubwordVocab(inner) => Ok(inner.ngram_indices(word)),
+            VocabWrap::FinalfusionSubwordVocab(inner) => Ok(inner.ngram_indices(word)),
+            VocabWrap::SimpleVocab(_) => Err(exceptions::ValueError::py_err(
+                "querying n-gram indices is not supported for this vocabulary",
+            )),
+        }
+    }
+
+    fn subword_indices(&self, word: &str) -> PyResult<Option<Vec<usize>>> {
+        let embeds = self.embeddings.borrow();
+        match embeds.vocab() {
+            VocabWrap::FastTextSubwordVocab(inner) => Ok(inner.subword_indices(word)),
+            VocabWrap::FinalfusionSubwordVocab(inner) => Ok(inner.subword_indices(word)),
+            VocabWrap::SimpleVocab(_) => Err(exceptions::ValueError::py_err(
+                "querying subwords' indices is not supported for this vocabulary",
+            )),
+        }
+    }
 }
 
 #[pyproto]

tests/conftest.py

@@ -25,6 +25,11 @@ def similarity_fifu(tests_root):
     yield finalfusion.Embeddings(os.path.join(tests_root, "similarity.fifu"))
 
 
+@pytest.fixture
+def subword_fifu(tests_root):
+    yield finalfusion.Embeddings(os.path.join(tests_root, "subword.fifu"))
+
+
 @pytest.fixture
 def embeddings_text_dims(tests_root):
     yield finalfusion.Embeddings.read_text_dims(os.path.join(tests_root, "embeddings.dims.txt"))

tests/subword.fifu

@@ -1,3 +1,74 @@
+TEST_NGRAM_INDICES = [
+    ('tüb',
+     14),
+    ('en>',
+     69),
+    ('übinge',
+     74),
+    ('gen',
+     124),
+    ('ing',
+     168),
+    ('ngen',
+     181),
+    ('bing',
+     197),
+    ('inge',
+     246),
+    ('übin',
+     250),
+    ('tübi',
+     276),
+    ('bingen',
+     300),
+    ('<tübin',
+     308),
+    ('bin',
+     325),
+    ('übing',
+     416),
+    ('gen>',
+     549),
+    ('ngen>',
+     590),
+    ('ingen>',
+     648),
+    ('tübing',
+     651),
+    ('übi',
+     707),
+    ('ingen',
+     717),
+    ('binge',
+     761),
+    ('<tübi',
+     817),
+    ('<tü',
+     820),
+    ('<tüb',
+     857),
+    ('nge',
+     860),
+    ('tübin',
+     1007)]
+
+
 def test_embeddings_with_norms_oov(embeddings_fifu):
     vocab = embeddings_fifu.vocab()
     assert vocab.item_to_indices("Something out of vocabulary") is None
+
+
+def test_ngram_indices(subword_fifu):
+    vocab = subword_fifu.vocab()
+    ngram_indices = sorted(vocab.ngram_indices("tübingen"), key=lambda tup: tup[1])
+    for ngram_index, test_ngram_index in zip(
+            ngram_indices, TEST_NGRAM_INDICES):
+        assert ngram_index == test_ngram_index
+
+
+def test_subword_indices(subword_fifu):
+    vocab = subword_fifu.vocab()
+    subword_indices = sorted(vocab.subword_indices("tübingen"))
+    for subword_index, test_ngram_index in zip(
+            subword_indices, TEST_NGRAM_INDICES):
+        assert subword_index == test_ngram_index[1]