| """Gemma 4 text decoder with LoRA and one linear head per field (needs torch; runs on the pod and on node12). |
| |
| Classifier-plan §1, §4. The checkpoints (``google/gemma-4-12B`` and ``-it``) are |
| ``Gemma4UnifiedForConditionalGeneration`` (model_type ``gemma4_unified``): text decoder plus light image/audio |
| projections. Only the text decoder is loaded, as ``Gemma4UnifiedForCausalLM`` with the text config and a |
| ``key_mapping`` that renames the checkpoint's ``model.language_model.*`` to ``model.*`` (without it every weight is |
| reported missing and the model would train from random weights; the yt sibling found this on 2026-10-18). Image and |
| audio weights are ignored, the tied LM head is dropped: ``HeadsModel`` keeps only ``lm.model`` (the final-norm |
| hidden states) and one ``Linear(hidden -> K_f)`` per taxonomy field. |
| |
| Pooling (the §4.4 screen): ``last`` = the hidden state of each row's last real token (the ``### LABELS`` suffix); |
| ``mean`` = the mean of final-norm states over the row's real tokens. Rows are right-padded with an attention mask. |
| Gemma needs its BOS token: inputs are built by ``training.formats.encode`` with ``bos_id(tok)``. |
| |
| The loss (``step_loss``) is normalised per optimiser step: per field, the summed cross-entropy over unmasked rows |
| divided by that step's unmasked count, fields weighted by ``Taxonomy.weights()`` (direction 2.0). Splitting a step |
| into micro-batches therefore leaves its gradient unchanged. Optional per-option class weights |
| (``Taxonomy.class_weight_vectors()``; nsb-v2: pnl breakeven 3.0) scale each row's cross-entropy by the weight of its |
| gold option; the divisor stays the unweighted count, so the split invariance holds. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import random |
| import shutil |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] |
| |
| BACKBONES = { |
| "gemma4_unified": dict(name="gemma4", cls="Gemma4UnifiedForCausalLM", |
| key_mapping={r"^model\.language_model\.": "model."}, bos=True), |
| "gemma4_unified_text": dict(name="gemma4", cls="Gemma4UnifiedForCausalLM", key_mapping=None, bos=True), |
| } |
| HUB_IDS = {"base": "google/gemma-4-12B", "it": "google/gemma-4-12B-it"} |
| POOLS = ("last", "mean") |
|
|
|
|
| def backbone(path) -> dict: |
| from transformers import AutoConfig |
|
|
| model_type = AutoConfig.from_pretrained(path).model_type |
| if model_type not in BACKBONES: |
| raise ValueError(f"unsupported backbone {model_type!r}; known: {sorted(BACKBONES)}") |
| return BACKBONES[model_type] |
|
|
|
|
| def text_config(path): |
| from transformers import AutoConfig |
|
|
| config = AutoConfig.from_pretrained(path) |
| return getattr(config, "text_config", config) |
|
|
|
|
| def load_backbone(path, device=None, dtype=torch.bfloat16, attn: str = "sdpa"): |
| """Text-only causal LM; refuses a checkpoint with missing text weights (they would be randomly initialised). |
| |
| Returns (lm, loading_info).""" |
| import transformers |
|
|
| spec = backbone(path) |
| kw = {} |
| if device is not None and torch.device(device).type == "cuda": |
| kw["device_map"] = {"": str(device)} |
| if spec["key_mapping"]: |
| kw["key_mapping"] = spec["key_mapping"] |
| cls = getattr(transformers, spec["cls"]) |
| lm, info = cls.from_pretrained(str(path), config=text_config(path), dtype=dtype, attn_implementation=attn, |
| output_loading_info=True, **kw) |
| missing = sorted(k for k in info["missing_keys"] if not k.startswith("lm_head")) |
| if missing: |
| raise RuntimeError(f"missing weights in {path}: {missing[:5]} ({len(missing)} total)") |
| return lm, info |
|
|
|
|
| def load_tokenizer(path): |
| from transformers import AutoTokenizer |
|
|
| tok = AutoTokenizer.from_pretrained(str(path)) |
| tok.padding_side = "right" |
| if tok.pad_token_id is None: |
| tok.pad_token = tok.eos_token |
| return tok |
|
|
|
|
| def tokenizer_fn(tok): |
| """Plain ids: no added special tokens, and special-token strings in post text stay ordinary text.""" |
| def tokenize(text: str) -> list[int]: |
| return tok(text, add_special_tokens=False, split_special_tokens=True)["input_ids"] |
| return tokenize |
|
|
|
|
| def bos_id(tok) -> int | None: |
| """Gemma degrades badly without BOS; ``encode`` prepends exactly one.""" |
| return tok.bos_token_id |
|
|
|
|
| def add_lora(module, r: int, alpha: int, dropout: float): |
| from peft import LoraConfig, get_peft_model |
|
|
| return get_peft_model(module, LoraConfig(r=r, lora_alpha=alpha, lora_dropout=dropout, |
| target_modules=LORA_TARGETS, bias="none")) |
|
|
|
|
| def pad(rows: list[list[int]], pad_id: int, device): |
| """Right-padded ids, attention mask and each row's last real index.""" |
| width = max(len(r) for r in rows) |
| ids = torch.full((len(rows), width), pad_id, dtype=torch.long) |
| mask = torch.zeros((len(rows), width), dtype=torch.long) |
| for i, r in enumerate(rows): |
| ids[i, :len(r)] = torch.tensor(r, dtype=torch.long) |
| mask[i, :len(r)] = 1 |
| last = mask.sum(1) - 1 |
| return ids.to(device), mask.to(device), last.to(device) |
|
|
|
|
| class HeadsModel(nn.Module): |
| """LoRA text decoder; pooled final-norm hidden state -> one linear head per field.""" |
|
|
| def __init__(self, backbone, hidden: int, sizes: list[int], pool: str = "last", dropout: float = 0.1): |
| super().__init__() |
| if pool not in POOLS: |
| raise ValueError(f"pool must be one of {POOLS}") |
| self.backbone = backbone |
| self.pool = pool |
| self.dropout = nn.Dropout(dropout) |
| self.heads = nn.ModuleList([nn.Linear(hidden, k) for k in sizes]) |
|
|
| def pooled(self, ids, mask, last) -> torch.Tensor: |
| h = self.backbone(input_ids=ids, attention_mask=mask, use_cache=False).last_hidden_state |
| if self.pool == "last": |
| return h[torch.arange(h.shape[0], device=h.device), last].float() |
| m = mask.unsqueeze(-1).to(h.dtype) |
| return ((h * m).sum(1) / m.sum(1).clamp_min(1)).float() |
|
|
| def forward(self, ids, mask, last) -> list[torch.Tensor]: |
| pooled = self.dropout(self.pooled(ids, mask, last)) |
| return [head(pooled) for head in self.heads] |
|
|
|
|
| def step_loss(logits: list[torch.Tensor], labels: torch.Tensor, counts: list[int], weights: list[float] | None = None, |
| class_weights: list[torch.Tensor | None] | None = None): |
| """Sum over fields of weight x (summed CE over this micro-batch's unmasked rows) / (step's unmasked count). |
| |
| ``labels`` is [rows, fields] with -1 for a masked target. Summing this over a step's micro-batches gives, per |
| field, the mean over the step's unmasked rows, whatever the split. ``class_weights`` (per field a [K] tensor or |
| None) multiplies each row's CE by its gold option's weight. Returns (loss, per-field CE sums).""" |
| weights = weights or [1.0] * len(logits) |
| class_weights = class_weights or [None] * len(logits) |
| total = logits[0].sum() * 0.0 |
| per_field = [] |
| for i, lg in enumerate(logits): |
| y = labels[:, i] |
| keep = y >= 0 |
| if bool(keep.any()): |
| cw = class_weights[i] |
| ce = F.cross_entropy(lg[keep].float(), y[keep], reduction="sum", |
| weight=None if cw is None else cw.to(device=lg.device, dtype=torch.float32)) |
| total = total + weights[i] * ce / counts[i] |
| per_field.append(float(ce.detach())) |
| else: |
| per_field.append(0.0) |
| return total, per_field |
|
|
|
|
| def build(cfg: dict, sizes: list[int], device, *, train: bool, adapter_dir=None, heads_path=None) -> HeadsModel: |
| """HeadsModel from a base checkpoint dir (``cfg["model"]``); a saved adapter and heads when given.""" |
| device = torch.device(device) |
| lm, _ = load_backbone(cfg["model"], device, dtype=getattr(torch, cfg.get("dtype", "bfloat16")), |
| attn=cfg.get("attn", "sdpa")) |
| text = lm.model |
| if train and cfg.get("gradient_checkpointing", True): |
| text.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) |
| text.config.use_cache = False |
| if adapter_dir: |
| from peft import PeftModel |
|
|
| text = PeftModel.from_pretrained(text, str(adapter_dir), is_trainable=train) |
| else: |
| text = add_lora(text, cfg["lora_r"], cfg["lora_alpha"], cfg["lora_dropout"]) |
| hidden = lm.config.hidden_size |
| del lm |
| model = HeadsModel(text, hidden, sizes, cfg.get("pool", "last"), cfg.get("head_dropout", 0.1)) |
| if heads_path: |
| load_heads(model, heads_path) |
| model = model.to(device) |
| |
| model.train(train) |
| return model |
|
|
|
|
| |
| def heads_state(model: HeadsModel, fields: list[str]) -> dict[str, torch.Tensor]: |
| sd = model.heads.state_dict() |
| out = {} |
| for i, f in enumerate(fields): |
| out[f"{f}.weight"] = sd[f"{i}.weight"].detach().cpu().contiguous() |
| out[f"{f}.bias"] = sd[f"{i}.bias"].detach().cpu().contiguous() |
| return out |
|
|
|
|
| def save_heads(model: HeadsModel, fields: list[str], path: Path) -> None: |
| from safetensors.torch import save_file |
|
|
| save_file(heads_state(model, fields), str(path), metadata={"fields": json.dumps(fields), "pool": model.pool}) |
|
|
|
|
| def load_heads(model: HeadsModel, path) -> list[str]: |
| """Load ``heads.safetensors``; its field list and pooling must match the model.""" |
| from safetensors import safe_open |
| from safetensors.torch import load_file |
|
|
| with safe_open(str(path), framework="pt") as fh: |
| meta = fh.metadata() or {} |
| fields = json.loads(meta["fields"]) |
| if meta.get("pool", model.pool) != model.pool: |
| raise ValueError(f"heads were trained with pool={meta.get('pool')}, model has pool={model.pool}") |
| if len(fields) != len(model.heads): |
| raise ValueError(f"{len(fields)} heads in {path}, model has {len(model.heads)}") |
| tensors = load_file(str(path)) |
| sd = {} |
| for i, f in enumerate(fields): |
| sd[f"{i}.weight"] = tensors[f"{f}.weight"] |
| sd[f"{i}.bias"] = tensors[f"{f}.bias"] |
| model.heads.load_state_dict(sd) |
| return fields |
|
|
|
|
| def rng_state() -> dict: |
| out = {"python": random.getstate(), "torch": torch.get_rng_state()} |
| if torch.cuda.is_available(): |
| out["cuda"] = torch.cuda.get_rng_state_all() |
| return out |
|
|
|
|
| def set_rng_state(state: dict) -> None: |
| random.setstate(state["python"]) |
| torch.set_rng_state(state["torch"]) |
| if "cuda" in state and torch.cuda.is_available(): |
| torch.cuda.set_rng_state_all(state["cuda"]) |
|
|
|
|
| def save_checkpoint(model: HeadsModel, target: Path, fields: list[str], state: dict | None = None, |
| trainer: dict | None = None) -> None: |
| """Atomic directory write: adapter/, heads.safetensors, state.json and (resume points) trainer.pt.""" |
| target = Path(target) |
| tmp = target.with_name(target.name + ".tmp") |
| shutil.rmtree(tmp, ignore_errors=True) |
| tmp.mkdir(parents=True) |
| model.backbone.save_pretrained(str(tmp / "adapter")) |
| save_heads(model, fields, tmp / "heads.safetensors") |
| if trainer is not None: |
| torch.save(trainer, tmp / "trainer.pt") |
| (tmp / "state.json").write_text(json.dumps(state or {}, indent=1)) |
| old = target.with_name(target.name + ".old") |
| shutil.rmtree(old, ignore_errors=True) |
| if target.exists(): |
| os.replace(target, old) |
| os.replace(tmp, target) |
| shutil.rmtree(old, ignore_errors=True) |
|
|
|
|
| @torch.no_grad() |
| def predict(model: HeadsModel, encoded: list[list[int]], pad_id: int, batches: list[list[int]], device, |
| field_names: list[str]) -> dict[int, dict[str, list[float]]]: |
| """Raw logits per field for each encoded row index (dropout off).""" |
| device = torch.device(device) |
| was_training = model.training |
| model.eval() |
| out = {} |
| amp = dict(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda") |
| for b in batches: |
| ids, mask, last = pad([encoded[i] for i in b], pad_id, device) |
| with torch.autocast(**amp): |
| logits = model(ids, mask, last) |
| values = [lg.float().cpu().tolist() for lg in logits] |
| for j, i in enumerate(b): |
| out[i] = {f: values[k][j] for k, f in enumerate(field_names)} |
| if was_training: |
| model.train() |
| return out |
|
|
