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