Update app.py

app.py

@@ -1,48 +1,35 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import BertTokenizer, BertForSequenceClassification
 import torch
+import gradio as gr
+
+app = FastAPI()
+
+# Your existing prediction code remains unchanged
+Label = ['Login Issue', 'Booking Issue', 'Delivery Issue', 'Laboratory Issue', 'Application Issue']
 
-# Load model and tokenizer from local directory
-MODEL_PATH = "./saved_model1"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
+# ... [Keep your existing CORS and model loading code] ...
 
-def predict(text):
-    # Preprocess input
-    inputs = tokenizer(
-        text,
-        padding=True,
-        truncation=True,
-        max_length=512,
-        return_tensors="pt"
-    )
-    
-    # Inference
+# Gradio Interface
+def gradio_predict(issue_text):
+    inputs = tokenizer(issue_text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         outputs = model(**inputs)
-    
-    # Postprocess output (modify based on your task)
-    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    predicted_class = torch.argmax(probabilities).item()
-    
+        predictions = torch.softmax(outputs.logits, dim=1)
+        label_idx = torch.argmax(predictions, dim=1).item()
     return {
-        "text": text,
-        "predicted_class": predicted_class,
-        "probabilities": probabilities.tolist()[0]
+        "Predicted Category": Label[label_idx],
+        "Confidence": f"{predictions[0][label_idx].item():.4f}"
     }
 
-# Create Gradio interface
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Textbox(label="Input Text", lines=3),
-    outputs=[
-        gr.Textbox(label="Processed Text"),
-        gr.Number(label="Predicted Class"),
-        gr.Label(label="Class Probabilities")
-    ],
-    title="BERT Model Deployment",
-    examples=[["Sample text 1"], ["Another example text"]]
+gradio_app = gr.Interface(
+    fn=gradio_predict,
+    inputs=gr.Textbox(label="Enter Issue Description"),
+    outputs=gr.JSON(label="Prediction Results"),
+    title="Issue Classifier",
+    description="BERT-based classification demo"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+# Mount Gradio to FastAPI
+app = gr.mount_gradio_app(app, gradio_app, path="/gradio")