Update app.py

app.py

@@ -37,8 +37,43 @@ model_race = RaceClassifier(n_classes=4)
 model_race.to(device)
 model_race.load_state_dict(torch.load('best_model_race.pt', map_location=torch.device('cpu')))
 
+def predict(*text):
+    tweets = [tweet for tweet in text if tweet]
+    print(tweets)
+    sentences = tweets
 
-max_textboxes = 10
+    tokenizer = AutoTokenizer.from_pretrained("vinai/bertweet-base", normalization=True)
+
+    encoded_sentences = tokenizer(
+        sentences,
+        padding=True,
+        truncation=True,
+        return_tensors='pt',
+        max_length=128,
+    )
+
+    input_ids = encoded_sentences["input_ids"].to(device)
+    attention_mask = encoded_sentences["attention_mask"].to(device)
+
+    model_race.eval()
+    with torch.no_grad():
+        outputs = model_race(input_ids, attention_mask)
+        probs = torch.nn.functional.softmax(outputs, dim=1)
+        predictions = torch.argmax(outputs, dim=1)
+        predictions = predictions.cpu().numpy()
+
+    output_string = "RACE\n Probabilities:\n"
+    for i, prob in enumerate(probs[0]):
+        print(f"{labels[i]} = {round(prob.item() * 100, 2)}%")
+        output_string += f"{labels[i]} = {round(prob.item() * 100, 2)}%\n"
+
+    print(labels[predictions[0]])
+    output_string += f"Predicted as: {labels[predictions[0]]}"
+
+    return output_string
+
+
+max_textboxes = 20
 
 
 def update_textboxes(k):