Fine-tune a small model

import hashlib, json, random, re
from collections import Counter

def normalise(s: str) -> str:
    "Lowercase, collapse whitespace, drop punctuation. Only for comparison."
    return re.sub(r"[^a-z0-9 ]", "", re.sub(r"\s+", " ", s.lower())).strip()

rows = [json.loads(line) for line in open("raw.jsonl")]

# One schema, enforced now rather than discovered by the trainer at 3am.
clean, seen = [], set()
for r in rows:
    if not r.get("input") or not r.get("output"):
        continue
    key = hashlib.sha1(normalise(r["input"]).encode()).hexdigest()
    if key in seen:            # exact-after-normalisation duplicate
        continue
    seen.add(key)
    clean.append({"input": r["input"].strip(), "output": r["output"].strip(), "key": key})

print(len(rows), "->", len(clean), "after dedupe")
print(Counter(r["output"] for r in clean).most_common())   # label balance

random.Random(0).shuffle(clean)
held_out, train = clean[:150], clean[150:]

# The check that matters: nothing in held-out may appear in train.
assert not ({r["key"] for r in held_out} & {r["key"] for r in train})

for name, part in (("train", train), ("held_out", held_out)):
    with open(name + ".jsonl", "w") as f:
        for r in part:
            f.write(json.dumps({"input": r["input"], "output": r["output"]}) + "\n")

import json, random

held_out = [json.loads(l) for l in open("held_out.jsonl")]

def grade(prediction: str, expected: str) -> float:
    "Exact match after the same normalisation used everywhere else."
    return float(normalise(prediction) == normalise(expected))

def score(predict, examples, n_boot=1000):
    marks = [grade(predict(e["input"]), e["output"]) for e in examples]
    mean = sum(marks) / len(marks)

    # Bootstrap the interval. Without it you cannot tell a real gain from noise,
    # and on 150 examples the noise is roughly plus or minus four points.
    rng = random.Random(0)
    boots = []
    for _ in range(n_boot):
        sample = [marks[rng.randrange(len(marks))] for _ in marks]
        boots.append(sum(sample) / len(sample))
    boots.sort()
    return mean, boots[int(0.025 * n_boot)], boots[int(0.975 * n_boot)]

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16,
                                             device_map="auto")

def ask(system: str, user: str, max_new_tokens: int = 64) -> str:
    msgs = [{"role": "system", "content": system}, {"role": "user", "content": user}]
    ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt")
    # Greedy: the baseline must be reproducible, so no sampling anywhere.
    out = model.generate(ids.to(model.device), max_new_tokens=max_new_tokens,
                         do_sample=False, pad_token_id=tok.eos_token_id)
    return tok.decode(out[0, ids.shape[-1]:], skip_special_tokens=True).strip()

MINIMAL = "Classify the support ticket."

STRONG = open("prompt_v3.txt").read()   # task statement + format + 5 examples

print("naive  ", score(lambda x: ask(MINIMAL, x), held_out))
print("prompted", score(lambda x: ask(STRONG, x), held_out))

from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=16,                # rank: the width of the bottleneck
    lora_alpha=32,       # the update is scaled by alpha / r, so this is 2x
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    # Attention projections are the usual choice. Adding the MLP projections
    # ("gate_proj", "up_proj", "down_proj") helps when the task needs new
    # knowledge rather than a new format, and roughly triples adapter size.
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
)

peft_model = get_peft_model(model, config)
peft_model.print_trainable_parameters()
# trainable params: 9,232,384 || all params: 1,552,946,688 || trainable%: 0.5945

# Check it by hand for one layer before you believe the number.
for name, p in peft_model.named_parameters():
    if p.requires_grad:
        print(name, tuple(p.shape))
        break
# base_model.model.model.layers.0.self_attn.q_proj.lora_A.default.weight (16, 1536)

from datasets import load_dataset
from trl import SFTTrainer, SFTConfig

data = load_dataset("json", data_files={"train": "train.jsonl",
                                        "eval": "held_out.jsonl"})

def to_chat(row):
    # The same three messages the baseline used. If this drifts from ask(),
    # the fine-tune is being trained for a format it will never be asked in.
    return {"messages": [
        {"role": "system", "content": STRONG_SYSTEM_ONLY},
        {"role": "user", "content": row["input"]},
        {"role": "assistant", "content": row["output"]},
    ]}

data = data.map(to_chat, remove_columns=["input", "output"])

args = SFTConfig(
    output_dir="runs/lora-r16-lr2e4",
    num_train_epochs=3,             # more than 3 on a small set memorises it
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,  # effective batch 16
    learning_rate=2e-4,             # 10-100x a full fine-tune: only A and B move
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    bf16=True,
    logging_steps=10,
    eval_strategy="epoch",
    save_strategy="epoch",
    report_to="wandb",
    seed=0,
)

trainer = SFTTrainer(model=peft_model, args=args,
                     train_dataset=data["train"], eval_dataset=data["eval"])
trainer.train()
trainer.save_model("adapter/")   # writes only the adapter, tens of MB

# evaluate.py --model base --prompt minimal
#             --model base --prompt strong
#             --model adapter/ --prompt strong
import argparse, json
from peft import PeftModel

ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)    # "base" or a path to an adapter
ap.add_argument("--prompt", choices=["minimal", "strong"], required=True)
a = ap.parse_args()

m = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16,
                                         device_map="auto")
if a.model != "base":
    m = PeftModel.from_pretrained(m, a.model)
m.eval()

system = MINIMAL if a.prompt == "minimal" else STRONG
mean, lo, hi = score(lambda x: ask(system, x), held_out)
print(json.dumps({"model": a.model, "prompt": a.prompt,
                  "score": round(mean, 3), "ci": [round(lo, 3), round(hi, 3)]}))

base + minimal prompt    0.41   [0.33, 0.49]
base + strong prompt     0.78   [0.71, 0.85]
LoRA fine-tune           0.83   [0.76, 0.89]

from peft import PeftModel

base = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16)
merged = PeftModel.from_pretrained(base, "adapter/").merge_and_unload()
merged.save_pretrained("merged-model/")
tok.save_pretrained("merged-model/")

# Then serve it like any other model:
#   vllm serve ./merged-model --port 8000
#
# Or keep the base weights shared and load the adapter by name:
#   vllm serve Qwen/Qwen2.5-1.5B-Instruct --enable-lora \
#     --lora-modules tickets=./adapter