| from fastapi import FastAPI |
| from pydantic import BaseModel |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification |
| import torch |
| import torch.nn.functional as F |
|
|
| app = FastAPI() |
|
|
| model_path = "results" |
|
|
| tokenizer = AutoTokenizer.from_pretrained(model_path) |
| model = AutoModelForSequenceClassification.from_pretrained(model_path) |
| model.eval() |
|
|
| class TextInput(BaseModel): |
| text: str |
|
|
| @app.post("/predict") |
| def predict(input: TextInput): |
| inputs = tokenizer( |
| input.text, |
| return_tensors="pt", |
| truncation=True, |
| padding=True, |
| max_length=256 |
| ) |
| with torch.no_grad(): |
| outputs = model(**inputs) |
| probs = F.softmax(outputs.logits, dim=1) |
| pred = torch.argmax(probs, dim=1).item() |
| confidence = (probs[0][pred].item())*100 |
| label = "POSITIVE" if pred == 1 else "NEGATIVE" |
| return { |
| "prediction": label, |
| "confidence": round(confidence, 4) |
| } |
