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

app = FastAPI()

model_name = "King-8/help-classifier"

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

labels = [
    "learning_help",
    "project_help",
    "attendance_issue",
    "technical_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()

    return {
        "label": labels[predicted_class]
    }