app.py

from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

app = FastAPI()

# Using v2 model (1000 examples, improved performance)
model_name = "King-8/help-classifier-v2"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

labels = [
    "learning_help",
    "project_help",
    "technical_issue",
    "attendance_issue",
    "general_guidance"
]

class InputText(BaseModel):
    text: str
    
@app.get("/")
def root():
    return {"message": "Help Classifier API is running"}

@app.post("/predict")
def predict(input: InputText):
    inputs = tokenizer(input.text, return_tensors="pt", truncation=True, padding=True)
    outputs = model(**inputs)

    probs = F.softmax(outputs.logits, dim=1)
    predicted_class = probs.argmax().item()
    confidence = probs.max().item()

    return {
        "label": labels[predicted_class],
        "confidence": round(confidence, 4)
    }