Spaces:
Runtime error
Runtime error
| from fastapi import FastAPI, HTTPException | |
| from pydantic import BaseModel | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| import torch | |
| import os | |
| from dotenv import load_dotenv | |
| app = FastAPI() | |
| load_dotenv() | |
| checkpoint = "GaaS-Team/DistilBERT-finetuned-GaaS" | |
| hf_token = os.getenv("HUGGINGFACE_TOKEN") | |
| tokenizer = AutoTokenizer.from_pretrained(checkpoint, use_auth_token=hf_token) | |
| model = AutoModelForSequenceClassification.from_pretrained(checkpoint, use_auth_token=hf_token) | |
| class TextInput(BaseModel): | |
| text: str | |
| async def predict_sentiment(input: TextInput): | |
| try: | |
| text = input.text | |
| inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt") | |
| with torch.no_grad(): | |
| logits = model(**inputs).logits | |
| predicted_class_id = logits.argmax(dim=-1).item() | |
| predicted_class = model.config.id2label[predicted_class_id] | |
| return {"sentiment": predicted_class} | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}") |