"""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()