Update app.py

app.py

@@ -31,7 +31,7 @@ async def predict(request: TextRequest):
         inputs = tokenizer.encode_plus(
             request.text,
             add_special_tokens=True,
-            max_length=64,
+            max_length=128,
             truncation=True,
             padding='max_length',
             return_attention_mask=True,
@@ -50,11 +50,11 @@ async def batch_predict(request: BatchTextRequest):
     try:
         model.eval()
         results = []
-        for text in request.texts:
+        for idx, text in enumerate(request.texts):
             inputs = tokenizer.encode_plus(
                 text,
                 add_special_tokens=True,
-                max_length=64,
+                max_length=128,
                 truncation=True,
                 padding='max_length',
                 return_attention_mask=True,
@@ -63,11 +63,12 @@ async def batch_predict(request: BatchTextRequest):
             with torch.no_grad():
                 logits = model(inputs['input_ids'], attention_mask=inputs['attention_mask']).logits
                 prediction = torch.argmax(logits, dim=1).item()
-                results.append({"text": text, "prediction": "Spam" if prediction == 1 else "Ham"})
+                results.append({"id": idx + 1, "text": text, "prediction": "Spam" if prediction == 1 else "Ham"})
         return {"results": results}
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Batch prediction failed: {e}")
 
+
 @app.get("/")
 async def root():
     return {"message": "Welcome to the MobileBERT API"}