Update app.py

app.py

@@ -58,16 +58,21 @@ problem_countries = [
 # === Prediction Function ===
 def predict_votes(resolution_text):
     vec = vectorizer.encode(resolution_text)
-    x_tensor = torch.tensor(vec, dtype=torch.float32).unsqueeze(0)  # batchify
+    x_tensor = torch.tensor(vec, dtype=torch.float32).unsqueeze(0)
 
     countries = []
     votes = []
 
     for country in country_encoder.classes_:
-        country_id = country_encoder.transform([country])[0]
-        c_tensor = torch.tensor([country_id], dtype=torch.long)
-
-        model = problem_model if country in problem_countries else main_model
+        is_problem = country in problem_countries
+        model = problem_model if is_problem else main_model
+
+        if is_problem:
+            problem_index = problem_countries.index(country)  # 0–45
+            c_tensor = torch.tensor([problem_index], dtype=torch.long)
+        else:
+            country_id = country_encoder.transform([country])[0]  # 0–192
+            c_tensor = torch.tensor([country_id], dtype=torch.long)
 
         with torch.no_grad():
             logit = model(x_tensor, c_tensor).squeeze()