fix prompt

lm_eval/tasks/ruler/niah_1.yaml

@@ -8,6 +8,7 @@ test_split: test
 download_dataset: !function utils.niah_single_1
 doc_to_text: "{{input}}"
 doc_to_target: "{{outputs[0]}}"
+gen_prefix: "{{gen_prefix}}"
 process_results: !function utils.process_results
 
 metric_list:

lm_eval/tasks/ruler/prepare.py

@@ -298,7 +298,9 @@ def generate_samples(
             "outputs": answer,
             "length": length,
             "max_length": max_seq_length,
-            "gen_prefix": f"The special magic {type_needle_v} for {query} mentioned in the provided text are",
+            "gen_prefix": f"The special magic {type_needle_v[:-1]} for {query} mentioned in the provided text is"
+            if num_needle_q * num_needle_v == 1
+            else f"The special magic {type_needle_v} for {query} mentioned in the provided text are",
         }
         if formatted_output["outputs"][0] not in formatted_output["input"]:
             assert (

lm_eval/tasks/ruler/utils.py

@@ -24,8 +24,9 @@ def get_tokenizer(
     return AutoTokenizer.from_pretrained(pretrained, trust_remote_code=True)
 
 
-# TOKENIZER = AutoTokenizer.from_pretrained(os.environ.get("TOKENIZER"))
-TEMPLATE = """Some special magic {type_needle_v} are hidden within the following text. Make sure to memorize it. I will quiz you about the {type_needle_v} afterwards.\n{context}\nWhat are all the special magic {type_needle_v} for {query} mentioned in the provided text? The special magic {type_needle_v} for {query} mentioned in the provided text are"""
+# TEMPLATE = """Some special magic {type_needle_v} are hidden within the following text. Make sure to memorize it. I will quiz you about the {type_needle_v} afterwards.\n{context}\nWhat are all the special magic {type_needle_v} for {query} mentioned in the provided text? The special magic {type_needle_v} for {query} mentioned in the provided text are"""
+TEMPLATE = """Some special magic {type_needle_v} are hidden within the following text. Make sure to memorize it. I will quiz you about the {type_needle_v} afterwards.\n{context}\nWhat are all the special magic {type_needle_v} for {query} mentioned in the provided text?"""
+
 
 SEQ_LENGTHS = (
     # 131072,