Stop returning tuples in iteration over embeddings

When iterating over embeddings with norms, we get 3-tuples. This means
that the user has to memorize which tuple element is what. So, instead
produce instances of a EmbeddingWithNorm class which has the
properties word/embedding/norm. For consistency, update iteration
without norms as well.

Author: danieldk

src/iter.rs

@@ -26,7 +26,7 @@ impl PyIterProtocol for PyEmbeddingIterator {
         Ok(slf.into())
     }
 
-    fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<(String, Py<PyArray1<f32>>)>> {
+    fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<PyEmbedding>> {
         let slf = &mut *slf;
 
         let embeddings = slf.embeddings.borrow();
@@ -39,16 +39,39 @@ impl PyIterProtocol for PyEmbeddingIterator {
             slf.idx += 1;
 
             let gil = pyo3::Python::acquire_gil();
-            Ok(Some((
+            Ok(Some(PyEmbedding {
                 word,
-                embed.into_owned().into_pyarray(gil.python()).to_owned(),
-            )))
+                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>>,
@@ -67,7 +90,7 @@ impl PyIterProtocol for PyEmbeddingWithNormIterator {
         Ok(slf.into())
     }
 
-    fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<(String, Py<PyArray1<f32>>, f32)>> {
+    fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<PyEmbeddingWithNorm>> {
         let slf = &mut *slf;
 
         let embeddings = slf.embeddings.borrow();
@@ -81,13 +104,43 @@ impl PyIterProtocol for PyEmbeddingWithNormIterator {
             slf.idx += 1;
 
             let gil = pyo3::Python::acquire_gil();
-            Ok(Some((
+            Ok(Some(PyEmbeddingWithNorm {
                 word,
-                embed.into_owned().into_pyarray(gil.python()).to_owned(),
+                embedding: embed.into_owned().into_pyarray(gil.python()).to_owned(),
                 norm,
-            )))
+            }))
         } else {
             Ok(None)
         }
     }
 }
+
+/// A word and its embedding and embedding norm.
+#[pyclass(name=EmbeddingWithNorm)]
+pub struct PyEmbeddingWithNorm {
+    embedding: Py<PyArray1<f32>>,
+    norm: f32,
+    word: String,
+}
+
+#[pymethods]
+impl PyEmbeddingWithNorm {
+    /// 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
+    }
+
+    /// Get the norm.
+    #[getter]
+    pub fn get_norm(&self) -> f32 {
+        self.norm
+    }
+}

src/lib.rs

@@ -9,7 +9,7 @@ mod embeddings_wrap;
 use embeddings_wrap::EmbeddingsWrap;
 
 mod iter;
-use iter::{PyEmbeddingIterator, PyEmbeddingWithNormIterator};
+use iter::{PyEmbedding, PyEmbeddingIterator, PyEmbeddingWithNorm, PyEmbeddingWithNormIterator};
 
 mod similarity;
 use similarity::PyWordSimilarity;
@@ -24,6 +24,8 @@ use vocab::PyVocab;
 #[pymodule]
 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

@@ -15,12 +15,10 @@
 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[1] * norm
-        test_embed = embeddings_text[embedding_with_norm[0]]
+        unnormed_embed = embedding_with_norm.embedding * norm
+        test_embed = embeddings_text[embedding_with_norm.word]
         assert numpy.allclose(
             unnormed_embed, test_embed), "Embedding from 'iter_with_norm()' fails to match!"
-        assert len(
-            embedding_with_norm) == 3, "The number of values returned by 'iter_with_norm()' does not match!"
 
 
 def test_embeddings_with_norms_oov(embeddings_fifu):
@@ -35,13 +33,11 @@ def test_indexing(embeddings_fifu):
 
 
 def test_embeddings(embeddings_fifu, embeddings_text):
-    for embedding_with_norm, norm in zip(embeddings_fifu, TEST_NORMS):
-        unnormed_embed = embedding_with_norm[1] * norm
-        test_embed = embeddings_text[embedding_with_norm[0]]
+    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!"
-        assert len(
-            embedding_with_norm) == 2, "The number of values returned by normal iterator does not match!"
 
 
 def test_embeddings_oov(embeddings_fifu):
@@ -52,4 +48,4 @@ def test_norms(embeddings_fifu):
     for embedding_with_norm, norm in zip(
             embeddings_fifu.iter_with_norm(), TEST_NORMS):
         assert pytest.approx(
-            embedding_with_norm[2]) == norm, "Norm fails to match!"
+            embedding_with_norm.norm) == norm, "Norm fails to match!"