Update app.py

app.py

@@ -73,22 +73,41 @@ class TransactionInput(BaseModel):
 async def root():
     return {"message": "Welcome to the Expense Categorization API. Use POST /predict to categorize expenses."}
 
-# Define predict endpoint
+# Define predict endpoint with confidence scores
 @app.post("/predict")
 async def predict_category(transaction: TransactionInput, request: Request):
     try:
         logger.info(f"Received request: {transaction.dict()}")
+        # Tokenize input
         tokens = tokenizer(transaction.description, return_tensors="pt", truncation=True, padding="max_length", max_length=64)
         input_ids = tokens["input_ids"].to(device)
         attention_mask = tokens["attention_mask"].to(device)
+        
+        # Get model predictions
         with torch.no_grad():
             main_logits, sub_logits = model(input_ids, attention_mask)
-        main_category = torch.argmax(main_logits, dim=1).cpu().item()
-        sub_category = torch.argmax(sub_logits, dim=1).cpu().item()
+        
+        # Compute softmax probabilities for main category
+        main_probs = torch.softmax(main_logits, dim=1)
+        main_category_idx = torch.argmax(main_probs, dim=1).cpu().item()
+        main_confidence = main_probs[0, main_category_idx].cpu().item()
+        
+        # Compute softmax probabilities for subcategory
+        sub_probs = torch.softmax(sub_logits, dim=1)
+        sub_category_idx = torch.argmax(sub_probs, dim=1).cpu().item()
+        sub_confidence = sub_probs[0, sub_category_idx].cpu().item()
+        
+        # Decode category labels
+        main_category = main_category_encoder.inverse_transform([main_category_idx])[0]
+        sub_category = sub_category_encoder.inverse_transform([sub_category_idx])[0]
+        
+        # Prepare response
         response = {
             "description": transaction.description,
-            "main_category": main_category_encoder.inverse_transform([main_category])[0],
-            "sub_category": sub_category_encoder.inverse_transform([sub_category])[0]
+            "main_category": main_category,
+            "main_confidence": round(main_confidence, 4),
+            "sub_category": sub_category,
+            "sub_confidence": round(sub_confidence, 4)
         }
         logger.info(f"Response: {response}")
         return response