Add embedding similarity queries.

Add a method to perform similarity queries based on an input
embedding rather than words.

Author: sebpuetz

Cargo.lock

@@ -1,4 +1,5 @@
 use std::cell::RefCell;
+use std::collections::HashSet;
 use std::fs::File;
 use std::io::{BufReader, BufWriter};
 use std::rc::Rc;
@@ -12,11 +13,11 @@ use finalfusion::prelude::*;
 use finalfusion::similarity::*;
 use itertools::Itertools;
 use ndarray::Array2;
-use numpy::{IntoPyArray, PyArray1, PyArray2};
+use numpy::{IntoPyArray, NpyDataType, PyArray1, PyArray2};
 use pyo3::class::iter::PyIterProtocol;
 use pyo3::prelude::*;
-use pyo3::types::PyTuple;
-use pyo3::{exceptions, PyMappingProtocol};
+use pyo3::types::{PyAny, PyList, PySet, PyTuple};
+use pyo3::{exceptions, PyMappingProtocol, PyTypeInfo};
 use toml::{self, Value};
 
 use crate::{EmbeddingsWrap, PyEmbeddingIterator, PyVocab, PyWordSimilarity};
@@ -143,18 +144,7 @@ impl PyEmbeddings {
                 exceptions::KeyError::py_err(format!("Unknown word or n-grams: {}", failed))
             })?;
 
-        let mut r = Vec::with_capacity(results.len());
-        for ws in results {
-            r.push(
-                Py::new(
-                    py,
-                    PyWordSimilarity::new(ws.word.to_owned(), ws.similarity.into_inner()),
-                )?
-                .into_object(py),
-            )
-        }
-
-        Ok(r)
+        Self::similarity_results(py, results)
     }
 
     /// Get the embedding for the given word.
@@ -258,7 +248,7 @@ impl PyEmbeddings {
 
     /// Perform a similarity query.
     #[args(limit = 10)]
-    fn similarity(&self, py: Python, word: &str, limit: usize) -> PyResult<Vec<PyObject>> {
+    fn word_similarity(&self, py: Python, word: &str, limit: usize) -> PyResult<Vec<PyObject>> {
         let embeddings = self.embeddings.borrow();
 
         let embeddings = embeddings.view().ok_or_else(|| {
@@ -271,18 +261,46 @@ impl PyEmbeddings {
             .word_similarity(word, limit)
             .ok_or_else(|| exceptions::KeyError::py_err("Unknown word and n-grams"))?;
 
-        let mut r = Vec::with_capacity(results.len());
-        for ws in results {
-            r.push(
-                Py::new(
-                    py,
-                    PyWordSimilarity::new(ws.word.to_owned(), ws.similarity.into_inner()),
-                )?
-                .into_object(py),
+        Self::similarity_results(py, results)
+    }
+
+    /// Perform a similarity query based on a query embedding.
+    #[args(limit = 10, skip = "None")]
+    fn embedding_similarity(
+        &self,
+        py: Python,
+        embedding: PyEmbedding,
+        skip: Option<Option<Skips>>,
+        limit: usize,
+    ) -> PyResult<Vec<PyObject>> {
+        let embeddings = self.embeddings.borrow();
+
+        let embeddings = embeddings.view().ok_or_else(|| {
+            exceptions::ValueError::py_err(
+                "Similarity queries are not supported for this type of embedding matrix",
             )
+        })?;
+
+        let embedding = embedding.0.as_array();
+
+        if embedding.shape()[0] != embeddings.storage().shape().1 {
+            return Err(exceptions::ValueError::py_err(format!(
+                "Incompatible embedding shapes: embeddings: ({},), query: ({},)",
+                embedding.shape()[0],
+                embeddings.storage().shape().1
+            )));
         }
 
-        Ok(r)
+        let results = if let Some(Some(skip)) = skip {
+            embeddings.embedding_similarity_masked(embedding, limit, &skip.0)
+        } else {
+            embeddings.embedding_similarity(embedding, limit)
+        };
+
+        Self::similarity_results(
+            py,
+            results.ok_or_else(|| exceptions::KeyError::py_err("Unknown word and n-grams"))?,
+        )
     }
 
     /// Write the embeddings to a finalfusion file.
@@ -304,6 +322,25 @@ impl PyEmbeddings {
     }
 }
 
+impl PyEmbeddings {
+    fn similarity_results(
+        py: Python,
+        results: Vec<WordSimilarityResult>,
+    ) -> PyResult<Vec<PyObject>> {
+        let mut r = Vec::with_capacity(results.len());
+        for ws in results {
+            r.push(
+                Py::new(
+                    py,
+                    PyWordSimilarity::new(ws.word.to_owned(), ws.similarity.into_inner()),
+                )?
+                .into_object(py),
+            )
+        }
+        Ok(r)
+    }
+}
+
 #[pyproto]
 impl PyMappingProtocol for PyEmbeddings {
     fn __getitem__(&self, word: &str) -> PyResult<Py<PyArray1<f32>>> {
@@ -372,3 +409,47 @@ where
         embeddings: Rc::new(RefCell::new(EmbeddingsWrap::View(embeddings.into()))),
     })
 }
+
+struct Skips<'a>(HashSet<&'a str>);
+
+impl<'a> FromPyObject<'a> for Skips<'a> {
+    fn extract(ob: &'a PyAny) -> Result<Self, PyErr> {
+        let mut set = ob
+            .len()
+            .map(|len| HashSet::with_capacity(len))
+            .unwrap_or_default();
+
+        let iter = if <PySet as PyTypeInfo>::is_instance(ob) {
+            ob.iter().unwrap()
+        } else if <PyList as PyTypeInfo>::is_instance(ob) {
+            ob.iter().unwrap()
+        } else {
+            return Err(exceptions::TypeError::py_err("Iterable expected"));
+        };
+
+        for el in iter {
+            set.insert(
+                el?.extract()
+                    .map_err(|_| exceptions::TypeError::py_err("Expected String"))?,
+            );
+        }
+        Ok(Skips(set))
+    }
+}
+
+struct PyEmbedding<'a>(&'a PyArray1<f32>);
+
+impl<'a> FromPyObject<'a> for PyEmbedding<'a> {
+    fn extract(ob: &'a PyAny) -> Result<Self, PyErr> {
+        let embedding = ob
+            .downcast_ref::<PyArray1<f32>>()
+            .map_err(|_| exceptions::TypeError::py_err("Expected array with dtype Float32"))?;
+        if embedding.data_type() != NpyDataType::Float32 {
+            return Err(exceptions::TypeError::py_err(format!(
+                "Expected dtype Float32, got {:?}",
+                embedding.data_type()
+            )));
+        };
+        Ok(PyEmbedding(embedding))
+    }
+}

src/embeddings.rs

@@ -1,3 +1,6 @@
+import pytest
+import numpy
+
 SIMILARITY_ORDER_STUTTGART_10 = [
     "Karlsruhe",
     "Mannheim",
@@ -57,10 +60,28 @@
 
 
 def test_similarity_berlin_40(similarity_fifu):
-    for idx, sim in enumerate(similarity_fifu.similarity("Berlin", 40)):
+    for idx, sim in enumerate(similarity_fifu.word_similarity("Berlin", 40)):
         assert SIMILARITY_ORDER[idx] == sim.word
 
 
 def test_similarity_stuttgart_10(similarity_fifu):
-    for idx, sim in enumerate(similarity_fifu.similarity("Stuttgart", 10)):
+    for idx, sim in enumerate(similarity_fifu.word_similarity("Stuttgart", 10)):
+        assert SIMILARITY_ORDER_STUTTGART_10[idx] == sim.word
+
+
+def test_embedding_similarity_stuttgart_10(similarity_fifu):
+    stuttgart = similarity_fifu.embedding("Stuttgart")
+    sims = similarity_fifu.embedding_similarity(stuttgart, limit=10)
+    assert sims[0].word == "Stuttgart"
+
+    for idx, sim in enumerate(sims[1:]):
         assert SIMILARITY_ORDER_STUTTGART_10[idx] == sim.word
+
+    for idx, sim in enumerate(similarity_fifu.embedding_similarity(stuttgart, skip={"Stuttgart"}, limit=10)):
+        assert SIMILARITY_ORDER_STUTTGART_10[idx] == sim.word
+
+
+def test_embedding_similarity_incompatible_shapes(similarity_fifu):
+    incompatible_embed = numpy.ones(1, dtype=numpy.float32)
+    with pytest.raises(ValueError):
+        similarity_fifu.embedding_similarity(incompatible_embed)