run via raw onnxruntime, drop torch/optimum (version conflict)

Dockerfile

@@ -8,27 +8,27 @@ ENV HOME=/home/user \
     HF_HOME=/home/user/.cache/huggingface
 WORKDIR /home/user/app
 
-# CPU-only torch first so transformers/optimum don't drag in the ~2 GB CUDA build.
-RUN pip install --user --no-cache-dir torch==2.* --index-url https://download.pytorch.org/whl/cpu
-
 COPY --chown=user requirements.txt .
 RUN pip install --user --no-cache-dir -r requirements.txt
 
 COPY --chown=user app.py .
 
-# Bake the int8 ONNX model into the image (no cold pull on first request) and
-# assert it loads, pools to 1024-dim, and outputs unit-norm vectors.
+# Bake the int8 ONNX model + tokenizer into the image (no cold pull on first
+# request) and assert it loads, pools to 1024-dim, and emits unit-norm vectors.
 RUN python -c "\
-import torch, torch.nn.functional as F; \
-from optimum.onnxruntime import ORTModelForFeatureExtraction; \
+import numpy as np, onnxruntime as ort; \
+from huggingface_hub import hf_hub_download; \
 from transformers import AutoTokenizer; \
 m='libryo-ai/BAAI-bge-m3-int8'; \
-tok=AutoTokenizer.from_pretrained(m); mdl=ORTModelForFeatureExtraction.from_pretrained(m); \
-e=tok(['ભાવ સમાચાર','mandi prices'],padding=True,truncation=True,max_length=512,return_tensors='pt'); \
-v=F.normalize(mdl(**e).last_hidden_state[:,0],p=2,dim=1); \
-assert v.shape==(2,1024), v.shape; \
-assert abs(float(v[0].norm())-1.0)<1e-3; \
-print('model ok', v.shape)"
+tok=AutoTokenizer.from_pretrained(m); \
+sess=ort.InferenceSession(hf_hub_download(m,'model.onnx'),providers=['CPUExecutionProvider']); \
+names={i.name for i in sess.get_inputs()}; \
+e=tok(['ભાવ સમાચાર','mandi prices'],padding=True,truncation=True,max_length=512,return_tensors='np'); \
+h=sess.run(None,{k:v for k,v in e.items() if k in names})[0]; \
+c=h[:,0]; c=c/np.linalg.norm(c,axis=1,keepdims=True); \
+assert c.shape==(2,1024), c.shape; \
+assert abs(float(np.linalg.norm(c[0]))-1.0)<1e-3; \
+print('model ok', c.shape)"
 
 EXPOSE 7860
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -12,8 +12,8 @@ pinned: false
 
 FastAPI `/embed` service hosting **`libryo-ai/BAAI-bge-m3-int8`** — the int8
 ONNX, dense-only build of bge-m3 (1024-dim, multilingual incl. Hindi/Gujarati).
-~3× faster + ~half the RAM of fp32, negligible accuracy loss. Served via
-optimum `ORTModelForFeatureExtraction` with **CLS pooling + L2-normalize**.
+~3× faster + ~half the RAM of fp32, negligible accuracy loss. Served via raw
+`onnxruntime` (no torch/optimum) with **CLS pooling + L2-normalize**.
 Offloads embedding off the NestJS API box (see `../INFRA-DEPLOYMENT-PLAN.md`).
 Deployed on a **free CPU Basic** Hugging Face Space (2 vCPU / 16 GB).
 

app.py

@@ -1,28 +1,33 @@
 import os
 from contextlib import asynccontextmanager
 
-import torch
-import torch.nn.functional as F
+import numpy as np
+import onnxruntime as ort
 from fastapi import FastAPI, Header, HTTPException
+from huggingface_hub import hf_hub_download
 from pydantic import BaseModel
-from optimum.onnxruntime import ORTModelForFeatureExtraction
 from transformers import AutoTokenizer
 
-# int8 ONNX bge-m3, dense-only. ~3x faster + ~half RAM vs fp32 sentence-transformers.
+# int8 ONNX bge-m3, dense-only. Run raw via onnxruntime — no torch/optimum, so
+# no export-path version conflicts. ~3x faster + ~half RAM vs fp32.
 MODEL_NAME = os.getenv("EMBED_MODEL", "libryo-ai/BAAI-bge-m3-int8")
+ONNX_FILE = os.getenv("EMBED_ONNX_FILE", "model.onnx")
 EMBED_TOKEN = os.getenv("EMBED_TOKEN")  # optional shared secret
 MAX_TOKENS = int(os.getenv("EMBED_MAX_TOKENS", "512"))
 BATCH = int(os.getenv("EMBED_BATCH", "32"))
 
-_model: ORTModelForFeatureExtraction | None = None
+_session: ort.InferenceSession | None = None
 _tokenizer = None
+_input_names: set[str] = set()
 
 
 @asynccontextmanager
 async def lifespan(_: FastAPI):
-    global _model, _tokenizer
+    global _session, _tokenizer, _input_names
     _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    _model = ORTModelForFeatureExtraction.from_pretrained(MODEL_NAME)
+    path = hf_hub_download(MODEL_NAME, ONNX_FILE)
+    _session = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
+    _input_names = {i.name for i in _session.get_inputs()}
     yield
 
 
@@ -41,7 +46,7 @@ class EmbedResponse(BaseModel):
 
 @app.get("/")
 def health():
-    return {"status": "ok", "model": MODEL_NAME, "ready": _model is not None}
+    return {"status": "ok", "model": MODEL_NAME, "ready": _session is not None}
 
 
 def _encode(texts: list[str]) -> list[list[float]]:
@@ -50,14 +55,15 @@ def _encode(texts: list[str]) -> list[list[float]]:
         chunk = texts[i : i + BATCH]
         enc = _tokenizer(
             chunk, padding=True, truncation=True,
-            max_length=MAX_TOKENS, return_tensors="pt",
+            max_length=MAX_TOKENS, return_tensors="np",
         )
-        with torch.no_grad():
-            hidden = _model(**enc).last_hidden_state
-        # bge-m3 dense embedding = CLS token (position 0), then L2-normalize
-        # so cosine == dot product downstream.
-        cls = F.normalize(hidden[:, 0], p=2, dim=1)
-        out.extend(cls.tolist())
+        # bge-m3 (xlm-roberta) has no token_type_ids; feed only what the graph wants.
+        feed = {k: v for k, v in enc.items() if k in _input_names}
+        hidden = _session.run(None, feed)[0]  # (B, T, 1024) last_hidden_state
+        # dense embedding = CLS token (position 0), then L2-normalize so cosine == dot.
+        cls = hidden[:, 0]
+        cls = cls / np.clip(np.linalg.norm(cls, axis=1, keepdims=True), 1e-12, None)
+        out.extend(cls.astype(np.float32).tolist())
     return out
 
 
@@ -67,7 +73,7 @@ def _encode(texts: list[str]) -> list[list[float]]:
 def embed(req: EmbedRequest, authorization: str | None = Header(default=None)):
     if EMBED_TOKEN and authorization != f"Bearer {EMBED_TOKEN}":
         raise HTTPException(status_code=401, detail="unauthorized")
-    if _model is None:
+    if _session is None:
         raise HTTPException(status_code=503, detail="model loading")
     if not req.texts:
         raise HTTPException(status_code=400, detail="texts required")

requirements.txt

@@ -1,4 +1,6 @@
 fastapi==0.115.*
 uvicorn[standard]==0.32.*
-optimum[onnxruntime]==1.*
+onnxruntime==1.*
 transformers==4.*
+huggingface_hub==0.*
+numpy<2