Initialize math_trainer Space for GPU SOTA training and eval.

.gitignore

@@ -0,0 +1,3 @@
+workspace/
+__pycache__/
+scripts/__pycache__/

README.md

@@ -1,12 +1,33 @@
 ---
-title: Math Trainer
-emoji: 🏆
-colorFrom: purple
-colorTo: indigo
+title: Math Conjecture Trainer
 sdk: gradio
-sdk_version: 6.6.0
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Math Conjecture Trainer Space
+
+This Space launches multi-stage `DeepSeek-Math` training for the conjecture
+project and pushes adapters to:
+
+- `NorthernTribe-Research/math-conjecture-model`
+
+## What this Space does
+
+1. Downloads dataset parquet splits from:
+   `NorthernTribe-Research/math-conjecture-training-corpus`
+2. Runs `scripts/train_sota.py` with staged curriculum config
+3. Optionally runs `scripts/eval_sota.py`
+4. Pushes training artifacts to your model repo on Hugging Face
+
+## Required Space secret
+
+Set this in Space Settings > Variables and Secrets:
+
+- `HF_TOKEN`: write-access token for your Hugging Face account
+
+## Default training config
+
+- `configs/deepseek_math_sota.yaml`
+- base model default: `deepseek-ai/deepseek-math-v2`
+- output root: `workspace/runs/math-conjecture-sota`

app.py

@@ -0,0 +1,272 @@
+#!/usr/bin/env python3
+"""Gradio app to run SOTA conjecture-model training on Hugging Face Space GPU."""
+
+from __future__ import annotations
+
+import datetime as dt
+import json
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+from typing import Generator, List, Tuple
+
+import gradio as gr
+import yaml
+from huggingface_hub import hf_hub_download
+
+
+ROOT = Path(__file__).resolve().parent
+WORKSPACE_DIR = ROOT / "workspace"
+DATA_DIR = WORKSPACE_DIR / "data" / "releases" / "v1"
+RUNTIME_DIR = WORKSPACE_DIR / "runtime"
+CONFIG_TEMPLATE = ROOT / "configs" / "deepseek_math_sota.yaml"
+TRAIN_SCRIPT = ROOT / "scripts" / "train_sota.py"
+EVAL_SCRIPT = ROOT / "scripts" / "eval_sota.py"
+
+
+def now_ts() -> str:
+    return dt.datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")
+
+
+def append_log(lines: List[str], message: str) -> str:
+    lines.append(f"[{now_ts()}] {message}")
+    text = "\n".join(lines)
+    if len(text) > 200_000:
+        text = text[-200_000:]
+    return text
+
+
+def ensure_workspace() -> None:
+    DATA_DIR.mkdir(parents=True, exist_ok=True)
+    RUNTIME_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def download_dataset(dataset_repo_id: str, token: str, log_lines: List[str]) -> Tuple[str, str]:
+    ensure_workspace()
+    for split in ("train", "validation", "test"):
+        cached_path = hf_hub_download(
+            repo_id=dataset_repo_id,
+            repo_type="dataset",
+            filename=f"{split}.parquet",
+            token=token,
+        )
+        out_path = DATA_DIR / f"{split}.parquet"
+        shutil.copy2(cached_path, out_path)
+        append_log(log_lines, f"Downloaded {split}.parquet to {out_path}")
+    return str(DATA_DIR / "train.parquet"), str(DATA_DIR / "validation.parquet")
+
+
+def write_runtime_config(
+    base_model_id: str,
+    model_repo_id: str,
+    train_file: str,
+    validation_file: str,
+) -> Path:
+    cfg = yaml.safe_load(CONFIG_TEMPLATE.read_text(encoding="utf-8"))
+    cfg["model"]["base_model"] = base_model_id
+    cfg["hub"]["repo_id"] = model_repo_id
+    cfg["data"]["default_train_file"] = train_file
+    cfg["data"]["default_validation_file"] = validation_file
+    cfg["global"]["output_root"] = str(WORKSPACE_DIR / "runs" / "math-conjecture-sota")
+    runtime_path = RUNTIME_DIR / "deepseek_math_sota.runtime.yaml"
+    runtime_path.write_text(yaml.safe_dump(cfg, sort_keys=False), encoding="utf-8")
+    return runtime_path
+
+
+def stream_subprocess(
+    cmd: List[str],
+    env: dict,
+    cwd: Path,
+    log_lines: List[str],
+    status_prefix: str,
+) -> Generator[Tuple[str, str], None, int]:
+    append_log(log_lines, f"Running command: {' '.join(cmd)}")
+    yield "\n".join(log_lines), f"{status_prefix}: running"
+    proc = subprocess.Popen(
+        cmd,
+        cwd=str(cwd),
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+    )
+    assert proc.stdout is not None
+    for line in proc.stdout:
+        line = line.rstrip()
+        if line:
+            append_log(log_lines, line)
+            yield "\n".join(log_lines), f"{status_prefix}: running"
+    ret = proc.wait()
+    append_log(log_lines, f"{status_prefix} finished with exit code {ret}")
+    yield "\n".join(log_lines), f"{status_prefix}: {'ok' if ret == 0 else 'failed'}"
+    return ret
+
+
+def run_pipeline(
+    hf_token: str,
+    dataset_repo_id: str,
+    model_repo_id: str,
+    base_model_id: str,
+    start_stage: int,
+    max_stages: int,
+    run_eval: bool,
+    eval_k: int,
+    eval_samples: int,
+) -> Generator[Tuple[str, str], None, None]:
+    log_lines: List[str] = []
+    try:
+        token = (hf_token or "").strip() or os.environ.get("HF_TOKEN", "").strip()
+        if not token:
+            raise ValueError("HF token is required. Set HF_TOKEN secret or provide it in the UI.")
+
+        dataset_repo_id = dataset_repo_id.strip()
+        model_repo_id = model_repo_id.strip()
+        base_model_id = base_model_id.strip()
+        if not dataset_repo_id or not model_repo_id or not base_model_id:
+            raise ValueError("dataset repo, model repo, and base model must be non-empty.")
+
+        append_log(log_lines, "Preparing local workspace.")
+        yield "\n".join(log_lines), "Preparing workspace"
+        train_file, validation_file = download_dataset(dataset_repo_id, token, log_lines)
+        yield "\n".join(log_lines), "Dataset ready"
+
+        runtime_cfg = write_runtime_config(
+            base_model_id=base_model_id,
+            model_repo_id=model_repo_id,
+            train_file=train_file,
+            validation_file=validation_file,
+        )
+        append_log(log_lines, f"Wrote runtime config: {runtime_cfg}")
+        yield "\n".join(log_lines), "Config ready"
+
+        env = os.environ.copy()
+        env["HF_TOKEN"] = token
+        env["HUGGINGFACE_HUB_TOKEN"] = token
+        env["PYTHONUNBUFFERED"] = "1"
+
+        train_cmd = [sys.executable, str(TRAIN_SCRIPT), "--config", str(runtime_cfg)]
+        train_cmd += ["--start-stage", str(int(start_stage))]
+        train_cmd += ["--max-stages", str(int(max_stages))]
+        train_gen = stream_subprocess(
+            cmd=train_cmd,
+            env=env,
+            cwd=ROOT,
+            log_lines=log_lines,
+            status_prefix="Training",
+        )
+        train_ret = None
+        while True:
+            try:
+                yield next(train_gen)
+            except StopIteration as stop:
+                train_ret = stop.value
+                break
+        if train_ret != 0:
+            return
+
+        if run_eval:
+            eval_report = WORKSPACE_DIR / "runs" / "latest_eval_report.json"
+            eval_cmd = [
+                sys.executable,
+                str(EVAL_SCRIPT),
+                "--config",
+                str(runtime_cfg),
+                "--base-model",
+                base_model_id,
+                "--adapter-path",
+                str(WORKSPACE_DIR / "runs" / "math-conjecture-sota" / "final_adapter"),
+                "--eval-file",
+                str(DATA_DIR / "test.parquet"),
+                "--k",
+                str(int(eval_k)),
+                "--max-samples",
+                str(int(eval_samples)),
+                "--output-json",
+                str(eval_report),
+            ]
+            eval_gen = stream_subprocess(
+                cmd=eval_cmd,
+                env=env,
+                cwd=ROOT,
+                log_lines=log_lines,
+                status_prefix="Evaluation",
+            )
+            eval_ret = None
+            while True:
+                try:
+                    yield next(eval_gen)
+                except StopIteration as stop:
+                    eval_ret = stop.value
+                    break
+            if eval_ret == 0 and eval_report.exists():
+                report = json.loads(eval_report.read_text(encoding="utf-8"))
+                summary = {
+                    "evaluated_rows": report.get("evaluated_rows"),
+                    "pass_at_1": report.get("pass_at_1"),
+                    "pass_at_k": report.get("pass_at_k"),
+                    "k": report.get("k"),
+                }
+                append_log(log_lines, f"Eval summary: {json.dumps(summary)}")
+
+        append_log(log_lines, "Pipeline completed.")
+        yield "\n".join(log_lines), "Completed"
+    except Exception as exc:
+        append_log(log_lines, f"Pipeline failed: {type(exc).__name__}: {exc}")
+        yield "\n".join(log_lines), "Failed"
+
+
+with gr.Blocks(title="Math Conjecture Trainer Space") as demo:
+    gr.Markdown(
+        """
+        # Math Conjecture Trainer
+        Launch multi-stage DeepSeek-Math fine-tuning on Space GPU and push adapters to your model repo.
+        """
+    )
+    with gr.Row():
+        hf_token = gr.Textbox(label="HF Token", type="password", placeholder="hf_xxx (or set HF_TOKEN secret)")
+        dataset_repo_id = gr.Textbox(
+            label="Dataset Repo",
+            value="NorthernTribe-Research/math-conjecture-training-corpus",
+        )
+    with gr.Row():
+        model_repo_id = gr.Textbox(
+            label="Model Repo",
+            value="NorthernTribe-Research/math-conjecture-model",
+        )
+        base_model_id = gr.Textbox(
+            label="Base Model",
+            value="deepseek-ai/deepseek-math-v2",
+        )
+    with gr.Row():
+        start_stage = gr.Slider(label="Start Stage", minimum=1, maximum=3, step=1, value=1)
+        max_stages = gr.Slider(label="How Many Stages", minimum=1, maximum=3, step=1, value=3)
+        run_eval = gr.Checkbox(label="Run Evaluation After Training", value=True)
+    with gr.Row():
+        eval_k = gr.Slider(label="Eval K", minimum=1, maximum=8, step=1, value=4)
+        eval_samples = gr.Slider(label="Eval Max Samples", minimum=50, maximum=1000, step=50, value=300)
+    run_button = gr.Button("Start Train + Push")
+    status = gr.Textbox(label="Status", value="Idle")
+    logs = gr.Textbox(label="Logs", lines=24, max_lines=30, show_copy_button=True)
+
+    run_button.click(
+        fn=run_pipeline,
+        inputs=[
+            hf_token,
+            dataset_repo_id,
+            model_repo_id,
+            base_model_id,
+            start_stage,
+            max_stages,
+            run_eval,
+            eval_k,
+            eval_samples,
+        ],
+        outputs=[logs, status],
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=1).launch()

configs/deepseek_math_sota.yaml

@@ -0,0 +1,113 @@
+global:
+  output_root: workspace/runs/math-conjecture-sota
+  seed: 17
+
+model:
+  base_model: deepseek-ai/deepseek-math-v2
+  trust_remote_code: true
+  use_bf16: true
+  load_in_4bit: true
+  bnb_4bit_quant_type: nf4
+  bnb_4bit_use_double_quant: true
+  attn_implementation: null
+  lora:
+    r: 96
+    alpha: 192
+    dropout: 0.05
+    bias: none
+    target_modules:
+      - q_proj
+      - k_proj
+      - v_proj
+      - o_proj
+      - gate_proj
+      - up_proj
+      - down_proj
+
+data:
+  default_train_file: workspace/data/releases/v1/train.parquet
+  default_validation_file: workspace/data/releases/v1/validation.parquet
+  prompt_field: prompt
+  target_field: target
+  final_answer_field: final_answer
+  proof_field: proof_formal
+  sample_weight_field: sample_weight
+  max_seq_length: 3072
+  min_loss_weight: 0.25
+  max_loss_weight: 6.0
+  family_boost:
+    conjecture_core: 2.5
+    formal_proof: 1.6
+    competition: 1.2
+    structured_reasoning: 1.0
+  system_prompt: |
+    You are a frontier mathematical reasoning model focused on unsolved
+    conjectures. Your outputs must be precise, technically coherent, and explicit
+    about uncertainty. Never claim a full proof unless it is derivable from given
+    assumptions or already established in cited prior results.
+
+training_defaults:
+  per_device_train_batch_size: 1
+  per_device_eval_batch_size: 1
+  gradient_accumulation_steps: 16
+  weight_decay: 0.01
+  warmup_ratio: 0.03
+  lr_scheduler_type: cosine
+  max_grad_norm: 1.0
+  gradient_checkpointing: true
+  logging_steps: 10
+  save_steps: 400
+  eval_steps: 400
+  save_total_limit: 3
+  dataloader_num_workers: 2
+
+stages:
+  - name: broad_math_bootstrap
+    max_train_samples: null
+    max_eval_samples: 3000
+    filters:
+      include_families:
+        - competition
+        - structured_reasoning
+        - formal_proof
+        - conjecture_core
+    training:
+      num_train_epochs: 1
+      learning_rate: 2.0e-5
+
+  - name: conjecture_specialization
+    max_train_samples: null
+    max_eval_samples: 2000
+    filters:
+      include_families:
+        - conjecture_core
+        - formal_proof
+      min_sample_weight: 2.0
+    training:
+      num_train_epochs: 2
+      learning_rate: 8.0e-6
+      save_steps: 250
+      eval_steps: 250
+
+  - name: conjecture_alignment
+    max_train_samples: null
+    max_eval_samples: null
+    filters:
+      include_families:
+        - conjecture_core
+      require_conjecture_id: true
+    training:
+      num_train_epochs: 3
+      learning_rate: 5.0e-6
+      save_steps: 100
+      eval_steps: 100
+
+hub:
+  push_to_hub: true
+  repo_id: NorthernTribe-Research/math-conjecture-model
+  private: false
+  upload_stage_checkpoints: true
+  commit_message: Train multi-stage SOTA curriculum for conjecture reasoning.
+
+credentials:
+  path: huggingface-api-key.json

requirements.txt

@@ -0,0 +1,10 @@
+gradio>=5.0.0
+torch>=2.3.0
+transformers>=4.48.0
+accelerate>=1.1.0
+datasets>=2.21.0
+peft>=0.14.0
+bitsandbytes>=0.45.0
+huggingface_hub>=0.26.0
+pyyaml>=6.0.2
+sentencepiece>=0.2.0

scripts/eval_sota.py

@@ -0,0 +1,299 @@
+#!/usr/bin/env python3
+"""Self-consistency evaluation for math-conjecture model checkpoints."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence
+
+import torch
+import yaml
+from datasets import load_dataset
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run pass@k-style evaluation on held-out split.")
+    parser.add_argument(
+        "--config",
+        type=Path,
+        default=Path("configs/deepseek_math_sota.yaml"),
+        help="Training config used for prompt formatting defaults.",
+    )
+    parser.add_argument(
+        "--base-model",
+        type=str,
+        default=None,
+        help="Override base model id from config.",
+    )
+    parser.add_argument(
+        "--adapter-path",
+        type=Path,
+        default=None,
+        help="Optional LoRA adapter path to load on top of base model.",
+    )
+    parser.add_argument(
+        "--eval-file",
+        type=Path,
+        default=Path("data/releases/v1/test.parquet"),
+        help="Parquet split used for evaluation.",
+    )
+    parser.add_argument("--max-samples", type=int, default=300, help="Maximum evaluation rows.")
+    parser.add_argument("--k", type=int, default=4, help="Number of sampled generations per prompt.")
+    parser.add_argument("--max-new-tokens", type=int, default=256, help="Generation length cap.")
+    parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature.")
+    parser.add_argument("--top-p", type=float, default=0.95, help="Nucleus sampling p.")
+    parser.add_argument("--seed", type=int, default=17, help="Random seed.")
+    parser.add_argument(
+        "--output-json",
+        type=Path,
+        default=Path("workspace/runs/latest_eval_report.json"),
+        help="Where to write evaluation report.",
+    )
+    return parser.parse_args()
+
+
+def as_text(value: Any) -> str:
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        return value.strip()
+    return str(value).strip()
+
+
+def load_config(path: Path) -> Dict[str, Any]:
+    cfg = yaml.safe_load(path.read_text(encoding="utf-8"))
+    if not isinstance(cfg, dict):
+        raise ValueError("Invalid YAML config.")
+    return cfg
+
+
+def normalize_answer(text: str) -> str:
+    text = text.strip().lower()
+    text = re.sub(r"\s+", " ", text)
+    text = text.replace("$", "")
+    return text
+
+
+def flatten_expected(row: Dict[str, Any], data_cfg: Dict[str, Any]) -> List[str]:
+    out: List[str] = []
+    final_field = as_text(data_cfg.get("final_answer_field")) or "final_answer"
+    target_field = as_text(data_cfg.get("target_field")) or "target"
+
+    final_answer = row.get(final_field)
+    if final_answer is not None:
+        txt = as_text(final_answer)
+        if txt:
+            out.append(txt)
+
+    target = row.get(target_field)
+    if target is None:
+        return out
+    if isinstance(target, str):
+        stripped = target.strip()
+        if not stripped:
+            return out
+        try:
+            target = json.loads(stripped)
+        except json.JSONDecodeError:
+            out.append(stripped)
+            return out
+
+    if isinstance(target, dict):
+        for value in target.values():
+            if isinstance(value, list):
+                for item in value:
+                    txt = as_text(item)
+                    if txt:
+                        out.append(txt)
+            else:
+                txt = as_text(value)
+                if txt:
+                    out.append(txt)
+    elif isinstance(target, list):
+        for item in target:
+            txt = as_text(item)
+            if txt:
+                out.append(txt)
+    else:
+        txt = as_text(target)
+        if txt:
+            out.append(txt)
+    return out
+
+
+def build_user_block(row: Dict[str, Any], data_cfg: Dict[str, Any]) -> str:
+    prompt_field = as_text(data_cfg.get("prompt_field")) or "prompt"
+    prompt = as_text(row.get(prompt_field))
+    if not prompt:
+        prompt = "Solve the math task."
+    meta_fields = [
+        ("task_type", "Task type"),
+        ("family", "Family"),
+        ("difficulty", "Difficulty"),
+        ("source_dataset", "Source"),
+        ("status_as_of", "Status as of"),
+    ]
+    lines = []
+    for key, label in meta_fields:
+        value = as_text(row.get(key))
+        if value:
+            lines.append(f"{label}: {value}")
+    if lines:
+        return f"{prompt}\n\nMetadata:\n" + "\n".join(lines)
+    return prompt
+
+
+def build_prompt_text(row: Dict[str, Any], tokenizer: AutoTokenizer, data_cfg: Dict[str, Any]) -> str:
+    system_prompt = as_text(data_cfg.get("system_prompt"))
+    if not system_prompt:
+        system_prompt = "You are a rigorous mathematical reasoning assistant."
+    user_block = build_user_block(row, data_cfg)
+    if getattr(tokenizer, "chat_template", None):
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_block},
+        ]
+        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    return f"System:\n{system_prompt}\n\nUser:\n{user_block}\n\nAssistant:\n"
+
+
+def extract_candidate_text(full_generation: str, prompt_text: str) -> str:
+    if full_generation.startswith(prompt_text):
+        return full_generation[len(prompt_text) :].strip()
+    return full_generation.strip()
+
+
+def is_match(candidate: str, expected_values: Sequence[str]) -> bool:
+    cand_norm = normalize_answer(candidate)
+    if not cand_norm:
+        return False
+    for expected in expected_values:
+        exp_norm = normalize_answer(expected)
+        if not exp_norm:
+            continue
+        if exp_norm in cand_norm or cand_norm in exp_norm:
+            return True
+        boxed = re.findall(r"\\boxed\{([^{}]+)\}", cand_norm)
+        if boxed and any(exp_norm in item for item in boxed):
+            return True
+    return False
+
+
+def load_model_and_tokenizer(
+    base_model: str,
+    adapter_path: Optional[Path],
+    trust_remote_code: bool,
+) -> tuple[Any, AutoTokenizer]:
+    tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=trust_remote_code, use_fast=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token or tokenizer.unk_token
+    if tokenizer.pad_token is None:
+        tokenizer.add_special_tokens({"pad_token": "<|pad|>"})
+
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto" if torch.cuda.is_available() else None,
+        trust_remote_code=trust_remote_code,
+    )
+    if adapter_path is not None:
+        model = PeftModel.from_pretrained(model, str(adapter_path))
+    model.eval()
+    return model, tokenizer
+
+
+def main() -> None:
+    args = parse_args()
+    cfg = load_config(args.config)
+    data_cfg = cfg.get("data", {})
+    model_cfg = cfg.get("model", {})
+    set_seed(args.seed)
+
+    base_model = args.base_model or as_text(model_cfg.get("base_model"))
+    if not base_model:
+        raise ValueError("Base model is required via --base-model or config.model.base_model.")
+
+    model, tokenizer = load_model_and_tokenizer(
+        base_model=base_model,
+        adapter_path=args.adapter_path,
+        trust_remote_code=bool(model_cfg.get("trust_remote_code", False)),
+    )
+
+    if not args.eval_file.exists():
+        raise FileNotFoundError(f"Evaluation file not found: {args.eval_file}")
+    ds = load_dataset("parquet", data_files={"eval": str(args.eval_file)})["eval"]
+
+    if args.max_samples > 0 and args.max_samples < len(ds):
+        ds = ds.select(range(args.max_samples))
+
+    total = 0
+    hit_at_1 = 0
+    hit_at_k = 0
+    records = []
+
+    for row in ds:
+        expected_values = flatten_expected(row, data_cfg)
+        if not expected_values:
+            continue
+        prompt_text = build_prompt_text(row, tokenizer, data_cfg)
+        inputs = tokenizer(prompt_text, return_tensors="pt", truncation=True, max_length=4096)
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            output_ids = model.generate(
+                **inputs,
+                do_sample=True,
+                temperature=args.temperature,
+                top_p=args.top_p,
+                num_return_sequences=args.k,
+                max_new_tokens=args.max_new_tokens,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+        generations = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+        candidates = [extract_candidate_text(text, prompt_text) for text in generations]
+        matches = [is_match(candidate, expected_values) for candidate in candidates]
+        total += 1
+        if matches and matches[0]:
+            hit_at_1 += 1
+        if any(matches):
+            hit_at_k += 1
+
+        records.append(
+            {
+                "uid": as_text(row.get("uid")),
+                "prompt": as_text(row.get(as_text(data_cfg.get("prompt_field")) or "prompt")),
+                "expected_values": expected_values[:5],
+                "candidates": candidates,
+                "matches": matches,
+            }
+        )
+
+    pass_at_1 = (hit_at_1 / total) if total else 0.0
+    pass_at_k = (hit_at_k / total) if total else 0.0
+    report = {
+        "base_model": base_model,
+        "adapter_path": str(args.adapter_path) if args.adapter_path is not None else None,
+        "eval_file": str(args.eval_file),
+        "evaluated_rows": total,
+        "k": args.k,
+        "pass_at_1": pass_at_1,
+        "pass_at_k": pass_at_k,
+        "temperature": args.temperature,
+        "top_p": args.top_p,
+        "max_new_tokens": args.max_new_tokens,
+        "samples": records[:30],
+    }
+    args.output_json.parent.mkdir(parents=True, exist_ok=True)
+    args.output_json.write_text(json.dumps(report, ensure_ascii=True, indent=2), encoding="utf-8")
+    print(json.dumps({k: report[k] for k in ("evaluated_rows", "pass_at_1", "pass_at_k", "k")}, indent=2))
+    print(f"Saved report to {args.output_json}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train_sota.py

@@ -0,0 +1,688 @@
+#!/usr/bin/env python3
+"""Multi-stage curriculum SFT for advancing the conjecture math model."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+import yaml
+from datasets import Dataset, DatasetDict, load_dataset
+from huggingface_hub import HfApi
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from torch.utils.data import WeightedRandomSampler
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    DataCollatorForSeq2Seq,
+    Trainer,
+    TrainingArguments,
+    set_seed,
+)
+
+DEFAULT_CONFIG_PATH = Path("configs/deepseek_math_sota.yaml")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Train DeepSeek-Math with a multi-stage SOTA curriculum recipe."
+    )
+    parser.add_argument(
+        "--config",
+        type=Path,
+        default=DEFAULT_CONFIG_PATH,
+        help="Path to multi-stage YAML config.",
+    )
+    parser.add_argument("--repo-id", type=str, default=None, help="Override hub.repo_id.")
+    parser.add_argument("--push-to-hub", action="store_true", help="Force push enabled.")
+    parser.add_argument("--no-push-to-hub", action="store_true", help="Force push disabled.")
+    parser.add_argument(
+        "--start-stage",
+        type=int,
+        default=1,
+        help="1-based stage index to start from.",
+    )
+    parser.add_argument(
+        "--max-stages",
+        type=int,
+        default=None,
+        help="Optional number of stages to run from --start-stage.",
+    )
+    parser.add_argument(
+        "--credentials-path",
+        type=Path,
+        default=None,
+        help="Override credentials.path.",
+    )
+    return parser.parse_args()
+
+
+def as_text(value: Any) -> str:
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        return value.strip()
+    return str(value).strip()
+
+
+def as_float(value: Any, default: float) -> float:
+    if value is None:
+        return default
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return default
+
+
+def as_int(value: Any, default: int) -> int:
+    if value is None:
+        return default
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return default
+
+
+def load_config(path: Path) -> Dict[str, Any]:
+    if not path.exists():
+        raise FileNotFoundError(f"Config not found: {path}")
+    cfg = yaml.safe_load(path.read_text(encoding="utf-8"))
+    if not isinstance(cfg, dict):
+        raise ValueError(f"Invalid config format: {path}")
+    for key in ("model", "data", "stages"):
+        if key not in cfg:
+            raise ValueError(f"Missing config section: {key}")
+    if not isinstance(cfg["stages"], list) or not cfg["stages"]:
+        raise ValueError("Config must contain at least one stage in stages[].")
+    cfg.setdefault("global", {})
+    cfg.setdefault("training_defaults", {})
+    cfg.setdefault("hub", {})
+    cfg.setdefault("credentials", {})
+    return cfg
+
+
+def apply_overrides(cfg: Dict[str, Any], args: argparse.Namespace) -> None:
+    if args.repo_id:
+        cfg.setdefault("hub", {})["repo_id"] = args.repo_id
+    if args.credentials_path is not None:
+        cfg.setdefault("credentials", {})["path"] = str(args.credentials_path)
+    if args.push_to_hub and args.no_push_to_hub:
+        raise ValueError("Cannot set both --push-to-hub and --no-push-to-hub.")
+    if args.push_to_hub:
+        cfg.setdefault("hub", {})["push_to_hub"] = True
+    if args.no_push_to_hub:
+        cfg.setdefault("hub", {})["push_to_hub"] = False
+
+
+def resolve_auth(cfg: Dict[str, Any]) -> Tuple[Optional[str], Optional[str]]:
+    token = as_text(os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")) or None
+    username = as_text(os.environ.get("HF_USERNAME")) or None
+    cred_path = as_text(cfg.get("credentials", {}).get("path"))
+    if cred_path:
+        path = Path(cred_path)
+        if path.exists():
+            data = json.loads(path.read_text(encoding="utf-8"))
+            if token is None:
+                token = as_text(data.get("key")) or None
+            if username is None:
+                username = as_text(data.get("username")) or None
+    return token, username
+
+
+def resolve_repo_id(cfg: Dict[str, Any], username: Optional[str], output_root: Path) -> Optional[str]:
+    repo_id = as_text(cfg.get("hub", {}).get("repo_id"))
+    if repo_id:
+        return repo_id
+    if not username:
+        return None
+    return f"{username}/{output_root.name}"
+
+
+def stringify_structured(value: Any) -> str:
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        text = value.strip()
+        if not text:
+            return ""
+        try:
+            parsed = json.loads(text)
+        except json.JSONDecodeError:
+            return text
+        return json.dumps(parsed, ensure_ascii=False, sort_keys=True)
+    return json.dumps(value, ensure_ascii=False, sort_keys=True)
+
+
+def build_user_block(row: Dict[str, Any], data_cfg: Dict[str, Any]) -> str:
+    prompt_field = as_text(data_cfg.get("prompt_field")) or "prompt"
+    prompt = as_text(row.get(prompt_field))
+    if not prompt:
+        prompt = "Solve the math task."
+    meta_fields = [
+        ("task_type", "Task type"),
+        ("family", "Family"),
+        ("difficulty", "Difficulty"),
+        ("source_dataset", "Source"),
+        ("status_as_of", "Status as of"),
+    ]
+    meta_lines = []
+    for key, label in meta_fields:
+        value = as_text(row.get(key))
+        if value:
+            meta_lines.append(f"{label}: {value}")
+    tags = row.get("topic_tags")
+    if isinstance(tags, list) and tags:
+        tag_text = ", ".join(as_text(tag) for tag in tags if as_text(tag))
+        if tag_text:
+            meta_lines.append(f"Tags: {tag_text}")
+    if not meta_lines:
+        return prompt
+    return f"{prompt}\n\nMetadata:\n" + "\n".join(meta_lines)
+
+
+def build_answer_block(row: Dict[str, Any], data_cfg: Dict[str, Any]) -> str:
+    target_field = as_text(data_cfg.get("target_field")) or "target"
+    final_answer_field = as_text(data_cfg.get("final_answer_field")) or "final_answer"
+    proof_field = as_text(data_cfg.get("proof_field")) or "proof_formal"
+
+    sections = []
+    target_text = stringify_structured(row.get(target_field))
+    if target_text:
+        sections.append(f"Structured target:\n{target_text}")
+
+    final_answer = stringify_structured(row.get(final_answer_field))
+    if final_answer:
+        sections.append(f"Final answer:\n{final_answer}")
+
+    proof_text = stringify_structured(row.get(proof_field))
+    if proof_text:
+        sections.append(f"Formal proof snippet:\n{proof_text}")
+
+    if not sections:
+        sections.append("No structured target provided.")
+    return "\n\n".join(sections).strip()
+
+
+def build_prompt_text(row: Dict[str, Any], tokenizer: AutoTokenizer, data_cfg: Dict[str, Any]) -> str:
+    system_prompt = as_text(data_cfg.get("system_prompt"))
+    if not system_prompt:
+        system_prompt = (
+            "You are a rigorous mathematical reasoning assistant focused on unsolved "
+            "conjectures. Produce checkable reasoning."
+        )
+    user_block = build_user_block(row, data_cfg)
+    if getattr(tokenizer, "chat_template", None):
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_block},
+        ]
+        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    return f"System:\n{system_prompt}\n\nUser:\n{user_block}\n\nAssistant:\n"
+
+
+def compute_loss_weight(row: Dict[str, Any], data_cfg: Dict[str, Any]) -> float:
+    sample_weight_field = as_text(data_cfg.get("sample_weight_field")) or "sample_weight"
+    base = as_float(row.get(sample_weight_field), 1.0)
+    family = as_text(row.get("family"))
+    family_boost = data_cfg.get("family_boost", {})
+    if isinstance(family_boost, dict):
+        base *= as_float(family_boost.get(family), 1.0)
+    min_w = as_float(data_cfg.get("min_loss_weight"), 0.1)
+    max_w = as_float(data_cfg.get("max_loss_weight"), 8.0)
+    if min_w > max_w:
+        min_w, max_w = max_w, min_w
+    return max(min_w, min(max_w, base))
+
+
+def stage_split_files(stage_cfg: Dict[str, Any], data_cfg: Dict[str, Any]) -> Dict[str, str]:
+    train_file = as_text(stage_cfg.get("train_file")) or as_text(data_cfg.get("default_train_file"))
+    valid_file = as_text(stage_cfg.get("validation_file")) or as_text(data_cfg.get("default_validation_file"))
+    train_path = Path(train_file)
+    valid_path = Path(valid_file)
+    if not train_path.exists():
+        raise FileNotFoundError(f"Missing train split for stage: {train_path}")
+    if not valid_path.exists():
+        raise FileNotFoundError(f"Missing validation split for stage: {valid_path}")
+    return {"train": str(train_path), "validation": str(valid_path)}
+
+
+def apply_filters(dataset: Dataset, filter_cfg: Dict[str, Any]) -> Dataset:
+    if not filter_cfg:
+        return dataset
+    include_families = set(filter_cfg.get("include_families", []) or [])
+    exclude_families = set(filter_cfg.get("exclude_families", []) or [])
+    include_task_types = set(filter_cfg.get("include_task_types", []) or [])
+    source_datasets = set(filter_cfg.get("source_datasets", []) or [])
+    require_conjecture_id = bool(filter_cfg.get("require_conjecture_id", False))
+    min_sample_weight = filter_cfg.get("min_sample_weight")
+    min_sample_weight = as_float(min_sample_weight, 0.0) if min_sample_weight is not None else None
+
+    def _keep(row: Dict[str, Any]) -> bool:
+        family = as_text(row.get("family"))
+        if include_families and family not in include_families:
+            return False
+        if exclude_families and family in exclude_families:
+            return False
+        if include_task_types:
+            task_type = as_text(row.get("task_type"))
+            if task_type not in include_task_types:
+                return False
+        if source_datasets:
+            source = as_text(row.get("source_dataset"))
+            if source not in source_datasets:
+                return False
+        if require_conjecture_id:
+            conjecture_id = as_text(row.get("conjecture_id"))
+            if not conjecture_id or conjecture_id.lower() == "null":
+                return False
+        if min_sample_weight is not None:
+            sample_weight = as_float(row.get("sample_weight"), 0.0)
+            if sample_weight < min_sample_weight:
+                return False
+        return True
+
+    return dataset.filter(_keep, desc="Applying stage filters")
+
+
+def maybe_select(dataset: Dataset, max_samples: Optional[int]) -> Dataset:
+    if max_samples is None:
+        return dataset
+    if max_samples <= 0:
+        raise ValueError("max_samples must be positive.")
+    if max_samples >= len(dataset):
+        return dataset
+    return dataset.select(range(max_samples))
+
+
+def tokenize_datasets(raw: DatasetDict, tokenizer: AutoTokenizer, data_cfg: Dict[str, Any]) -> DatasetDict:
+    max_len = as_int(data_cfg.get("max_seq_length"), 2048)
+    if max_len < 64:
+        raise ValueError("data.max_seq_length must be >= 64")
+    eos = tokenizer.eos_token or ""
+    remove_columns = raw["train"].column_names
+
+    def _tokenize(row: Dict[str, Any]) -> Dict[str, Any]:
+        prompt_text = build_prompt_text(row, tokenizer, data_cfg)
+        answer_text = build_answer_block(row, data_cfg)
+        full_text = f"{prompt_text}{answer_text}{eos}"
+        prompt_ids = tokenizer(prompt_text, add_special_tokens=False)["input_ids"]
+        full_enc = tokenizer(
+            full_text,
+            add_special_tokens=False,
+            truncation=True,
+            max_length=max_len,
+        )
+        input_ids = full_enc["input_ids"]
+        attention_mask = full_enc["attention_mask"]
+        if not input_ids:
+            fallback = tokenizer.eos_token_id
+            if fallback is None:
+                fallback = tokenizer.pad_token_id
+            if fallback is None:
+                fallback = 0
+            input_ids = [fallback]
+            attention_mask = [1]
+            labels = [fallback]
+        else:
+            prompt_len = min(len(prompt_ids), len(input_ids))
+            labels = [-100] * prompt_len + input_ids[prompt_len:]
+            if prompt_len >= len(input_ids):
+                labels[-1] = input_ids[-1]
+        loss_weight = compute_loss_weight(row, data_cfg)
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+            "loss_weight": float(loss_weight),
+        }
+
+    tokenized = raw.map(
+        _tokenize,
+        remove_columns=remove_columns,
+        desc="Tokenizing prompt/answer pairs",
+    )
+    tokenized = tokenized.filter(
+        lambda row: any(token != -100 for token in row["labels"]),
+        desc="Dropping prompt-only rows",
+    )
+    return tokenized
+
+
+def build_model_and_tokenizer(model_cfg: Dict[str, Any], training_defaults: Dict[str, Any]) -> Tuple[Any, AutoTokenizer]:
+    base_model = as_text(model_cfg.get("base_model"))
+    if not base_model:
+        raise ValueError("model.base_model is required.")
+
+    use_bf16 = bool(model_cfg.get("use_bf16", True))
+    dtype = torch.bfloat16 if use_bf16 else torch.float16
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        base_model,
+        trust_remote_code=bool(model_cfg.get("trust_remote_code", False)),
+        use_fast=True,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token or tokenizer.unk_token
+    if tokenizer.pad_token is None:
+        tokenizer.add_special_tokens({"pad_token": "<|pad|>"})
+
+    model_kwargs: Dict[str, Any] = {
+        "trust_remote_code": bool(model_cfg.get("trust_remote_code", False)),
+        "torch_dtype": dtype,
+    }
+    attn_impl = as_text(model_cfg.get("attn_implementation"))
+    if attn_impl:
+        model_kwargs["attn_implementation"] = attn_impl
+
+    load_in_4bit = bool(model_cfg.get("load_in_4bit", True))
+    if load_in_4bit:
+        if not torch.cuda.is_available():
+            raise RuntimeError("4-bit loading requested but CUDA is not available.")
+        model_kwargs["quantization_config"] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type=as_text(model_cfg.get("bnb_4bit_quant_type")) or "nf4",
+            bnb_4bit_use_double_quant=bool(model_cfg.get("bnb_4bit_use_double_quant", True)),
+            bnb_4bit_compute_dtype=dtype,
+        )
+        model_kwargs["device_map"] = "auto"
+
+    model = AutoModelForCausalLM.from_pretrained(base_model, **model_kwargs)
+    if tokenizer.pad_token_id is not None:
+        model.config.pad_token_id = tokenizer.pad_token_id
+    model.config.use_cache = False
+
+    if load_in_4bit:
+        model = prepare_model_for_kbit_training(
+            model,
+            use_gradient_checkpointing=bool(training_defaults.get("gradient_checkpointing", True)),
+        )
+
+    lora_cfg = model_cfg.get("lora", {})
+    peft_cfg = LoraConfig(
+        r=as_int(lora_cfg.get("r"), 64),
+        lora_alpha=as_int(lora_cfg.get("alpha"), 128),
+        lora_dropout=as_float(lora_cfg.get("dropout"), 0.05),
+        bias=as_text(lora_cfg.get("bias")) or "none",
+        task_type="CAUSAL_LM",
+        target_modules=lora_cfg.get("target_modules"),
+    )
+    model = get_peft_model(model, peft_cfg)
+    model.print_trainable_parameters()
+    return model, tokenizer
+
+
+class WeightedLossCollator:
+    def __init__(self, tokenizer: AutoTokenizer, model: Any) -> None:
+        self.base = DataCollatorForSeq2Seq(
+            tokenizer=tokenizer,
+            model=model,
+            label_pad_token_id=-100,
+            pad_to_multiple_of=8,
+        )
+
+    def __call__(self, features: list[Dict[str, Any]]) -> Dict[str, Any]:
+        weights = [float(feature.pop("loss_weight", 1.0)) for feature in features]
+        batch = self.base(features)
+        batch["loss_weight"] = torch.tensor(weights, dtype=torch.float32)
+        return batch
+
+
+class WeightedLossTrainer(Trainer):
+    def _get_train_sampler(self):
+        if self.train_dataset is None:
+            return None
+        if "loss_weight" not in self.train_dataset.column_names:
+            return super()._get_train_sampler()
+        weights = self.train_dataset["loss_weight"]
+        if not weights:
+            return super()._get_train_sampler()
+        weight_tensor = torch.tensor(weights, dtype=torch.double)
+        return WeightedRandomSampler(
+            weights=weight_tensor,
+            num_samples=len(weight_tensor),
+            replacement=True,
+        )
+
+    def compute_loss(
+        self,
+        model: Any,
+        inputs: Dict[str, Any],
+        return_outputs: bool = False,
+        num_items_in_batch: Optional[torch.Tensor] = None,
+    ):
+        loss_weight = inputs.pop("loss_weight", None)
+        labels = inputs.get("labels")
+        if labels is None:
+            return super().compute_loss(
+                model=model,
+                inputs=inputs,
+                return_outputs=return_outputs,
+                num_items_in_batch=num_items_in_batch,
+            )
+
+        model_inputs = {k: v for k, v in inputs.items() if k != "labels"}
+        outputs = model(**model_inputs)
+        logits = outputs.logits
+
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+        token_losses = torch.nn.functional.cross_entropy(
+            shift_logits.view(-1, shift_logits.size(-1)),
+            shift_labels.view(-1),
+            ignore_index=-100,
+            reduction="none",
+        ).view(shift_labels.size())
+        token_mask = shift_labels.ne(-100).float()
+        seq_den = token_mask.sum(dim=1).clamp(min=1.0)
+        seq_loss = (token_losses * token_mask).sum(dim=1) / seq_den
+
+        if loss_weight is not None:
+            normalized = loss_weight.to(seq_loss.device).float().clamp(min=0.05)
+            loss = (seq_loss * normalized).sum() / normalized.sum()
+        else:
+            loss = seq_loss.mean()
+
+        if return_outputs:
+            return loss, outputs
+        return loss
+
+
+def build_training_args(
+    output_dir: Path,
+    training_cfg: Dict[str, Any],
+    use_bf16: bool,
+    has_eval_split: bool,
+) -> TrainingArguments:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    return TrainingArguments(
+        output_dir=str(output_dir),
+        num_train_epochs=as_float(training_cfg.get("num_train_epochs"), 1.0),
+        per_device_train_batch_size=as_int(training_cfg.get("per_device_train_batch_size"), 1),
+        per_device_eval_batch_size=as_int(training_cfg.get("per_device_eval_batch_size"), 1),
+        gradient_accumulation_steps=as_int(training_cfg.get("gradient_accumulation_steps"), 1),
+        learning_rate=as_float(training_cfg.get("learning_rate"), 2e-5),
+        weight_decay=as_float(training_cfg.get("weight_decay"), 0.0),
+        warmup_ratio=as_float(training_cfg.get("warmup_ratio"), 0.0),
+        lr_scheduler_type=as_text(training_cfg.get("lr_scheduler_type")) or "cosine",
+        max_grad_norm=as_float(training_cfg.get("max_grad_norm"), 1.0),
+        gradient_checkpointing=bool(training_cfg.get("gradient_checkpointing", True)),
+        logging_steps=as_int(training_cfg.get("logging_steps"), 10),
+        save_steps=as_int(training_cfg.get("save_steps"), 500),
+        save_total_limit=as_int(training_cfg.get("save_total_limit"), 3),
+        dataloader_num_workers=as_int(training_cfg.get("dataloader_num_workers"), 0),
+        seed=as_int(training_cfg.get("seed"), 17),
+        bf16=use_bf16,
+        fp16=not use_bf16,
+        remove_unused_columns=False,
+        report_to="none",
+        evaluation_strategy="steps" if has_eval_split else "no",
+        eval_steps=as_int(training_cfg.get("eval_steps"), 500) if has_eval_split else None,
+    )
+
+
+def push_folder(
+    api: HfApi,
+    repo_id: str,
+    folder_path: Path,
+    commit_message: str,
+    path_in_repo: Optional[str] = None,
+) -> None:
+    kwargs: Dict[str, Any] = {
+        "repo_id": repo_id,
+        "repo_type": "model",
+        "folder_path": str(folder_path),
+        "commit_message": commit_message,
+    }
+    if path_in_repo:
+        kwargs["path_in_repo"] = path_in_repo
+    api.upload_folder(**kwargs)
+
+
+def main() -> None:
+    args = parse_args()
+    cfg = load_config(args.config)
+    apply_overrides(cfg, args)
+
+    seed = as_int(cfg.get("global", {}).get("seed"), 17)
+    set_seed(seed)
+
+    output_root = Path(as_text(cfg.get("global", {}).get("output_root")) or "model_development/runs/math-conjecture-sota")
+    output_root.mkdir(parents=True, exist_ok=True)
+
+    token, username = resolve_auth(cfg)
+    repo_id = resolve_repo_id(cfg, username=username, output_root=output_root)
+    push_to_hub = bool(cfg.get("hub", {}).get("push_to_hub", False))
+    if push_to_hub:
+        if token is None:
+            raise ValueError("Hub push requested but token is missing.")
+        if repo_id is None:
+            raise ValueError("Hub push requested but repo_id is missing.")
+
+    model, tokenizer = build_model_and_tokenizer(cfg["model"], cfg.get("training_defaults", {}))
+    data_cfg = cfg["data"]
+    stage_reports = []
+
+    start_stage = max(1, args.start_stage)
+    stages = cfg["stages"]
+    end_stage = len(stages)
+    if args.max_stages is not None:
+        if args.max_stages <= 0:
+            raise ValueError("--max-stages must be positive.")
+        end_stage = min(end_stage, start_stage + args.max_stages - 1)
+
+    for index in range(start_stage, end_stage + 1):
+        stage = stages[index - 1]
+        stage_name = as_text(stage.get("name")) or f"stage_{index:02d}"
+        stage_slug = f"{index:02d}_{stage_name.replace(' ', '_')}"
+        stage_output_dir = output_root / stage_slug
+
+        split_files = stage_split_files(stage, data_cfg)
+        raw = load_dataset("parquet", data_files=split_files)
+        filters = stage.get("filters", {})
+        raw["train"] = apply_filters(raw["train"], filters)
+        raw["validation"] = apply_filters(raw["validation"], filters)
+        raw["train"] = maybe_select(raw["train"], stage.get("max_train_samples"))
+        raw["validation"] = maybe_select(raw["validation"], stage.get("max_eval_samples"))
+        if len(raw["train"]) == 0:
+            raise ValueError(f"Stage {stage_slug} has zero train rows after filtering.")
+
+        tokenized = tokenize_datasets(raw, tokenizer, data_cfg)
+        train_dataset = tokenized["train"]
+        eval_dataset = tokenized["validation"] if len(tokenized["validation"]) > 0 else None
+
+        merged_training = dict(cfg.get("training_defaults", {}))
+        merged_training.update(stage.get("training", {}))
+        merged_training["seed"] = seed
+        training_args = build_training_args(
+            output_dir=stage_output_dir,
+            training_cfg=merged_training,
+            use_bf16=bool(cfg["model"].get("use_bf16", True)),
+            has_eval_split=eval_dataset is not None,
+        )
+        collator = WeightedLossCollator(tokenizer=tokenizer, model=model)
+        trainer = WeightedLossTrainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            tokenizer=tokenizer,
+            data_collator=collator,
+        )
+
+        train_result = trainer.train()
+        trainer.log_metrics("train", train_result.metrics)
+        trainer.save_metrics("train", train_result.metrics)
+        trainer.save_state()
+        if eval_dataset is not None:
+            eval_metrics = trainer.evaluate()
+            trainer.log_metrics("eval", eval_metrics)
+            trainer.save_metrics("eval", eval_metrics)
+        trainer.save_model(str(stage_output_dir))
+        tokenizer.save_pretrained(str(stage_output_dir))
+
+        report = {
+            "stage_index": index,
+            "stage_name": stage_name,
+            "output_dir": str(stage_output_dir),
+            "train_rows": len(train_dataset),
+            "eval_rows": len(eval_dataset) if eval_dataset is not None else 0,
+            "train_metrics": train_result.metrics,
+        }
+        stage_reports.append(report)
+
+    final_dir = output_root / "final_adapter"
+    final_dir.mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(str(final_dir))
+    tokenizer.save_pretrained(str(final_dir))
+
+    summary = {
+        "config_path": str(args.config),
+        "repo_id": repo_id,
+        "seed": seed,
+        "stages_ran": stage_reports,
+        "final_adapter_dir": str(final_dir),
+    }
+    summary_path = output_root / "training_summary.json"
+    summary_path.write_text(json.dumps(summary, ensure_ascii=True, indent=2), encoding="utf-8")
+
+    if push_to_hub and repo_id is not None and token is not None:
+        api = HfApi(token=token)
+        api.create_repo(
+            repo_id=repo_id,
+            repo_type="model",
+            private=bool(cfg.get("hub", {}).get("private", False)),
+            exist_ok=True,
+        )
+        commit_message = as_text(cfg.get("hub", {}).get("commit_message")) or "Upload SOTA curriculum adapter."
+        push_folder(api, repo_id, final_dir, commit_message=commit_message)
+        if bool(cfg.get("hub", {}).get("upload_stage_checkpoints", False)):
+            for report in stage_reports:
+                stage_dir = Path(report["output_dir"])
+                path_in_repo = f"checkpoints/{Path(report['output_dir']).name}"
+                push_folder(
+                    api,
+                    repo_id,
+                    stage_dir,
+                    commit_message=f"Upload stage checkpoint {report['stage_name']}",
+                    path_in_repo=path_in_repo,
+                )
+        api.upload_file(
+            path_or_fileobj=str(summary_path),
+            path_in_repo="training_summary.json",
+            repo_id=repo_id,
+            repo_type="model",
+            commit_message="Upload training summary for SOTA curriculum run.",
+        )
+        print(f"Pushed training artifacts to https://huggingface.co/{repo_id}")
+
+    print(f"Training complete. Final adapter: {final_dir}")
+    print(f"Training summary: {summary_path}")
+
+
+if __name__ == "__main__":
+    main()