Update script.py

script.py

@@ -1,111 +1,38 @@
-"""
-IOL-AI 2024 submission — self-contained model script.
+import subprocess, sys
+subprocess.run([sys.executable, "-m", "pip", "install", "-q",
+                "transformers>=4.43", "accelerate>=0.30", "torch>=2.2", "pandas"], check=True)
 
-Paste this as `script.py` into an empty HF model repo and submit that repo.
-It installs its own deps, downloads the model from the Hub, reads the test
-set from /tmp/data/test.csv, and writes submission.csv. No weights or
-requirements.txt needed in the repo — just this one file.
-
-Pick a model that fits T4 medium (16 GB VRAM) and finishes 10 rows in < 1 h:
-
-  Thinking (set THINKING = True):
-    Qwen/Qwen3-14B                            ~9 GB
-    Qwen/Qwen3-8B                             ~5 GB
-    deepseek-ai/DeepSeek-R1-Distill-Qwen-14B  ~9 GB
-
-  Standard (set THINKING = False):
-    Qwen/Qwen2.5-14B-Instruct                 ~9 GB
-    google/gemma-3-12b-it                     ~8 GB
-    Qwen/Qwen2.5-7B-Instruct                  ~5 GB
-"""
-import subprocess
-import sys
-
-subprocess.run([
-    sys.executable, "-m", "pip", "install", "-q",
-    "transformers>=4.51", "bitsandbytes>=0.43", "accelerate>=0.30",
-    "torch>=2.2", "pandas",
-], check=True)
-
-import re
-import torch
+import json
 import pandas as pd
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# ── Configuration ─────────────────────────────────────────────────────────────
-MODEL_ID       = "Qwen/Qwen3-14B"
-THINKING       = True
-MAX_NEW_TOKENS = 4096 if THINKING else 512
-TEST_PATH      = "/tmp/data/test.csv"
+MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
 
-# ── Load model in 4-bit ───────────────────────────────────────────────────────
-bnb = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+tok = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    quantization_config=bnb,
-    device_map="auto",
-    trust_remote_code=True,
-)
-model.eval()
-
-# ── Prompt ────────────────────────────────────────────────────────────────────
-SYSTEM = (
-    "You are solving problems from the International Linguistics Olympiad. "
-    "Each problem gives you linguistic data and asks you to find patterns. "
-    "Provide one answer per line in the order the items appear. "
-    "Be concise — output only the answers, no commentary."
-)
-
-def build_prompt(row):
-    return row["context"].strip() + "\n\n" + row["query"].strip()
+    MODEL_ID, torch_dtype=torch.float16, device_map="auto"
+).eval()
 
-def strip_thinking(text):
-    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
+df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
 
-def run(row):
+rows = []
+for _, r in df.iterrows():
     messages = [
-        {"role": "system", "content": SYSTEM},
-        {"role": "user",   "content": build_prompt(row)},
+        {"role": "system", "content":
+            "You solve International Linguistics Olympiad problems. Answer every numbered "
+            "item. Put each answer on its own line, in order, with no numbering and no extra text."},
+        {"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"},
     ]
-    kwargs = {"enable_thinking": True} if THINKING else {}
-    text = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True, **kwargs
-    )
-    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+    ids = tok.apply_chat_template(
+        messages, add_generation_prompt=True, return_tensors="pt",
+    ).to(model.device)
     with torch.no_grad():
-        out = model.generate(
-            **inputs,
-            max_new_tokens=MAX_NEW_TOKENS,
-            do_sample=False,
-            pad_token_id=tokenizer.eos_token_id,
-        )
-    gen = out[0][inputs["input_ids"].shape[-1]:]
-    decoded = tokenizer.decode(gen, skip_special_tokens=True).strip()
-    return strip_thinking(decoded) if THINKING else decoded
-
-# ── Run inference ─────────────────────────────────────────────────────────────
-test_df = pd.read_csv(TEST_PATH, dtype=str).fillna("")
-
-preds = []
-for i, (_, row) in enumerate(test_df.iterrows()):
-    try:
-        pred = run(row)
-    except Exception as e:
-        print(f"[{i}] failed: {e}")
-        pred = ""
-    preds.append(pred)
-    print(f"  {i+1}/{len(test_df)} done")
-
-# ── Write submission ──────────────────────────────────────────────────────────
-pd.DataFrame({
-    "id":          test_df["id"],
-    "pred":        preds,
-    "explanation": [""] * len(preds),
-}).to_csv("submission.csv", index=False)
-print(f"Saved submission.csv — {len(preds)} rows")
+        out = model.generate(ids, max_new_tokens=512, do_sample=False)
+    text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
+    answers = [ln.strip() for ln in text.splitlines() if ln.strip()]
+    rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
+    print(f"{len(rows)}/{len(df)} done", flush=True)
+
+pd.DataFrame(rows).to_csv("submission.csv", index=False)
+print("wrote submission.csv", flush=True)