Correct the model card (pooling is at the LABELS suffix's last token); add code/classifier from project commit 4ced7554f
3d7c238 verified | """Qwen3.5 text backbone with LoRA and one linear head per field (needs torch; runs on the pod). | |
| Adapted from the read-only reference ~/yt_live_analysis/training/model.py (arm H). Only the text decoder is | |
| loaded: the checkpoints are Qwen3_5ForConditionalGeneration, and loading Qwen3_5ForCausalLM with the | |
| text config maps the language-model weights and ignores the vision tower and the MTP layer. The LM head is | |
| dropped. Rows are right-padded (left padding would run pad tokens through the Gated DeltaNet state before | |
| the post) and each row is pooled at its last real token. The loss is normalised per optimiser step | |
| (``data.step_counts``), so splitting a step into micro-batches leaves its gradient unchanged. | |
| """ | |
| from __future__ import annotations | |
| from pathlib import Path | |
| import torch | |
| from torch import nn | |
| import torch.nn.functional as F | |
| LORA_TARGETS = ['q_proj', 'k_proj', 'v_proj', 'o_proj', 'in_proj_qkv', 'in_proj_z', 'out_proj', | |
| 'gate_proj', 'up_proj', 'down_proj'] | |
| def text_config(path): | |
| from transformers import AutoConfig | |
| config = AutoConfig.from_pretrained(path) | |
| return getattr(config, 'text_config', config) | |
| def load_backbone(path, device, dtype=torch.bfloat16, attn='sdpa'): | |
| """Text-only load; refuses a checkpoint with missing weights (they would be randomly initialised).""" | |
| from transformers import Qwen3_5ForCausalLM | |
| kw = {'device_map': {'': str(device)}} if device.type == 'cuda' else {} | |
| lm, info = Qwen3_5ForCausalLM.from_pretrained(path, config=text_config(path), dtype=dtype, attn_implementation=attn, | |
| output_loading_info=True, **kw) | |
| if info['missing_keys']: | |
| raise RuntimeError('Missing weights in %s: %s' % (path, sorted(info['missing_keys'])[:5])) | |
| return lm | |
| def load_tokenizer(path): | |
| from transformers import AutoTokenizer | |
| tok = AutoTokenizer.from_pretrained(path, split_special_tokens=True) | |
| tok.padding_side = 'right' | |
| if tok.pad_token_id is None: | |
| tok.pad_token = tok.eos_token | |
| return tok | |
| def tokenizer_fn(tok): | |
| """Plain token ids: no added special tokens, and special-token strings in the 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 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 backbone; the final-norm hidden state at the last real token feeds one linear head per field.""" | |
| def __init__(self, backbone, hidden: int, sizes: list[int], dropout: float = 0.1): | |
| super().__init__() | |
| self.backbone = backbone | |
| self.dropout = nn.Dropout(dropout) | |
| self.heads = nn.ModuleList([nn.Linear(hidden, k) for k in sizes]) | |
| def forward(self, ids, mask, last) -> list[torch.Tensor]: | |
| h = self.backbone(input_ids=ids, attention_mask=mask, use_cache=False).last_hidden_state | |
| pooled = h[torch.arange(h.shape[0], device=h.device), last].float() | |
| pooled = self.dropout(pooled) | |
| return [head(pooled) for head in self.heads] | |
| def step_loss(logits: list[torch.Tensor], labels: torch.Tensor, counts: list[int]): | |
| """Sum over fields of (summed cross-entropy over this micro-batch's non-null rows) / (step's non-null count). | |
| ``labels`` is [rows, fields] with -1 for a null target. Summing this over a step's micro-batches gives, | |
| per field, the mean over the step's non-null rows, whatever the split. Returns (loss, per-field sums). | |
| """ | |
| total = logits[0].sum() * 0.0 | |
| per_field = [] | |
| for i, lg in enumerate(logits): | |
| y = labels[:, i] | |
| keep = y >= 0 | |
| if bool(keep.any()): | |
| ce = F.cross_entropy(lg[keep].float(), y[keep], reduction='sum') | |
| total = total + 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): | |
| lm = load_backbone(cfg['model'], device, dtype=getattr(torch, cfg.get('dtype', 'bfloat16'))) | |
| text = lm.model | |
| if train: | |
| 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']) | |
| lm.lm_head = None | |
| model = HeadsModel(text, lm.config.hidden_size, sizes, cfg['head_dropout']) | |
| if heads_path: | |
| model.heads.load_state_dict(torch.load(heads_path, map_location='cpu')) | |
| model = model.to(device) | |
| # from_pretrained returns modules in eval mode, and gradient checkpointing only runs in training mode. | |
| model.train(train) | |
| return model | |
| def save_checkpoint(model: HeadsModel, target: Path, extra: dict | None = None, trainer=None) -> None: | |
| """Atomic directory write: adapter/, heads.pt, state.json and (for resume points) trainer.pt.""" | |
| import json | |
| import shutil | |
| tmp = target.with_name(target.name + '.tmp') | |
| shutil.rmtree(tmp, ignore_errors=True) | |
| tmp.mkdir(parents=True) | |
| model.backbone.save_pretrained(str(tmp / 'adapter')) | |
| torch.save(model.heads.state_dict(), tmp / 'heads.pt') | |
| if trainer is not None: | |
| torch.save(trainer, tmp / 'trainer.pt') | |
| (tmp / 'state.json').write_text(json.dumps(extra or {})) | |
| shutil.rmtree(target, ignore_errors=True) | |
| tmp.rename(target) | |
| 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).""" | |
| 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 | |