File size: 7,684 Bytes
618e931
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
"""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()