Train Qwen3.6-35B-A3B on forecast prefs with Tinker

"""Reinforcement.tech compiled Tinker loop

H1: Reinforcement: Build Your Own Reward Model

INPUT
  Signal : Preference pairs (Pairwise preference)
  In     : Preferred / rejected pairs from real traces.
  Task   : Calibrated Forecasting — Scored predictions over time
  Model  : Qwen/Qwen3.6-35B-A3B (MOE, 35B-A3B + vision)

OUTPUT
  Loop         : Runnable preference loop.
  Loss         : pairwise preference (no reference policy)
  LoRA rank    : 32
  Steps        : 50

Tinker primitives used in this file:
  sample                   preview the adapter
  forward_backward_custom  pairwise preference loss
  optim_step               Adam update on the adapter
  save_state               checkpoint weights + optimizer

Requires:
  uv pip install tinker torch
  export TINKER_API_KEY=...
  Docs: https://tinker-docs.thinkingmachines.ai/tinker/quickstart/

Dataset JSONL (one record per line). Pass --data PATH.
  environment : {"prompt": str, "metadata": {...}}
  dpo         : {"prompt": str, "chosen": str, "rejected": str}
  sdft        : {"prompt": str, "completion": str}
  metadata is passed through to the verifier (expected answer, test file, Lean goal).
  --eval-data uses the same schema on a held-out set (never optim_step).
  Without --data this file runs one built-in example (a smoke test, not training).

CLI (overrides the module constants; --seed is only an args field):
  --data --eval-data --eval-every --log --steps --rank --lr --out --seed
  --resume --group-size --prompts-per-step --max-tokens --temperature
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import random
from pathlib import Path

try:
    import tinker
    from tinker import types
except ImportError:  # --help and load_dataset work without Tinker installed
    tinker = None
    types = None

BASE_MODEL = "Qwen/Qwen3.6-35B-A3B"
LORA_RANK = 32
LEARNING_RATE = 1e-5
STEPS = 50
MAX_TOKENS = 256
TEMPERATURE = 0.4
GROUP_SIZE = 8
PROMPTS_PER_STEP = 1
EVAL_EVERY = 10
SAVE_EVERY = 10
RUN_NAME = Path(__file__).stem
DATASET_KEYS = ("prompt", "chosen", "rejected",)

LARGE_MODEL_NOTE = ""


def parse_args():
    parser = argparse.ArgumentParser(description="Reinforcement.tech compiled Tinker loop")
    parser.add_argument("--data", help="JSONL dataset path")
    parser.add_argument("--eval-data", help="Held-out JSONL. Same schema as --data. Never used for optim_step.")
    parser.add_argument("--eval-every", type=int, default=EVAL_EVERY)
    parser.add_argument("--log", help="Append per-step JSONL: {step, mean_reward, n_datums, loss}")
    parser.add_argument("--steps", type=int, default=STEPS)
    parser.add_argument("--rank", type=int, default=LORA_RANK)
    parser.add_argument("--lr", type=float, default=float(LEARNING_RATE))
    parser.add_argument("--out", default=RUN_NAME, help="Run name / checkpoint prefix")
    parser.add_argument(
        "--seed",
        type=int,
        default=0,
        help="Shuffle seed for --data / --eval-data. Not copied in apply_args; rows_for_run reads args.seed.",
    )
    parser.add_argument("--resume", help="tinker:// path printed by save_state (weights + optimizer)")
    parser.add_argument("--group-size", type=int, default=GROUP_SIZE)
    parser.add_argument("--prompts-per-step", type=int, default=PROMPTS_PER_STEP)
    parser.add_argument("--max-tokens", type=int, default=MAX_TOKENS)
    parser.add_argument("--temperature", type=float, default=TEMPERATURE)
    return parser.parse_args()


def apply_args(args) -> None:
    global LORA_RANK, STEPS, LEARNING_RATE, RUN_NAME
    global GROUP_SIZE, PROMPTS_PER_STEP, MAX_TOKENS, TEMPERATURE, EVAL_EVERY
    # --seed is intentionally not a module constant. rows_for_run(args) and the
    # eval shuffle read args.seed so two shuffles stay independent of apply_args.
    LORA_RANK = args.rank
    STEPS = args.steps
    LEARNING_RATE = args.lr
    RUN_NAME = args.out
    GROUP_SIZE = args.group_size
    PROMPTS_PER_STEP = args.prompts_per_step
    MAX_TOKENS = args.max_tokens
    TEMPERATURE = args.temperature
    EVAL_EVERY = args.eval_every


def require_key() -> None:
    if not os.environ.get("TINKER_API_KEY"):
        raise SystemExit("Set TINKER_API_KEY before running this loop.")


def load_dataset(path: str) -> list[dict]:
    """Read JSONL. One record per line. Schema: see the module docstring."""
    rows: list[dict] = []
    with open(path, encoding="utf-8") as handle:
        for line_no, line in enumerate(handle, 1):
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            missing = [key for key in DATASET_KEYS if key not in row]
            if missing:
                raise SystemExit(f"{path}:{line_no} missing {missing}")
            rows.append(row)
    if not rows:
        raise SystemExit(f"{path} had no JSONL rows.")
    return rows


def rows_for_run(args) -> list[dict]:
    if args.data:
        rows = load_dataset(args.data)
    else:
        print("No --data given; running on 1 built-in example. This is a smoke test, not training.")
        rows = list(EXAMPLE_ROWS)
    rng = random.Random(args.seed)
    rng.shuffle(rows)
    needed = max(args.steps, 1) * max(args.prompts_per_step, 1)
    if needed > len(rows):
        print(f"warning: --steps/--prompts-per-step need {needed} rows; dataset has {len(rows)}; cycling.")
    return rows


def eval_rows_for_run(args) -> list[dict]:
    if not args.eval_data:
        return []
    rows = load_dataset(args.eval_data)
    rng = random.Random(args.seed)
    rng.shuffle(rows)
    return rows


def print_config(args, rows: list[dict], eval_rows: list[dict]) -> None:
    print("=== run config ===")
    print(f"model={BASE_MODEL}")
    print(f"rank={LORA_RANK}  steps={STEPS}  lr={LEARNING_RATE}  seed={args.seed}")
    print(f"group_size={GROUP_SIZE}  prompts_per_step={PROMPTS_PER_STEP}")
    print(f"max_tokens={MAX_TOKENS}  temperature={TEMPERATURE}")
    print(f"data={args.data or '(EXAMPLE_ROWS)'}  n_train={len(rows)}")
    print(f"eval_data={args.eval_data or '(none)'}  n_eval={len(eval_rows)}  eval_every={EVAL_EVERY}")
    print(f"log={args.log or '(none)'}  resume={args.resume or '(none)'}  out={RUN_NAME}")
    print("==================")
    if LARGE_MODEL_NOTE:
        print(LARGE_MODEL_NOTE)


def append_metrics(path: str | None, record: dict) -> None:
    if not path:
        return
    out = Path(path)
    out.parent.mkdir(parents=True, exist_ok=True)
    with out.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record) + "\n")


def metric_loss(result) -> float | None:
    metrics = getattr(result, "metrics", None) or {}
    if hasattr(metrics, "get"):
        for key in ("loss", "dpo_loss"):
            if metrics.get(key) is not None:
                return metrics[key]
    loss = getattr(result, "loss", None)
    return float(loss) if loss is not None else None


async def connect(args):
    if tinker is None or types is None:
        raise SystemExit("uv pip install tinker  — then rerun.")
    service = tinker.ServiceClient()
    if args.resume:
        training = await service.create_training_client_from_state_with_optimizer_async(args.resume)
        print(f"resumed weights+optimizer from {args.resume}")
    else:
        training = await service.create_lora_training_client_async(
            base_model=BASE_MODEL,
            rank=LORA_RANK,
            user_metadata={"product": "reinforcement.tech", "run": RUN_NAME},
        )
    tokenizer = training.get_tokenizer()
    return service, training, tokenizer


async def sampling_client(training):
    """Ephemeral on-policy sampler. Do not pass name= — it is deprecated and ignored."""
    return await training.save_weights_and_get_sampling_client_async()


async def checkpoint(training, step: int) -> None:
    if STEPS == 0:
        return
    if step % SAVE_EVERY != 0 and step != STEPS - 1:
        return
    saved = await training.save_state_async(name=f"{RUN_NAME}-step-{step}")
    result = await saved.result_async()
    path = getattr(result, "path", None)
    print(f"checkpoint name={RUN_NAME}-step-{step} path={path}")
    print("resume later with --resume PATH")

BETA = 0.1
LENGTH_NORMALIZE = True  # set False to compare raw summed logprobs (length-biased)
# Built-in smoke-test pair. Used only when --data is absent.
PREFERRED = "0.31 — base rate plus one independent signal; do not round to 50%."
REJECTED = "0.50 — no new signal, so I rounded to even odds and stopped."
PROMPT = "Complete the task. Prefer the trace a senior builder would keep."
EXAMPLE_ROWS = [{"prompt": PROMPT, "chosen": PREFERRED, "rejected": REJECTED}]


def as_sft_datum(tokenizer, prompt: str, completion: str) -> types.Datum:
    prompt_tokens = tokenizer.encode(prompt)
    completion_tokens = tokenizer.encode(completion)
    full = prompt_tokens + completion_tokens
    n_prefix = max(len(prompt_tokens) - 1, 0)
    return types.Datum(
        model_input=types.ModelInput.from_ints(tokens=full[:-1]),
        loss_fn_inputs=dict(
            target_tokens=full[1:],
            weights=[0.0] * n_prefix + [1.0] * len(completion_tokens),
        ),
    )


def dpo_loss(data, logprobs_list):
    """Pairwise logistic preference. No frozen reference policy (closer to SLiC than textbook DPO).
    LENGTH_NORMALIZE divides each sum by completion-token count so the loss does not
    prefer shorter strings. Turn it off only if you want the raw-sum (length-biased) objective.
    Batches are (chosen, rejected) pairs; loss is the mean over pairs.
    """
    import torch
    import torch.nn.functional as F

    def scored_tokens(datum, logprobs):
        weights = datum.loss_fn_inputs.get("weights")
        if weights is None:
            return max(len(logprobs), 1)
        try:
            return max(int((weights > 0).sum()), 1)
        except TypeError:
            return max(sum(1 for weight in weights if weight), 1)

    pair_losses = []
    for index in range(0, len(logprobs_list), 2):
        preferred_lp = logprobs_list[index].sum()
        rejected_lp = logprobs_list[index + 1].sum()
        if LENGTH_NORMALIZE:
            preferred_lp = preferred_lp / scored_tokens(data[index], logprobs_list[index])
            rejected_lp = rejected_lp / scored_tokens(data[index + 1], logprobs_list[index + 1])
        pair_losses.append(-F.logsigmoid(BETA * (preferred_lp - rejected_lp)))
    loss = torch.stack(pair_losses).mean()
    return loss, {"dpo_loss": float(loss.detach())}


async def forward_ce(training, datums) -> float:
    future = await training.forward_async(datums, loss_fn="cross_entropy")
    result = await future.result_async()
    loss = metric_loss(result)
    return float(loss) if loss is not None else 0.0


async def evaluate(eval_rows, training, tokenizer) -> float:
    margins: list[float] = []
    for row in eval_rows:
        chosen = as_sft_datum(tokenizer, str(row["prompt"]), str(row["chosen"]))
        rejected = as_sft_datum(tokenizer, str(row["prompt"]), str(row["rejected"]))
        chosen_ce = await forward_ce(training, [chosen])
        rejected_ce = await forward_ce(training, [rejected])
        margins.append(rejected_ce - chosen_ce)
    return sum(margins) / max(len(margins), 1)


async def train(args) -> None:
    require_key()
    rows = rows_for_run(args)
    eval_rows = eval_rows_for_run(args)
    print_config(args, rows, eval_rows)
    _service, training, tokenizer = await connect(args)
    preview_prompt_text = str(rows[0]["prompt"]) if rows else PROMPT

    for step in range(STEPS):
        datums = []
        for offset in range(PROMPTS_PER_STEP):
            row = rows[(step * PROMPTS_PER_STEP + offset) % len(rows)]
            datums.append(as_sft_datum(tokenizer, str(row["prompt"]), str(row["chosen"])))
            datums.append(as_sft_datum(tokenizer, str(row["prompt"]), str(row["rejected"])))
        fwdbwd = await training.forward_backward_custom_async(datums, dpo_loss)
        optim = await training.optim_step_async(types.AdamParams(learning_rate=LEARNING_RATE))
        result = await fwdbwd.result_async()
        await optim.result_async()
        await checkpoint(training, step)
        loss = metric_loss(result)
        append_metrics(args.log, {"step": step, "mean_reward": None, "n_datums": len(datums), "loss": loss})
        print(f"step {step:03d}  {result.metrics}")
        if eval_rows and step % EVAL_EVERY == 0:
            eval_margin = await evaluate(eval_rows, training, tokenizer)
            append_metrics(
                args.log,
                {"step": step, "mean_reward": eval_margin, "n_datums": len(eval_rows), "loss": None, "split": "eval"},
            )
            print(f"step {step:03d}  eval_margin={eval_margin:.3f}  n={len(eval_rows)}")

    sampling = await sampling_client(training)
    preview_prompt = types.ModelInput.from_ints(tokenizer.encode(preview_prompt_text))
    preview = await sampling.sample_async(
        prompt=preview_prompt,
        num_samples=1,
        sampling_params=types.SamplingParams(max_tokens=MAX_TOKENS, temperature=TEMPERATURE),
    )
    print("sample:", tokenizer.decode(preview.sequences[0].tokens))


if __name__ == "__main__":
    args = parse_args()
    apply_args(args)
    asyncio.run(train(args))