Add masked analogy functionality.

Author: sebpuetz

Cargo.lock

@@ -20,7 +20,7 @@ features = ["extension-module"]
 
 [dependencies]
 failure = "0.1"
-finalfusion = "0.6"
+finalfusion = "0.7"
 libc = "0.2"
 ndarray = "0.12"
 numpy = "0.5"

Cargo.toml

@@ -63,14 +63,15 @@ impl PyEmbeddings {
     ///
     /// This returns words for the analogy query *w1* is to *w2*
     /// as *w3* is to ?.
-    #[args(limit = 10)]
+    #[args(limit = 10, mask = "(true, true, true)")]
     fn analogy(
         &self,
         py: Python,
         word1: &str,
         word2: &str,
         word3: &str,
         limit: usize,
+        mask: (bool, bool, bool),
     ) -> PyResult<Vec<PyObject>> {
         use EmbeddingsWrap::*;
         let embeddings = self.embeddings.borrow();
@@ -83,10 +84,11 @@ impl PyEmbeddings {
             }
         };
 
-        let results = match embeddings.analogy(word1, word2, word3, limit) {
-            Some(results) => results,
-            None => return Err(exceptions::KeyError::py_err("Unknown word or n-grams")),
-        };
+        let results =
+            match embeddings.analogy_masked(word1, word2, word3, limit, [mask.0, mask.1, mask.2]) {
+                Some(results) => results,
+                None => return Err(exceptions::KeyError::py_err("Unknown word or n-grams")),
+            };
 
         let mut r = Vec::with_capacity(results.len());
         for ws in results {

src/embeddings.rs

@@ -0,0 +1,64 @@
+import finalfusion
+
+ANALOGY_ORDER = [
+    "Deutschland",
+    "Westdeutschland",
+    "Sachsen",
+    "Mitteldeutschland",
+    "Brandenburg",
+    "Polen",
+    "Norddeutschland",
+    "Dänemark",
+    "Schleswig-Holstein",
+    "Österreich",
+    "Bayern",
+    "Thüringen",
+    "Bundesrepublik",
+    "Ostdeutschland",
+    "Preußen",
+    "Deutschen",
+    "Hessen",
+    "Potsdam",
+    "Mecklenburg",
+    "Niedersachsen",
+    "Hamburg",
+    "Süddeutschland",
+    "Bremen",
+    "Russland",
+    "Deutschlands",
+    "BRD",
+    "Litauen",
+    "Mecklenburg-Vorpommern",
+    "DDR",
+    "West-Berlin",
+    "Saarland",
+    "Lettland",
+    "Hannover",
+    "Rostock",
+    "Sachsen-Anhalt",
+    "Pommern",
+    "Schweden",
+    "Deutsche",
+    "deutschen",
+    "Westfalen",
+]
+
+def test_analogies():
+    embeds = finalfusion.Embeddings('tests/analogy.fifu')
+    for idx, analogy in enumerate(embeds.analogy("Paris", "Frankreich", "Berlin", 40)):
+        assert ANALOGY_ORDER[idx] == analogy.word
+
+    assert embeds.analogy("Paris", "Frankreich", "Paris", 1, (True, False, True))[0].word == "Frankreich"
+    assert embeds.analogy("Paris", "Frankreich", "Paris", 1, (True, True, True))[0].word != "Frankreich"
+    assert embeds.analogy("Frankreich", "Frankreich", "Frankreich", 1, (False, False, False))[0].word == "Frankreich"
+    assert embeds.analogy("Frankreich", "Frankreich", "Frankreich", 1, (False, False, True))[0].word != "Frankreich"
+    try:
+        embeds.analogy("Paris", "Frankreich", "Paris", 1, (True, True))
+        assert True == False
+    except:
+        ()
+    try:
+        embeds.analogy("Paris", "Frankreich", "Paris", 1, (True, True, True, True))
+        assert True == False
+    except:
+        ()