diff --git a/python/python/lancedb/embeddings/instructor.py b/python/python/lancedb/embeddings/instructor.py index 675a0139c..37ae1c296 100644 --- a/python/python/lancedb/embeddings/instructor.py +++ b/python/python/lancedb/embeddings/instructor.py @@ -101,8 +101,7 @@ class InstructorEmbeddingFunction(TextEmbeddingFunction): @weak_lru(maxsize=1) def ndims(self): - model = self.get_model() - return model.encode("foo").shape[0] + return len(self.generate_embeddings([[self.source_instruction, "foo"]])[0]) def compute_query_embeddings(self, query: str, *args, **kwargs) -> List[np.array]: return self.generate_embeddings([[self.query_instruction, query]]) diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index f2d8971c3..678270f19 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -64,6 +64,23 @@ def test_embedding_function(tmp_path): assert np.allclose(actual, expected) +def test_instructor_ndims_uses_instruction(): + instructor = get_registry().get("instructor").create() + model = MagicMock() + model.encode.return_value = np.zeros((1, 384)) + + with patch.object(type(instructor), "get_model", return_value=model): + assert instructor.ndims() == 384 + + model.encode.assert_called_once_with( + [[instructor.source_instruction, "foo"]], + batch_size=instructor.batch_size, + show_progress_bar=instructor.show_progress_bar, + normalize_embeddings=instructor.normalize_embeddings, + device=instructor.device, + ) + + def test_embedding_function_variables(): @register("variable-testing") class VariableTestingFunction(TextEmbeddingFunction):