add war domain

interfaces/illframes.py

@@ -8,7 +8,7 @@ from transformers import AutoModelForSequenceClassification
 from transformers import AutoTokenizer
 from huggingface_hub import HfApi
 
-from label_dicts import ILLFRAMES_MIGRATION_LABEL_NAMES, ILLFRAMES_COVID_LABEL_NAMES
+from label_dicts import ILLFRAMES_MIGRATION_LABEL_NAMES, ILLFRAMES_COVID_LABEL_NAMES, ILLFRAMES_WAR_LABEL_NAMES
 
 HF_TOKEN = os.environ["hf_read"]
 
@@ -18,7 +18,8 @@ languages = [
 
 domains = {
     "Covid": "covid",
-    "Migration": "migration"
+    "Migration": "migration",
+    "War": "war"
 }
 
 
@@ -55,16 +56,6 @@ def build_huggingface_path(domain: str):
 
 def predict(text, model_id, tokenizer_id, label_names):
     device = torch.device("cpu")
-
-    # --- DEBUG ---
-    
-    disk_space = get_disk_space('/data/')
-    print("Disk Space Info:")
-    for key, value in disk_space.items():
-        print(f"{key}: {value}")
-        
-    # ---
-
     try:
         model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, offload_folder="offload", device_map="auto", token=HF_TOKEN)
     except:
@@ -101,8 +92,10 @@ def predict_illframes(text, language, domain):
 
     if domain == "migration":
         label_names = ILLFRAMES_MIGRATION_LABEL_NAMES
-    else:
+    elif domain == "covid":
         label_names = ILLFRAMES_COVID_LABEL_NAMES
+    elif domain == "war":
+        label_names = ILLFRAMES_WAR_LABEL_NAMES
 
     return predict(text, model_id, tokenizer_id, label_names)