Fix for embed-multi bug, closes #3

simonw · Oct 26, 2023 · caf8716 · caf8716
1 parent 2905d86
commit caf8716
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/llm_embed_jina.py b/llm_embed_jina.py
@@ -22,5 +22,5 @@ def embed_batch(self, texts):
             self._model = AutoModel.from_pretrained(
                 "jinaai/{}".format(self.model_id), trust_remote_code=True
             )
-        results = self._model.encode(texts)
+        results = self._model.encode(list(texts))
         return (list(map(float, result)) for result in results)
diff --git a/tests/test_embed_jina.py b/tests/test_embed_jina.py
@@ -1,3 +1,5 @@
+from click.testing import CliRunner
+from llm.cli import cli
 import llm
 
 
@@ -6,3 +8,23 @@ def test_jina_embed_small():
     floats = model.embed("hello world")
     assert len(floats) == 512
     assert all(isinstance(f, float) for f in floats)
+
+
+def test_jina_embed_multi(tmpdir):
+    db_path = str(tmpdir / "test.db")
+    runner = CliRunner()
+    result = runner.invoke(
+        cli,
+        [
+            "embed-multi",
+            "-m",
+            "jina-embeddings-v2-small-en",
+            "test",
+            "-",
+            "-d",
+            db_path,
+        ],
+        input='[{"id": "a", "text": "abc"}]',
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 0