Update app.py

app.py

@@ -6,6 +6,8 @@ from transformers import DistilBertTokenizer, DistilBertModel, AutoModel, AutoTo
 from langdetect import detect
 from huggingface_hub import snapshot_download
 import os
+from typing import List
+
 
 # Device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -84,34 +86,44 @@ app.add_middleware(
 
 
 class TextIn(BaseModel):
-    text: str
+    texts: List[str]
 
 
+@app.post("/api/predict")
 @app.post("/api/predict")
 def predict(data: TextIn):
-    text = data.text
-    try:
-        lang = detect(text)
-    except:
-        lang = "unknown"
-
-    if lang == "en":
-        tokenizer = english_tokenizer
-        model = english_model
-        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
-        with torch.no_grad():
-            outputs = model(**inputs)
-            probs = torch.sigmoid(outputs).squeeze().cpu().tolist()
-        return {"language": "English", "predictions": dict(zip(english_labels, probs))}
-
-    else:
-        tokenizer = hinglish_tokenizer
-        model = hinglish_model
-        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
-        with torch.no_grad():
-            outputs = model(**inputs)
-            probs = torch.softmax(outputs, dim=1).squeeze().cpu().tolist()
-        return {"language": "Hinglish", "predictions": dict(zip(hinglish_labels, probs))}
+    results = []
+
+    for text in data.texts:
+        try:
+            lang = detect(text)
+        except:
+            lang = "unknown"
+
+        if lang == "en":
+            tokenizer = english_tokenizer
+            model = english_model
+            inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
+            with torch.no_grad():
+                outputs = model(**inputs)
+                probs = torch.sigmoid(outputs).squeeze().cpu().tolist()
+            predictions = dict(zip(english_labels, probs))
+        else:
+            tokenizer = hinglish_tokenizer
+            model = hinglish_model
+            inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
+            with torch.no_grad():
+                outputs = model(**inputs)
+                probs = torch.softmax(outputs, dim=1).squeeze().cpu().tolist()
+            predictions = dict(zip(hinglish_labels, probs))
+
+        results.append({
+            "text": text,
+            "language": lang if lang in ["en", "hi"] else "unknown",
+            "predictions": predictions
+        })
+
+    return {"results": results}
 
 
 @app.get("/")