Update app.py

app.py

@@ -53,11 +53,42 @@ async def predict(request: TextRequest):
         raise HTTPException(status_code=500, detail=f"Prediction failed: {e}")
 
 
+# @app.post("/batch_predict/")
+# async def batch_predict(request: BatchTextRequest):
+#     try:
+#         model.eval()
+#         logger.info(f"Received batch prediction request for {len(request.texts)} texts")
+#         inputs = tokenizer(
+#             request.texts,
+#             add_special_tokens=True,
+#             max_length=64,
+#             truncation=True,
+#             padding='max_length',
+#             return_attention_mask=True,
+#             return_tensors='pt'
+#         )
+
+#         with torch.no_grad():
+#             logits = model(inputs['input_ids'], attention_mask=inputs['attention_mask']).logits
+#             predictions = torch.argmax(logits, dim=1).tolist()
+
+#         results = [
+#             {"id": idx + 1, "text": text, "prediction": "Spam" if pred == 1 else "Ham"}
+#             for idx, (text, pred) in enumerate(zip(request.texts, predictions))
+#         ]
+#         logger.info(f"Batch prediction results: {results}")
+#         return {"results": results}
+#     except Exception as e:
+#         logger.error(f"Batch prediction failed: {e}")
+#         raise HTTPException(status_code=500, detail=f"Batch prediction failed: {e}")
+
+
 @app.post("/batch_predict/")
 async def batch_predict(request: BatchTextRequest):
     try:
         model.eval()
-        logger.info(f"Received batch prediction request for {len(request.texts)} texts")
+        
+        # Batch encode all texts in the request at once
         inputs = tokenizer(
             request.texts,
             add_special_tokens=True,
@@ -68,20 +99,22 @@ async def batch_predict(request: BatchTextRequest):
             return_tensors='pt'
         )
 
+        # Run batch inference
         with torch.no_grad():
             logits = model(inputs['input_ids'], attention_mask=inputs['attention_mask']).logits
             predictions = torch.argmax(logits, dim=1).tolist()
-
+        
+        # Format results
         results = [
             {"id": idx + 1, "text": text, "prediction": "Spam" if pred == 1 else "Ham"}
             for idx, (text, pred) in enumerate(zip(request.texts, predictions))
         ]
-        logger.info(f"Batch prediction results: {results}")
+        
         return {"results": results}
-    except Exception as e:
-        logger.error(f"Batch prediction failed: {e}")
-        raise HTTPException(status_code=500, detail=f"Batch prediction failed: {e}")
 
+    except Exception as e:
+        logging.error(f"Batch prediction failed: {e}")
+        raise HTTPException(status_code=500, detail="Batch prediction failed. Please try again.")
 
 @app.get("/")
 async def root():