Add a test for custom metric on BruteKNN

msluszniak · May 14, 2024 · 582b220 · 582b220
1 parent f358b24
commit 582b220
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/test/scholar/neighbors/brute_knn_test.exs b/test/scholar/neighbors/brute_knn_test.exs
@@ -119,5 +119,10 @@ defmodule Scholar.Neighbors.BruteKNNTest do
 
       assert distances_pred == distances_true
     end
+
+    test "custom metric" do
+      model = BruteKNN.fit(data(), num_neighbors: 3, batch_size: 1, metric: :cosine)
+      assert {_, _} = BruteKNN.predict(model, query())
+    end
   end
 end