import os os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" import re import json import time import sys import subprocess try: import bitsandbytes # noqa: F401 except ImportError: subprocess.run([sys.executable, "-m", "pip", "install", "-q", "bitsandbytes"], check=True) import bitsandbytes import torch import pandas as pd from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig # Use the local model directory (uploaded to HF repo root) MODEL_ID = "." TIME_LIMIT_S = 28 * 60 START = time.time() def time_left(): return TIME_LIMIT_S - (time.time() - START) def load_model(): """Load model with 4-bit quantization""" bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True, ) tok = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, quantization_config=bnb_config, device_map="auto", torch_dtype=torch.float16, local_files_only=True, ).eval() return tok, model def generate(messages, max_new_tokens=3500): ids = tok.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt" ).to(model.device) with torch.no_grad(): out = model.generate( ids, max_new_tokens=max_new_tokens, do_sample=False, temperature=None, top_p=None, pad_token_id=tok.eos_token_id, ) return tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip() def count_items(query: str) -> int: nums = re.findall(r"(?m)^\s*(\d+)\s*[.)]\s", query) if nums: return max(int(n) for n in nums) lines = [l for l in query.splitlines() if l.strip()] return max(1, len(lines)) def extract_json_list(text: str): text = text.strip() try: val = json.loads(text) if isinstance(val, list): return [str(x) for x in val] except Exception: pass m = re.search(r"\[.*\]", text, re.DOTALL) if m: try: val = json.loads(m.group(0)) if isinstance(val, list): return [str(x) for x in val] except Exception: pass return None def extract_answers_fallback(text: str, n: int): lines = [l.strip() for l in text.splitlines() if l.strip()] cleaned = [] for l in lines: l = re.sub(r"^\s*\d+[.)]\s*", "", l) l = l.strip(" -\t") if l: cleaned.append(l) if len(cleaned) >= n: return cleaned[:n] parts = [p.strip() for p in text.replace("\n", ",").split(",") if p.strip()] if len(parts) >= n: return parts[:n] cleaned = cleaned or parts or [text.strip()] while len(cleaned) < n: cleaned.append(cleaned[-1] if cleaned else "") return cleaned[:n] def fit_answers(answers, n): answers = list(answers) if len(answers) < n: answers = answers + [answers[-1] if answers else ""] * (n - len(answers)) return answers[:n] REASONING_SYSTEM = ( "You are an expert at International Linguistics Olympiad (IOL) problems. " "You will be given a self-contained puzzle about an unfamiliar language: some example " "data (context) and a query asking you to translate, match, fill in blanks, or convert " "numbers. Work out the underlying grammar/vocabulary rules carefully and systematically " "from the given examples only. Show your step-by-step reasoning: identify morphemes, " "patterns, correspondences, and test your hypothesis against every example before " "answering the query." ) EXTRACT_SYSTEM = ( "You convert a linguistics-puzzle solution into a strict output format. " "Given the original problem and a reasoning trace, output ONLY a single JSON object " "with exactly two keys:\n" ' "answers": a JSON list of strings, one per numbered item in the query, in order.\n' ' "explanation": a short (3-6 bullet points, plain text with \n between bullets) ' "human-readable summary of the key rules used to derive the answers. This is NOT the " "full reasoning trace, just a concise justification a person can read in under a minute.\n" "Output nothing except that JSON object." ) def solve_row(context: str, query: str, n_items: int): problem_text = f"{context.strip()}\n\n{query.strip()}" reasoning = "" if time_left() > 60: try: reasoning = generate( [ {"role": "system", "content": REASONING_SYSTEM}, {"role": "user", "content": problem_text}, ], max_new_tokens=3500, ) except Exception as e: reasoning = f"(reasoning pass failed: {e})" answers, explanation = None, "" if time_left() > 30: extract_prompt = ( f"PROBLEM:\n{problem_text}\n\n" f"REASONING TRACE:\n{reasoning}\n\n" f"The query has {n_items} numbered item(s). Now output the JSON object as instructed." ) try: raw = generate( [ {"role": "system", "content": EXTRACT_SYSTEM}, {"role": "user", "content": extract_prompt}, ], max_new_tokens=800, ) try: obj = json.loads(raw.strip()) except Exception: m = re.search(r"\{.*\}", raw, re.DOTALL) obj = json.loads(m.group(0)) if m else None if obj and isinstance(obj, dict): a = obj.get("answers") if isinstance(a, list): answers = [str(x) for x in a] explanation = str(obj.get("explanation", "")).strip() except Exception: pass if answers is None: candidate = extract_json_list(reasoning) answers = candidate if candidate is not None else extract_answers_fallback(reasoning, n_items) if not explanation: snippet = reasoning.strip().replace("\n", " ") explanation = (snippet[:400] + "...") if len(snippet) > 400 else snippet answers = fit_answers(answers, n_items) return answers, explanation def main(): # Load model once global tok, model tok, model = load_model() df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("") out_rows = [] for i, r in df.iterrows(): rid = r["id"] context, query = r.get("context", ""), r.get("query", "") n_items = count_items(query) if time_left() < 45: answers = [""] * n_items explanation = "" else: try: answers, explanation = solve_row(context, query, n_items) except Exception as e: answers = [""] * n_items explanation = f"(error: {e})" out_rows.append( { "id": rid, "pred": json.dumps(answers, ensure_ascii=False), "explanation": explanation, } ) print(f"[{i + 1}/{len(df)}] id={rid} n_items={n_items} time_left={time_left():.0f}s", flush=True) pd.DataFrame(out_rows, columns=["id", "pred", "explanation"]).to_csv( "submission.csv", index=False ) print("wrote submission.csv", flush=True) if __name__ == "__main__": main()