Update app.py

app.py

@@ -3,11 +3,18 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 import numpy as np
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
 MODEL_PATH = 'finiteautomata/bertweet-base-sentiment-analysis'
 tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
 model = model.to(device)
 
+inputs = tokenizer(query, return_tensors='pt', truncation=True)
+inputs = inputs.to(device)
+outputs = model(**inputs)
+label2id = model.config.label2id
 
 logits = outputs.logits
 sigmoid = torch.nn.Sigmoid()
@@ -19,7 +26,6 @@ for i, k in enumerate(label2id.keys()):
 
 
 label2id = {k: v for k, v in sorted(label2id.items(), key=lambda item: item[1], reverse=True)}
-label2id
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
@@ -37,8 +43,8 @@ def get_predictions(input_text: str) -> dict:
     for i, k in enumerate(label2id.keys()):
         label2id[k] = probs[i]
     label2id = {k: float(v) for k, v in sorted(label2id.items(), key=lambda item: item[1].item(), reverse=True)}
-    print(label2id)
-
+    return label2id
+    
 import gradio as gr
 gr.Interface(
     fn=get_predictions,