ProbioticFarmer's picture
Upload script.py from A100-80GB 4-tier competition evaluation
618e931 verified
Raw History Blame Contribute Delete
7.68 kB
"""Hugging Face Script Competition Entrypoint (`huggingface/competitions`) for
`ProbioticFarmer/Qwen3.8-27B-BarronPolyStep-RealCoder-v3`.
When executed inside a Hugging Face Competition GPU container (`a100-large`):
1. Reads the hidden/private test dataset from `COMPETITION_DATA_DIR` (or `--data-dir`),
supporting `test.jsonl` or `test.csv`.
2. Loads the 64-layer `Qwen/Qwen3.8-27B` hybrid backbone (`48 GatedDeltaNet` Triton +
`16 FlashAttention-2` layers) and applies the Layer-63 FP32 Barron PolyStep PEFT LoRA
adapter (`adapter_model.safetensors`, `lora_scale=2.0`).
3. Scores candidate code completions (or generates free-form solutions) using exact
multi-token sequence log-probability under the Barron PolyStep head.
4. Writes `submission.csv` and `submission.jsonl` for automated private-test grading.
"""
from __future__ import annotations
import argparse
import csv
import json
import os
import sys
import time
from typing import Any, Dict, List, Tuple
import numpy as np
import torch
import torch.nn.functional as F
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
if SCRIPT_DIR not in sys.path:
sys.path.insert(0, SCRIPT_DIR)
from cuda_qwen38_gpu import MLXQwen38Engine
def find_lora_path() -> str:
candidates = [
os.path.join(SCRIPT_DIR, "adapter_model.safetensors"),
os.path.join(
SCRIPT_DIR,
"checkpoints",
"Qwen3.8-27B-BarronPolyStep-RealCoder-v3",
"adapter_model.safetensors",
),
]
for p in candidates:
if os.path.exists(p):
return p
raise FileNotFoundError(f"Could not find adapter_model.safetensors in {candidates}")
@torch.inference_mode()
def score_problem_candidates(
engine: MLXQwen38Engine,
prefix: str,
candidates: List[str],
) -> Tuple[List[float], str]:
"""Score candidate completions using 64-layer forward pass + chunked candidate evaluation."""
p_ids = engine.tokenizer.encode(prefix).ids
engine.reset_cache()
logits_p = engine._prefill_chunk(p_ids, offset=0)
engine.cached_tokens = list(p_ids)
engine.snapshot_delta_state(len(p_ids), pinned=True)
lp0 = F.log_softmax(logits_p[0].float(), dim=-1)
scores: List[float] = []
for cand in candidates:
c_ids = engine.tokenizer.encode(cand).ids
if not c_ids:
scores.append(-1e9)
continue
total_lp = float(lp0[c_ids[0]].item())
if len(c_ids) > 1:
engine.restore_to_arbitrary_index(len(p_ids))
cur_off = len(p_ids)
for idx in range(len(c_ids) - 1):
lg = engine._prefill_chunk([c_ids[idx]], offset=cur_off)
cur_off += 1
lpf = F.log_softmax(lg[0].float(), dim=-1)
total_lp += float(lpf[c_ids[idx + 1]].item())
scores.append(total_lp / len(c_ids))
best_idx = int(np.argmax(scores))
return scores, candidates[best_idx]
def load_blind_test_items(data_dir: str) -> List[Dict[str, Any]]:
jsonl_path = os.path.join(data_dir, "test.jsonl")
csv_path = os.path.join(data_dir, "test.csv")
items: List[Dict[str, Any]] = []
if os.path.exists(jsonl_path):
with open(jsonl_path, encoding="utf-8") as f:
for line in f:
if line.strip():
items.append(json.loads(line))
return items
if os.path.exists(csv_path):
with open(csv_path, encoding="utf-8") as f:
reader = csv.DictReader(f)
for row in reader:
cands = json.loads(row["candidates"]) if "candidates" in row and row["candidates"] else []
items.append(
{
"id": row.get("id", ""),
"task_id": row.get("task_id", row.get("id", "")),
"prefix": row.get("prefix", row.get("prompt", "")),
"candidates": cands,
"suffix": row.get("suffix", ""),
}
)
return items
raise FileNotFoundError(f"No test.jsonl or test.csv found in {data_dir}")
def main() -> None:
parser = argparse.ArgumentParser(description="HF Script Competition Runner for PolyStep-RealCoder-v3")
parser.add_argument(
"--data-dir",
type=str,
default=os.environ.get("COMPETITION_DATA_DIR", "/tmp/data"),
help="Path to mounted private test dataset directory",
)
parser.add_argument(
"--output",
type=str,
default="submission.csv",
help="Output CSV path expected by Hugging Face Competitions evaluator",
)
parser.add_argument(
"--model-mode",
type=str,
choices=["polystep", "base"],
default="polystep",
help="Run with Barron PolyStep LoRA ('polystep') or unmodified Qwen3.8-27B ('base')",
)
args = parser.parse_args()
t0 = time.time()
lora_path = find_lora_path() if args.model_mode == "polystep" else None
engine = MLXQwen38Engine(lora_path=lora_path, lora_scale=2.0 if lora_path else 0.0)
if args.model_mode == "base":
engine.active_mlp_lora.clear()
engine.active_attn_lora.clear()
items = load_blind_test_items(args.data_dir)
print(
f"[HF Competition Runner] Loaded {len(items)} blind test problems from {args.data_dir} "
f"(mode={args.model_mode})",
flush=True,
)
predictions: List[Dict[str, Any]] = []
for idx, item in enumerate(items, 1):
prob_id = item.get("id", idx)
task_id = item.get("task_id", f"task_{prob_id}")
prefix = item["prefix"]
suffix = item.get("suffix", "")
candidates = item["candidates"]
scores, chosen_code = score_problem_candidates(engine, prefix, candidates)
pred_idx = int(np.argmax(scores))
sorted_sc = sorted(scores, reverse=True)
margin = float(sorted_sc[0] - sorted_sc[1]) if len(sorted_sc) > 1 else 0.0
predictions.append(
{
"id": prob_id,
"task_id": task_id,
"prediction_idx": pred_idx,
"predicted_code": chosen_code,
"confidence_margin_nats": round(margin, 4),
"solution": prefix + chosen_code + suffix,
}
)
print(
f" [Blind Test {idx:02d}/{len(items):02d}] id={prob_id} ({task_id}) -> "
f"pred_idx={pred_idx} ({chosen_code.strip()!r}, margin={margin:+.4f} nats)",
flush=True,
)
os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
with open(args.output, "w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(
f,
fieldnames=["id", "task_id", "prediction_idx", "predicted_code", "confidence_margin_nats"],
)
writer.writeheader()
for p in predictions:
writer.writerow(
{
"id": p["id"],
"task_id": p["task_id"],
"prediction_idx": p["prediction_idx"],
"predicted_code": p["predicted_code"],
"confidence_margin_nats": p["confidence_margin_nats"],
}
)
jsonl_out = os.path.splitext(args.output)[0] + ".jsonl"
with open(jsonl_out, "w", encoding="utf-8") as f:
for p in predictions:
f.write(json.dumps(p) + "\n")
print(
f"[HF Competition Runner] Wrote {len(predictions)} predictions to {args.output} "
f"and {jsonl_out} in {time.time() - t0:.2f}s.",
flush=True,
)
if __name__ == "__main__":
main()