Bug in binary

INGEOTEC · Jan 31, 2024 · 6159d51 · 6159d51
1 parent be9e76b
commit 6159d51
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/EvoMSA/back_prop.py b/EvoMSA/back_prop.py
@@ -185,7 +185,11 @@ def decision_function(self, D: List[Union[dict, list]]) -> np.ndarray:
 
     def predict(self, D: List[Union[dict, list]]) -> np.ndarray:
         df = self.decision_function(D)
-        return self.classes_[df.argmax(axis=1)]
+        if df.shape[1] == 1:
+            index = np.where(df.flatten() > 0, 1, 0)
+        else:
+            index = df.argmax(axis=1)
+        return self.classes_[index]
 
     @staticmethod
     def array(data):

diff --git a/EvoMSA/tests/test_back_prop.py b/EvoMSA/tests/test_back_prop.py
@@ -48,8 +48,15 @@ def test_binary():
     D = list(tweet_iterator(TWEETS))
     D = [x for x in D if x['klass'] in {'N', 'P'}]
     bow = BoWBP(lang='es').fit(D)
+    hy = bow.predict(D)
+    acc = (np.array([x['klass'] for x in D]) == hy).mean()
+    assert acc > 0.9
     dense = DenseBoWBP(lang='es',
                        voc_size_exponent=13).fit(D)
+    hy = dense.predict(D)
+    acc = (np.array([x['klass'] for x in D]) == hy).mean()
+    assert acc > 0.85
+
 
 
 # def test_BoWBP_validation_set():