service-internal's picture
Update main.py
6644ffe verified
raw
history blame
1.23 kB
import os
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf-cache"
os.environ["HF_HOME"] = "/tmp/hf-home"
from fastapi import FastAPI, Request
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
from scipy.special import softmax
import numpy as np
app = FastAPI()
MODEL = "cardiffnlp/twitter-roberta-base-sentiment-latest"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
config = AutoConfig.from_pretrained(MODEL)
model = AutoModelForSequenceClassification.from_pretrained(MODEL)
def preprocess(text):
tokens = []
for t in text.split():
if t.startswith("@") and len(t) > 1:
t = "@user"
elif t.startswith("http"):
t = "http"
tokens.append(t)
return " ".join(tokens)
@app.post("/analyze")
async def analyze(request: Request):
data = await request.json()
text = preprocess(data.get("text", ""))
encoded_input = tokenizer(text, return_tensors='pt')
output = model(**encoded_input)
scores = output[0][0].detach().numpy()
scores = softmax(scores)
result = [
{"label": config.id2label[i], "score": round(float(scores[i]), 4)}
for i in scores.argsort()[::-1]
]
return {"result": result}