Fine-tune an LLM and decide whether to ship it

def gates(before, after, guidelines, min_drop=0.2, max_regression=0.05, min_clean=0.9):
    return {
        "improved": after["violations per answer"] <= (1 - min_drop) * before["violations per answer"],
        "no_regression": all(after[g] - before[g] <= max_regression for g in guidelines),
        "clean": after["ends cleanly"] >= min_clean,
    }

import requests
import torch
import torch.nn as nn
from tinygpt import load_pretrained

baseline, tok = load_pretrained("tiny-gpt-chat")
candidate, _ = load_pretrained("tiny-gpt-chat")
end = tok.token_id("<end>")
pairs = requests.get("https://api.gradient.test/v1/preferences").json()["data"]
kinds = sorted({" ".join(p["prompt"].split()[3:-1]) for p in pairs})
test_names = {"Kestrel", "Lark", "River", "Sierra", "Summit", "Willow"}
held_back = test_names | {"Alder", "Basin", "Cove", "Ember"}        # test and validation names

def to_tensors(rows):
    enc = []
    for prompt, answer in rows:
        p, a = tok.encode(f"<user> {prompt} <assistant>"), tok.encode(answer) + [end]
        enc.append(((p + a)[:65], ([-100] * len(p) + a)[:65]))
    T = max(len(ids) for ids, _ in enc)
    return (torch.tensor([ids + [end] * (T - len(ids)) for ids, _ in enc]),
            torch.tensor([lab + [-100] * (T - len(lab)) for _, lab in enc]))

def finetune(model, rows, steps=80):
    X, Y = to_tensors(rows)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
    model.train()
    torch.manual_seed(1)
    for step in range(steps):
        b = torch.randint(0, len(X), (16,))
        logits = model(X[b][:, :-1])
        loss = nn.functional.cross_entropy(logits.reshape(-1, logits.shape[-1]), Y[b][:, 1:].reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
    model.eval()

@torch.no_grad()
def evaluate(model):
    torch.manual_seed(0)
    items, clean = [], 0
    for product in [f"{n} {k}" for n in sorted(test_names) for k in kinds]:
        ids = torch.tensor([tok.encode(f"<user> Describe the {product} . <assistant>")] * 3)   # three answers each
        for row in model.generate(ids, 60, temperature=0.8, stop_id=end)[:, ids.shape[1]:].tolist():
            reply = tok.decode(row)
            clean += reply.endswith("<end>")
            items.append({"product": product, "text": reply.split("<end>")[0].strip()})
    flaws = [r["flaws"] for r in requests.post("https://api.gradient.test/v1/brand-check", json={"items": items}).json()["results"]]
    share = lambda test: sum(any(test(f) for f in fs) for fs in flaws) / len(flaws)
    return {"violations per answer": sum(map(len, flaws)) / len(flaws),
            "names another product": share(lambda f: f.startswith("doesn't name")),
            "repeats a word": share(lambda f: f.startswith("repeats")),
            "wrong or missing material": share(lambda f: "material" in f or "isn't used" in f),
            "missing repair promise": share(lambda f: f == "no repair promise"),
            "ends cleanly": clean / len(items)}

curated = [(p["prompt"], p["chosen"]) for p in pairs if not p["chosen_flaws"] and p["prompt"].split()[2] not in held_back]
finetune(candidate, curated)
before, after = evaluate(baseline), evaluate(candidate)
print(f"candidate A: {len(curated)} curated answers\n")
print(f"{'':<28} {'baseline':>9} {'candidate':>10}")
for metric in before:
    print(f"{metric:<28} {before[metric]:>9.2f} {after[metric]:>10.2f}")
print()
for gate, ok in gates(before, after).items():
    print(f"{'PASS' if ok else 'FAIL'}  {gate}")

product = re.compile(r"^(?:Meet the|Our|The) ([A-Z][a-z]+) (.+?)(?:,| is| are) made")
rows = []
for d in docs:                                         # the catalog's real descriptions...
    name, kind = product.match(d).groups()
    if name in held_back:
        continue
    answer = d if "repair promise" in d else d + " Backed by our repair promise."   # ...plus the brand rule
    rows.append((f"Describe the {name} {kind} .", answer))

finetune(candidate, rows)
before, after = evaluate(baseline), evaluate(candidate)
print(f"candidate B: {len(rows)} catalog descriptions with the repair promise\n")
print(f"{'':<28} {'baseline':>9} {'candidate':>10}")
for metric in before:
    print(f"{metric:<28} {before[metric]:>9.2f} {after[metric]:>10.2f}")
print()
results = gates(before, after)
for gate, ok in results.items():
    print(f"{'PASS' if ok else 'FAIL'}  {gate}")
print("\nDecision:", "ship candidate B" if all(results.values()) else "do not ship")

def decision(gate_results):
    failed = [name for name, ok in gate_results.items() if not ok]
    return "ship" if not failed else "do not ship: " + ", ".join(failed)

def model_card(name, base, metrics_before, metrics_after):
    return {"model": name, "base_model": base,
            "evaluation": {m: {"before": metrics_before[m], "after": metrics_after[m]} for m in metrics_before}}

before = {"violations per answer": 1.44, "wrong material": 0.22, "ends cleanly": 1.0}
after = {"violations per answer": 1.14, "wrong material": 0.32, "ends cleanly": 0.99}
g = gates(before, after, ["wrong material"])
print(g)
print(decision(g))