classifier / sentimental.py
narutoSiskovich's picture
Create sentimental.py
9380e7e verified
raw
history blame
1.16 kB
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}