Add a prompting-based method

background_inference.py

@@ -4,7 +4,7 @@ import utils
 import datasets
 import eval_utils
 from constants import DIALECTS_WITH_LABELS
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
 
 model_name = sys.argv[1]
 commit_id = sys.argv[2]
@@ -19,10 +19,17 @@ utils.update_model_queue(
 )
 
 try:
-    tokenizer = AutoTokenizer.from_pretrained(model_name, revision=commit_id)
-    model = AutoModelForSequenceClassification.from_pretrained(
-        model_name, revision=commit_id
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name, revision=commit_id, access_token=os.environ["HF_TOKEN"]
     )
+    if inference_function == "prompt_chat_LLM":
+        model = AutoModel.from_pretrained(
+            model_name, revision=commit_id, access_token=os.environ["HF_TOKEN"]
+        )
+    else:
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name, revision=commit_id, access_token=os.environ["HF_TOKEN"]
+        )
 
     # Load the dataset
     dataset_name = os.environ["DATASET_NAME"]

constants.py

@@ -18,6 +18,28 @@ DIALECTS = [
     "UAE",
     "Yemen",
 ]
+
+DIALECT_IN_ARABIC = {
+    "Algeria": "الجزائر",
+    "Bahrain": "البحرين",
+    "Egypt": "مصر",
+    "Iraq": "العراق",
+    "Jordan": "الأردن",
+    "Kuwait": "الكويت",
+    "Lebanon": "لبنان",
+    "Libya": "ليبيا",
+    "Morocco": "المغرب",
+    "Oman": "عمان",
+    "Palestine": "فلسطين",
+    "Qatar": "قطر",
+    "Saudi_Arabia": "المملكة العربية السعودية",
+    "Sudan": "السودان",
+    "Syria": "سوريا",
+    "Tunisia": "تونس",
+    "UAE": "الإمارات",
+    "Yemen": "اليمن",
+}
+
 assert len(DIALECTS) == 18
 
 DIALECTS_WITH_LABELS = [

eval_utils.py

@@ -1,5 +1,5 @@
 import torch
-from constants import DIALECTS, DIALECTS_WITH_LABELS
+from constants import DIALECTS, DIALECTS_WITH_LABELS, DIALECT_IN_ARABIC
 
 
 def predict_top_p(model, tokenizer, text, P=0.9):
@@ -25,3 +25,26 @@ def predict_top_p(model, tokenizer, text, P=0.9):
             break
 
     return [DIALECTS[i] for i, p in enumerate(predictions) if p == 1]
+
+
+def prompt_chat_LLM(model, tokenizer, text):
+    """Prompt the user to determine whether the input text is acceptable in each of the 11 dialects."""
+    predicted_dialects = []
+    for dialect in DIALECTS_WITH_LABELS:
+        messages = [
+            {
+                "role": "user",
+                "content": f"حدد إذا كانت الجملة الأتية مقبولة في أحد اللهجات المستخدمة في {DIALECT_IN_ARABIC[dialect]}. أجب ب 'نعم' أو 'لا' فقط."
+                + "\n"
+                + f'الجملة: "{text}"',
+            },
+        ]
+        input_ids = tokenizer.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
+        )
+        gen_tokens = model.generate(input_ids, max_new_tokens=20)
+        gen_text = tokenizer.decode(gen_tokens[0], skip_special_tokens=True)
+        # TODO: Add a condition for the case of "لا" and other responses (e.g., refuse to answer)
+        if "نعم" in gen_text:
+            predicted_dialects.append(dialect)
+    return predicted_dialects