| 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) |
| } |
