Update app.py

app.py

@@ -1,16 +1,24 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from sentence_transformers import SentenceTransformer
-import os
+from transformers import AutoTokenizer, AutoModel
+import torch
+import torch.nn.functional as F
 
 app = FastAPI()
 
-model = SentenceTransformer("thenlper/gte-small")
+# Charger le modèle depuis HF sans passer par SentenceTransformer
+MODEL_NAME = "thenlper/gte-small"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModel.from_pretrained(MODEL_NAME)
 
 class EmbedInput(BaseModel):
     text: str
 
 @app.post("/embed")
 async def embed_text(payload: EmbedInput):
-    embedding = model.encode(payload.text, normalize_embeddings=True)
-    return {"embedding": embedding.tolist()}
+    inputs = tokenizer(payload.text, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        embeddings = outputs.last_hidden_state[:, 0]  # CLS token
+        normalized = F.normalize(embeddings, p=2, dim=1)
+    return {"embedding": normalized[0].tolist()}