fix bug

app.py

@@ -1,15 +1,28 @@
 import gradio as gr
+import numpy as np
 import torch
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-device = 0 if torch.cuda.is_available() else -1
-
-sentiment_pipeline = pipeline("text-classification", 
-                              model="winain7788/bert-finetuned-sem_eval-english", 
-                              device=device)
+id2label = {0: 'anger', 1: 'anticipation', 2: 'disgust', 3: 'fear', 4: 'joy', 5: 'love', 6: 'optimism', 7: 'pessimism', 8: 'sadness', 9: 'surprise', 10: 'trust'}
+tokenizer = AutoTokenizer.from_pretrained("winain7788/bert-finetuned-sem_eval-english")
+model = AutoModelForSequenceClassification.from_pretrained("winain7788/bert-finetuned-sem_eval-english")
 
 async def get_sentiment(text):
-    return sentiment_pipeline([text])
+    encoding = tokenizer(text, return_tensors="pt")
+    encoding = {k: v.to(model.device) for k,v in encoding.items()}
+
+    outputs = model(**encoding)
+    logits = outputs.logits
+    logits.shape
+    # apply sigmoid + threshold
+    sigmoid = torch.nn.Sigmoid()
+    probs = sigmoid(logits.squeeze().cpu())
+    predictions = np.zeros(probs.shape)
+    predictions[np.where(probs >= 0.5)] = 1
+
+    # turn predicted id's into actual label names
+    predicted_labels = [id2label[idx] for idx, label in enumerate(predictions) if label == 1.0]
+    return predicted_labels
 
 demo = gr.Interface(fn=get_sentiment, inputs="text", outputs="json")