SDPO self-distillation examples for badlogicgames/pi-mono: full HF Jobs + Trackio script and minimal educational script

train_sdpo_pi_mono_minimal.py

# /// script
# dependencies = [
#   "datasets>=3.0.0",
#   "peft>=0.13.0",
#   "torch",
#   "transformers",
#   "trl>=1.5.0",
# ]
# ///

"""Small SDPO example using filtered badlogicgames/pi-mono traces.

This file is intentionally direct:

1. Load the Hub parquet export for badlogicgames/pi-mono.
2. Filter traces down to useful errors or user corrections.
3. Print the prepared training-data summary.
4. Train TRL's SDPOTrainer on the prepared rows.

Edit the constants below, then run:

  uv run train_sdpo_pi_mono_minimal.py

Or on Hugging Face Jobs:

  hf jobs uv run <raw-gist-url-for-this-file> \
    --flavor t4-small \
    --timeout 1h \
    --secrets HF_TOKEN \
    --env HUB_MODEL_ID=username/qwen2.5-0.5b-sdpo-pi-mono-minimal
"""

from __future__ import annotations

import json
import os
import re
import tempfile
import urllib.parse
import urllib.request
from typing import Any

from datasets import Dataset, load_dataset


MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
DATASET = "badlogicgames/pi-mono"
HUB_MODEL_ID = os.environ.get("HUB_MODEL_ID", "burtenshaw/qwen2.5-0.5b-sdpo-pi-mono-minimal")
OUTPUT_DIR = "outputs/sdpo-pi-mono-minimal"

SYSTEM_PROMPT = (
    "You are a careful coding agent. Diagnose failures using concrete evidence "
    "from commands, tests, compiler output, and file paths."
)

ERROR_RE = re.compile(
    r"(pytest|jest|vitest|assert|expected|actual|traceback|exception|syntaxerror|"
    r"typeerror|build|compile|lint|ts\d{4}|mypy|ruff|no such file|not found|"
    r"command not found|permission denied|unauthorized|validation failed|"
    r"must have required properties|exit code [1-9]|command exited with code [1-9])",
    re.I,
)
USER_CORRECTION_RE = re.compile(
    r"\b(wrong|fails?|failed|error|bug|not quite|instead|should|missing|broken|fix|actually)\b",
    re.I,
)


def text_from_content(content: Any) -> str:
    if isinstance(content, str):
        return content
    if not isinstance(content, list):
        return ""
    parts: list[str] = []
    for item in content:
        if isinstance(item, str):
            parts.append(item)
        elif isinstance(item, dict):
            value = item.get("text") or item.get("content") or item.get("thinking")
            if isinstance(value, str):
                parts.append(value)
    return "\n".join(parts)


def clean(text: str) -> str:
    text = re.sub(r"/Users/[^/\s]+/workspaces/pi-mono", "$WORKSPACE", text)
    text = re.sub(r"/private/tmp/[^\s]+", "$TMP", text)
    text = re.sub(r"\x1b\[[0-9;?]*[A-Za-z]", "", text)
    return re.sub(r"\n{3,}", "\n\n", text).strip()


def trim(text: str) -> str:
    text = clean(text)
    if len(text) <= 1200:
        return text
    return f"{text[:600].rstrip()}\n\n[... trimmed ...]\n\n{text[-600:].lstrip()}"


def hint_score(text: str, role: str) -> float:
    text_lower = text.lower()
    score = 2.0 if role == "user" else 1.0
    if "pytest" in text_lower or "assert" in text_lower:
        score += 3.0
    if re.search(r"(build|compile|lint|ts\d{4}|mypy|ruff)", text, re.I):
        score += 2.5
    if re.search(r"(traceback|exception|syntaxerror|typeerror)", text, re.I):
        score += 2.0
    if re.search(r"(no such file|not found|command not found)", text, re.I):
        score += 1.5
    if len(text) >= 180:
        score += 0.5
    return score


def make_example(row: dict[str, Any]) -> dict[str, Any] | None:
    prompt_text = str(row.get("prompt") or "").strip()
    if len(prompt_text) < 30:
        return None

    saw_assistant = False
    hints: list[tuple[float, str, str]] = []
    for event in row.get("traces") or []:
        message = event.get("message") if isinstance(event, dict) else None
        if not isinstance(message, dict):
            continue

        role = str(message.get("role") or "")
        text = text_from_content(message.get("content"))
        if role == "assistant":
            saw_assistant = True
        elif role == "toolResult" and message.get("isError") is True and ERROR_RE.search(text):
            hints.append((hint_score(text, role), "tool error", text))
        elif role == "user" and saw_assistant and USER_CORRECTION_RE.search(text):
            hints.append((hint_score(text, role), "user correction", text))

    if not hints:
        return None

    hints.sort(reverse=True, key=lambda item: item[0])
    best_hints = hints[:2]
    privileged_context = "\n\n".join(f"{label}:\n{trim(text)}" for _, label, text in best_hints)
    return {
        "prompt": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt_text},
        ],
        "privileged_context": privileged_context,
        "filter_score": round(sum(score for score, _, _ in best_hints), 3),
        "prompt_text": prompt_text,
    }


def load_source_dataset() -> Dataset:
    query = urllib.parse.urlencode({"dataset": DATASET})
    with urllib.request.urlopen(f"https://datasets-server.huggingface.co/parquet?{query}", timeout=120) as response:
        payload = json.load(response)

    urls = [
        item["url"]
        for item in payload.get("parquet_files", [])
        if item.get("config") == "default" and item.get("split") == "train" and item.get("url")
    ]
    if not urls:
        return load_dataset(DATASET, split="train")

    cache_dir = os.path.join(tempfile.gettempdir(), "pi_mono_sdpo_minimal")
    os.makedirs(cache_dir, exist_ok=True)
    local_paths: list[str] = []
    for index, url in enumerate(urls):
        local_path = os.path.join(cache_dir, f"default-train-{index}.parquet")
        if not os.path.exists(local_path):
            tmp_path = f"{local_path}.tmp"
            with urllib.request.urlopen(url, timeout=300) as response, open(tmp_path, "wb") as handle:
                while chunk := response.read(1024 * 1024):
                    handle.write(chunk)
            os.replace(tmp_path, local_path)
        local_paths.append(local_path)

    return load_dataset("parquet", data_files={"train": local_paths}, split="train")


def prepare_dataset() -> Dataset:
    raw = load_source_dataset()
    raw = raw.select(range(min(400, len(raw))))

    examples = [example for row in raw if (example := make_example(row)) is not None]
    examples.sort(key=lambda item: item["filter_score"], reverse=True)
    examples = examples[:16]
    if not examples:
        raise RuntimeError("No rows survived filtering.")
    return Dataset.from_list(examples)


def completion_text(completion: Any) -> str:
    if isinstance(completion, list) and completion and isinstance(completion[0], dict):
        return str(completion[0].get("content") or "")
    return str(completion or "")


def context_terms(context: str) -> list[str]:
    terms: list[str] = []
    stop = {"that", "this", "with", "from", "have", "error", "failed", "expected", "actual"}
    for match in re.findall(r"[A-Za-z_][A-Za-z0-9_./@:-]{3,}", context.lower()):
        term = match.strip(".,;:()[]{}'\"`")
        if "/" in term:
            term = term.rsplit("/", 1)[-1]
        if len(term) >= 4 and term not in stop and term not in terms:
            terms.append(term)
        if len(terms) == 8:
            break
    return terms


def trace_grounding_reward(
    completions: list[Any],
    privileged_context: list[str] | None = None,
    **_: Any,
) -> list[float]:
    contexts = privileged_context or ["" for _ in completions]
    rewards: list[float] = []
    for completion, context in zip(completions, contexts, strict=False):
        text = completion_text(completion).lower()
        words = set(re.findall(r"[a-zA-Z_][a-zA-Z0-9_/-]*", text))
        score = 0.1 if len(words) >= 20 else 0.0
        if words.intersection({"fix", "test", "build", "run", "error", "file", "change"}):
            score += 0.2
        score += min(sum(1 for term in context_terms(context) if term in text), 3) * 0.2
        rewards.append(max(0.0, min(1.0, score)))
    return rewards


def print_summary(dataset: Dataset) -> None:
    scores = [row["filter_score"] for row in dataset]
    print(
        json.dumps(
            {
                "selected_samples": len(dataset),
                "min_score": min(scores),
                "max_score": max(scores),
                "avg_score": round(sum(scores) / len(scores), 3),
                "first_sample": {
                    "score": dataset[0]["filter_score"],
                    "prompt": dataset[0]["prompt_text"][:240],
                    "privileged_context": dataset[0]["privileged_context"][:600],
                },
            },
            indent=2,
        )
    )


def train(dataset: Dataset) -> None:
    from peft import LoraConfig
    from trl.experimental.sdpo import SDPOConfig, SDPOTrainer

    trainer = SDPOTrainer(
        model=MODEL,
        reward_funcs=trace_grounding_reward,
        args=SDPOConfig(
            output_dir=OUTPUT_DIR,
            hub_model_id=HUB_MODEL_ID,
            push_to_hub=bool(os.environ.get("HF_TOKEN")),
            report_to="none",
            max_steps=2,
            learning_rate=5e-5,
            per_device_train_batch_size=1,
            gradient_accumulation_steps=2,
            num_generations=2,
            generation_batch_size=2,
            max_prompt_length=768,
            max_completion_length=160,
            include_environment_feedback=True,
            sdpo_policy_loss_mode="distillation_only",
            distillation_alpha=1.0,
            full_logit_distillation=False,
            teacher_regularization="ema",
            logging_steps=1,
            save_strategy="no",
            eval_strategy="no",
            remove_unused_columns=False,
            fp16=True,
        ),
        train_dataset=dataset,
        peft_config=LoraConfig(
            r=8,
            lora_alpha=16,
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        ),
    )
    trainer.train()
    trainer.save_model(OUTPUT_DIR)
    if os.environ.get("HF_TOKEN"):
        trainer.push_to_hub()


def main() -> None:
    dataset = prepare_dataset()
    print_summary(dataset)
    train(dataset)


if __name__ == "__main__":
    main()