Pull Ollama models. Try to make fastembed more efficient.

Author: double16

requirements.txt

@@ -1,16 +1,17 @@
-haystack-ai~=2.21.0
-qdrant-haystack~=9.4.0
-fastembed-haystack~=1.5.0
+haystack-ai~=2.22.0
+qdrant-haystack~=10.0.0
+fastembed-haystack~=2.0.0
 numpy<2
-ollama-haystack~=5.3.0
-amazon-bedrock-haystack~=5.3.1
-openai~=2.14.0
-google-genai-haystack~=2.3.0
+ollama-haystack~=6.0.0
+amazon-bedrock-haystack~=6.2.0
+openai~=2.15.0
+google-genai-haystack~=3.1.0
+haystack-experimental~=0.16.0
 mcp-haystack~=1.1.0
 chardet~=5.2.0
 sentence-transformers~=5.2.0
 prompt_toolkit~=3.0.51
-mcp[cli]~=1.23.1
+mcp[cli]~=1.25.0
 httpx~=0.28.1
 uv~=0.9.15
 tldextract~=5.3.0
@@ -26,7 +27,7 @@ lxml~=6.0.0
 tinycss2~=1.5.1
 json5~=0.13.0
 html5lib~=1.1
-optimum~=2.1.0
+optimum[onnxruntime]~=2.1.0
 ddgs~=9.10.0
 pycryptodome~=3.23.0
 psutil~=7.2.1

shyhurricane/generator_config.py

@@ -1,6 +1,7 @@
 import argparse
 import logging
 import os
+from math import ceil
 from typing import Optional, Dict, Any, Union, List
 
 import requests
@@ -30,6 +31,7 @@
 from pydantic import BaseModel, Field
 
 from shyhurricane.doc_type_model_map import ModelConfig
+from shyhurricane.utils import process_cpu_count
 
 logger = logging.getLogger(__name__)
 
@@ -235,6 +237,14 @@ def from_env():
         )
         return generator_config
 
+    def ollama_url(self) -> str:
+        return "http://" + (self.ollama_host or OLLAMA_HOST_DEFAULT)
+
+    def ollama_pull(self, model_id: str):
+        model, tag = model_id.rsplit(":", maxsplit=1)
+        r = requests.post(f"{self.ollama_url()}/api/pull", json={"model": model, "tag": tag, "force": False})
+        r.raise_for_status()
+
     def apply_reasoning_default(self):
         self.ollama_host = self.ollama_host or OLLAMA_HOST_DEFAULT
         if self.ollama_model or self.gemini_model or self.openai_model or self.bedrock_model:
@@ -325,8 +335,9 @@ def create_chat_generator(self,
                 # https://huggingface.co/docs/inference-providers/guides/gpt-oss
                 _generation_kwargs["effort"] = "high"
             logger.info("Using Ollama chat with model %s at %s", self.ollama_model, self.ollama_host)
+            self.ollama_pull(self.ollama_model)
             return OllamaChatGenerator(
-                url="http://" + (self.ollama_host or OLLAMA_HOST_DEFAULT),
+                url=self.ollama_url(),
                 model=self.ollama_model,
                 timeout=ollama_timeout,
                 generation_kwargs=_generation_kwargs | (generation_kwargs or {}),
@@ -374,8 +385,9 @@ def create_generator(self,
                 "temperature": temperature or self.temperature,
             }
             logger.info("Using Ollama generator with model %s at %s", self.ollama_model, self.ollama_host)
+            self.ollama_pull(self.ollama_model)
             return OllamaGenerator(
-                url="http://" + (self.ollama_host or OLLAMA_HOST_DEFAULT),
+                url=self.ollama_url(),
                 model=self.ollama_model,
                 generation_kwargs=_generation_kwargs | (generation_kwargs or {}),
             )
@@ -397,7 +409,7 @@ def _embedder_enable_ollama(self) -> bool:
         # v0.12.11, v0.13.0 - macos has use after free failures
         # v0.14.0 - macos embedding is working
         try:
-            resp_version = requests.get("http://" + (self.ollama_host or OLLAMA_HOST_DEFAULT) + "/api/version")
+            resp_version = requests.get(self.ollama_url() + "/api/version")
             resp_version.raise_for_status()
             version = float(".".join(resp_version.json()["version"].split(".")[0:2]))
             return version >= 0.14
@@ -466,9 +478,10 @@ def create_document_embedder(self, model_config: ModelConfig):
             )
         elif self.ollama_model and self._embedder_enable_ollama():
             logger.info("Using Ollama document embedder with model %s at %s", model_path, self.ollama_host)
+            self.ollama_pull(model_path)
             return OllamaDocumentEmbedder(
                 model=model_path,
-                url="http://" + (self.ollama_host or OLLAMA_HOST_DEFAULT),
+                url=self.ollama_url(),
                 progress_bar=False,
             )
 
@@ -501,9 +514,10 @@ def create_text_embedder(self, model_config: ModelConfig):
             )
         elif self.ollama_model and self._embedder_enable_ollama():
             logger.info("Using Ollama text embedder with model %s at %s", model_path, self.ollama_host)
+            self.ollama_pull(model_path)
             return OllamaTextEmbedder(
                 model=model_path,
-                url="http://" + (self.ollama_host or OLLAMA_HOST_DEFAULT),
+                url=self.ollama_url(),
             )
 
         logger.info("Using local text embedder with model %s", model_path)
@@ -528,14 +542,18 @@ def create_sparse_document_embedder(self, model_config: ModelConfig):
         return FastembedSparseDocumentEmbedder(
             model=model_config.model_name,
             cache_dir=self._fastembed_cache_dir(),
-            batch_size=1,
+            threads=max(1, ceil(process_cpu_count() / 2)),
+            batch_size=32,
+            parallel=0,
             progress_bar=False,
         )
 
     def create_sparse_text_embedder(self, model_config: ModelConfig):
         return FastembedSparseTextEmbedder(
             model=model_config.model_name,
             cache_dir=self._fastembed_cache_dir(),
+            threads=max(1, ceil(process_cpu_count() / 2)),
+            parallel=0,
             progress_bar=False,
         )
 

shyhurricane/utils.py

@@ -721,3 +721,15 @@ def coerce_to_dict(value: Any, kv_sep: str = None, element_sep: str = None) -> D
         it = iter(value)
         return dict(zip_longest(it, it, fillvalue=None))
     return {str(value): ""}
+
+
+def process_cpu_count() -> int:
+    try:
+        return os.process_cpu_count()
+    except Exception:
+        pass
+    try:
+        return len(os.sched_getaffinity(0))
+    except Exception:
+        pass
+    return os.cpu_count()