Update app.py

app.py

@@ -34,6 +34,24 @@ race_labels = {
     3: "White"
 }
 
+age_labels = {
+    0: "Adult",
+    1: "Elderly",
+    2: "Young"
+}
+
+education_labels = {
+    0: "Educated",
+    1: "Uneducated"
+}
+
+gender_labels = {
+    0: "Female",
+    1: "Male",
+    2: "Non-Binary",
+    3: "Transgender"
+}
+
 orientation_labels = {
     0: "Heterosexual",
     1: "LGBTQ"
@@ -43,6 +61,18 @@ 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_age = RaceClassifier(n_classes=3)
+model_age.to(device)
+model_age.load_state_dict(torch.load('best_model_age_last.pt', map_location=torch.device('cpu')))
+
+model_education = RaceClassifier(n_classes=2)
+model_education.to(device)
+model_education.load_state_dict(torch.load('best_model_education_last.pt', map_location=torch.device('cpu')))
+
+model_gender = RaceClassifier(n_classes=4)
+model_gender.to(device)
+model_gender.load_state_dict(torch.load('best_model_gender_last.pt', map_location=torch.device('cpu')))
+
 model_orientation = RaceClassifier(n_classes=2)
 model_orientation.to(device)
 model_orientation.load_state_dict(torch.load('best_model_orientation_last.pt', map_location=torch.device('cpu')))
@@ -88,13 +118,21 @@ def predict(*text):
     attention_mask = encoded_sentences["attention_mask"].to(device)
 
     race_probs, race_predictions = evaluate(model_race, input_ids, attention_mask)
+    age_probs, age_predictions = evaluate(model_age, input_ids, attention_mask)
+    education_probs, education_predictions = evaluate(model_education, input_ids, attention_mask)
+    gender_probs, gender_predictions = evaluate(model_gender, input_ids, attention_mask)
     orientation_probs, orientation_predictions = evaluate(model_orientation, input_ids, attention_mask)
 
     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 += write_output(age_probs, age_predictions,"age",age_labels)
     final_output += "\n"
+    final_output += write_output(education_probs,education_predictions,"education", education_labels)
+    final_output += "\n"
+    final_output += write_output(gender_probs, gender_predictions, "gender", gender_labels)
+    final_output += "\n"
+    final_output += write_output(orientation_probs, orientation_predictions, "sexual orientation", orientation_labels)
 
     return final_output