Upload script.py with huggingface_hub

script.py

@@ -93,7 +93,7 @@ SYSTEM = (
 )
 
 SYSTEM_POST_EXPLAIN = (
-    "You are a helpful assistant that explains the reasoning behind the answers to the International Linguistics Olympiad problems.
+    "You are a helpful assistant that explains the reasoning behind the answers to the International Linguistics Olympiad problems.\n"
     "You are given the following information:\n"
     "- The context of the problem\n"
     "- The task type\n"
@@ -101,69 +101,36 @@ SYSTEM_POST_EXPLAIN = (
     "- The answer\n"
     "- The reasoning\n"
     "You need to explain the reasoning behind the answer in a way that is easy to understand and concisely focused on the key insights and rules deduced and applied.\n"
-    "Do not include any other text, do not includethe answer in the explanation, "
-    " and do not invent any new information."
+    "Do not include any other text, do not include the answer in the explanation, "
+    "and do not invent any new information."
 )
-tok = load_tokenizer(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID, torch_dtype=torch.float16, device_map="auto"
-).eval()
 
-with open("/tmp/data/test.csv", encoding="utf-8", newline="") as f:
-    test_rows = list(csv.DictReader(f))
+# Prefer a dedicated header line; also allow same-line answers after the colon.
+# Fall back to the last "final answer(s):" substring if no clean header line exists.
+FINAL_ANSWERS_LINE_RE = re.compile(
+    # Header alone on the line (markdown / punctuation allowed), OR
+    # header with first answer on the same line after a colon.
+    r"(?im)^[^\w\n]*final answers?[^\w\n]*:?[ \t]*(?=\n|$)|"
+    r"(?im)^[^\w\n]*final answers?\s*:\s*"
+)
+FINAL_ANSWERS_INLINE_RE = re.compile(
+    r"(?is)\bfinal answers?\s*:\s*"
+)
 
-outputs_queries_types = []
-for r in test_rows:
 
-    # Create the prompt.
-    messages = [
-        {"role": "system", "content": SYSTEM},
-        {"role": "user", "content": 
-        f"CONTEXT:{r['context'].strip()}\nTASK TYPE:`{r['task_type']}`\n\nQUERY:{r['query'].strip()}"},
-    ]
-    ids = tok.apply_chat_template(
-        messages, add_generation_prompt=True, return_tensors="pt",
-    ).to(model.device)
+def extract_raw_final(text: str) -> str:
+    """Return text after the last final-answers marker, or '' if none found."""
+    line_matches = list(FINAL_ANSWERS_LINE_RE.finditer(text))
+    if line_matches:
+        return text[line_matches[-1].end() :]
 
-    # Generate the answer.
-    done = False
-    attempts = 0
-    while not done and attempts < MAX_ATTEMPTS:
-        with torch.no_grad():
-            out = model.generate(
-                ids, 
-                max_new_tokens=MAX_NEW_TOKENS, 
-                do_sample=True,
-                temperature=TEMPERATURE,
-                top_p=TOP_P,)
-        text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
-        # IF NO FINAL ANSWER keyword is used, try again.
-        attempts += 1
-        if "final answer" not in text.lower() or text.lower().split('final answer')[1].split('\n')==0:
-            print(f'TRYING AGAIN...attempts #{attempts+1}/{MAX_ATTEMPTS}')
-        else:
-            done = True
+    inline_matches = list(FINAL_ANSWERS_INLINE_RE.finditer(text))
+    if inline_matches:
+        return text[inline_matches[-1].end() :]
 
-    # Generate the explanation.
-    messages_post_explain = [
-        {"role": "system", "content": SYSTEM_POST_EXPLAIN},
-        {"role": "user", "content": 
-        f"CONTEXT:{r['context'].strip()}\nTASK TYPE:`{r['task_type']}`\n\nQUERY:{r['query'].strip()}\n\nANSWER:{text.strip()}"},
-    ]
-    ids_post_explain = tok.apply_chat_template(
-        messages_post_explain, add_generation_prompt=True, return_tensors="pt",
-    ).to(model.device)
-    with torch.no_grad():
-        out_post_explain = model.generate(
-            ids_post_explain, 
-            max_new_tokens=MAX_NEW_TOKENS, 
-            do_sample=False)
-    text_post_explain = tok.decode(out_post_explain[0][ids_post_explain.shape[-1]:], skip_special_tokens=True).strip()
+    return ""
 
-    outputs_queries_types.append((text, text_post_explain, r['id'], r['query'], r['task_type']))
-    print(f"{len(outputs_queries_types)}/{len(test_rows)} done", flush=True)
 
-# Postprocess and store the answers.
 def expected_answer_count(query: str, task_type: str) -> int:
     if task_type == "match_letters":
         numbered = re.findall(r"^\s*\d+\.", query, re.MULTILINE)
@@ -209,46 +176,32 @@ def split_single_line_answer(text: str, expected: int, task_type: str) -> list[s
     return [text]
 
 
-def postprocess_answer(text, query, task_type):
-    """Keep only the lines after the last 'FINAL ANSWERS:' marker, one answer per line,
-    stopping at the first empty line. If eval_type is multiple and only one line as answer, split at whitespace."""
-    # Updated regex to be more flexible with surrounding characters
-    marker_match = list(re.finditer(r"(?im)^[^\w\n]*final answers?[^\w\n]*:?\s*$", text))
-    if marker_match:
-        text_after_marker = text[marker_match[-1].end():]
-        #print('FOUND FINAL ANSWER', text_after_marker)
-    else:
-        #print("No 'FINAL ANSWERS:' marker found")
-        return []
-
+def parse_answer_lines(text_after_marker: str, query: str, task_type: str) -> list[str]:
+    """Parse cleaned answer lines from the raw final-answers section."""
     answers = []
-    answer_lines = text_after_marker.splitlines()
-    for i, line in enumerate(answer_lines):
-        stripped_line = line.strip('`').strip()
-
-        # Stop processing if an empty line is encountered (not as first line)
-        if stripped_line=='':
-          continue
+    for line in text_after_marker.splitlines():
+        stripped_line = line.strip("`").strip()
+        if stripped_line == "":
+            continue
 
-        # Use a more precise regex to only remove numbering if it's a prefix to other text
-        # This ensures that lines which are just numbers (e.g., '1') are not stripped.
         match_numbered_prefix = re.match(r"^\s*\d+[.)]\s+(.*)", stripped_line)
         if match_numbered_prefix:
             cleaned_line = match_numbered_prefix.group(1).strip()
         else:
             cleaned_line = stripped_line
 
-        # Remove any bold markdown '**'
         cleaned_line = re.sub(r"\*\*", "", cleaned_line).strip()
 
-        # Specific handling for 'match_letters' task type to strip extra words
-        if task_type == 'match_letters':
+        if task_type == "match_letters":
             parts = [
                 part.strip("().[]")
                 for part in re.split(r"[\s,;]+", cleaned_line)
                 if part.strip()
             ]
-            if not (len(parts) > 1 and all(re.fullmatch(r"[A-Za-z]", part) for part in parts)):
+            if not (
+                len(parts) > 1
+                and all(re.fullmatch(r"[A-Za-z]", part) for part in parts)
+            ):
                 match_letter_word = re.match(
                     r"^\s*(?:\(([A-Za-z])\)|\[([A-Za-z])\]|([A-Za-z]))\.?:?\s*(.*)$",
                     cleaned_line,
@@ -261,21 +214,83 @@ def postprocess_answer(text, query, task_type):
                     )
                     cleaned_line = letter.upper()
 
-        # Append the cleaned, non-empty line
         if cleaned_line:
             answers.append(cleaned_line)
 
-    #print('PARSED ANSWERS', answers)
-
-    # Compare against QUERY length: sometimes model forgets newlines
-    #print('QUERY', query)
     expected = expected_answer_count(query, task_type)
-    query_len = len(query.splitlines()) - 2
-    #print(query_len)
     if len(answers) == 1 and expected > 1:
         answers = split_single_line_answer(answers[0], expected, task_type)
     return answers
 
+
+def postprocess_answer(text, query, task_type):
+    """Keep only the content after the last 'FINAL ANSWERS' marker, one answer per line."""
+    text_after_marker = extract_raw_final(text)
+    if not text_after_marker.strip():
+        return []
+    return parse_answer_lines(text_after_marker, query, task_type)
+
+
+tok = load_tokenizer(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID, torch_dtype=torch.float16, device_map="auto"
+).eval()
+
+with open("/tmp/data/test.csv", encoding="utf-8", newline="") as f:
+    test_rows = list(csv.DictReader(f))
+
+outputs_queries_types = []
+for r in test_rows:
+
+    # Create the prompt.
+    messages = [
+        {"role": "system", "content": SYSTEM},
+        {"role": "user", "content":
+        f"CONTEXT:{r['context'].strip()}\nTASK TYPE:`{r['task_type']}`\n\nQUERY:{r['query'].strip()}"},
+    ]
+    ids = tok.apply_chat_template(
+        messages, add_generation_prompt=True, return_tensors="pt",
+    ).to(model.device)
+
+    # Generate the answer.
+    done = False
+    attempts = 0
+    text = ""
+    while not done and attempts < MAX_ATTEMPTS:
+        with torch.no_grad():
+            out = model.generate(
+                ids,
+                max_new_tokens=MAX_NEW_TOKENS,
+                do_sample=True,
+                temperature=TEMPERATURE,
+                top_p=TOP_P,)
+        text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
+        # Retry until we can extract a non-empty final-answers section.
+        attempts += 1
+        if extract_raw_final(text).strip():
+            done = True
+        else:
+            print(f'TRYING AGAIN...attempts #{attempts}/{MAX_ATTEMPTS}')
+
+    # Generate the explanation.
+    messages_post_explain = [
+        {"role": "system", "content": SYSTEM_POST_EXPLAIN},
+        {"role": "user", "content":
+        f"CONTEXT:{r['context'].strip()}\nTASK TYPE:`{r['task_type']}`\n\nQUERY:{r['query'].strip()}\n\nANSWER:{text.strip()}"},
+    ]
+    ids_post_explain = tok.apply_chat_template(
+        messages_post_explain, add_generation_prompt=True, return_tensors="pt",
+    ).to(model.device)
+    with torch.no_grad():
+        out_post_explain = model.generate(
+            ids_post_explain,
+            max_new_tokens=MAX_NEW_TOKENS,
+            do_sample=False)
+    text_post_explain = tok.decode(out_post_explain[0][ids_post_explain.shape[-1]:], skip_special_tokens=True).strip()
+
+    outputs_queries_types.append((text, text_post_explain, r['id'], r['query'], r['task_type']))
+    print(f"{len(outputs_queries_types)}/{len(test_rows)} done", flush=True)
+
 rows = []
 for answer, explanation, row_id, query, task_type in outputs_queries_types:
     answers = postprocess_answer(answer, query, task_type)
@@ -284,4 +299,4 @@ with open("submission.csv", "w", encoding="utf-8", newline="") as f:
     writer = csv.DictWriter(f, fieldnames=["id", "pred", "explanation"])
     writer.writeheader()
     writer.writerows(rows)
-print("wrote submission.csv", flush=True)
+print("wrote submission.csv", flush=True)