fix: use onnxruntime directly without torch dependency

app/pipeline.py

@@ -1,4 +1,4 @@
-from haystack import Document, Pipeline
+from haystack import Document, Pipeline, component
 from haystack.document_stores.in_memory import InMemoryDocumentStore
 from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
 from haystack.components.builders import PromptBuilder
@@ -7,86 +7,103 @@ from typing import Optional, List, Union, Dict
 import numpy as np
 import logging
 import os
+import onnxruntime as ort
+from transformers import AutoTokenizer
 
 from .config import DatasetConfig, DATASET_CONFIGS, MODEL_CONFIG
 
 logger = logging.getLogger(__name__)
 
 
+@component
 class ONNXTextEmbedder:
-    """ONNX-based text embedder using optimum."""
+    """ONNX-based text embedder using onnxruntime directly."""
     
     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"
+        # Download and cache the ONNX model
+        from huggingface_hub import hf_hub_download
+        onnx_path = hf_hub_download(
+            repo_id=model_name,
+            filename="onnx/model.onnx"
+        )
+        
+        # Create ONNX Runtime session
+        self.session = ort.InferenceSession(
+            onnx_path,
+            providers=['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 _mean_pooling(self, token_embeddings: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
+        """Perform mean pooling on token embeddings using numpy."""
+        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        sum_mask = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
+        return sum_embeddings / sum_mask
     
+    @component.output_types(embedding=List[float])
     def run(self, text: str):
-        import torch
-        
         encoded = self.tokenizer(
             text, 
             padding=True, 
             truncation=True, 
             max_length=512, 
-            return_tensors='pt'
+            return_tensors='np'
         )
         
-        with torch.no_grad():
-            model_output = self.model(**encoded)
+        # Run ONNX inference
+        inputs = {
+            'input_ids': encoded['input_ids'].astype(np.int64),
+            'attention_mask': encoded['attention_mask'].astype(np.int64)
+        }
         
-        embeddings = self._mean_pooling(model_output, encoded['attention_mask'])
-        embedding = embeddings[0].numpy().tolist()
+        if 'token_type_ids' in encoded:
+            inputs['token_type_ids'] = encoded['token_type_ids'].astype(np.int64)
+        
+        outputs = self.session.run(None, inputs)
+        token_embeddings = outputs[0]
+        
+        embeddings = self._mean_pooling(token_embeddings, encoded['attention_mask'])
+        embedding = embeddings[0].tolist()
         
         return {"embedding": embedding}
-    
-    def warm_up(self):
-        pass
 
 
+@component
 class ONNXDocumentEmbedder:
-    """ONNX-based document embedder using optimum."""
+    """ONNX-based document embedder using onnxruntime directly."""
     
     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"
+        # Download and cache the ONNX model
+        from huggingface_hub import hf_hub_download
+        onnx_path = hf_hub_download(
+            repo_id=model_name,
+            filename="onnx/model.onnx"
+        )
+        
+        # Create ONNX Runtime session
+        self.session = ort.InferenceSession(
+            onnx_path,
+            providers=['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 _mean_pooling(self, token_embeddings: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
+        """Perform mean pooling on token embeddings using numpy."""
+        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        sum_mask = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
+        return sum_embeddings / sum_mask
     
-    def run(self, documents):
-        import torch
-        from haystack import Document
-        
+    @component.output_types(documents=List[Document])
+    def run(self, documents: List[Document]):
         texts = [doc.content for doc in documents]
         
         encoded = self.tokenizer(
@@ -94,23 +111,33 @@ class ONNXDocumentEmbedder:
             padding=True, 
             truncation=True, 
             max_length=512, 
-            return_tensors='pt'
+            return_tensors='np'
         )
         
-        with torch.no_grad():
-            model_output = self.model(**encoded)
+        # Run ONNX inference
+        inputs = {
+            'input_ids': encoded['input_ids'].astype(np.int64),
+            'attention_mask': encoded['attention_mask'].astype(np.int64)
+        }
+        
+        if 'token_type_ids' in encoded:
+            inputs['token_type_ids'] = encoded['token_type_ids'].astype(np.int64)
         
-        embeddings = self._mean_pooling(model_output, encoded['attention_mask'])
+        outputs = self.session.run(None, inputs)
+        token_embeddings = outputs[0]
+        
+        embeddings = self._mean_pooling(token_embeddings, 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())
+        for doc, embedding in zip(documents, embeddings):
+            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:
@@ -173,10 +200,6 @@ class RAGPipeline:
         
         self.retriever = InMemoryEmbeddingRetriever(self.document_store)
         
-        # Warm up the embedders
-        self.doc_embedder.warm_up()
-        self.text_embedder.warm_up()
-        
         # Initialize prompt template
         self.prompt_builder = PromptBuilder(template=self.config.prompt_template or """
         Given the following context, please answer the question.

requirements.txt

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