File size: 3,987 Bytes
0753b25 c773f1d 0753b25 | 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 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI(
title="Support Ticket Classifier API",
description="Classify customer support tickets into 5 categories using BERTweet",
version="1.0.0",
docs_url="/docs",
redoc_url="/redoc"
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
logger.info("Loading support ticket classifier model...")
try:
model = AutoModelForSequenceClassification.from_pretrained("Adedayo2000/BERTweet-Support-Ticket-Classifier")
tokenizer = AutoTokenizer.from_pretrained("Adedayo2000/BERTweet-Support-Ticket-Classifier")
logger.info("Model loaded successfully!")
except Exception as e:
logger.error(f"Failed to load model: {e}")
model = None
tokenizer = None
class TicketInput(BaseModel):
text: str
class Config:
json_schema_extra = {
"example": {
"text": "I can't login to my account, password reset not working"
}
}
@app.get("/")
def home():
return {
"message": "Support Ticket Classifier API",
"status": "live",
"model": "Adedayo2000/BERTweet-Support-Ticket-Classifier",
"accuracy": "85%",
"categories": [
"Account Access",
"Billing/Payment",
"Feature Request",
"General Inquiry",
"Technical Support"
],
"endpoints": {
"GET /": "API information",
"POST /classify": "Classify support ticket",
"GET /categories": "List all categories",
"GET /docs": "Interactive API documentation",
"GET /health": "Health check"
}
}
@app.get("/health")
def health_check():
return {
"status": "healthy",
"model_loaded": model is not None and tokenizer is not None
}
@app.get("/categories")
def get_categories():
if model is None:
return {"error": "Model not loaded"}
return {
"categories": list(model.config.id2label.values()),
"count": len(model.config.id2label)
}
@app.post("/classify")
def classify_ticket(input: TicketInput):
if not input.text or not input.text.strip():
return {
"error": "Ticket text cannot be empty",
"example": {"text": "I can't login to my account"}
}
if model is None or tokenizer is None:
return {
"error": "Model not loaded",
"status": "unavailable"
}
try:
inputs = tokenizer(
input.text,
return_tensors='pt',
truncation=True,
max_length=128,
padding=True
)
with torch.no_grad():
outputs = model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
pred_idx = torch.argmax(probs, dim=-1).item()
category = model.config.id2label[pred_idx]
confidence_score = probs[0][pred_idx].item()
confidence_percentage = f"{confidence_score * 100:.2f}%"
all_predictions = {
model.config.id2label[i]: f"{probs[0][i].item() * 100:.2f}%"
for i in range(len(probs[0]))
}
return {
"text": input.text,
"category": category,
"confidence": confidence_percentage,
"confidence_score": round(confidence_score, 4),
#"all_categories": all_predictions
}
except Exception as e:
logger.error(f"Classification error: {e}")
return {
"error": "Classification failed",
"details": str(e)
} |