diff --git a/docs/web/integrations/embedding/openai.mdx b/docs/web/integrations/embedding/openai.mdx index fbfa62d9d..4dd1ee8ba 100644 --- a/docs/web/integrations/embedding/openai.mdx +++ b/docs/web/integrations/embedding/openai.mdx @@ -10,7 +10,7 @@ LanceDB registers the OpenAI embeddings function in the registry by default, as | Parameter | Type | Default Value | Description | |---|---|---|---| | `name` | `str` | `"text-embedding-ada-002"` | The name of the model. | -| `dim` | `int` | Model default | Requested output dimensions for models that support shortening embeddings. Use `func.ndims()` when defining the vector field. | +| `dim` | `int` | Model default | Requested output dimensions for models that support shortening embeddings. Required for model names other than the OpenAI ones listed by `model_names()` (for example a model served behind an OpenAI-compatible `base_url`), where it is the model's embedding dimension. Use `func.ndims()` when defining the vector field. | | `use_azure` | bool | `False` | Set true to use Azure OpenAI SDK | Set your OpenAI API key before running the example: diff --git a/python/python/lancedb/embeddings/openai.py b/python/python/lancedb/embeddings/openai.py index d4ecbdf3f..8592598cc 100644 --- a/python/python/lancedb/embeddings/openai.py +++ b/python/python/lancedb/embeddings/openai.py @@ -36,6 +36,7 @@ class OpenAIEmbeddings(TextEmbeddingFunction): embedding_function = openai.create( name="", base_url="http://localhost:11434", + dim=768, # the model's embedding dimension; required for unlisted names ) ``` @@ -77,8 +78,15 @@ class OpenAIEmbeddings(TextEmbeddingFunction): return self.dim or 3072 elif self.name == "text-embedding-3-small": return self.dim or 1536 + elif self.dim: + # A name not listed above, e.g. a model served behind an + # OpenAI-compatible `base_url`: `dim` is its embedding size + return self.dim else: - raise ValueError(f"Unknown model name {self.name}") + raise ValueError( + f"Unknown model name {self.name}: set `dim` to the model's " + "embedding dimension" + ) def generate_embeddings( self, texts: Union[List[str], "np.ndarray"] diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 32503c04f..c18ddd644 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -602,6 +602,51 @@ def test_openai_propagates_api_key(monkeypatch): assert len(actual.text) > 0 +@pytest.mark.parametrize("base_url", [None, "http://localhost:11434/v1"]) +def test_openai_compatible_model_uses_dim_for_ndims(base_url): + # Unlisted model names use the explicit dimension independently of base_url. + openai = get_registry().get("openai") + func = openai.create( + name="nomic-embed-text", + base_url=base_url, + dim=768, + ) + assert func.ndims() == 768 + + class Schema(LanceModel): + text: str = func.SourceField() + vector: Vector(func.ndims()) = func.VectorField() + + schema = Schema.to_arrow_schema() + assert schema.field("vector").type == pa.list_(pa.float32(), 768) + + +def test_openai_unknown_model_without_dim_raises_clear_error(): + openai = get_registry().get("openai") + func = openai.create( + name="nomic-embed-text", + base_url="http://localhost:11434/v1", + ) + with pytest.raises(ValueError, match="dim"): + func.ndims() + + +@pytest.mark.parametrize( + ("name", "dim", "expected_dim"), + [ + ("text-embedding-ada-002", None, 1536), + ("text-embedding-ada-002", 512, 1536), + ("text-embedding-3-large", None, 3072), + ("text-embedding-3-large", 1024, 1024), + ("text-embedding-3-small", None, 1536), + ("text-embedding-3-small", 512, 512), + ], +) +def test_openai_known_model_ndims(name, dim, expected_dim): + func = get_registry().get("openai").create(name=name, dim=dim) + assert func.ndims() == expected_dim + + @patch("time.sleep") def test_openai_no_retry_on_401(mock_sleep): """