Upload app.py with huggingface_hub

app.py

@@ -13,12 +13,9 @@ app = FastAPI()
 
 
 class EmbedRequest(BaseModel):
-    text: str | list[str]
-    normalize: bool = True
-
-
-class BatchEmbedRequest(BaseModel):
-    texts: list[str]
+    text: str | list[str] | None = None
+    texts: list[str] | None = None
+    model: str | None = None
     normalize: bool = True
 
 
@@ -29,20 +26,34 @@ def health():
 
 @app.post("/embed")
 def embed(req: EmbedRequest):
-    texts = [req.text] if isinstance(req.text, str) else req.text
-    embeddings = model.encode(texts, normalize_embeddings=req.normalize)
+    # Accept both "text" (single/list) and "texts" (list) fields
+    if req.texts:
+        input_texts = req.texts
+    elif req.text:
+        input_texts = [req.text] if isinstance(req.text, str) else req.text
+    else:
+        return {"error": "Provide 'text' or 'texts' field"}, 400
+    embeddings = model.encode(input_texts, normalize_embeddings=req.normalize)
     return {
         "embeddings": embeddings.tolist(),
         "model": MODEL_NAME,
         "dimensions": embeddings.shape[1],
+        "tokens": len(input_texts) * 32,
     }
 
 
 @app.post("/embed_batch")
-def embed_batch(req: BatchEmbedRequest):
-    embeddings = model.encode(req.texts, normalize_embeddings=req.normalize)
+def embed_batch(req: EmbedRequest):
+    if req.texts:
+        input_texts = req.texts
+    elif req.text:
+        input_texts = [req.text] if isinstance(req.text, str) else req.text
+    else:
+        return {"error": "Provide 'text' or 'texts' field"}, 400
+    embeddings = model.encode(input_texts, normalize_embeddings=req.normalize)
     return {
         "embeddings": embeddings.tolist(),
         "model": MODEL_NAME,
         "dimensions": embeddings.shape[1],
+        "tokens": len(input_texts) * 32,
     }