Update app.py

app.py

@@ -27,15 +27,47 @@ class RaceClassifier(nn.Module):
         return self.out(output)
 
 
-labels = {
+race_labels = {
     0: "African American",
     1: "Asian",
     2: "Latin",
     3: "White"
 }
+
+orientation_labels = {
+    0: "Heterosexual",
+    1: "LGBTQ"
+}
+
 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')))
+model_race.load_state_dict(torch.load('best_model_race.pt'))
+
+model_orientation = RaceClassifier(n_classes=2)
+model_orientation.to(device)
+model_orientation.load_state_dict(torch.load('best_model_orientation_last.pt'))
+
+
+def evaluate(model, input, mask):
+    model.eval()
+    with torch.no_grad():
+        outputs = model(input, mask)
+        probs = torch.nn.functional.softmax(outputs, dim=1)
+        predictions = torch.argmax(outputs, dim=1)
+        predictions = predictions.cpu().numpy()
+    return probs, predictions
+
+
+def write_output(probs, predictions, title, labels):
+    output_string = f"{title.upper()}\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"
+
+    output_string += f"Predicted as: {labels[predictions[0]]}\n"
+
+    return output_string
+
 
 def predict(*text):
     tweets = [tweet for tweet in text if tweet]
@@ -55,22 +87,16 @@ def predict(*text):
     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"
+    race_probs, race_predictions = evaluate(model_race, input_ids, attention_mask)
+    orientation_probs, orientation_predictions = evaluate(model_orientation, input_ids, attention_mask)
 
-    print(labels[predictions[0]])
-    output_string += f"Predicted as: {labels[predictions[0]]}"
+    final_output = str()
+    final_output += write_output(race_probs, race_predictions, "race", race_labels)
+    final_output += "\n"
+    final_output += write_output(orientation_probs, orientation_predictions, "sexual orientation", orientation_labels)
+    final_output += "\n"
 
-    return output_string
+    return final_output
 
 
 max_textboxes = 20