import torch
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel

BASE    = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
ADAPTER = "bmarlin96/meridian-router-lora"
LABELS  = ["shortage", "delivery", "recall", "billing", "marketing", "confirmation"]

tok = AutoTokenizer.from_pretrained(BASE)
tok.pad_token = tok.eos_token
base = AutoModelForSequenceClassification.from_pretrained(BASE, num_labels=len(LABELS))
base.config.pad_token_id = tok.pad_token_id
model = PeftModel.from_pretrained(base, ADAPTER)
model.eval()

app = FastAPI()

class Message(BaseModel):
    text: str

@app.get("/")
def health():
    return {"status": "ok"}

@app.post("/classify")
def classify(msg: Message):
    inputs = tok(msg.text, return_tensors="pt", truncation=True, max_length=96)
    with torch.no_grad():
        probs = torch.softmax(model(**inputs).logits, dim=-1)[0]
    i = int(probs.argmax())
    return {"label": LABELS[i], "confidence": round(float(probs[i]), 3)}