Less indirection and fixed tests

ankane · Oct 27, 2019 · 45fd9ae · 45fd9ae
1 parent 9114ebb
commit 45fd9ae
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 12 deletions.
diff --git a/lib/fasttext/classifier.rb b/lib/fasttext/classifier.rb
@@ -32,11 +32,17 @@ def fit(x, y = nil)
 
     # TODO predict multiple in C++ for performance
     def predict(text, k: 1, threshold: 0.0)
-      if text.is_a?(Array)
-        text.map { |t| predict_one(t, k: k, threshold: threshold) }
-      else
-        predict_one(text, k: k, threshold: threshold)
-      end
+      multiple = text.is_a?(Array)
+      text = [text] unless multiple
+
+      result =
+        text.map do |t|
+          m.predict(prep_text(t), k, threshold).map do |v|
+            [remove_prefix(v[1]), v[0]]
+          end.to_h
+        end
+
+      multiple ? result : result.first
     end
 
     def test(x, y = nil, k: 1)
@@ -66,12 +72,6 @@ def labels(include_freq: false)
 
     private
 
-    def predict_one(text, k:, threshold:)
-      m.predict(prep_text(text), k, threshold).map do |v|
-        [remove_prefix(v[1]), v[0]]
-      end.to_h
-    end
-
     def input_path(x, y)
       if x.is_a?(String)
         raise ArgumentError, "Cannot pass y with file" if y

diff --git a/test/classifier_test.rb b/test/classifier_test.rb
@@ -26,7 +26,7 @@ def test_works
     assert model.sentence_vector("first document")
 
     assert model.predict("First document")
-    asert model.predict(["First document", "Second document"], k: 3)
+    assert model.predict(["First document", "Second document"], k: 3)
 
     # TODO fix flaky test
     # pred = model.predict("First document").first