Update app.py

app.py

@@ -15,6 +15,15 @@ from huggingface_hub import hf_hub_download, login
 label_to_int = pd.read_pickle('label_to_int.pkl')
 int_to_label = {v: k for k, v in label_to_int.items()}
 
+# Update labels based on the given conditions
+for k, v in int_to_label.items():
+    if "KOREA" in v:
+        int_to_label[k] = "KOREA"
+    elif "KINGDOM" in v:
+        int_to_label[k] = "UK"
+    elif "RUSSIAN" in v:
+        int_to_label[k] = "RUSSIA"
+
 class LogisticRegressionTorch(nn.Module):
     def __init__(self, input_dim: int, output_dim: int):
         super(LogisticRegressionTorch, self).__init__()