Upload script.py with huggingface_hub

script.py

@@ -0,0 +1,204 @@
+import os
+import re
+import json
+import pandas as pd
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# The repo is the working directory at run time, and there is no network.
+os.environ["HF_HUB_OFFLINE"] = "1"
+os.environ["TRANSFORMERS_OFFLINE"] = "1"
+MODEL_ID = "."
+MAX_NEW_TOKENS = 2000
+TEMPERATURE = 0.8
+TOP_P = 0.95
+MAX_ATTEMPTS = 5
+
+SYSTEM = (
+    "You solve International Linguistics Olympiad problems by reasoning from the "
+    "data in CONTEXT you are given to solve the problems in QUERY. \n"
+    "There are common TASK TYPES that we specify below, but "
+    "you may meet a TASK TYPE you have never seen: read the "
+    "instruction and the examples, and answer the QUERY in the same form they use.\n\n"
+    "Common TASK TYPES and what to return: \n"
+    "`translation`: return the translated form only, in the language the task asks for; \n"
+    "`fill_blanks`: return only the missing form for each indicated blank "
+    "(beware: this could be many different things: a word, a part of a word or a phonetic transcription---pay close attention to what part of the CONTEXT is missing in QUERY); \n"
+    "`match_letters`: return only the option letter (for example A, B, C); \n"
+    "`text_to_num`: return the number in digits; \n"
+    "`num_to_text`: return the number written out in words, in the language asked; \n"
+    "any other type: return exactly what the instruction asks for, nothing else. \n\n"
+    "As the first part of your answer, reason step by step about (1) the linguistic "
+    "rules that can be deduced from the given examples in CONTEXT, and (2) "
+    "how to apply them to the given problems in QUERY, and (3) in what format answers need to be returned (words, numbers, phonetic transcriptions, ...). \n"
+    "Then write a draft of the final answer. "
+    "Subsequently, compare it with the format requirements again, "
+    "and verify it's compliant with the deduced rules, and it is complete, i.e. has an answer for each element in QUERY. "
+    "If necessary, correct and refine."
+    "Finally, write a line that says exactly `FINAL ANSWERS:` "
+    "and, below it, write the answers to the items requested in QUERY (not those in CONTEXT),"
+    "one answer per line (separated by \n) in the order the items are asked for in the QUERY -- the "
+    "bare answer only, no numbering, no quotes, no extra text, according to the given TASK TYPE."
+)
+
+tok = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID, torch_dtype=torch.float16, device_map="auto"
+).eval()
+
+df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
+
+outputs_queries_types = []
+for _, r in df.iterrows():
+
+    # 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
+    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
+
+    outputs_queries_types.append((text, r['id'], r['query'], r['task_type']))
+    print(f"{len(outputs_queries_types)}/{len(df)} 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)
+        return len(numbered) or 1
+
+    if "blanks" in query.lower():
+        range_match = re.search(r"\((\d+)-(\d+)\)", query)
+        if range_match:
+            return int(range_match.group(2)) - int(range_match.group(1)) + 1
+        return len(re.findall(r"\(\d+\)", query)) or 1
+
+    numbered = re.findall(r"^\s*\d+[.)]", query, re.MULTILINE)
+    return len(numbered) or 1
+
+
+def split_single_line_answer(text: str, expected: int, task_type: str) -> list[str]:
+    text = text.strip()
+    if expected <= 1:
+        return [text]
+
+    def try_split(pattern: str) -> list[str] | None:
+        parts = [part.strip() for part in re.split(pattern, text) if part.strip()]
+        return parts if len(parts) == expected else None
+
+    if task_type == "match_letters":
+        for pattern in (r"\s+", r",\s*", r";\s*"):
+            if result := try_split(pattern):
+                return result
+        letters = re.findall(r"[A-Za-z]", text)
+        if len(letters) == expected:
+            return [letter.upper() for letter in letters]
+        return [text]
+
+    if task_type in ("text_to_num", "num_to_text"):
+        for pattern in (r",\s*", r";\s*", r"\s+"):
+            if result := try_split(pattern):
+                return result
+        return [text]
+
+    for pattern in (r";\s*", r",\s*"):
+        if result := try_split(pattern):
+            return result
+    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 []
+
+    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
+
+        # 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':
+            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)):
+                match_letter_word = re.match(
+                    r"^\s*(?:\(([A-Za-z])\)|\[([A-Za-z])\]|([A-Za-z]))\.?:?\s*(.*)$",
+                    cleaned_line,
+                )
+                if match_letter_word:
+                    letter = (
+                        match_letter_word.group(1)
+                        or match_letter_word.group(2)
+                        or match_letter_word.group(3)
+                    )
+                    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
+
+rows = []
+for answer, row_id, query, task_type in outputs_queries_types:
+    answers = postprocess_answer(answer, query, task_type)
+    rows.append({"id": row_id, "pred": json.dumps(answers, ensure_ascii=False)})
+pd.DataFrame(rows).to_csv("submission.csv", index=False)
+print("wrote submission.csv", flush=True)