from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# === Загрузка модели ===
MODEL_NAME = "nlptown/bert-base-multilingual-uncased-sentiment"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
model.eval()

# === FastAPI приложение ===
app = FastAPI(title="Sentiment Analysis API")

# === Схема запроса ===
class TextRequest(BaseModel):
    text: str

# === Логика сентимента ===
def analyze_sentiment(message: str) -> float:
    inputs = tokenizer(message, return_tensors="pt", truncation=True)
    with torch.no_grad():
        logits = model(**inputs).logits
    probs = torch.softmax(logits, dim=-1)
    stars = torch.argmax(probs, dim=-1).item() + 1  # от 1 до 5
    sentiment = (stars - 3) * 2.5  # нормируем -5..+5
    return round(sentiment, 2)

# === Эндпоинт API ===
@app.post("/sentiment")
def sentiment(request: TextRequest):
    score = analyze_sentiment(request.text)
    return {"sentiment_score": score}