feat/ ollama embedder

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.3.7-dev1
+
+### Features
+
+* **Add Ollama embedder** Adds support for creating embeddings vector using Ollama
+
 ## 0.3.6
 
 ### Fixes

requirements/embed/ollama.in

@@ -0,0 +1,3 @@
+-c ../common/constraints.txt
+
+ollama

requirements/embed/ollama.txt

@@ -0,0 +1,37 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile ./embed/ollama.in --output-file ./embed/ollama.txt --no-strip-extras --python-version 3.9
+annotated-types==0.7.0
+    # via pydantic
+anyio==4.6.2.post1
+    # via httpx
+certifi==2024.8.30
+    # via
+    #   httpcore
+    #   httpx
+exceptiongroup==1.2.2
+    # via anyio
+h11==0.14.0
+    # via httpcore
+httpcore==1.0.7
+    # via httpx
+httpx==0.27.2
+    # via ollama
+idna==3.10
+    # via
+    #   anyio
+    #   httpx
+ollama==0.4.2
+    # via -r ./embed/ollama.in
+pydantic==2.10.3
+    # via ollama
+pydantic-core==2.27.1
+    # via pydantic
+sniffio==1.3.1
+    # via
+    #   anyio
+    #   httpx
+typing-extensions==4.12.2
+    # via
+    #   anyio
+    #   pydantic
+    #   pydantic-core

test/integration/embedders/test_ollama.py

@@ -0,0 +1,43 @@
+import json
+import os
+from pathlib import Path
+
+from test.integration.embedders.utils import validate_embedding_output, validate_raw_embedder
+from test.integration.utils import requires_env
+from unstructured_ingest.embed.ollama import (
+    OllamaEmbeddingConfig,
+    OllamaEmbeddingEncoder,
+)
+from unstructured_ingest.v2.processes.embedder import Embedder, EmbedderConfig
+
+API_KEY = "OLLAMA_API_KEY"
+
+
+def get_api_key() -> str:
+    api_key = os.getenv(API_KEY, None)
+    assert api_key
+    return api_key
+
+
+@requires_env(API_KEY)
+def test_ollama_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder_config = EmbedderConfig(embedding_provider="ollama", embedding_api_key=api_key)
+    embedder = Embedder(config=embedder_config)
+    results = embedder.run(elements_filepath=embedder_file)
+    assert results
+    with embedder_file.open("r") as f:
+        original_elements = json.load(f)
+    validate_embedding_output(original_elements=original_elements, output_elements=results)
+
+
+@requires_env(API_KEY)
+def test_raw_ollama_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder = OllamaEmbeddingEncoder(config=OllamaEmbeddingConfig(api_key=api_key))
+    validate_raw_embedder(
+        embedder=embedder,
+        embedder_file=embedder_file,
+        expected_dimensions=(768,),
+        expected_is_unit_vector=False,
+    )

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.6"  # pragma: no cover
+__version__ = "0.3.7-dev1"  # pragma: no cover

unstructured_ingest/cli/interfaces.py

@@ -421,6 +421,7 @@ def get_cli_options() -> t.List[click.Option]:
             "vertexai",
             "voyageai",
             "octoai",
+            "ollama",
         ]
         options = [
             click.Option(

unstructured_ingest/embed/ollama.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import numpy as np
+from pydantic import Field
+
+from unstructured_ingest.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from ollama import embed as OllamaClient
+
+
+class OllamaEmbeddingConfig(EmbeddingConfig):
+    embedder_model_name: Optional[str] = Field(default="all-minilm", alias="model_name")
+
+    @requires_dependencies(
+        ["ollama"],
+        extras="embed-ollama",
+    )
+    def get_client(self) -> "OllamaClient":
+        from ollama import embed as OllamaClient
+
+        return OllamaClient
+
+
+@dataclass
+class OllamaEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: OllamaEmbeddingConfig
+
+    def get_exemplary_embedding(self) -> list[float]:
+        return self.embed_query(query="Q")
+
+    def num_of_dimensions(self) -> tuple[int, ...]:
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.shape(exemplary_embedding)
+
+    def is_unit_vector(self) -> bool:
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
+
+    def embed_query(self, query: str) -> list[float]:
+        return self._embed_documents(texts=[query])[0]
+
+    def _embed_documents(self, texts: list[str]) -> list[list[float]]:
+        client = self.config.get_client()
+        _r = client(model=self.config.embedder_model_name, input=texts)
+        return _r["embeddings"]
+
+    def embed_documents(self, elements: list[dict]) -> list[dict]:
+        embeddings = self._embed_documents([e.get("text", "") for e in elements])
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings
+
+    def _add_embeddings_to_elements(self, elements: list[dict], embeddings: list) -> list[dict]:
+        assert len(elements) == len(embeddings)
+        elements_w_embedding = []
+
+        for i, element in enumerate(elements):
+            element["embeddings"] = embeddings[i]
+            elements_w_embedding.append(element)
+        return elements

unstructured_ingest/interfaces.py

@@ -212,20 +212,23 @@ def get_embedder(self) -> "BaseEmbeddingEncoder":
             )
 
             return OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig(**kwargs))
+
         elif self.provider == "huggingface":
             from unstructured_ingest.embed.huggingface import (
                 HuggingFaceEmbeddingConfig,
                 HuggingFaceEmbeddingEncoder,
             )
 
             return HuggingFaceEmbeddingEncoder(config=HuggingFaceEmbeddingConfig(**kwargs))
+
         elif self.provider == "octoai":
             from unstructured_ingest.embed.octoai import (
                 OctoAiEmbeddingConfig,
                 OctoAIEmbeddingEncoder,
             )
 
             return OctoAIEmbeddingEncoder(config=OctoAiEmbeddingConfig(**kwargs))
+
         elif self.provider == "aws-bedrock":
             from unstructured_ingest.embed.bedrock import (
                 BedrockEmbeddingConfig,
@@ -239,20 +242,31 @@ def get_embedder(self) -> "BaseEmbeddingEncoder":
                     region_name=self.aws_region,
                 )
             )
+
         elif self.provider == "vertexai":
             from unstructured_ingest.embed.vertexai import (
                 VertexAIEmbeddingConfig,
                 VertexAIEmbeddingEncoder,
             )
 
             return VertexAIEmbeddingEncoder(config=VertexAIEmbeddingConfig(**kwargs))
+
         elif self.provider == "voyageai":
             from unstructured_ingest.embed.voyageai import (
                 VoyageAIEmbeddingConfig,
                 VoyageAIEmbeddingEncoder,
             )
 
             return VoyageAIEmbeddingEncoder(config=VoyageAIEmbeddingConfig(**kwargs))
+
+        elif self.provider == "ollama":
+            from unstructured_ingest.embed.ollama import (
+                OllamaEmbeddingConfig,
+                OllamaEmbeddingEncoder,
+            )
+
+            return OllamaEmbeddingEncoder(config=OllamaEmbeddingConfig(**kwargs))
+
         else:
             raise ValueError(f"{self.provider} not a recognized encoder")
 

unstructured_ingest/v2/processes/embedder.py

@@ -24,6 +24,7 @@ class EmbedderConfig(BaseModel):
             "octoai",
             "mixedbread-ai",
             "togetherai",
+            "ollama",
         ]
     ] = Field(default=None, description="Type of the embedding class to be used.")
     embedding_api_key: Optional[SecretStr] = Field(
@@ -146,6 +147,11 @@ def get_togetherai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncod
             config=TogetherAIEmbeddingConfig.model_validate(embedding_kwargs)
         )
 
+    def get_ollama_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured_ingest.embed.ollama import OllamaEmbeddingConfig, OllamaEmbeddingEncoder
+
+        return OllamaEmbeddingEncoder(config=OllamaEmbeddingConfig.model_validate(embedding_kwargs))
+
     def get_embedder(self) -> "BaseEmbeddingEncoder":
         kwargs: dict[str, Any] = {}
         if self.embedding_api_key:
@@ -176,6 +182,8 @@ def get_embedder(self) -> "BaseEmbeddingEncoder":
             return self.get_togetherai_embedder(embedding_kwargs=kwargs)
         if self.embedding_provider == "azure-openai":
             return self.get_azure_openai_embedder(embedding_kwargs=kwargs)
+        if self.embedding_provider == "ollama":
+            return self.get_ollama_embedder(embedding_kwargs=kwargs)
 
         raise ValueError(f"{self.embedding_provider} not a recognized encoder")