Spaces:
Runtime error
Runtime error
File size: 2,736 Bytes
58ed1c6 f27ac44 58ed1c6 921649c 58ed1c6 f27ac44 921649c f27ac44 921649c 58ed1c6 f27ac44 921649c f27ac44 921649c f27ac44 921649c f27ac44 58ed1c6 f27ac44 58ed1c6 f27ac44 921649c 58ed1c6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import BertTokenizer, BertForSequenceClassification
import torch
import gradio as gr
app = FastAPI()
# Prediction labels
LABELS = [
'Login Issue',
'Booking Issue',
'Delivery Issue',
'Laboratory Issue',
'Application Issue'
]
# 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_classifier(text):
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = model(**inputs)
probs = torch.softmax(outputs.logits, dim=1)
pred_idx = torch.argmax(probs).item()
return {
"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_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 system for customer support issues"
)
app = gr.mount_gradio_app(app, gradio_app, path="/gradio")
|