Remove the specialized embedding iterator with norms

Instead, let the normal embeddings iterator return norms. Getting the
norms is a simple, predictable, memory traversal, and thus very cheap.

Author: danieldk

src/embeddings.rs

@@ -15,9 +15,7 @@ use pyo3::types::PyTuple;
 use pyo3::{exceptions, PyMappingProtocol};
 use toml::{self, Value};
 
-use crate::{
-    EmbeddingsWrap, PyEmbeddingIterator, PyEmbeddingWithNormIterator, PyVocab, PyWordSimilarity,
-};
+use crate::{EmbeddingsWrap, PyEmbeddingIterator, PyVocab, PyWordSimilarity};
 
 /// finalfusion embeddings.
 #[pyclass(name=Embeddings)]
@@ -257,10 +255,6 @@ impl PyEmbeddings {
                 .map_err(|err| exceptions::IOError::py_err(err.to_string())),
         }
     }
-
-    fn iter_with_norm(&self) -> PyResult<PyEmbeddingWithNormIterator> {
-        Ok(PyEmbeddingWithNormIterator::new(self.embeddings.clone(), 0))
-    }
 }
 
 #[pyproto]

src/iter.rs

@@ -32,70 +32,6 @@ impl PyIterProtocol for PyEmbeddingIterator {
         let embeddings = slf.embeddings.borrow();
         let vocab = embeddings.vocab();
 
-        if slf.idx < vocab.len() {
-            let word = vocab.words()[slf.idx].to_string();
-            let embed = embeddings.storage().embedding(slf.idx);
-
-            slf.idx += 1;
-
-            let gil = pyo3::Python::acquire_gil();
-            Ok(Some(PyEmbedding {
-                word,
-                embedding: embed.into_owned().into_pyarray(gil.python()).to_owned(),
-            }))
-        } else {
-            Ok(None)
-        }
-    }
-}
-
-/// A word and its embedding.
-#[pyclass(name=Embedding)]
-pub struct PyEmbedding {
-    embedding: Py<PyArray1<f32>>,
-    word: String,
-}
-
-#[pymethods]
-impl PyEmbedding {
-    /// Get the embedding.
-    #[getter]
-    pub fn get_embedding(&self) -> Py<PyArray1<f32>> {
-        let gil = Python::acquire_gil();
-        self.embedding.clone_ref(gil.python())
-    }
-
-    /// Get the word.
-    #[getter]
-    pub fn get_word(&self) -> &str {
-        &self.word
-    }
-}
-
-#[pyclass(name=EmbeddingWithNormIterator)]
-pub struct PyEmbeddingWithNormIterator {
-    embeddings: Rc<RefCell<EmbeddingsWrap>>,
-    idx: usize,
-}
-
-impl PyEmbeddingWithNormIterator {
-    pub fn new(embeddings: Rc<RefCell<EmbeddingsWrap>>, idx: usize) -> Self {
-        PyEmbeddingWithNormIterator { embeddings, idx }
-    }
-}
-
-#[pyproto]
-impl PyIterProtocol for PyEmbeddingWithNormIterator {
-    fn __iter__(slf: PyRefMut<Self>) -> PyResult<Py<PyEmbeddingWithNormIterator>> {
-        Ok(slf.into())
-    }
-
-    fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<PyEmbeddingWithNorm>> {
-        let slf = &mut *slf;
-
-        let embeddings = slf.embeddings.borrow();
-        let vocab = embeddings.vocab();
-
         if slf.idx < vocab.len() {
             let word = vocab.words()[slf.idx].to_string();
             let embed = embeddings.storage().embedding(slf.idx);
@@ -104,7 +40,7 @@ impl PyIterProtocol for PyEmbeddingWithNormIterator {
             slf.idx += 1;
 
             let gil = pyo3::Python::acquire_gil();
-            Ok(Some(PyEmbeddingWithNorm {
+            Ok(Some(PyEmbedding {
                 word,
                 embedding: embed.into_owned().into_pyarray(gil.python()).to_owned(),
                 norm,
@@ -116,15 +52,15 @@ impl PyIterProtocol for PyEmbeddingWithNormIterator {
 }
 
 /// A word and its embedding and embedding norm.
-#[pyclass(name=EmbeddingWithNorm)]
-pub struct PyEmbeddingWithNorm {
+#[pyclass(name=Embedding)]
+pub struct PyEmbedding {
     embedding: Py<PyArray1<f32>>,
     norm: f32,
     word: String,
 }
 
 #[pymethods]
-impl PyEmbeddingWithNorm {
+impl PyEmbedding {
     /// Get the embedding.
     #[getter]
     pub fn get_embedding(&self) -> Py<PyArray1<f32>> {

src/lib.rs

@@ -9,7 +9,7 @@ mod embeddings_wrap;
 use embeddings_wrap::EmbeddingsWrap;
 
 mod iter;
-use iter::{PyEmbedding, PyEmbeddingIterator, PyEmbeddingWithNorm, PyEmbeddingWithNormIterator};
+use iter::{PyEmbedding, PyEmbeddingIterator};
 
 mod similarity;
 use similarity::PyWordSimilarity;
@@ -25,7 +25,6 @@ use vocab::PyVocab;
 fn finalfusion(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<PyEmbeddings>()?;
     m.add_class::<PyEmbedding>()?;
-    m.add_class::<PyEmbeddingWithNorm>()?;
     m.add_class::<PyWordSimilarity>()?;
     m.add_class::<PyVocab>()?;
     Ok(())

tests/test_embedding.py

@@ -12,11 +12,11 @@
 ]
 
 
-def test_embeddings_with_norms(embeddings_fifu, embeddings_text):
-    for embedding_with_norm, norm in zip(
-            embeddings_fifu.iter_with_norm(), TEST_NORMS):
-        unnormed_embed = embedding_with_norm.embedding * norm
-        test_embed = embeddings_text[embedding_with_norm.word]
+def test_embeddings(embeddings_fifu, embeddings_text):
+    for embedding, norm in zip(
+            embeddings_fifu, TEST_NORMS):
+        unnormed_embed = embedding.embedding * embedding.norm
+        test_embed = embeddings_text[embedding.word]
         assert numpy.allclose(
             unnormed_embed, test_embed), "Embedding from 'iter_with_norm()' fails to match!"
 
@@ -32,20 +32,12 @@ def test_indexing(embeddings_fifu):
         embeddings_fifu["Something out of vocabulary"]
 
 
-def test_embeddings(embeddings_fifu, embeddings_text):
-    for embedding, norm in zip(embeddings_fifu, TEST_NORMS):
-        unnormed_embed = embedding.embedding * norm
-        test_embed = embeddings_text[embedding.word]
-        assert numpy.allclose(
-            unnormed_embed, test_embed), "Embedding from normal iterator fails to match!"
-
-
 def test_embeddings_oov(embeddings_fifu):
     assert embeddings_fifu.embedding("Something out of vocabulary") is None
 
 
 def test_norms(embeddings_fifu):
-    for embedding_with_norm, norm in zip(
-            embeddings_fifu.iter_with_norm(), TEST_NORMS):
+    for embedding, norm in zip(
+            embeddings_fifu, TEST_NORMS):
         assert pytest.approx(
-            embedding_with_norm.norm) == norm, "Norm fails to match!"
+            embedding.norm) == norm, "Norm fails to match!"