Add default value to embedding lookup.

Allow specification of a default return value for embedding
lookups. Iterable, Array, Scalar and None are permitted values.

Author: sebpuetz

src/embeddings.rs

@@ -11,18 +11,19 @@ use finalfusion::io as ffio;
 use finalfusion::prelude::*;
 use finalfusion::similarity::*;
 use itertools::Itertools;
+use ndarray::Array1;
 use numpy::{IntoPyArray, NpyDataType, PyArray1};
 use pyo3::class::iter::PyIterProtocol;
 use pyo3::prelude::*;
-use pyo3::types::{PyAny, PyTuple};
+use pyo3::types::{PyAny, PyIterator, PyTuple};
 use pyo3::{exceptions, PyMappingProtocol};
 use toml::{self, Value};
 
 use crate::storage::PyStorage;
 use crate::{EmbeddingsWrap, PyEmbeddingIterator, PyVocab, PyWordSimilarity};
 
 /// finalfusion embeddings.
-#[pyclass(name=Embeddings)]
+#[pyclass(name = Embeddings)]
 pub struct PyEmbeddings {
     // The use of Rc + RefCell should be safe in this crate:
     //
@@ -179,17 +180,49 @@ impl PyEmbeddings {
         Self::similarity_results(py, results)
     }
 
+    /// embedding(word,/, default)
+    /// --
+    ///
     /// Get the embedding for the given word.
     ///
     /// If the word is not known, its representation is approximated
-    /// using subword units.
-    fn embedding(&self, word: &str) -> Option<Py<PyArray1<f32>>> {
+    /// using subword units. #
+    ///
+    /// If no representation can be calculated:
+    ///  - `None` if `default` is `None`
+    ///  - an array filled with `default` if `default` is a scalar
+    ///  - an array if `default` is a 1-d array
+    ///  - an array filled with values from `default` if it is an iterator over floats.
+    #[args(default = "PyEmbeddingDefault::default()")]
+    fn embedding(
+        &self,
+        word: &str,
+        default: PyEmbeddingDefault,
+    ) -> PyResult<Option<Py<PyArray1<f32>>>> {
         let embeddings = self.embeddings.borrow();
+        let gil = pyo3::Python::acquire_gil();
+        if let PyEmbeddingDefault::Embedding(array) = &default {
+            if array.as_ref(gil.python()).shape()[0] != embeddings.storage().shape().1 {
+                return Err(exceptions::ValueError::py_err(format!(
+                    "Invalid shape of default embedding: {}",
+                    array.as_ref(gil.python()).shape()[0]
+                )));
+            }
+        }
 
-        embeddings.embedding(word).map(|e| {
-            let gil = pyo3::Python::acquire_gil();
-            e.into_owned().into_pyarray(gil.python()).to_owned()
-        })
+        if let Some(embedding) = embeddings.embedding(word) {
+            return Ok(Some(
+                embedding.into_owned().into_pyarray(gil.python()).to_owned(),
+            ));
+        };
+        match default {
+            PyEmbeddingDefault::Constant(constant) => {
+                let nd_arr = Array1::from_elem([embeddings.storage().shape().1], constant);
+                Ok(Some(nd_arr.into_pyarray(gil.python()).to_owned()))
+            }
+            PyEmbeddingDefault::Embedding(array) => Ok(Some(array)),
+            PyEmbeddingDefault::None => Ok(None),
+        }
     }
 
     fn embedding_with_norm(&self, word: &str) -> Option<Py<PyTuple>> {
@@ -415,6 +448,58 @@ where
     })
 }
 
+pub enum PyEmbeddingDefault {
+    Embedding(Py<PyArray1<f32>>),
+    Constant(f32),
+    None,
+}
+
+impl<'a> Default for PyEmbeddingDefault {
+    fn default() -> Self {
+        PyEmbeddingDefault::None
+    }
+}
+
+impl<'a> FromPyObject<'a> for PyEmbeddingDefault {
+    fn extract(ob: &'a PyAny) -> Result<Self, PyErr> {
+        if ob.is_none() {
+            return Ok(PyEmbeddingDefault::None);
+        }
+        if let Ok(emb) = ob
+            .extract()
+            .map(|e: &PyArray1<f32>| PyEmbeddingDefault::Embedding(e.to_owned()))
+        {
+            return Ok(emb);
+        }
+
+        if let Ok(constant) = ob.extract().map(PyEmbeddingDefault::Constant) {
+            return Ok(constant);
+        }
+        if let Ok(embed) = ob
+            .iter()
+            .and_then(|iter| collect_array_from_py_iter(iter, ob.len().ok()))
+            .map(PyEmbeddingDefault::Embedding)
+        {
+            return Ok(embed);
+        }
+
+        Err(exceptions::TypeError::py_err(
+            "failed to construct default value.",
+        ))
+    }
+}
+
+fn collect_array_from_py_iter(iter: PyIterator, len: Option<usize>) -> PyResult<Py<PyArray1<f32>>> {
+    let mut embed_vec = len.map(Vec::with_capacity).unwrap_or_default();
+    for item in iter {
+        let item = item.and_then(|item| item.extract())?;
+        embed_vec.push(item);
+    }
+    let gil = Python::acquire_gil();
+    let embed = PyArray1::from_vec(gil.python(), embed_vec).to_owned();
+    Ok(embed)
+}
+
 struct Skips<'a>(HashSet<&'a str>);
 
 impl<'a> FromPyObject<'a> for Skips<'a> {

tests/test_embeddings.py

@@ -22,7 +22,7 @@ def test_embeddings(embeddings_fifu, embeddings_text, embeddings_text_dims):
     # The correct dimensionality of the other embedding types is asserted
     # in the pairwise comparisons below.
     assert fifu_storage.shape() == (7, 10)
-    
+
     for embedding, storage_row in zip(embeddings_fifu, fifu_storage):
         assert numpy.allclose(
             embedding.embedding, embeddings_text[embedding.word]), "FiFu and text embedding mismatch"
@@ -32,6 +32,32 @@ def test_embeddings(embeddings_fifu, embeddings_text, embeddings_text_dims):
             embedding.embedding, storage_row), "FiFu and storage row  mismatch"
 
 
+def test_unknown_embeddings(embeddings_fifu):
+    assert embeddings_fifu.embedding("OOV") is None, "Unknown lookup with no default failed"
+    assert embeddings_fifu.embedding(
+        "OOV", default=None) is None, "Unknown lookup with 'None' default failed"
+    assert numpy.allclose(embeddings_fifu.embedding(
+        "OOV", default=[10]*10), numpy.array([10.]*10)), "Unknown lookup with 'list' default failed"
+    assert numpy.allclose(embeddings_fifu.embedding("OOV", default=numpy.array(
+        [10.]*10)), numpy.array([10.]*10)), "Unknown lookup with array default failed"
+    assert numpy.allclose(embeddings_fifu.embedding(
+        "OOV", default=10), numpy.array([10.]*10)), "Unknown lookup with 'int' scalar default failed"
+    assert numpy.allclose(embeddings_fifu.embedding(
+        "OOV", default=10.), numpy.array([10.]*10)), "Unknown lookup with 'float' scalar default failed"
+    with pytest.raises(TypeError):
+        embeddings_fifu.embedding(
+            "OOV", default="not working"), "Unknown lookup with 'str' default succeeded"
+    with pytest.raises(ValueError):
+        embeddings_fifu.embedding(
+            "OOV", default=[10.]*5), "Unknown lookup with incorrectly shaped 'list' default succeeded"
+    with pytest.raises(ValueError):
+        embeddings_fifu.embedding(
+            "OOV", default=numpy.array([10.]*5)), "Unknown lookup with incorrectly shaped array default succeeded"
+    with pytest.raises(ValueError):
+        embeddings_fifu.embedding(
+            "OOV", default=range(7)), "Unknown lookup with iterable default with incorrect number succeeded"
+
+
 def test_embeddings_pq(similarity_fifu, similarity_pq):
     for embedding in similarity_fifu:
         embedding_pq = similarity_pq.embedding("Berlin")