service-internal commited on
Commit
a2004de
·
verified ·
1 Parent(s): 5cd8347

Upload main.py

Browse files
Files changed (1) hide show
  1. main.py +44 -0
main.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, Request
2
+ from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
3
+ from scipy.special import softmax
4
+ import numpy as np
5
+ import uvicorn
6
+
7
+ app = FastAPI()
8
+
9
+ # Load model and tokenizer
10
+ MODEL = "cardiffnlp/twitter-roberta-base-sentiment-latest"
11
+ tokenizer = AutoTokenizer.from_pretrained(MODEL)
12
+ config = AutoConfig.from_pretrained(MODEL)
13
+ model = AutoModelForSequenceClassification.from_pretrained(MODEL)
14
+
15
+ # Preprocessing function
16
+ def preprocess(text):
17
+ tokens = []
18
+ for t in text.split():
19
+ if t.startswith("@") and len(t) > 1:
20
+ t = "@user"
21
+ elif t.startswith("http"):
22
+ t = "http"
23
+ tokens.append(t)
24
+ return " ".join(tokens)
25
+
26
+ # Inference route
27
+ @app.post("/analyze")
28
+ async def analyze(request: Request):
29
+ data = await request.json()
30
+ text = preprocess(data.get("text", ""))
31
+
32
+ encoded_input = tokenizer(text, return_tensors='pt')
33
+ output = model(**encoded_input)
34
+ scores = output[0][0].detach().numpy()
35
+ scores = softmax(scores)
36
+
37
+ ranking = np.argsort(scores)[::-1]
38
+ result = []
39
+ for i in ranking:
40
+ label = config.id2label[i]
41
+ score = round(float(scores[i]), 4)
42
+ result.append({"label": label, "score": score})
43
+
44
+ return {"result": result}