train_head.py

"""Qwen3-Embedding-0.6B + LoRA + prototype head — faithful example of the jev-topic-head spec.

Mirrors: LoRA r16/a32/dropout .05 on attention+FFN (rest frozen); head = 56 anchors seeded
from topic-description embeddings + one learned scale (init 20); logits = scale * cos(v, P);
CE w/ label smoothing .05; two AdamW groups (lora 3e-4 wd .01 / head 2e-3 no wd); warmup
10% -> linear decay to 5%; WeightedRandomSampler 1/sqrt(count), hard pairs x2; early stop on
val macro-F1; seeds; ship nothing, just save runs/seed{S}/.

Deviations from the (partly inferred) spec, on purpose and documented:
  - last-token pooling per the official Qwen3-Embedding recipe
  - "bf16" read as bf16 base weights, fp32 head; fp32 end to end if dtype: fp32
  - data is synthetic (make_toy_data.py), so scores validate plumbing, not quality
"""
import argparse, copy, json, math, os, random

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import yaml
from sklearn.metrics import f1_score
from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler
from transformers import AutoModel, AutoTokenizer
from peft import LoraConfig, get_peft_model


def set_seed(s):
    random.seed(s)
    np.random.seed(s)
    torch.manual_seed(s)
    torch.cuda.manual_seed_all(s)


def read_jsonl(path):
    rows = []
    with open(path) as f:
        for line in f:
            rows.append(json.loads(line))
    return rows


class ItemDataset(Dataset):
    def __init__(self, path, topic2id):
        self.rows = read_jsonl(path)
        self.t2i = topic2id

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, i):
        r = self.rows[i]
        return r["text"], self.t2i[r["jev_topic"]]


def collate(batch):
    texts, ys = zip(*batch)
    return list(texts), torch.tensor(ys, dtype=torch.long)


def encode(model, tok, texts, max_length, device):
    """Last-token pooling (Qwen3-Embedding recipe), fp32 output."""
    enc = tok(texts, padding=True, truncation=True, max_length=max_length,
              return_tensors="pt").to(device)
    out = model(**enc).last_hidden_state
    mask = enc["attention_mask"]
    last = mask.sum(dim=1) - 1
    v = out[torch.arange(out.size(0), device=device), last]
    return v.float()


class PrototypeHead(nn.Module):
    def __init__(self, seed_protos, scale_init):
        super().__init__()
        self.protos = nn.Parameter(seed_protos.clone().float())
        self.scale = nn.Parameter(torch.tensor(float(scale_init)))

    def logits(self, v):
        v = F.normalize(v, dim=-1)
        p = F.normalize(self.protos, dim=-1)
        return self.scale * (v @ p.t())


@torch.no_grad()
def predict_logits(model, head, tok, texts, max_length, device, batch=32):
    model.eval()
    all_logits = []
    for i in range(0, len(texts), batch):
        v = encode(model, tok, texts[i:i + batch], max_length, device)
        all_logits.append(head.logits(v).float().cpu())
    return torch.cat(all_logits)


def macro_f1(y_true, y_pred, n_topics):
    return f1_score(y_true, y_pred, average="macro", zero_division=0,
                    labels=list(range(n_topics)))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="config.yaml")
    ap.add_argument("--seed", type=int, default=None)
    args = ap.parse_args()
    cfg = yaml.safe_load(open(args.config))
    seed = args.seed if args.seed is not None else 0
    set_seed(seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16 if cfg["train"]["dtype"] == "bf16" else torch.float32

    data_dir = cfg["data"]["dir"]
    taxonomy = json.load(open(os.path.join(data_dir, cfg["data"]["taxonomy"])))
    topics = sorted(taxonomy)                      # fixed, seed-independent order
    topic2id = {t: i for i, t in enumerate(topics)}
    K = len(topics)

    tok = AutoTokenizer.from_pretrained(cfg["base_model"])
    tok.padding_side = "right"
    base = AutoModel.from_pretrained(cfg["base_model"], torch_dtype=dtype).to(device)
    base.config.use_cache = False

    # ---- seed prototypes from topic-description embeddings (one forward pass) ----
    descs = [taxonomy[t] for t in topics]
    seed_protos = encode(base, tok, descs, cfg["max_length"], device)
    print(f"seeded protos {tuple(seed_protos.shape)}")

    # ---- LoRA: the only thing that changes the encoder ----
    lcfg = LoraConfig(r=cfg["lora"]["r"], lora_alpha=cfg["lora"]["alpha"],
                      lora_dropout=cfg["lora"]["dropout"],
                      target_modules=cfg["lora"]["target_modules"])
    model = get_peft_model(base, lcfg)
    model.print_trainable_parameters()

    head = PrototypeHead(seed_protos, cfg["head"]["scale_init"]).to(device)

    lora_params = [p for p in model.parameters() if p.requires_grad]
    n_lora = sum(p.numel() for p in lora_params)
    n_head = sum(p.numel() for p in head.parameters())
    print(f"trainable: {n_lora:,} lora + {n_head:,} head, rest frozen")

    # ---- data + class-balanced sampling (1/sqrt(count), hard pairs x2) ----
    train_rows = read_jsonl(os.path.join(data_dir, cfg["data"]["train"]))
    counts = {}
    for r in train_rows:
        counts[r["jev_topic"]] = counts.get(r["jev_topic"], 0) + 1
    hard = set()
    for a, b in cfg["data"].get("hard_pairs", []):
        hard.update([a, b])
    weights = [1.0 / math.sqrt(counts.get(r["jev_topic"], 1)) * (2.0 if r["jev_topic"] in hard else 1.0)
               for r in train_rows]
    sampler = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)

    ds = ItemDataset(os.path.join(data_dir, cfg["data"]["train"]), topic2id)
    dl = DataLoader(ds, batch_size=cfg["train"]["batch_size"], sampler=sampler,
                    collate_fn=collate)
    val_ds = ItemDataset(os.path.join(data_dir, cfg["data"]["validation"]), topic2id)
    val_texts = [r["text"] for r in val_ds.rows]
    val_y = [topic2id[r["jev_topic"]] for r in val_ds.rows]

    # ---- two AdamW groups + warmup->linear-decay-to-5% schedule ----
    o = cfg["optimizer"]
    opt = torch.optim.AdamW([
        {"params": lora_params, "lr": o["lora_lr"], "weight_decay": o["lora_wd"]},
        {"params": list(head.parameters()), "lr": o["head_lr"], "weight_decay": o["head_wd"]},
    ])
    epochs = cfg["train"]["epochs"]
    spe = math.ceil(len(ds) / cfg["train"]["batch_size"])
    total_steps = epochs * spe
    warmup = max(1, int(o["warmup_frac"] * total_steps))
    floor = o["final_lr_frac"]

    def lr_lambda(step):
        if step < warmup:
            return (step + 1) / warmup
        prog = (step - warmup) / max(1, total_steps - warmup)
        return floor + (1.0 - floor) * (1.0 - min(1.0, prog))

    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)

    out_dir = os.path.join(cfg["out_root"], f"seed{seed}")
    os.makedirs(out_dir, exist_ok=True)
    history, best_f1, best_state = [], -1.0, None
    patience, bad = cfg["train"]["patience"], 0

    for ep in range(epochs):
        model.train()
        head.train()
        losses = []
        for texts, y in dl:
            y = y.to(device)
            v = encode(model, tok, texts, cfg["max_length"], device)
            logits = head.logits(v)
            loss = F.cross_entropy(logits, y,
                                   label_smoothing=cfg["loss"]["label_smoothing"])
            opt.zero_grad(set_to_none=True)
            loss.backward()
            opt.step()
            sched.step()
            losses.append(loss.item())

        lg = predict_logits(model, head, tok, val_texts, cfg["max_length"], device)
        pred = lg.argmax(-1).numpy()
        f1 = macro_f1(val_y, pred, K)
        acc = float((pred == np.array(val_y)).mean())
        print(f"epoch {ep:2d} | loss {np.mean(losses):.3f} | val acc {acc:.3f} | val macro-F1 {f1:.4f}")
        history.append({"epoch": ep, "loss": float(np.mean(losses)),
                        "val_acc": acc, "val_macro_f1": float(f1)})

        if f1 > best_f1:
            best_f1, bad = float(f1), 0
            best_state = {"head": copy.deepcopy(head.state_dict()),
                          "lora": {k: v.detach().cpu().clone()
                                   for k, v in model.state_dict().items()
                                   if "lora_" in k and "requires_grad" not in k}}
        else:
            bad += 1
            if bad >= patience:
                print(f"early stop at epoch {ep} (patience {patience})")
                break

    # ---- save best (the shipped artifact is seed 0, like the spec) ----
    model.load_state_dict(best_state["lora"], strict=False)
    model.save_pretrained(os.path.join(out_dir, "lora"))
    head.load_state_dict(best_state["head"])
    head_dir = os.path.join(out_dir, "head")
    os.makedirs(head_dir, exist_ok=True)
    torch.save({"protos": head.protos.detach().cpu(),
                "scale": head.scale.detach().cpu(),
                "dim": head.protos.shape[1], "topics": topics,
                "base_model": cfg["base_model"]},
               os.path.join(head_dir, "head.pt"))
    json.dump({"seed": seed, "best_val_macro_f1": best_f1, "history": history},
              open(os.path.join(out_dir, "history.json"), "w"), indent=2)
    print(f"saved {out_dir} | best val macro-F1 {best_f1:.4f}")


if __name__ == "__main__":
    main()