Update app.py

app.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from transformers import BertTokenizer, BertForSequenceClassification
 import torch
@@ -6,30 +7,92 @@ import gradio as gr
 
 app = FastAPI()
 
-# Your existing prediction code remains unchanged
-Label = ['Login Issue', 'Booking Issue', 'Delivery Issue', 'Laboratory Issue', 'Application Issue']
+# Prediction labels
+LABELS = [
+    'Login Issue', 
+    'Booking Issue', 
+    'Delivery Issue', 
+    'Laboratory Issue', 
+    'Application Issue'
+]
 
-# ... [Keep your existing CORS and model loading code] ...
+# CORS Configuration
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Load model and tokenizer
+try:
+    model = BertForSequenceClassification.from_pretrained("./saved_model1")
+    tokenizer = BertTokenizer.from_pretrained("./saved_model1")
+    model.eval()
+except Exception as e:
+    raise RuntimeError(f"Model loading failed: {str(e)}")
+
+# Request Model
+class PredictionRequest(BaseModel):
+    issue: str
+
+# FastAPI Endpoint
+@app.post("/predict")
+async def predict(request: PredictionRequest):
+    try:
+        inputs = tokenizer(
+            request.issue, 
+            return_tensors="pt", 
+            truncation=True, 
+            padding=True,
+            max_length=512
+        )
+        
+        with torch.no_grad():
+            outputs = model(**inputs)
+            probabilities = torch.softmax(outputs.logits, dim=1)
+            label_idx = torch.argmax(probabilities).item()
+            
+        return {
+            "category": LABELS[label_idx],
+            "confidence": round(probabilities[0][label_idx].item(), 4)
+        }
+        
+    except Exception as e:
+        raise HTTPException(
+            status_code=500, 
+            detail=f"Prediction error: {str(e)}"
+        )
 
 # Gradio Interface
-def gradio_predict(issue_text):
-    inputs = tokenizer(issue_text, return_tensors="pt", truncation=True, padding=True)
+def gradio_classifier(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         outputs = model(**inputs)
-        predictions = torch.softmax(outputs.logits, dim=1)
-        label_idx = torch.argmax(predictions, dim=1).item()
+        probs = torch.softmax(outputs.logits, dim=1)
+        pred_idx = torch.argmax(probs).item()
+    
     return {
-        "Predicted Category": Label[label_idx],
-        "Confidence": f"{predictions[0][label_idx].item():.4f}"
+        "Prediction": LABELS[pred_idx],
+        "Confidence Score": float(probs[0][pred_idx].item()),
+        "All Probabilities": {
+            label: round(float(probs[0][i]), 4) 
+            for i, label in enumerate(LABELS)
+        }
     }
 
+# Mount Gradio interface
 gradio_app = gr.Interface(
-    fn=gradio_predict,
-    inputs=gr.Textbox(label="Enter Issue Description"),
-    outputs=gr.JSON(label="Prediction Results"),
+    fn=gradio_classifier,
+    inputs=gr.Textbox(lines=3, placeholder="Enter issue description...", label="Issue"),
+    outputs=[
+        gr.Label(label="Predicted Category"),
+        gr.Number(label="Confidence Score"),
+        gr.JSON(label="Class Probabilities")
+    ],
     title="Issue Classifier",
-    description="BERT-based classification demo"
+    description="BERT-based classification system for customer support issues"
 )
 
-# Mount Gradio to FastAPI
 app = gr.mount_gradio_app(app, gradio_app, path="/gradio")