Spaces:
Runtime error
Runtime error
| from fastapi import FastAPI, Request | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
| import torch | |
| app = FastAPI() | |
| # Allow calls from MERN stack | |
| origins = ["http://localhost:3000", "http://localhost:5000"] | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=origins, | |
| allow_methods=["*"], | |
| allow_headers=["*"] | |
| ) | |
| model_dir = "./spam_detector_model" | |
| model = AutoModelForSequenceClassification.from_pretrained(model_dir) | |
| tokenizer = AutoTokenizer.from_pretrained(model_dir) | |
| labels = list(model.config.id2label.values()) | |
| async def predict(req: Request): | |
| data = await req.json() | |
| inputs = tokenizer(data["text"], return_tensors="pt", truncation=True, padding=True) | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| probs = torch.nn.functional.softmax(outputs.logits, dim=-1) | |
| pred = torch.argmax(probs).item() | |
| confidence = probs[0][pred].item() | |
| return { | |
| "label": labels[pred], | |
| "confidence": round(confidence, 3) | |
| } | |