Update app.py

app.py

@@ -81,21 +81,35 @@ app = FastAPI()
 def greet_json():
     return {"Hello": "World!"}
 
-#
-# # Predict endpoint for JSON input
-# @app.post("/predict")
-# async def predict_image(file: UploadFile = File(...)):
-#     try:
-#         # Read and preprocess the uploaded image
-#         image = Image.open(file.file)
-#         image = preprocess_image(image)
-#
-#         # Make prediction
-#         model.eval()
-#         with torch.no_grad():
-#             output = model(image)
-#             prediction = output.argmax(dim=1).item()
-#
-#         return JSONResponse(content={"prediction": f"The digit is {prediction}"})
-#     except Exception as e:
-#         return JSONResponse(content={"error": str(e)}, status_code=500)
+# Input model for request validation
+class TextAspectInput(BaseModel):
+    text: str
+    aspect: str
+
+
+# Sentiment labels
+SENTIMENT_LABELS = {0: "Negative", 1: "Neutral", 2: "Positive"}
+
+
+# Predict endpoint
+@app.post("/predict")
+async def predict_sentiment(input_data: TextAspectInput):
+    print(input_data)
+    try:
+        # Extract text and aspect
+        text = input_data.text
+        aspect = input_data.aspect
+
+        # Process input
+        input_ids = process_text(text, aspect)
+
+        # Make prediction
+        with torch.no_grad():
+            logits = model(input_ids)
+            probabilities = torch.softmax(logits, dim=-1)
+            prediction = probabilities.argmax(dim=-1).item()
+            sentiment = SENTIMENT_LABELS[prediction]
+
+        return {"sentiment": sentiment, "probabilities": probabilities.squeeze().tolist()}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))