Upload script.py with huggingface_hub

script.py

@@ -40,11 +40,12 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 os.environ["HF_HUB_OFFLINE"] = "1"
 os.environ["TRANSFORMERS_OFFLINE"] = "1"
 MODEL_ID = "."
-MAX_NEW_TOKENS = 2000
+MAX_NEW_TOKENS = 1200  # was 2000; cut to fit T4 wall clock
+MAX_NEW_TOKENS_EXPLAIN = 600
 TEMPERATURE = 0.8
 TOP_P = 0.95
-NUM_SAMPLES = 5  # independent rollouts per problem
-MAX_ATTEMPTS = 2  # parse-retries per sample; ensemble covers the rest
+NUM_SAMPLES = 3  # independent rollouts per problem (was 5; timeout)
+MAX_ATTEMPTS = 1  # no per-sample retry; majority vote absorbs misses
 
 
 def load_tokenizer(model_id: str = "."):
@@ -357,7 +358,7 @@ def generate_explanation(
     with torch.no_grad():
         out = model.generate(
             ids,
-            max_new_tokens=MAX_NEW_TOKENS,
+            max_new_tokens=MAX_NEW_TOKENS_EXPLAIN,
             do_sample=False,
         )
     return tok.decode(out[0][ids.shape[-1] :], skip_special_tokens=True).strip()
@@ -371,52 +372,54 @@ model = AutoModelForCausalLM.from_pretrained(
 with open("/tmp/data/test.csv", encoding="utf-8", newline="") as f:
     test_rows = list(csv.DictReader(f))
 
-rows = []
-for idx, r in enumerate(test_rows, start=1):
-    messages = [
-        {"role": "system", "content": SYSTEM},
-        {
-            "role": "user",
-            "content": (
-                f"CONTEXT:{r['context'].strip()}\n"
-                f"TASK TYPE:`{r['task_type']}`\n\n"
-                f"QUERY:{r['query'].strip()}"
-            ),
-        },
-    ]
-    ids = tok.apply_chat_template(
-        messages, add_generation_prompt=True, return_tensors="pt",
-    ).to(model.device)
-
-    print(f"{idx}/{len(test_rows)} sampling {NUM_SAMPLES}x", flush=True)
-    raw_texts: list[str] = []
-    rollouts: list[list[str]] = []
-    for s in range(1, NUM_SAMPLES + 1):
-        text = generate_sample(tok, model, ids)
-        answers = postprocess_answer(text, r["query"], r["task_type"])
-        raw_texts.append(text)
-        rollouts.append(answers)
-        print(f"  sample {s}/{NUM_SAMPLES} parsed={answers!r}", flush=True)
-
-    expected = expected_answer_count(r["query"], r["task_type"])
-    voted = majority_vote(rollouts, expected, r["task_type"])
-    source_i = select_source_sample(rollouts, voted, r["task_type"])
-    print(f"  vote -> {voted!r} (explain from sample {source_i + 1})", flush=True)
-
-    explanation = generate_explanation(
-        tok, model, r["context"], r["task_type"], r["query"], raw_texts[source_i],
-    )
-
-    rows.append(
-        {
-            "id": r["id"],
-            "pred": json.dumps(voted, ensure_ascii=False),
-            "explanation": explanation,
-        }
-    )
-
+# Write incrementally so a wall-clock kill still leaves a partial submission.csv.
 with open("submission.csv", "w", encoding="utf-8", newline="") as f:
     writer = csv.DictWriter(f, fieldnames=["id", "pred", "explanation"])
     writer.writeheader()
-    writer.writerows(rows)
+    f.flush()
+
+    for idx, r in enumerate(test_rows, start=1):
+        messages = [
+            {"role": "system", "content": SYSTEM},
+            {
+                "role": "user",
+                "content": (
+                    f"CONTEXT:{r['context'].strip()}\n"
+                    f"TASK TYPE:`{r['task_type']}`\n\n"
+                    f"QUERY:{r['query'].strip()}"
+                ),
+            },
+        ]
+        ids = tok.apply_chat_template(
+            messages, add_generation_prompt=True, return_tensors="pt",
+        ).to(model.device)
+
+        print(f"{idx}/{len(test_rows)} sampling {NUM_SAMPLES}x", flush=True)
+        raw_texts: list[str] = []
+        rollouts: list[list[str]] = []
+        for s in range(1, NUM_SAMPLES + 1):
+            text = generate_sample(tok, model, ids)
+            answers = postprocess_answer(text, r["query"], r["task_type"])
+            raw_texts.append(text)
+            rollouts.append(answers)
+            print(f"  sample {s}/{NUM_SAMPLES} parsed={answers!r}", flush=True)
+
+        expected = expected_answer_count(r["query"], r["task_type"])
+        voted = majority_vote(rollouts, expected, r["task_type"])
+        source_i = select_source_sample(rollouts, voted, r["task_type"])
+        print(f"  vote -> {voted!r} (explain from sample {source_i + 1})", flush=True)
+
+        explanation = generate_explanation(
+            tok, model, r["context"], r["task_type"], r["query"], raw_texts[source_i],
+        )
+
+        writer.writerow(
+            {
+                "id": r["id"],
+                "pred": json.dumps(voted, ensure_ascii=False),
+                "explanation": explanation,
+            }
+        )
+        f.flush()
+
 print("wrote submission.csv", flush=True)