File size: 1,055 Bytes
fc05405 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | 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)} |