feat: switch to ONNX runtime for embeddings

app/config.py

@@ -73,7 +73,8 @@ DATASET_CONFIGS = {
     ),
 }
 
-# Default configuration for embedding model
+# Default configuration for embedding model (ONNX)
 MODEL_CONFIG = {
-    "embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
+    "embedding_model": "Xenova/all-MiniLM-L6-v2",  # Pre-exported ONNX model
+    "use_onnx": True,
 }

app/pipeline.py

@@ -1,13 +1,118 @@
 from haystack import Document, Pipeline
 from haystack.document_stores.in_memory import InMemoryDocumentStore
-from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder
 from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
 from haystack.components.builders import PromptBuilder
 from datasets import load_dataset
-from haystack.dataclasses import ChatMessage
 from typing import Optional, List, Union, Dict
+import numpy as np
+import logging
+import os
+
 from .config import DatasetConfig, DATASET_CONFIGS, MODEL_CONFIG
 
+logger = logging.getLogger(__name__)
+
+
+class ONNXTextEmbedder:
+    """ONNX-based text embedder using optimum."""
+    
+    def __init__(self, model_name: str):
+        from optimum.onnxruntime import ORTModelForFeatureExtraction
+        from transformers import AutoTokenizer
+        
+        logger.info(f"Loading ONNX model: {model_name}")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        # Load pre-exported ONNX model (no export needed)
+        self.model = ORTModelForFeatureExtraction.from_pretrained(
+            model_name, 
+            file_name="onnx/model.onnx",
+            provider="CPUExecutionProvider"
+        )
+        logger.info("ONNX model loaded successfully")
+    
+    def _mean_pooling(self, model_output, attention_mask):
+        token_embeddings = model_output[0]
+        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+    
+    def run(self, text: str):
+        import torch
+        
+        encoded = self.tokenizer(
+            text, 
+            padding=True, 
+            truncation=True, 
+            max_length=512, 
+            return_tensors='pt'
+        )
+        
+        with torch.no_grad():
+            model_output = self.model(**encoded)
+        
+        embeddings = self._mean_pooling(model_output, encoded['attention_mask'])
+        embedding = embeddings[0].numpy().tolist()
+        
+        return {"embedding": embedding}
+    
+    def warm_up(self):
+        pass
+
+
+class ONNXDocumentEmbedder:
+    """ONNX-based document embedder using optimum."""
+    
+    def __init__(self, model_name: str, progress_bar: bool = False):
+        from optimum.onnxruntime import ORTModelForFeatureExtraction
+        from transformers import AutoTokenizer
+        
+        logger.info(f"Loading ONNX model: {model_name}")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        # Load pre-exported ONNX model (no export needed)
+        self.model = ORTModelForFeatureExtraction.from_pretrained(
+            model_name, 
+            file_name="onnx/model.onnx",
+            provider="CPUExecutionProvider"
+        )
+        self.progress_bar = progress_bar
+        logger.info("ONNX model loaded successfully")
+    
+    def _mean_pooling(self, model_output, attention_mask):
+        token_embeddings = model_output[0]
+        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+    
+    def run(self, documents):
+        import torch
+        from haystack import Document
+        
+        texts = [doc.content for doc in documents]
+        
+        encoded = self.tokenizer(
+            texts, 
+            padding=True, 
+            truncation=True, 
+            max_length=512, 
+            return_tensors='pt'
+        )
+        
+        with torch.no_grad():
+            model_output = self.model(**encoded)
+        
+        embeddings = self._mean_pooling(model_output, encoded['attention_mask'])
+        
+        docs_with_embeddings = []
+        for doc, embedding in zip(documents, embeddings.numpy()):
+            doc_with_emb = Document(content=doc.content, meta=doc.meta, embedding=embedding.tolist())
+            docs_with_embeddings.append(doc_with_emb)
+        
+        return {"documents": docs_with_embeddings}
+    
+    def warm_up(self):
+        pass
+
+
 class RAGPipeline:
     def __init__(
         self,
@@ -53,22 +158,19 @@ class RAGPipeline:
                 )
                 self.documents.append(document)
 
- # Documents loaded silently - remove verbose output
-        
-        # Initialize components
+        # Initialize components with ONNX
         self.document_store = InMemoryDocumentStore()
-        self.doc_embedder = SentenceTransformersDocumentEmbedder(
-            model=embedding_model or MODEL_CONFIG["embedding_model"],
-            progress_bar=False
-        )
-        self.text_embedder = SentenceTransformersTextEmbedder(
-            model=embedding_model or MODEL_CONFIG["embedding_model"],
+        
+        model_name = embedding_model or MODEL_CONFIG["embedding_model"]
+        
+        self.doc_embedder = ONNXDocumentEmbedder(
+            model_name=model_name,
             progress_bar=False
         )
-        self.text_embedder = SentenceTransformersTextEmbedder(
-            model=embedding_model or MODEL_CONFIG["embedding_model"],
-            progress_bar=False
+        self.text_embedder = ONNXTextEmbedder(
+            model_name=model_name
         )
+        
         self.retriever = InMemoryEmbeddingRetriever(self.document_store)
         
         # Warm up the embedders

requirements.txt

@@ -1,6 +1,9 @@
 haystack-ai==2.10.3
 datasets==3.3.2
-sentence-transformers==3.4.1
+optimum[onnxruntime]==1.23.3
+onnxruntime==1.20.1
+onnxscript>=0.1.0
+transformers==4.46.3
 fastapi==0.115.4
 uvicorn==0.31.0
 openai==1.57.0