🔬Modal Research Exchange

{
  "base_model": "Qwen/Qwen3.5-4B",
  "baseline": "Base Qwen3.5-4B with no LoRA, prompted zero-shot with the identical grading prompt, renderer, decoding and 1000 eval pairs.",
  "batch_size": 64,
  "budget": {
    "gpu_hours": 1.5,
    "max_wall_minutes": 90
  },
  "command": "python - <<'EOF'\n# Custom part: PubMedQA -> grading-pair data builder; training loop is stock tinker SDK\nimport os, random, tinker\nfrom datasets import load_dataset\nfrom tinker_cookbook import renderers\nfrom tinker_cookbook.tokenizer_utils import get_tokenizer\nfrom tinker_cookbook.supervised.data import conversation_to_datum\nM='Qwen/Qwen3.5-4B'\nrenderer=renderers.get_renderer('qwen3', get_tokenizer(M))\ndef pair(ex):\n    gold=ex['final_decision']; ok=random.random()<0.5\n    cand=gold if ok else random.choice([l for l in ['yes','no','maybe'] if l!=gold])\n    ctx=' '.join(ex['context']['contexts'])\n    return [{'role':'user','content':f'Context: {ctx}\\nQuestion: {ex[\"question\"]}\\nCandidate answer: {cand}\\nIs the candidate answer correct? Reply correct or incorrect.'},\n            {'role':'assistant','content':'correct' if ok else 'incorrect'}]\ntrain=load_dataset('qiaojin/PubMedQA','pqa_artificial',split='train').shuffle(seed=0).select(range(9600))\nsc=tinker.ServiceClient(base_url=os.environ['TINKER_BASE_URL'])\ntc=sc.create_lora_training_client(base_model=M, rank=32)\nfor step in range(150):\n    batch=[conversation_to_datum(pair(train[step*64+i]), renderer, 1024) for i in range(64)]\n    tc.forward_backward(batch, loss_fn='cross_entropy').result()\n    tc.optim_step(tinker.AdamParams(learning_rate=1e-4)).result()\nprint(tc.save_weights_for_sampler(name='pubmedqa-judge').result().path)\nEOF\n# Eval: same grading prompt on 1000 pairs from load_dataset('qiaojin/PubMedQA','pqa_labeled'), sampled greedily from the saved LoRA and from the base model",
  "dataset": {
    "hf_path": "qiaojin/PubMedQA",
    "n_examples": 9600,
    "name": "PubMedQA (pqa_artificial for train, pqa_labeled for eval)",
    "split": "pqa_artificial/train"
  },
  "eval": {
    "benchmark": "PubMedQA pqa_labeled grading pairs (balanced correct/incorrect)",
    "metric": "judge accuracy vs gold verdict",
    "n_examples": 1000
  },
  "hypothesis": "A Qwen3.5-4B LoRA judge fine-tuned on PubMedQA pqa_artificial grades candidate answers on expert-labeled pqa_labeled at least 10 points more accurately than the same base model used zero-shot.",
  "lora_rank": 32,
  "max_seq_len": 1024,
  "method": "Turn PubMedQA into a grading task. Each example is (question, abstract context, candidate answer yes/no/maybe), and the target output is 'correct' or 'incorrect'. Build the pairs from gold labels, half with the gold answer and half with a randomly chosen wrong label. The treatment arm trains a rank-32 LoRA on about 9.6k such pairs from pqa_artificial (150 steps x batch 64) with the standard SFT cross-entropy, computed on the assistant turn only. Both arms are evaluated on 1000 balanced grading pairs built from the expert-labeled pqa_labeled split, using the same prompt and renderer. Greedy decoding goes through a sampling client, and the first output token is parsed as the verdict. This is the minimal check that Spindle can produce a usable bio grader before moving to Latch Bio's private data.",
  "optimizer": {
    "lr": 0.0001,
    "name": "adam"
  },
  "steps": 150,
  "success_criterion": "Validated if LoRA judge accuracy minus base judge accuracy is at least 10 points. The standard error is about 1.5 points per arm at n=1000, so gaps under 3 points count as noise. Failed if the gap is under 3 points.",
  "title": "LoRA SFT judge vs zero-shot base judge on PubMedQA answer grading",
  "training_type": "sft"
}

"""LoRA SFT judge vs zero-shot base judge on PubMedQA answer grading (Spindle, stock Tinker SDK).

Reads TINKER_BASE_URL, TINKER_API_KEY, MODAL_PROXY_TOKEN_ID and
MODAL_PROXY_TOKEN_SECRET. Turns PubMedQA into a grading task: given
(question, abstract context, candidate yes/no/maybe), output "correct" or
"incorrect". Trains a LoRA on grading pairs from pqa_artificial, then compares
greedy-decoded judge accuracy of the LoRA and the zero-shot base model on 1000
balanced grading pairs from the expert-labeled pqa_labeled split. Both arms use
the same prompt, chat template and decoding. Prints one final line
`RESULT {...}` with the metrics.
"""

import json
import math
import os
import random
import re
import time

import tinker
from datasets import load_dataset
from tinker import types

# CONFIG-START
CONFIG = {
    "base_model": "Qwen/Qwen3.5-4B",
    "lora_rank": 32,
    "dataset": "qiaojin/PubMedQA",
    "split": "pqa_artificial/train",      # train pairs: <config>/<split>
    "eval_split": "pqa_labeled/train",     # eval pairs: expert-labeled, 1000 rows
    "n_examples": 9600,
    "n_eval": 1000,
    "steps": 150,
    "batch_size": 64,
    "max_seq_len": 1024,
    "lr": 0.0001,
    "optimizer": "adam",
    "max_wall_seconds": 2700,
    "max_new_tokens": 8,
    "seed": 0
}
# CONFIG-END

LABELS = ["yes", "no", "maybe"]
PROMPT = (
    "Context: {ctx}\nQuestion: {question}\nCandidate answer: {cand}\n"
    "Is the candidate answer correct? Reply correct or incorrect."
)

STARTED = time.monotonic()


def elapsed():
    return time.monotonic() - STARTED


def chat_prompt_ids(tokenizer, text):
    # Thinking is disabled so the first generated word is the verdict; the same
    # template is used for training and for both eval arms.
    ids = tokenizer.apply_chat_template(
        [{"role": "user", "content": text}],
        tokenize=True,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    if hasattr(ids, "keys"):
        ids = ids["input_ids"]
    return list(ids)


def build_prompt_ids(tokenizer, question, contexts, cand):
    """Prompt token ids; the abstract context is truncated so that prompt + verdict fit max_seq_len."""
    ctx = " ".join(contexts)
    limit = CONFIG["max_seq_len"] - 16
    ids = chat_prompt_ids(tokenizer, PROMPT.format(ctx=ctx, question=question, cand=cand))
    if len(ids) <= limit:
        return ids
    overhead = len(chat_prompt_ids(tokenizer, PROMPT.format(ctx="", question=question, cand=cand)))
    budget = max(limit - overhead - 4, 32)
    ctx_ids = tokenizer.encode(ctx, add_special_tokens=False)[:budget]
    ctx = tokenizer.decode(ctx_ids)
    ids = chat_prompt_ids(tokenizer, PROMPT.format(ctx=ctx, question=question, cand=cand))
    return ids[-limit:] if len(ids) > limit else ids


def load_split(spec):
    config_name, split = spec.split("/")
    return load_dataset(CONFIG["dataset"], config_name, split=split)


def load_examples(tokenizer, spec, n, rng, balanced):
    """Return n grading pairs: (prompt_ids, verdict_text, gold_label, candidate).

    balanced=True gives exactly half correct / half incorrect (used for eval);
    otherwise each pair is correct with probability 0.5 (training, as in the spec)."""
    ds = load_split(spec).shuffle(seed=CONFIG["seed"])
    n = min(n, len(ds))
    ds = ds.select(range(n))
    if balanced:
        flags = [i < n // 2 for i in range(n)]
        rng.shuffle(flags)
    else:
        flags = [rng.random() < 0.5 for _ in range(n)]
    examples = []
    for row, ok in zip(ds, flags):
        gold = row["final_decision"].strip().lower()
        cand = gold if ok else rng.choice([l for l in LABELS if l != gold])
        ids = build_prompt_ids(tokenizer, row["question"], row["context"]["contexts"], cand)
        examples.append((ids, "correct" if ok else "incorrect", gold, cand))
    return examples


def make_datum(tokenizer, prompt_ids, response):
    """Loss weights cover only the assistant (verdict) tokens.

    Returns the Datum and its number of loss tokens."""
    response_ids = tokenizer.encode(response, add_special_tokens=False) + [tokenizer.eos_token_id]
    tokens = (prompt_ids + response_ids)[: CONFIG["max_seq_len"] + 1]
    weights = ([0.0] * len(prompt_ids) + [1.0] * len(response_ids))[: len(tokens)]
    datum = types.Datum(
        model_input=types.ModelInput.from_ints(tokens[:-1]),
        loss_fn_inputs={"target_tokens": tokens[1:], "weights": weights[1:]},
    )
    return datum, sum(weights[1:])


def mean_loss(output, batch):
    loss_tokens = sum(n for _, n in batch)
    return output.metrics["loss:sum"] / max(loss_tokens, 1.0)


def eval_loss(training, batch, chunk=64):
    total, tokens = 0.0, 0.0
    for i in range(0, len(batch), chunk):
        part = batch[i : i + chunk]
        output = training.forward([d for d, _ in part], "cross_entropy").result()
        total += output.metrics["loss:sum"]
        tokens += sum(n for _, n in part)
    return total / max(tokens, 1.0)


def optim_params():
    if CONFIG["optimizer"] == "dimuon":
        return types.DimuonParams(learning_rate=CONFIG["lr"])
    return types.AdamParams(learning_rate=CONFIG["lr"])


def parse_verdict(text):
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.S)
    text = text.replace("<think>", "").replace("</think>", "").strip().lower()
    m = re.match(r"[^a-z]*([a-z]+)", text)
    if not m:
        return None
    word = m.group(1)
    return word if word in ("correct", "incorrect") else None


def judge_eval(sampler, tokenizer, examples, chunk=200):
    """Greedy-decode each grading prompt and parse the first word as the verdict."""
    params = types.SamplingParams(max_tokens=CONFIG["max_new_tokens"], temperature=0.0)
    verdicts = []
    for i in range(0, len(examples), chunk):
        part = examples[i : i + chunk]
        futures = [
            sampler.sample(
                prompt=types.ModelInput.from_ints(ids), num_samples=1, sampling_params=params
            )
            for ids, _, _, _ in part
        ]
        for fut in futures:
            seq = fut.result().sequences[0]
            verdicts.append(parse_verdict(tokenizer.decode(seq.tokens, skip_special_tokens=True)))
    n = len(examples)
    correct = sum(v == ex[1] for v, ex in zip(verdicts, examples))
    invalid = sum(v is None for v in verdicts)
    pos = [(v, ex) for v, ex in zip(verdicts, examples) if ex[1] == "correct"]
    neg = [(v, ex) for v, ex in zip(verdicts, examples) if ex[1] == "incorrect"]
    acc = correct / max(n, 1)
    return {
        "judge_acc": acc,
        "judge_acc_se": math.sqrt(acc * (1 - acc) / max(n, 1)),
        "acc_on_correct_pairs": sum(v == "correct" for v, _ in pos) / max(len(pos), 1),
        "acc_on_incorrect_pairs": sum(v == "incorrect" for v, _ in neg) / max(len(neg), 1),
        "invalid_rate": invalid / max(n, 1),
        "frac_pred_correct": sum(v == "correct" for v in verdicts) / max(n, 1),
        "n": n,
    }


def main():
    rng = random.Random(CONFIG["seed"])
    service = tinker.ServiceClient(
        base_url=os.environ["TINKER_BASE_URL"],
        api_key=os.environ["TINKER_API_KEY"],
        default_headers_overrides={
            "Modal-Key": os.environ["MODAL_PROXY_TOKEN_ID"],
            "Modal-Secret": os.environ["MODAL_PROXY_TOKEN_SECRET"],
        },
    )
    training = service.create_lora_training_client(
        base_model=CONFIG["base_model"], rank=CONFIG["lora_rank"]
    )
    tokenizer = training.get_tokenizer()
    print(f"training client ready for {CONFIG['base_model']} after {elapsed():.0f}s", flush=True)

    eval_examples = load_examples(tokenizer, CONFIG["eval_split"], CONFIG["n_eval"], rng, balanced=True)
    train_examples = load_examples(tokenizer, CONFIG["split"], CONFIG["n_examples"], rng, balanced=False)
    train_data = [make_datum(tokenizer, ids, v) for ids, v, _, _ in train_examples]
    eval_data = [make_datum(tokenizer, ids, v) for ids, v, _, _ in eval_examples]
    print(
        f"loaded {len(train_data)} train / {len(eval_data)} eval grading pairs after {elapsed():.0f}s "
        f"(train frac correct {sum(v == 'correct' for _, v, _, _ in train_examples) / max(len(train_examples), 1):.3f})",
        flush=True,
    )

    eval_before = eval_loss(training, eval_data) if eval_data else None
    print(f"eval loss before training: {eval_before}", flush=True)

    # Baseline arm: zero-shot base model, identical prompt/template/decoding/eval pairs.
    t0 = time.monotonic()
    base_sampler = service.create_sampling_client(base_model=CONFIG["base_model"])
    base_metrics = judge_eval(base_sampler, tokenizer, eval_examples)
    base_eval_seconds = time.monotonic() - t0
    print(f"base judge: {json.dumps(base_metrics)} ({base_eval_seconds:.0f}s)", flush=True)

    # Reserve time for the LoRA eval inside the wall-time budget.
    reserve = 1.5 * base_eval_seconds + 120
    losses = []
    for step in range(CONFIG["steps"]):
        if elapsed() > CONFIG["max_wall_seconds"] - reserve:
            print(f"stopping at step {step}: wall-time budget reached", flush=True)
            break
        start = step * CONFIG["batch_size"] % len(train_data)
        batch = (train_data + train_data)[start : start + CONFIG["batch_size"]]
        step_started = time.monotonic()
        output = training.forward_backward([d for d, _ in batch], "cross_entropy").result()
        training.optim_step(optim_params()).result()
        loss = mean_loss(output, batch)
        if not math.isfinite(loss):
            raise RuntimeError(f"step {step}: loss is {loss}")
        losses.append(loss)
        print(f"step {step}: train loss {loss:.4f} ({time.monotonic() - step_started:.1f}s)", flush=True)

    eval_after = eval_loss(training, eval_data) if eval_data else None
    print(f"eval loss after training: {eval_after}", flush=True)

    lora_sampler = training.save_weights_and_get_sampling_client(name="pubmedqa-judge")
    lora_metrics = judge_eval(lora_sampler, tokenizer, eval_examples)
    print(f"LoRA judge: {json.dumps(lora_metrics)}", flush=True)

    gap = 100.0 * (lora_metrics["judge_acc"] - base_metrics["judge_acc"])
    if gap >= 10.0:
        verdict = "validated"
    elif gap < 3.0:
        verdict = "failed"
    else:
        verdict = "inconclusive"

    result = {
        "train_loss_first": losses[0] if losses else None,
        "train_loss_last": losses[-1] if losses else None,
        "train_loss_min": min(losses) if losses else None,
        "eval_loss_before": eval_before,
        "eval_loss_after": eval_after,
        "steps": len(losses),
        "n_examples": len(train_data),
        "n_eval": len(eval_data),
        "base_model": CONFIG["base_model"],
        "lora_rank": CONFIG["lora_rank"],
        "lr": CONFIG["lr"],
        "optimizer": CONFIG["optimizer"],
        "dataset": CONFIG["dataset"],
        "judge_acc": lora_metrics["judge_acc"],
        "judge_acc_se": lora_metrics["judge_acc_se"],
        "lora_judge": lora_metrics,
        "baseline": {
            "arm": "zero-shot base model (no LoRA)",
            "eval_loss": eval_before,
            **base_metrics,
        },
        "judge_acc_gap_points": gap,
        "verdict": verdict,
        "wall_seconds": round(elapsed(), 1),
    }
    print("RESULT " + json.dumps(result), flush=True)


if __name__ == "__main__":
    main()

train_loss_first

0.8029

train_loss_last

0.01785

train_loss_min

0.0001048

eval_loss_before

0.7103

eval_loss_after

0.2683

steps

150

n_examples

9,600

n_eval

1,000

lora_rank

32

lr

0.0001

judge_acc

0.85

judge_acc_se

0.01129

judge_acc_gap_points

30.9

wall_seconds

1,084.2

base_model

Qwen/Qwen3.5-4B

optimizer

adam

dataset

qiaojin/PubMedQA

lora_judge

{'judge_acc': 0.85, 'judge_acc_se': 0.011291589790636216, 'acc_on_correct_pairs': 0.81, 'acc_on_incorrect_pairs': 0.89, 'invalid_rate': 0.0, 'frac_pred_correct': 0.46, 'n': 1000}

baseline

{'arm': 'zero-shot base model (no LoRA)', 'eval_loss': 0.7102990943645127, 'judge_acc': 0.541, 'judge_acc_se': 0.015758140753274162, 'acc_on_correct_pairs': 0.982, 'acc_on_incorrect_pairs': 0.1, 'invalid_rate': 0.0, 'frac_pred_correct': 0.941, 'n': 1000}

verdict

validated

command

python experiment_run_04595af10e.py

config

{'base_model': 'Qwen/Qwen3.5-4B', 'lora_rank': 32, 'dataset': 'qiaojin/PubMedQA', 'split': 'pqa_artificial/train', 'n_examples': 4096, 'n_eval': 256, 'steps': 150, 'batch_size': 64, 'max_seq_len': 1024, 'lr': 0.0001, 'optimizer': 'adam', 'max_wall_seconds': 2700}

modal_call_id

fc-01M6FN0T0YRW78MF8YRMM5RDTX

script_source

Claude (claude-opus-5-5) adapted the reference script