Post-train Qwen3.5-4B into a Wikipedia research agent, one command at a time

mkdir wiki-researcher && cd wiki-researcher
uv venv --python 3.12 .venv
source .venv/bin/activate

tinker==0.36.0
transformers==5.22.1
jinja2==3.1.6
datasets==5.1.0
bm25s==0.3.16
PyStemmer==3.1.0

uv pip install -r requirements.txt

export TINKER_API_KEY="tml-..."   # add this line to ~/.zshrc so you don't retype it

"""Shared settings. Every script imports from here, so change a value once."""

MODEL = "Qwen/Qwen3.5-4B"          # Tinker ID of the model we post-train
LOCAL_MODEL = "Qwen/Qwen3.5-0.8B"  # same chat template; runs free on a laptop or Colab

# Tinker prices for MODEL in USD per 1M tokens (tinker-docs Models & Pricing, 2026-11-01).
PRICE_PREFILL = 0.33
PRICE_SAMPLE = 1.005
PRICE_TRAIN = 0.737

TOP_K = 3           # paragraphs returned by one search
MAX_SEARCHES = 5    # research budget: after this many searches the model must answer
MAX_TOKENS = 512    # max tokens the model may write in one turn
LORA_RANK = 32

"""Step 0: check that your key works and the model answers."""
import asyncio

import tinker

from config import MODEL


async def main():
    service = tinker.ServiceClient()
    client = await service.create_sampling_client_async(base_model=MODEL)
    tok = client.get_tokenizer()
    prompt = tok.apply_chat_template([{"role": "user", "content": "In one sentence: who wrote Hamlet?"}],
                                     add_generation_prompt=True, enable_thinking=False,
                                     tokenize=True, return_dict=False)
    result = await client.sample_async(prompt=tinker.ModelInput.from_ints(prompt), num_samples=1,
                                       sampling_params=tinker.SamplingParams(max_tokens=50, temperature=0.0))
    print(tok.decode(result.sequences[0].tokens, skip_special_tokens=True))


asyncio.run(main())

python hello_tinker.py

"""Step 1: download HotpotQA, write the question splits and the paragraph corpus.

data/train.jsonl  90k questions to train on
data/val.jsonl    200 questions to pick checkpoints with
data/test.jsonl   500 questions you only open in Step 5
data/corpus.jsonl every Wikipedia paragraph that appears in HotpotQA (one per title)
"""
import hashlib
import json
import os
import random

from datasets import load_dataset


def write(path, rows):
    with open(path, "w") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def question(row):
    return {"id": row["id"], "question": row["question"], "answers": [row["answer"]],
            "type": row["type"], "level": row["level"]}


os.makedirs("data", exist_ok=True)
train = load_dataset("hotpotqa/hotpot_qa", "distractor", split="train")
dev = load_dataset("hotpotqa/hotpot_qa", "distractor", split="validation")

corpus, seen = [], set()
for split in (train, dev):
    for context in split["context"]:
        for title, sentences in zip(context["title"], context["sentences"]):
            if title not in seen:
                seen.add(title)
                corpus.append({"id": len(corpus), "title": title, "text": "".join(sentences).strip()})
write("data/corpus.jsonl", corpus)

rng = random.Random(0)
dev_rows = [question(r) for r in dev]
rng.shuffle(dev_rows)
write("data/test.jsonl", dev_rows[:500])
write("data/val.jsonl", dev_rows[500:700])
train_rows = [question(r) for r in train]
rng.shuffle(train_rows)
write("data/train.jsonl", train_rows)

digest = hashlib.sha256(open("data/test.jsonl", "rb").read()).hexdigest()[:16]
print(f"corpus {len(corpus):,} paragraphs | train {len(train_rows):,} | val 200 | test 500")
print(f"test.jsonl sha256 {digest}  <- write this down; it must never change")

python prep_data.py

"""Exact Match and F1, normalized the way SQuAD and HotpotQA do it."""
import re
import string
from collections import Counter


def normalize(text):
    text = text.lower()
    text = "".join(ch for ch in text if ch not in set(string.punctuation))
    text = re.sub(r"\b(a|an|the)\b", " ", text)
    return " ".join(text.split())


def exact_match(prediction, answers):
    return float(any(normalize(prediction) == normalize(a) for a in answers))


def f1(prediction, answers):
    best = 0.0
    for answer in answers:
        pred, gold = normalize(prediction).split(), normalize(answer).split()
        # yes/no/noanswer must match exactly, as in the official HotpotQA script
        if {" ".join(pred), " ".join(gold)} & {"yes", "no", "noanswer"} and pred != gold:
            continue
        common = sum((Counter(pred) & Counter(gold)).values())
        if common == 0:
            continue
        precision, recall = common / len(pred), common / len(gold)
        best = max(best, 2 * precision * recall / (precision + recall))
    return best

import unittest

from score import exact_match, f1, normalize


class ScoreTest(unittest.TestCase):
    def test_articles_and_case(self):
        self.assertEqual(exact_match("The United States", ["United States"]), 1.0)

    def test_punctuation(self):
        self.assertEqual(normalize("Arthur's Magazine."), "arthurs magazine")

    def test_partial_credit(self):
        self.assertAlmostEqual(f1("President Richard Nixon", ["Richard Nixon"]), 0.8)

    def test_yes_no_is_all_or_nothing(self):
        self.assertEqual(f1("yes", ["no"]), 0.0)
        self.assertEqual(f1("no it was not", ["no"]), 0.0)
        self.assertEqual(f1("Yes", ["yes"]), 1.0)

    def test_numbers(self):
        self.assertEqual(exact_match("1,844", ["1844"]), 1.0)

    def test_best_of_several_answers(self):
        self.assertEqual(f1("NYC", ["New York City", "NYC"]), 1.0)

    def test_empty_prediction(self):
        self.assertEqual(f1("", ["Paris"]), 0.0)


if __name__ == "__main__":
    unittest.main()

python -m unittest test_score

"""The agent harness: question -> searches -> answer, recorded token by token.

The same code runs evaluation, collects SFT data and produces RL rollouts, so the
model is trained in exactly the environment it is tested in.
"""
import re
from dataclasses import dataclass, field

from config import MAX_SEARCHES, MAX_TOKENS, TOP_K
from score import f1

SYSTEM = """You are a research agent. You answer questions by searching Wikipedia.

Each reply contains exactly one action:
- To search, write a short query inside <search></search> and stop. The results arrive in the next message.
- When you know the answer, write only the answer inside <answer></answer>: a name, date, number, or yes/no.

Example. Question: Who directed the film in which Tom Hanks plays a FedEx employee stranded on an island?
Reply 1: <search>Tom Hanks FedEx employee stranded island film</search>
Reply 2: <search>Cast Away 2000 film director</search>
Reply 3: <answer>Robert Zemeckis</answer>

Most questions need two facts: find the first, then search for the second. Never answer before you have searched."""

SYSTEM_CLOSED_BOOK = """Answer the question from memory.
Write only the final answer inside <answer></answer>: a name, date, number, or yes/no."""


@dataclass
class Sample:
    tokens: list
    logprobs: list
    stop_reason: str  # "stop" (end of turn) or "length" (hit MAX_TOKENS)


@dataclass
class Episode:
    id: str
    question: str
    answers: list
    prompt: list                                  # chat-template tokens of system + question
    tokens: list = field(default_factory=list)    # everything after the prompt
    sampled: list = field(default_factory=list)   # 1 = the model wrote this token, 0 = the harness did
    logprobs: list = field(default_factory=list)  # sampler logprob per token (0.0 for harness tokens)
    searches: list = field(default_factory=list)
    answer: str = ""
    end: str = ""          # answer | no_action | length | budget
    prefill: int = 0       # tokens sent to the sampler, summed over turns (for cost)
    generated: int = 0     # tokens the sampler produced

    @property
    def reward(self):
        return f1(self.answer, self.answers) if self.end == "answer" else 0.0


class LocalPolicy:
    """Runs the model with transformers on your own machine: free, slow, no logprobs."""

    def __init__(self, name):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        self.torch = torch
        self.tokenizer = AutoTokenizer.from_pretrained(name)
        device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
        self.model = AutoModelForCausalLM.from_pretrained(name, dtype="auto").to(device)
        self.end = self.tokenizer.convert_tokens_to_ids("<|im_end|>")

    async def sample(self, ids, temperature):
        x = self.torch.tensor([ids], device=self.model.device)
        out = self.model.generate(x, attention_mask=self.torch.ones_like(x), max_new_tokens=MAX_TOKENS,
                                  do_sample=temperature > 0, temperature=temperature or None,
                                  top_p=None, top_k=None, eos_token_id=self.end)
        new = out[0, len(ids):].tolist()
        return Sample(new, [0.0] * len(new), "stop" if new and new[-1] == self.end else "length")


class TinkerPolicy:
    """Samples from a model on Tinker: a base model or one of your checkpoints."""

    def __init__(self, sampling_client):
        import tinker
        self.tinker = tinker
        self.client = sampling_client
        self.tokenizer = sampling_client.get_tokenizer()
        self.end = self.tokenizer.convert_tokens_to_ids("<|im_end|>")

    async def sample(self, ids, temperature):
        params = self.tinker.SamplingParams(max_tokens=MAX_TOKENS, temperature=temperature, stop=[self.end])
        result = await self.client.sample_async(prompt=self.tinker.ModelInput.from_ints(ids),
                                                num_samples=1, sampling_params=params)
        seq = result.sequences[0]
        if seq.logprobs is None:
            raise RuntimeError("the sampler returned no logprobs; RL needs them")
        return Sample(list(seq.tokens), list(seq.logprobs), seq.stop_reason)


def chat_prompt(tokenizer, system, question):
    return tokenizer.apply_chat_template(
        [{"role": "system", "content": system}, {"role": "user", "content": question}],
        add_generation_prompt=True, enable_thinking=False, tokenize=True, return_dict=False)


def user_turn(tokenizer, content):
    """Tokens that close the model's turn, add a user message and open the next model turn.

    Built by hand instead of re-rendering the whole chat: the Qwen template drops the
    empty <think></think> from earlier turns, which would change tokens the model
    already produced. test_harness.py checks this string against the real template.
    """
    text = f"\n<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
    return tokenizer.encode(text, add_special_tokens=False)


def format_results(docs):
    if not docs:
        return "<result>\nNo results. Try other keywords.\n</result>"
    return "<result>\n" + "\n\n".join(f"[{d['title']}] {d['text']}" for d in docs) + "\n</result>"


async def run_episode(policy, searcher, item, temperature=0.0, closed_book=False,
                      top_k=TOP_K, max_searches=MAX_SEARCHES):
    tok = policy.tokenizer
    system = SYSTEM_CLOSED_BOOK if closed_book else SYSTEM
    ep = Episode(item["id"], item["question"], item["answers"], chat_prompt(tok, system, item["question"]))

    def add(tokens, logprobs, sampled):
        ep.tokens += tokens
        ep.logprobs += logprobs
        ep.sampled += [sampled] * len(tokens)

    while True:
        ids = ep.prompt + ep.tokens
        out = await policy.sample(ids, temperature)
        ep.prefill += len(ids)
        ep.generated += len(out.tokens)
        add(out.tokens, out.logprobs, 1)
        if out.stop_reason == "length":
            ep.end = "length"
            return ep
        if not out.tokens or out.tokens[-1] != policy.end:  # some samplers leave out the stop token
            add([policy.end], [0.0], 0)
        reply = tok.decode(out.tokens, skip_special_tokens=True)

        answer = re.search(r"<answer>(.*?)</answer>", reply, re.S)
        if answer:
            ep.answer, ep.end = answer.group(1).strip(), "answer"
            return ep
        query = re.search(r"<search>(.*?)</search>", reply, re.S)
        if closed_book or not query:
            ep.end = "no_action"
            return ep
        if len(ep.searches) == max_searches:
            ep.end = "budget"
            return ep

        ep.searches.append(query.group(1).strip())
        content = format_results(searcher.search(ep.searches[-1], top_k))
        if len(ep.searches) == max_searches:
            content += "\nThat was your last search. Reply with <answer></answer> now."
        env = user_turn(tok, content)
        add(env, [0.0] * len(env), 0)

"""Run the agent on a split, score every answer, write runs/<name>/.

python eval.py --name base                          base model on Tinker, val split
python eval.py --name closed --closed-book          no search: the floor
python eval.py --name sft-e2 --model-path tinker://...   one of your checkpoints
python eval.py --name local --local --limit 10      free, on your own machine
"""
import argparse
import asyncio
import json
import os
import time

from config import LOCAL_MODEL, MAX_SEARCHES, MODEL, PRICE_PREFILL, PRICE_SAMPLE, TOP_K
from harness import LocalPolicy, TinkerPolicy, run_episode
from score import exact_match, f1


async def make_policy(args):
    if args.local:
        return LocalPolicy(LOCAL_MODEL)
    import tinker
    service = tinker.ServiceClient()
    if args.model_path:
        return TinkerPolicy(await service.create_sampling_client_async(model_path=args.model_path))
    return TinkerPolicy(await service.create_sampling_client_async(base_model=MODEL))


def cost(ep):
    return (ep.prefill * PRICE_PREFILL + ep.generated * PRICE_SAMPLE) / 1e6


def record(ep, tokenizer, free=False):
    return {"id": ep.id, "question": ep.question, "answers": ep.answers, "prediction": ep.answer,
            "em": exact_match(ep.answer, ep.answers), "f1": f1(ep.answer, ep.answers),
            "searches": ep.searches, "end": ep.end, "prefill": ep.prefill,
            "generated": ep.generated, "usd": 0.0 if free else cost(ep),
            "trace": tokenizer.decode(ep.tokens, skip_special_tokens=False)}


def summarize(rows):
    n = len(rows)
    mean = lambda key: sum(r[key] for r in rows) / n
    return {"n": n, "em": mean("em"), "f1": mean("f1"),
            "searches": sum(len(r["searches"]) for r in rows) / n,
            "answered": sum(r["end"] == "answer" for r in rows) / n,
            "usd_per_q": mean("usd"), "usd_total": sum(r["usd"] for r in rows)}


async def main(args):
    items = [json.loads(line) for line in open(f"data/{args.split}.jsonl")][: args.limit]
    policy = await make_policy(args)
    searcher = None
    if not args.closed_book:
        from search import Searcher  # built in Step 2
        searcher = Searcher()
    gate = asyncio.Semaphore(1 if args.local else args.concurrency)
    out_dir = f"runs/{args.name}"
    os.makedirs(out_dir, exist_ok=True)
    started, rows = time.time(), []

    async def one(item):
        async with gate:
            ep = await run_episode(policy, searcher, item, temperature=0.0, closed_book=args.closed_book,
                                   top_k=args.top_k, max_searches=args.max_searches)
        rows.append(record(ep, policy.tokenizer, free=args.local))
        if len(rows) % 25 == 0 or args.local:
            print(f"  {len(rows)}/{len(items)}  F1 so far {summarize(rows)['f1']:.3f}", flush=True)

    await asyncio.gather(*(one(item) for item in items))
    rows.sort(key=lambda r: r["id"])
    with open(f"{out_dir}/episodes.jsonl", "w") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
    summary = summarize(rows) | {"split": args.split, "model_path": args.model_path or MODEL,
                                 "top_k": args.top_k, "max_searches": args.max_searches,
                                 "seconds": round(time.time() - started)}
    json.dump(summary, open(f"{out_dir}/summary.json", "w"), indent=2)
    print(f"\n{args.name}: EM {summary['em']:.3f}  F1 {summary['f1']:.3f}  "
          f"searches {summary['searches']:.2f}  answered {summary['answered']:.0%}  "
          f"${summary['usd_per_q']:.4f}/q  (total ${summary['usd_total']:.2f}, {summary['seconds']}s)")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--name", required=True)
    p.add_argument("--split", default="val", choices=["val", "test", "train"])
    p.add_argument("--model-path", help="tinker://... checkpoint; default is the base model")
    p.add_argument("--local", action="store_true", help="run LOCAL_MODEL with transformers")
    p.add_argument("--closed-book", action="store_true")
    p.add_argument("--limit", type=int)
    p.add_argument("--top-k", type=int, default=TOP_K)
    p.add_argument("--max-searches", type=int, default=MAX_SEARCHES)
    p.add_argument("--concurrency", type=int, default=32)
    asyncio.run(main(p.parse_args()))

python eval.py --name closed --closed-book

"""BM25 search over the paragraph corpus.

python search.py                      build the index into data/bm25 (once, ~2 min)
python search.py "Ed Wood director"   try a query
"""
import json
import sys

import bm25s
import Stemmer

INDEX = "data/bm25"
stemmer = Stemmer.Stemmer("english")


def build():
    corpus = [json.loads(line) for line in open("data/corpus.jsonl")]
    texts = [f"{doc['title']}. {doc['text']}" for doc in corpus]
    index = bm25s.BM25()
    index.index(bm25s.tokenize(texts, stopwords="en", stemmer=stemmer))
    index.save(INDEX, corpus=corpus)
    print(f"indexed {len(corpus):,} paragraphs into {INDEX}")


class Searcher:
    def __init__(self):
        self.index = bm25s.BM25.load(INDEX, load_corpus=True)

    def search(self, query, k=3):
        tokens = bm25s.tokenize([query], stopwords="en", stemmer=stemmer, show_progress=False)
        if not tokens.ids[0]:  # the query was only stopwords
            return []
        docs, _ = self.index.retrieve(tokens, k=k, show_progress=False)
        return list(docs[0])


if __name__ == "__main__":
    if len(sys.argv) == 1:
        build()
    else:
        for doc in Searcher().search(" ".join(sys.argv[1:])):
            print(f"[{doc['title']}] {doc['text'][:200]}\n")

python search.py                                   # build, once
python search.py "Scott Derrickson nationality"    # try a query

import asyncio
import unittest

from transformers import AutoTokenizer

from config import MODEL
from harness import SYSTEM, Sample, chat_prompt, run_episode, user_turn

tok = AutoTokenizer.from_pretrained(MODEL)
END = tok.convert_tokens_to_ids("<|im_end|>")


class ScriptedPolicy:
    """Replays fixed replies, so the harness can be tested without a model."""

    def __init__(self, replies):
        self.tokenizer, self.end, self.replies = tok, END, list(replies)

    async def sample(self, ids, temperature):
        tokens = tok.encode(self.replies.pop(0), add_special_tokens=False) + [END]
        return Sample(tokens, [-0.5] * len(tokens), "stop")


class FakeSearcher:
    def search(self, query, k=3):
        return [{"title": "Arthur's Magazine", "text": "Arthur's Magazine was published from 1844 to 1846."}]


ITEM = {"id": "q1", "question": "Which magazine started first?", "answers": ["Arthur's Magazine"]}


def run(replies, **kw):
    return asyncio.run(run_episode(ScriptedPolicy(replies), FakeSearcher(), ITEM, **kw))


class HarnessTest(unittest.TestCase):
    def test_user_turn_matches_chat_template(self):
        # Our hand-built tokens must equal what the template renders for the next turn.
        msgs = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": "Q"}]
        prompt = chat_prompt(tok, SYSTEM, "Q")
        reply = tok.encode("<search>x</search>", add_special_tokens=False) + [END]
        ours = prompt + reply + user_turn(tok, "<result>R</result>")
        rendered = tok.apply_chat_template(
            msgs + [{"role": "assistant", "content": "<search>x</search>"},
                    {"role": "user", "content": "<result>R</result>"}],
            add_generation_prompt=True, enable_thinking=False, tokenize=False)
        # The template drops the old turn's empty <think> block; everything else must match.
        expected = rendered.replace("<|im_start|>assistant\n<search>", "<|im_start|>assistant\n<think>\n\n</think>\n\n<search>")
        self.assertEqual(tok.decode(ours), expected)

    def test_search_then_answer(self):
        ep = run(["<search>Arthur's Magazine</search>", "<answer>Arthur's Magazine</answer>"])
        self.assertEqual(ep.end, "answer")
        self.assertEqual(ep.searches, ["Arthur's Magazine"])
        self.assertEqual(ep.reward, 1.0)

    def test_masks_line_up(self):
        ep = run(["<search>a</search>", "<answer>b</answer>"])
        self.assertEqual(len(ep.tokens), len(ep.sampled))
        self.assertEqual(len(ep.tokens), len(ep.logprobs))
        self.assertIn("1844", tok.decode([t for t, s in zip(ep.tokens, ep.sampled) if s == 0]))
        self.assertNotIn("1844", tok.decode([t for t, s in zip(ep.tokens, ep.sampled) if s == 1]))

    def test_no_action_scores_zero(self):
        ep = run(["I think it is Arthur's Magazine."])
        self.assertEqual((ep.end, ep.reward), ("no_action", 0.0))

    def test_budget_forces_answer(self):
        ep = run(["<search>a</search>"] * 3 + ["<answer>Arthur's Magazine</answer>"], max_searches=3)
        self.assertEqual((len(ep.searches), ep.end), (3, "answer"))
        self.assertIn("last search", tok.decode(ep.tokens))

    def test_budget_exceeded(self):
        ep = run(["<search>a</search>"] * 3, max_searches=2)
        self.assertEqual(ep.end, "budget")


if __name__ == "__main__":
    unittest.main()

python -m unittest test_score test_harness

uv pip install torch
python eval.py --name local --local --limit 10
python show.py runs/local --n 3

python eval.py --name base
python show.py runs/base --n 10

"""Step 3a: build SFT data by rejection sampling (STaR).

The base model tries each train question several times in the harness. We keep the
shortest attempt that searched at least once and got the answer right.

python make_sft.py --questions 1500 --tries 4
"""
import argparse
import asyncio
import json

import tinker

from config import MODEL, PRICE_PREFILL, PRICE_SAMPLE
from harness import TinkerPolicy, run_episode
from search import Searcher


async def main(args):
    items = [json.loads(line) for line in open("data/train.jsonl")][: args.questions]
    service = tinker.ServiceClient()
    client = await (service.create_sampling_client_async(model_path=args.model_path) if args.model_path
                    else service.create_sampling_client_async(base_model=MODEL))
    policy, searcher = TinkerPolicy(client), Searcher()
    gate = asyncio.Semaphore(64)
    kept, solved, usd, done = [], 0, 0.0, 0

    async def attempt(item):
        async with gate:
            return await run_episode(policy, searcher, item, temperature=args.temperature)

    async def one(item):
        nonlocal solved, usd, done
        episodes = await asyncio.gather(*(attempt(item) for _ in range(args.tries)))
        usd += sum(e.prefill * PRICE_PREFILL + e.generated * PRICE_SAMPLE for e in episodes) / 1e6
        good = [e for e in episodes if e.reward >= args.min_f1 and e.searches]
        solved += any(e.reward >= args.min_f1 for e in episodes)
        if good:
            best = min(good, key=lambda e: len(e.tokens))
            kept.append({"id": best.id, "question": best.question, "answers": best.answers,
                         "prompt": best.prompt, "tokens": best.tokens, "sampled": best.sampled,
                         "searches": best.searches, "reward": best.reward})
        done += 1
        if done % 100 == 0:
            print(f"  {done}/{len(items)}  kept {len(kept)}  ${usd:.2f}", flush=True)

    await asyncio.gather(*(one(item) for item in items))
    with open(args.out, "w") as f:
        for row in kept:
            f.write(json.dumps(row) + "\n")
    print(f"\nkept {len(kept)}/{len(items)} questions  |  solved@{args.tries} {solved / len(items):.1%}  |  "
          f"avg searches {sum(len(r['searches']) for r in kept) / max(len(kept), 1):.2f}  |  ${usd:.2f}")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--questions", type=int, default=1500)
    p.add_argument("--tries", type=int, default=4)
    p.add_argument("--temperature", type=float, default=0.8)
    p.add_argument("--min-f1", type=float, default=0.8)
    p.add_argument("--model-path", help="sample from a checkpoint instead of the base model")
    p.add_argument("--out", default="data/sft.jsonl")
    asyncio.run(main(p.parse_args()))

python make_sft.py --questions 1500 --tries 4

python -c "
import json; from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained('Qwen/Qwen3.5-4B')
r = json.loads(open('data/sft.jsonl').readline())
print(tok.decode(r['prompt'] + r['tokens']))"

"""Turn episodes into Tinker training examples (Datum).

A model predicts the next token, so the input is every token except the last and the
target is every token except the first. Masks are shifted the same way.
"""
import tinker


def sft_datum(prompt, tokens, sampled):
    """Cross-entropy only on tokens the model wrote; prompt and search results get weight 0."""
    full = prompt + tokens
    weights = [0.0] * len(prompt) + [float(s) for s in sampled]
    return tinker.Datum(model_input=tinker.ModelInput.from_ints(full[:-1]),
                        loss_fn_inputs={"target_tokens": full[1:], "weights": weights[1:]})


def rl_datum(ep, advantage):
    """Importance-sampling loss: push up the model's own tokens in proportion to the advantage."""
    full = ep.prompt + ep.tokens
    n = len(ep.prompt) - 1
    return tinker.Datum(model_input=tinker.ModelInput.from_ints(full[:-1]),
                        loss_fn_inputs={"target_tokens": full[1:],
                                        "logprobs": [0.0] * n + ep.logprobs,
                                        "advantages": [0.0] * n + [advantage * s for s in ep.sampled]})

"""Step 3b: LoRA SFT on Tinker, loss only on tokens the model wrote.

python train_sft.py --epochs 3
Prints a sampler path (for eval.py) and a state path (for train_grpo.py) after each epoch.
"""
import argparse
import asyncio
import json
import os
import random

import tinker

from config import LORA_RANK, MODEL, PRICE_TRAIN
from datums import sft_datum


async def main(args):
    rows = [json.loads(line) for line in open(args.data)]
    data = [sft_datum(r["prompt"], r["tokens"], r["sampled"]) for r in rows]
    tokens_per_epoch = sum(d.model_input.length for d in data)
    print(f"{len(data)} examples, {tokens_per_epoch:,} tokens/epoch, "
          f"≈ ${tokens_per_epoch * args.epochs * PRICE_TRAIN / 1e6:.2f} to train")

    service = tinker.ServiceClient()
    trainer = await service.create_lora_training_client_async(base_model=MODEL, rank=LORA_RANK)
    steps_per_epoch = (len(data) + args.batch - 1) // args.batch
    total, step = steps_per_epoch * args.epochs, 0
    rng = random.Random(0)
    os.makedirs("runs", exist_ok=True)

    for epoch in range(1, args.epochs + 1):
        rng.shuffle(data)
        for i in range(0, len(data), args.batch):
            batch = data[i : i + args.batch]
            lr = args.lr * (1 - step / total)  # linear decay to 0
            fwd = await trainer.forward_backward_async(batch, "cross_entropy")
            opt = await trainer.optim_step_async(tinker.AdamParams(learning_rate=lr))
            result = await fwd.result_async()
            await opt.result_async()
            step += 1
            logprobs = [lp for out in result.loss_fn_outputs for lp in out["logprobs"].tolist()]
            weights = [w for d in batch for w in d.loss_fn_inputs["weights"].tolist()]
            loss = -sum(lp * w for lp, w in zip(logprobs, weights)) / sum(weights)
            if step % 5 == 0 or step == total:
                print(f"epoch {epoch}  step {step}/{total}  loss {loss:.4f}  lr {lr:.2e}", flush=True)

        name = f"{args.name}-e{epoch}"
        sampler = await (await trainer.save_weights_for_sampler_async(name=name)).result_async()
        state = await (await trainer.save_state_async(name=name + "-state")).result_async()
        print(f"saved {name}\n  sampler: {sampler.path}\n  state:   {state.path}")
        with open("runs/checkpoints.jsonl", "a") as f:
            f.write(json.dumps({"name": name, "sampler": sampler.path, "state": state.path}) + "\n")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--data", default="data/sft.jsonl")
    p.add_argument("--name", default="sft")
    p.add_argument("--epochs", type=int, default=3)
    p.add_argument("--batch", type=int, default=32)
    p.add_argument("--lr", type=float, default=2e-4)
    asyncio.run(main(p.parse_args()))

"""Print a checkpoint path saved in runs/checkpoints.jsonl.

python ckpt.py sft-e2          sampler path, for eval.py --model-path
python ckpt.py sft-e2 state    state path, for train_grpo.py --init-state
"""
import json
import sys

name, kind = sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else "sampler"
paths = [json.loads(line)[kind] for line in open("runs/checkpoints.jsonl") if json.loads(line)["name"] == name]
if not paths:
    sys.exit(f"no checkpoint named {name} in runs/checkpoints.jsonl")
print(paths[-1])

python train_sft.py --epochs 3

for e in 1 2 3; do
  python eval.py --name sft-e$e --model-path $(python ckpt.py sft-e$e)
done
python ablate.py runs/base runs/sft-e1 runs/sft-e2 runs/sft-e3

"""Step 4: online GRPO in the agent harness.

Each step: take Q new train questions, run G episodes per question with the current
weights, reward = F1, advantage = reward - mean reward of the same question's group.
Groups where every episode scored the same teach nothing and are skipped.

python train_grpo.py --init-state tinker://.../sft-e2-state --steps 30
"""
import argparse
import asyncio
import json
import os

import tinker

from config import LORA_RANK, MODEL, PRICE_PREFILL, PRICE_SAMPLE, PRICE_TRAIN
from datums import rl_datum
from harness import TinkerPolicy, run_episode
from search import Searcher


async def main(args):
    service = tinker.ServiceClient()
    if args.init_state:
        trainer = await service.create_training_client_from_state_async(args.init_state)
    else:
        trainer = await service.create_lora_training_client_async(base_model=MODEL, rank=LORA_RANK)
    items = [json.loads(line) for line in open("data/train.jsonl")][args.offset :]
    searcher, gate = Searcher(), asyncio.Semaphore(128)
    os.makedirs(f"runs/{args.name}", exist_ok=True)
    spent = 0.0

    for step in range(1, args.steps + 1):
        policy = TinkerPolicy(await trainer.save_weights_and_get_sampling_client_async())
        batch = items[(step - 1) * args.questions : step * args.questions]

        async def attempt(item):
            async with gate:
                return await run_episode(policy, searcher, item, temperature=args.temperature)

        groups = await asyncio.gather(*(asyncio.gather(*(attempt(it) for _ in range(args.group)))
                                        for it in batch))
        data, skipped = [], 0
        for group in groups:
            rewards = [e.reward for e in group]
            if max(rewards) - min(rewards) < 1e-6:
                skipped += 1
                continue
            mean = sum(rewards) / len(rewards)
            data += [rl_datum(e, e.reward - mean) for e in group]

        if data:
            fwd = await trainer.forward_backward_async(data, "importance_sampling")
            opt = await trainer.optim_step_async(tinker.AdamParams(learning_rate=args.lr))
            await fwd.result_async()
            await opt.result_async()

        episodes = [e for g in groups for e in g]
        spent += sum(e.prefill * PRICE_PREFILL + e.generated * PRICE_SAMPLE for e in episodes) / 1e6
        spent += sum(d.model_input.length for d in data) * PRICE_TRAIN / 1e6
        log = {"step": step, "reward": sum(e.reward for e in episodes) / len(episodes),
               "skipped_groups": skipped / len(groups),
               "searches": sum(len(e.searches) for e in episodes) / len(episodes),
               "answered": sum(e.end == "answer" for e in episodes) / len(episodes),
               "budget_hit": sum(e.end in ("budget", "length") for e in episodes) / len(episodes),
               "tokens": sum(len(e.tokens) for e in episodes) / len(episodes), "usd": round(spent, 2)}
        print("  ".join(f"{k} {v:.3f}" if isinstance(v, float) else f"{k} {v}" for k, v in log.items()), flush=True)
        with open(f"runs/{args.name}/train_log.jsonl", "a") as f:
            f.write(json.dumps(log) + "\n")

        if step % args.save_every == 0 or step == args.steps:
            name = f"{args.name}-s{step}"
            sampler = await (await trainer.save_weights_for_sampler_async(name=name)).result_async()
            state = await (await trainer.save_state_async(name=name + "-state")).result_async()
            print(f"saved {name}\n  sampler: {sampler.path}\n  state:   {state.path}")
            with open("runs/checkpoints.jsonl", "a") as f:
                f.write(json.dumps({"name": name, "sampler": sampler.path, "state": state.path}) + "\n")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--init-state", help="tinker:// state path from train_sft.py; omit to start from the base model")
    p.add_argument("--name", default="grpo")
    p.add_argument("--steps", type=int, default=30)
    p.add_argument("--questions", type=int, default=32, help="questions per step")
    p.add_argument("--group", type=int, default=8, help="episodes per question")
    p.add_argument("--lr", type=float, default=4e-5)
    p.add_argument("--temperature", type=float, default=1.0)
    p.add_argument("--offset", type=int, default=1500, help="skip the questions used for SFT")
    p.add_argument("--save-every", type=int, default=10)
    asyncio.run(main(p.parse_args()))

python train_grpo.py --init-state $(python ckpt.py sft-e2 state) --steps 30   # change e2 to the epoch you picked

for s in 10 20 30; do
  python eval.py --name grpo-s$s --model-path $(python ckpt.py grpo-s$s)
done
python ablate.py runs/sft-e2 runs/grpo-s10 runs/grpo-s20 runs/grpo-s30

"""Step 5: compare runs on the same questions with 95% bootstrap confidence intervals.

python ablate.py runs/test-base runs/test-sft runs/test-grpo
The first run is the baseline; every other run gets a paired difference against it.
"""
import json
import random
import sys


def load(run):
    rows = {}
    for line in open(f"{run}/episodes.jsonl"):
        r = json.loads(line)
        rows[r["id"]] = r
    return rows


def ci(values, rng, n=1000):
    means = sorted(sum(rng.choices(values, k=len(values))) / len(values) for _ in range(n))
    return means[int(0.025 * n)], means[int(0.975 * n)]


runs = sys.argv[1:]
data = [load(r) for r in runs]
ids = sorted(set.intersection(*(set(d) for d in data)))
rng = random.Random(0)
base = [data[0][i]["f1"] for i in ids]
print(f"{len(ids)} shared questions\n")
print(f"{'run':<22}{'F1':>7}{'95% CI':>16}{'EM':>7}{'searches':>10}{'$/q':>9}{'Δ F1 vs first':>24}")
for run, d in zip(runs, data):
    f1s = [d[i]["f1"] for i in ids]
    lo, hi = ci(f1s, rng)
    mean = lambda key: sum(d[i][key] for i in ids) / len(ids)
    searches = sum(len(d[i]["searches"]) for i in ids) / len(ids)
    delta = ""
    if d is not data[0]:
        diffs = [a - b for a, b in zip(f1s, base)]
        dlo, dhi = ci(diffs, rng)
        delta = f"{sum(diffs) / len(diffs):+.3f} [{dlo:+.3f}, {dhi:+.3f}]"
    print(f"{run:<22}{mean('f1'):>7.3f}{f'[{lo:.3f}, {hi:.3f}]':>16}{mean('em'):>7.3f}"
          f"{searches:>10.2f}{mean('usd'):>9.4f}{delta:>24}")
print("\nA difference is real only if its interval does not contain 0.")

python eval.py --split test --name test-closed --closed-book
python eval.py --split test --name test-base
python eval.py --split test --name test-sft  --model-path $(python ckpt.py sft-e2)
python eval.py --split test --name test-grpo --model-path $(python ckpt.py grpo-s30)
python ablate.py runs/test-base runs/test-closed runs/test-sft runs/test-grpo

python eval.py --split test --name test-grpo-k1 --model-path $(python ckpt.py grpo-s30) --top-k 1
python eval.py --split test --name test-grpo-b2 --model-path $(python ckpt.py grpo-s30) --max-searches 2
python ablate.py runs/test-grpo runs/test-grpo-k1 runs/test-grpo-b2