diff --git a/README.md b/README.md index 6f4b05a03..97c7c3005 100644 --- a/README.md +++ b/README.md @@ -406,6 +406,7 @@ foundation for building customizable, enterprise-grade training services. | Component Type | Component Link | Component Function | Author | | -------------- | -------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------- | ------------------- | | Patch | [qwen3_moe_transformers4_patch](https://www.modelscope.cn/models/twinkle-kit/qwen3_moe_transformers4_patch) | Fixes Qwen3 MoE model hang issue during FSDP2 training, effective for transformers==4.x | ModelScope Official | +| Decision | [Typed decision training](cookbook/transformers/decision/README.md) | Candidate scoring, soft-label losses, LoRA head persistence and question accuracy | Community | ## Contributions diff --git a/README_ZH.md b/README_ZH.md index 3d5f05fac..1d0112761 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -389,6 +389,7 @@ for epoch in range(3): | 组件类型 | 组件链接 | 组件功能 | 作者 | | -------- | -------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- | ----------------- | | Patch | [qwen3_moe_transformers4_patch](https://www.modelscope.cn/models/twinkle-kit/qwen3_moe_transformers4_patch) | 修复 Qwen3 MoE 模型在 FSDP2 训练期间挂起的问题，适用于 transformers==4.x | ModelScope 官方 | +| Decision | [决策模型训练](cookbook/transformers/decision/README.md) | 候选打分、软标签损失、LoRA 打分头保存与问题级精度 | 社区 | ## 开源贡献 diff --git a/cookbook/transformers/decision/README.md b/cookbook/transformers/decision/README.md new file mode 100644 index 000000000..a472f2723 --- /dev/null +++ b/cookbook/transformers/decision/README.md @@ -0,0 +1,136 @@ +# Typed decision training + +Twinkle can train and score typed decisions through its existing TransformersModel, InputProcessor, +Loss and PEFT interfaces. Predictions are probabilities over each question's legal options; +there is no autoregressive answer-token generation. A probability is not automatically calibrated: +retain the checkpoint's calibration settings, or fit new settings on a disjoint calibration split. + +## 1. Model and task + +| Family | Execution | Head | Loss preset | +| --- | --- | --- | --- | +| generic CausalLM | One question per prompt | Tokenizer-derived verbalizer | scoring | +| JEV | One question per prompt; text/image input | Released verbalizer and bias | distill_rps | +| Clef | All questions in a record jointly; text input | Span pooling and joint schema decoder | smooth_brier | +| OmniJev | Cached prefix and separate option branches; text/image/precomposed video | FiLM and ordinal head | ce_rps | +| StartLux | One question per prompt; text input | Last-position FP32 A-Z projection | scoring | + +`noul`, `choice`, and `score` use different option orders across families. The preprocessor preserves +those orders and remaps targets. Do not substitute one family's prompt or calibration for another. +Generic/JEV verbalizers require distinct single-token symbols and allow up to 16 choice options. +StartLux allows 2–26 options; wide choices must be explicitly decomposed outside this recipe. + +## 2. Validation and scope + +| Check | Evidence / scope | +| --- | --- | +| Loss and gradients | Four presets compared against independent unpadded formulas at 1e-12 | +| Accumulation | Uneven question counts match the full-batch normalized gradient | +| LoRA persistence | Tiny CausalLM updates both adapter and head; save/reload logits match exactly on CPU | +| Head contracts | Clef joint spans/padding, OmniJev FiLM/ordinal gradients, JEV image payload | +| Regression suite | Decision loss, encoding, augmentation and checkpoint regression tests; full repository pre-commit checks | +| Native training API | Tiny model forward/backward, optimizer step, save, resume from step 3 to 6 and reload; StartLux-4B two-step training, save and adapter reload/scoring on one Ascend NPU | +| Released-model quality | Full checkpoint-specific accuracy and visual benchmark validation pending | + +Tests use synthetic fixtures and do not establish released-checkpoint quality. Clef image span +alignment, packing, sequence/context parallelism and multi-adapter serving are not supported here. +The example uses one device; multi-rank training requires separate sharding and checkpoint validation. + +## 3. Installation and reproducibility + +Develop from a checkout and install its declared dependencies without changing repository bounds: + +```bash +python -m pip install -e . pytest +python -m pytest tests/decision -q --confcutdir=tests/decision +pre-commit run --all-files +``` + +An Ascend validation environment uses the base image below (no image export required), Python 3.12, +Torch 2.10.0, torch-npu 2.10.0.post4, CANN 9.1.0, Transformers 5.14.1 and PEFT 0.19.0. +Use a new container, mount only assigned devices and a task directory, and install Twinkle into an +isolated virtual environment. The image also contains unrelated inference packages; installing +training dependencies does not validate those packages' inference environments. + +```text +quay.io/ascend/vllm-ascend@sha256:8a397e82a22f55de25ae355a957b91e843a16ce704a03badc0a8e9915f68d84f +``` + +If the container exposes one physical NPU, its runtime-visible ordinal is `0`, even when its host +card number differs. Record both device identities in private deployment notes. + +## 4. Train, resume and score + +No model or dataset downloads are needed for this synthetic smoke test; dependencies must already +be installed. Use a fresh output directory for every run. + +```bash +python cookbook/transformers/decision/make_tiny.py --output /tmp/decision-tiny +python cookbook/transformers/decision/decision.py train \ + --model /tmp/decision-tiny --data /tmp/decision-tiny/mini.jsonl \ + --output /tmp/decision-train --steps 3 --dtype float32 +python cookbook/transformers/decision/decision.py infer \ + --model /tmp/decision-tiny --adapter /tmp/decision-train/final \ + --data /tmp/decision-tiny/mini.jsonl --output /tmp/decision-infer --dtype float32 +python cookbook/transformers/decision/decision.py train \ + --model /tmp/decision-tiny --resume /tmp/decision-train/final \ + --data /tmp/decision-tiny/mini.jsonl --output /tmp/decision-resume --steps 6 --dtype float32 +``` + +`--adapter` initializes weights; `--resume` restores training state. They are mutually exclusive. +`steps` is the total target number of updates. The recipe records consumed examples with the +checkpoint; preserve data ordering and augmentation settings when resuming. + +| Released family | Arguments to add to the same train/infer command | +| --- | --- | +| JEV | `--family jev --model --head-dir --adapter /adapter_vllm` | +| Clef | `--family clef --model --head-dir ` | +| OmniJev | `--family omnijev --model --head-dir --adapter ` | +| StartLux | `--family startlux --model ` | + +For JEV/OmniJev, keep the released adapter: omitting it discards the released decision capability. +Use the matching tokenizer, head artifacts and calibration. StartLux also requires the official +`startlux_decision` Python package on PYTHONPATH. The validated source revision is +`StartLuxLabs/StartLux-Decision@0e7a2e81b9c92756e26d8edd843a44d50e362669`. +For StartLux 4B/35B, select language attention targets and exclude the vision tower and vocabulary +head; do not claim training benefit from a smoke test. + +Training accepts hard labels or finite nonnegative distributions summing to one. Targets are never +rendered into model inputs. `--shuffle-fields`, `--drop-instructions` and `--subsample-options` +are explicit training-only augmentation options for generic/JEV/Clef/OmniJev; StartLux rejects these flags. Disable augmentation for evaluation. +Inference writes `predictions.jsonl` and `metrics.json` with per-question `acc` and per-encoded-record +`record_acc`. A JEV/OmniJev source record is expanded before this accounting. + +Open-Jev is an explicit online download, separate from the offline smoke dataset: + +```bash +python cookbook/transformers/decision/fetch_open_jev.py \ + --dataset ZefanCai/Open-Jev --revision \ + --subset release-v2-redistributable --split train --output /tmp/open-jev-train.jsonl +# Use --open-jev-format --data /tmp/open-jev-train.jsonl with train or infer. +``` + +Respect the dataset's license and preserve separate training, validation, calibration and test splits. +Image rows accept local images. OmniJev videos must first be converted to a frame mosaic with +`twinkle.model.decision.omnijev_media.video_state`; this needs ffmpeg/ffprobe. Passing raw video paths +into the training encoder raises an error. + +## 5. Acceptance before publication + +| Gate | Required evidence | +| --- | --- | +| Input | Exact prompt IDs, option order, target exclusion, image grids and branch spans | +| Numeric | Head logits/probabilities and loss/gradient comparisons against each released reference | +| Training | Finite updates, adapter/head changes, frozen-backbone audit, checkpoint recovery | +| Quality | Independent held-out scoring; report each suite and its denominator | +| Performance | Same device, precision, input and warmup; keep performance artifacts local | + +The scoring head is saved through PEFT `modules_to_save`; vocabulary embeddings are weak references +and are not duplicated into head state. FP32 activation must happen after adapter creation and before +the optimizer. Source attribution is retained in reused numerical modules. + +## 6. Summary + +Use the common Twinkle training API while retaining each model's decision format and head semantics. +Passing the included contract tests is an engineering milestone; it does not replace full released +model accuracy, multimodal, distributed or performance acceptance. diff --git a/cookbook/transformers/decision/decision.py b/cookbook/transformers/decision/decision.py new file mode 100644 index 000000000..2736f708a --- /dev/null +++ b/cookbook/transformers/decision/decision.py @@ -0,0 +1,152 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Train and evaluate typed decisions through Twinkle's Transformers backend.""" +import argparse +import json +import math +from pathlib import Path + +import torch +from transformers import AutoProcessor, AutoTokenizer + +import twinkle +from twinkle.loss.decision import DecisionLoss +from twinkle.metric.decision import DecisionAccuracy +from twinkle.model import TransformersModel +from twinkle.model.decision.encoding import DecisionEncoder +from twinkle.model.decision.model import activate_decision_head, attach_decision_head, decision_lora_config, initialize_missing_adapter_head +from twinkle.preprocessor.decision import (ClefPreprocessor, DecisionAugmentConfig, JevPreprocessor, + OmniJevPreprocessor, OpenJevPreprocessor, StartLuxPreprocessor) +from twinkle.processor.decision import DecisionInputProcessor + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('action', choices=['train', 'infer']) + parser.add_argument('--model', required=True, help='Local base or merged checkpoint directory') + parser.add_argument('--family', choices=['generic', 'jev', 'clef', 'omnijev', 'startlux'], default='generic') + parser.add_argument('--head-dir', help='Original decision head artifacts, required for released models') + parser.add_argument('--adapter', help='Released or trained PEFT adapter directory') + parser.add_argument('--resume', help='Resume optimizer and model state from a training checkpoint') + parser.add_argument('--data', required=True, help='Local JSONL; use disjoint files for training and testing') + parser.add_argument('--open-jev-format', action='store_true') + parser.add_argument('--output', required=True) + parser.add_argument('--loss', choices=['scoring', 'distill_rps', 'smooth_brier', 'ce_rps']) + parser.add_argument('--steps', type=int, default=10) + parser.add_argument('--lr', type=float, default=1e-4) + parser.add_argument('--max-length', type=int, default=8192) + parser.add_argument('--lora-targets', nargs='+', default=['q_proj', 'v_proj']) + parser.add_argument('--dtype', choices=['float32', 'bfloat16'], default='bfloat16') + parser.add_argument('--shuffle-fields', action='store_true') + parser.add_argument('--drop-instructions', type=float, default=0) + parser.add_argument('--subsample-options', type=int, default=0) + args = parser.parse_args() + if args.resume and (args.action != 'train' or args.adapter): + raise ValueError('Use --resume for training recovery; --adapter is a separate weights-only initialization') + if args.family == 'startlux' and (args.shuffle_fields or args.drop_instructions or args.subsample_options): + raise ValueError('StartLux augmentation is not supported; preserve the official renderer') + output = Path(args.output) + if output.exists(): + raise FileExistsError('Choose a new output directory') + if args.steps < 1: + raise ValueError('steps must be positive') + twinkle.initialize(mode='local') + if args.family in ('omnijev', 'jev'): + processor = AutoProcessor.from_pretrained(args.model) + else: + processor = AutoTokenizer.from_pretrained(args.model) + tokenizer = getattr(processor, 'tokenizer', processor) + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + if tokenizer.pad_token_id is None: + raise ValueError('Tokenizer needs a pad or EOS token') + wrapper = TransformersModel(model_id=args.model, dtype=getattr(torch, args.dtype), + mixed_precision='no' if args.dtype == 'float32' else 'bf16', + attn_implementation='sdpa') + attach_decision_head(wrapper.model, processor, args.family, args.head_dir or args.model) + if args.adapter: + from peft import PeftConfig, set_peft_model_state_dict + from peft.utils.save_and_load import load_peft_weights + adapter = PeftConfig.from_pretrained(args.adapter) + adapter.inference_mode = args.action != 'train' + factory_head = 'scoring_head' not in (adapter.modules_to_save or []) + adapter.modules_to_save = list(dict.fromkeys(list(adapter.modules_to_save or []) + ['scoring_head'])) + wrapper.add_adapter_to_model('decision', adapter) + state = load_peft_weights(args.adapter, device='cpu') + if factory_head: + state = initialize_missing_adapter_head(wrapper.model, state, 'decision') + restored = set_peft_model_state_dict(wrapper.model, state, adapter_name='decision') + missing_lora = [key for key in restored.missing_keys if 'lora_' in key] + if missing_lora: + raise ValueError(f'Missing adapter weights: {missing_lora}') + if restored.unexpected_keys: + raise ValueError(f'Unexpected adapter keys: {restored.unexpected_keys}') + else: + wrapper.add_adapter_to_model('decision', decision_lora_config( + r=8, lora_alpha=16, target_modules=args.lora_targets)) + activate_decision_head(wrapper.model) + wrapper._default_tokenizer = processor + wrapper.set_processor(DecisionInputProcessor, pad_token_id=tokenizer.pad_token_id, + padding_side='right' if args.action == 'train' else 'left') + mode = args.loss or {'generic': 'scoring', 'jev': 'distill_rps', 'clef': 'smooth_brier', + 'omnijev': 'ce_rps', 'startlux': 'scoring'}[args.family] + wrapper.set_loss(DecisionLoss, mode=mode) + if args.action == 'train': + wrapper.set_optimizer('AdamW', lr=args.lr, foreach=False) + start_step = 0 + if args.resume: + progress = wrapper.resume_from_checkpoint(args.resume, adapter_name='decision') + start_step = int(progress['consumed_train_samples']) + if start_step >= args.steps: + raise ValueError('steps must exceed the already completed step count') + # Token-level accuracy is not defined for a distribution-valued label. + group = wrapper.optimizer_group['decision'] + for state in (group.train_status, group.eval_status): + state.metrics = [DecisionAccuracy(wrapper.device_mesh, group._dp_group)] + config = None + if args.action == 'train': + config = DecisionAugmentConfig(drop_instructions=args.drop_instructions, + subsample_options=args.subsample_options, shuffle_fields=args.shuffle_fields) + pre_cls = {'generic': JevPreprocessor, 'jev': JevPreprocessor, 'clef': ClefPreprocessor, + 'omnijev': OmniJevPreprocessor, 'startlux': StartLuxPreprocessor}[args.family] + source = [json.loads(line) for line in Path(args.data).read_text().splitlines() if line.strip()] + preprocessor = pre_cls(augment_config=config) + if args.open_jev_format: + if args.family == 'startlux': + raise ValueError('Open-Jev normalization supports generic/JEV/Clef/OmniJev schemas') + preprocessor = OpenJevPreprocessor(pre_cls(), augment_config=config) + normalized, _ = preprocessor(source) + encoder = DecisionEncoder(processor, args.family, args.max_length) + rows = [encoder(row) for row in normalized] + if not rows: + raise ValueError('Dataset has no decision records') + output.mkdir(parents=True) + (output / 'run.json').write_text(json.dumps(vars(args), indent=2)) + if args.action == 'train': + if any('target_probs' not in row for row in rows): + raise ValueError('All training records require targets') + for step in range(start_step, args.steps): + result = wrapper.forward_backward(inputs=[rows[step % len(rows)]]) + if not math.isfinite(float(result['loss'])): + raise RuntimeError('Non-finite decision loss') + norm = wrapper.clip_grad_norm(max_grad_norm=1.) + if norm is not None and not math.isfinite(float(norm)): + raise RuntimeError('Non-finite decision gradient') + wrapper.step(); wrapper.zero_grad() + print(json.dumps({'step': step + 1, 'loss': float(result['loss'])}), flush=True) + wrapper.save(name='final', output_dir=str(output), adapter_name='decision', save_optimizer=True, + consumed_train_samples=args.steps) + else: + metric = DecisionAccuracy() + with (output / 'predictions.jsonl').open('w') as stream: + for row in rows: + result = wrapper.forward_only(inputs=[row], return_logits=True) + status = group.eval_status + metric.accumulate(status.inputs, status.outputs) + scores = result['logits'].masked_fill(~result['option_mask'], -torch.inf) + probs = scores.float().softmax(-1).cpu().tolist() + stream.write(json.dumps({'probabilities': probs}) + '\n') + (output / 'metrics.json').write_text(json.dumps(metric.calculate(), indent=2)) + + +if __name__ == '__main__': + main() diff --git a/cookbook/transformers/decision/fetch_open_jev.py b/cookbook/transformers/decision/fetch_open_jev.py new file mode 100644 index 000000000..ae7b38db5 --- /dev/null +++ b/cookbook/transformers/decision/fetch_open_jev.py @@ -0,0 +1,30 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Download an explicitly selected Open-Jev split as local JSONL.""" +import argparse +import json +from pathlib import Path + + +def main(): + from modelscope.msdatasets import MsDataset + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--dataset', choices=['ZefanCai/Open-Jev', 'ZefanCai/Open-Jev-v1.1'], + default='ZefanCai/Open-Jev') + parser.add_argument('--revision', required=True) + parser.add_argument('--subset', required=True) + parser.add_argument('--split', required=True) + parser.add_argument('--output', required=True) + args = parser.parse_args() + output = Path(args.output) + if output.exists(): + raise FileExistsError(output) + dataset = MsDataset.load(args.dataset, subset_name=args.subset, split=args.split, + version=args.revision).to_hf_dataset() + with output.open('x') as stream: + for row in dataset: + stream.write(json.dumps(row, ensure_ascii=False) + '\n') + print(json.dumps({'rows': len(dataset), **vars(args)})) + + +if __name__ == '__main__': + main() diff --git a/cookbook/transformers/decision/make_tiny.py b/cookbook/transformers/decision/make_tiny.py new file mode 100644 index 000000000..b50da0c4c --- /dev/null +++ b/cookbook/transformers/decision/make_tiny.py @@ -0,0 +1,37 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Create a random tiny checkpoint and synthetic decisions, with no downloads.""" +import argparse +import json +from pathlib import Path + +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--output', required=True) + args = parser.parse_args() + output = Path(args.output) + output.mkdir(parents=True, exist_ok=False) + words = ['[PAD]', '[UNK]', '[EOS]', 'false', 'true'] + list('012345') + list('ABCDEFGHIJKLMNOP') + tokenizer = Tokenizer(models.WordLevel({s: i for i, s in enumerate(words)}, unk_token='[UNK]')) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + fast = PreTrainedTokenizerFast(tokenizer_object=tokenizer, pad_token='[PAD]', unk_token='[UNK]', eos_token='[EOS]') + fast.save_pretrained(output) + config = LlamaConfig(vocab_size=len(words), hidden_size=32, intermediate_size=64, + num_hidden_layers=1, num_attention_heads=4, num_key_value_heads=2, + pad_token_id=0, eos_token_id=2) + LlamaForCausalLM(config).save_pretrained(output) + records = [ + {'state': 'one plus one equals two', 'questions': [{'kind': 'noul', 'question': 'Correct?', 'gold': 1}]}, + {'state': 'cold water', 'questions': [{'kind': 'choice', 'question': 'Temperature?', + 'options': ['cold', 'hot'], 'gold': 0}]}, + {'state': 'example rating', 'questions': [{'kind': 'score', 'question': 'Rate it', + 'gold': [.1, .1, .1, .2, .4, .1]}]}, + ] + (output / 'mini.jsonl').write_text(''.join(json.dumps(row) + '\n' for row in records)) + + +if __name__ == '__main__': + main() diff --git a/src/twinkle/loss/__init__.py b/src/twinkle/loss/__init__.py index 926dcb8a3..863e69cb8 100644 --- a/src/twinkle/loss/__init__.py +++ b/src/twinkle/loss/__init__.py @@ -2,6 +2,7 @@ from .base import Loss from .chunked_cross_entropy import ChunkedCrossEntropyLoss from .cross_entropy import CrossEntropyLoss +from .decision import DecisionLoss from .dpo import CPOLoss, DPOLoss, ORPOLoss, SimPOLoss from .gkd import GKDLoss from .grpo import BNPOLoss, CISPOLoss, DRGRPOLoss, GRPOLoss, GSPOLoss, PPOLoss, SAPOLoss @@ -13,6 +14,7 @@ torch_loss_mapping = { 'mse': MSELoss, + 'decision': DecisionLoss, 'chunked_cross_entropy': ChunkedCrossEntropyLoss, 'cross_entropy': CrossEntropyLoss, 'liger_fused_linear_cross_entropy': LigerFusedLinearCrossEntropyLoss, diff --git a/src/twinkle/loss/decision.py b/src/twinkle/loss/decision.py new file mode 100644 index 000000000..aead7aff6 --- /dev/null +++ b/src/twinkle/loss/decision.py @@ -0,0 +1,70 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import math +import torch + +from .base import Loss + + +class DecisionLoss(Loss): + """Question-normalized proper scoring losses for ragged option sets. + + Return a sum and a question count so Twinkle normalizes over all data + ranks and accumulated micro-batches, including uneven question counts. + """ + + require_logits = True + require_logps = False + + def __init__(self, mode='scoring', rps_weight=0.5, brier_weight=1.0, label_smoothing=0.1): + if mode not in ('scoring', 'distill_rps', 'smooth_brier', 'ce_rps'): + raise ValueError(f'Unknown decision loss: {mode}') + if not all(math.isfinite(x) and x >= 0 for x in (rps_weight, brier_weight)): + raise ValueError('Loss weights must be finite and nonnegative') + if not 0 <= label_smoothing < 1: + raise ValueError('label_smoothing must be in [0, 1)') + self.mode = mode + self.rps_weight = rps_weight + self.brier_weight = brier_weight + self.label_smoothing = label_smoothing + + def micro_batch_scale(self, inputs, indices): + return 1.0 + + def __call__(self, inputs, outputs, **kwargs): + raw_logits = outputs['logits'] + # Preserve double precision for numerical/gradient reference tests. + logits = raw_logits if raw_logits.dtype == torch.float64 else raw_logits.float() + mask = outputs['option_mask'].bool() + target = inputs['labels'].to(device=logits.device, dtype=logits.dtype) + if logits.ndim != 2 or logits.shape != mask.shape or logits.shape != target.shape: + raise ValueError('Expected matching [questions, options] logits, mask and targets') + if not logits.shape[0] or not bool(mask.any(-1).all()): + raise ValueError('Every decision must have at least one legal option') + if not bool(torch.isfinite(logits[mask]).all()): + raise ValueError('Non-finite legal option logits') + if not bool(torch.isfinite(target).all()) or bool((target < 0).any()): + raise ValueError('Targets must be finite and nonnegative') + if bool((target[~mask] != 0).any()) or not torch.allclose( + target.sum(-1), torch.ones_like(target[:, 0]), atol=1e-6, rtol=0): + raise ValueError('Target distributions must sum to one over legal options') + n_options = mask.sum(-1) + log_probs = torch.log_softmax(logits.masked_fill(~mask, -torch.inf), -1).masked_fill(~mask, 0) + probs = log_probs.exp().masked_fill(~mask, 0) + smooth = target + if self.mode == 'smooth_brier': + smooth = (1 - self.label_smoothing) * target + self.label_smoothing * mask.to( + logits.dtype) / n_options[:, None] + per_question = -(smooth * log_probs).sum(-1) + if self.mode == 'smooth_brier': + per_question = per_question + self.brier_weight * ((probs - target).square() * mask).sum(-1) / n_options + if self.mode in ('distill_rps', 'ce_rps'): + kinds = outputs.get('kinds') + if kinds is None or kinds.shape != (logits.shape[0], ): + raise ValueError('RPS requires one kind per question') + # Include only CDF boundaries preceding the final real option. + boundaries = torch.arange(logits.shape[1], device=logits.device)[None] < n_options[:, None] - 1 + rps = ((probs.cumsum(-1) - target.cumsum(-1)).square() * boundaries).sum(-1) + if self.mode == 'distill_rps': + rps = rps / (n_options - 1).clamp_min(1) + per_question = per_question + self.rps_weight * rps * (kinds == 2) + return {'loss': per_question.sum(), 'num_tokens': logits.new_tensor(logits.shape[0])} diff --git a/src/twinkle/metric/decision.py b/src/twinkle/metric/decision.py new file mode 100644 index 000000000..8ebe50ef9 --- /dev/null +++ b/src/twinkle/metric/decision.py @@ -0,0 +1,46 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .base import Metric + + +class DecisionAccuracy(Metric): + """Count questions and jointly correct records, rather than language tokens.""" + + def __init__(self, device_mesh=None, process_group=None, **kwargs): + super().__init__(device_mesh, process_group, **kwargs) + self.reset() + + def accumulate(self, inputs, outputs, **kwargs): + logits, target = outputs.get('logits'), inputs.get('labels') + if logits is None or target is None: + return + mask = outputs['option_mask'] + correct = logits.masked_fill(~mask, -float('inf')).argmax(-1) == target.argmax(-1) + counts = inputs['decision_meta']['counts'] + if sum(counts) != correct.numel(): + raise ValueError('Question counts do not cover decision predictions') + self.correct += int(correct.sum()) + self.questions += correct.numel() + offset = 0 + for count in counts: + self.record_correct += int(correct[offset:offset + count].all()) + self.records += 1 + offset += count + + def reset(self): + self.correct = self.questions = self.record_correct = self.records = 0 + + def calculate(self): + rows = self.gather_results([{ + 'correct': self.correct, + 'questions': self.questions, + 'record_correct': self.record_correct, + 'records': self.records + }]) + total = {key: sum(row[key] for row in rows) for key in rows[0]} + self.reset() + if not total['questions']: + return {} + return { + **total, 'acc': total['correct'] / total['questions'], + 'record_acc': total['record_correct'] / total['records'] + } diff --git a/src/twinkle/model/decision/__init__.py b/src/twinkle/model/decision/__init__.py new file mode 100644 index 000000000..85b3e739d --- /dev/null +++ b/src/twinkle/model/decision/__init__.py @@ -0,0 +1 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. diff --git a/src/twinkle/model/decision/encoding.py b/src/twinkle/model/decision/encoding.py new file mode 100644 index 000000000..bba3afdf8 --- /dev/null +++ b/src/twinkle/model/decision/encoding.py @@ -0,0 +1,248 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# Clef span encoding adapted from ms-swift fbc61cf2 (Apache-2.0). +import json +import string +from types import SimpleNamespace +from typing import Any, Dict, List, Tuple + +from .heads import QUESTION_TYPES + +_CLEF_SYSTEM_PROMPT = ('Read the complete state and schema. Decide every field jointly. Each answer ' + 'must be exactly one of that field\'s allowed options.') +_CLEF_PREFIX_TMPL = '<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\nSTATE:\n' +_CLEF_SUFFIX = '\n<|im_end|>\n<|im_start|>assistant\n\n\n\n\nJOINT SCHEMA DECISIONS:' +_CLEF_NOUL_CRITERIA = { + 'true': 'The proposition is true or the answer is yes.', + 'false': 'The proposition is false or the answer is no.', +} + + +class DecisionEncoder: + """Encode normalized decision rows without inserting answer labels into prompts.""" + + def __init__(self, processor, family='generic', max_length=8192): + if family not in ('generic', 'jev', 'clef', 'omnijev', 'startlux'): + raise ValueError(f'Unknown decision family {family}') + self.processor = processor + self.tokenizer = getattr(processor, 'tokenizer', processor) + self.family = family + self.max_length = max_length + + def _tokenize(self, text): + return self.tokenizer.encode(text, add_special_tokens=False) + + @staticmethod + def _clef_fields(inputs): + row = inputs.extra_kwargs + state = row.get('state', row['messages'][0]['content']) + if isinstance(state, str): + try: + state = json.loads(state) + except ValueError: + pass + criteria = [json.loads(c) if isinstance(c, str) else c for c in row['criteria']] + fields = list(zip(row['question_ids'], row['kinds'], row['questions'], criteria)) + if not fields or any(len(row[k]) != len(fields) for k in ('question_ids', 'kinds', 'questions', 'criteria')): + raise ValueError('Inconsistent Clef schema fields') + return state, fields + + def __call__(self, row): + if row.get('videos'): + raise ValueError('Precompose videos with video_state before encoding') + inputs = SimpleNamespace(extra_kwargs=row, images=row.get('images'), videos=row.get('videos')) + if self.family == 'startlux': + from startlux_decision import jevfmt + data = json.loads(row['startlux_row']) + order = [o['id'] for o in data['options']] + ids, rendered = jevfmt.render_ids(data, self.tokenizer, order, max_length=self.max_length or 8192) + if rendered != order: + raise ValueError('Renderer changed option order') + encoded = { + 'input_ids': ids, + 'decision_meta': { + 'n_options': [len(order)], + 'kinds': [QUESTION_TYPES[data['type']]] + } + } + elif self.family == 'clef': + encoded = self._encode_clef(inputs) + elif self.family == 'omnijev': + encoded = self._encode_omnijev(row) + else: + if len(row['questions']) != 1: + raise ValueError('JEV/generic requires one question per encoded row') + kind, options = row['kinds'][0], row['options'][0] + if kind == 'choice' and len(options) > 16: + raise ValueError('This verbalizer supports at most 16 choice options') + lines = [f'{string.ascii_uppercase[i]}) {o}' + for i, o in enumerate(options)] if kind == 'choice' else options + prompt = (f'[kind] {kind}\n[state] {row["messages"][0]["content"]}\n' + f'[question] {row["questions"][0]}\n[options]\n' + '\n'.join(lines) + '\n[decision]:') + if row.get('images'): + if self.family != 'jev' or not hasattr(self.processor, 'image_processor'): + raise ValueError('Image decisions require a JEV vision processor') + from PIL import Image + frames = [] + for source in row['images']: + if isinstance(source, Image.Image): + frames.append(source.convert('RGB')) + else: + with Image.open(source) as image: + frames.append(image.convert('RGB')) + marker = '<|vision_start|><|image_pad|><|vision_end|>' + prompt = prompt.replace('[state] ', '[state] ' + marker * len(frames), 1) + encoded = dict(self.processor(text=[prompt], images=frames, return_tensors='pt')) + encoded['input_ids'] = encoded['input_ids'][0].tolist() + encoded.pop('attention_mask', None) + if 'mm_token_type_ids' in encoded: + encoded['mm_token_type_ids'] = encoded['mm_token_type_ids'][0].tolist() + else: + encoded = {'input_ids': self._tokenize(prompt)} + encoded['decision_meta'] = { + 'n_options': [len(options)], + 'kinds': [QUESTION_TYPES[kind]], + 'mm_counts': { + 'image': len(row.get('images') or []), + 'video': 0 + } + } + if self.max_length is not None and len(encoded['input_ids']) > self.max_length: + raise ValueError('Decision prompt exceeds max_length; answer positions must not be truncated') + if 'target_probs' in row: + encoded['target_probs'] = row['target_probs'] + return encoded + + def _encode_omnijev(self, row): + from . import omnijev_media as media + from .omnijev_branch import OPT_CLOSE, OPT_OPEN + instruction = row['messages'][0]['content'] + images = row.get('images') or [] + video = row.get('omnijev_video') + if images: + frames, content = media.video_content(media.compose_state(images, video), video, instruction) + else: + frames, content = [], [{'type': 'text', 'text': instruction}] + messages = [{ + 'role': 'user', + 'content': content + }, { + 'role': 'assistant', + 'content': '\n' + media.option_body(row['options'][0]) + }] + text = self.processor.apply_chat_template(messages, tokenize=False, continue_final_message=True) + if frames: + encoded = dict(self.processor(text=[text], images=frames, return_tensors='pt')) + encoded['input_ids'] = encoded['input_ids'][0].tolist() + encoded.pop('attention_mask', None) + if 'mm_token_type_ids' in encoded: + encoded['mm_token_type_ids'] = encoded['mm_token_type_ids'][0].tolist() + else: + encoded = {'input_ids': self._tokenize(text)} + ids = encoded['input_ids'] + + def positions(token): + sub = self._tokenize(token) + if not sub: + raise ValueError('Missing option marker tokens') + return [i for i in range(len(ids) - len(sub) + 1) if ids[i:i + len(sub)] == sub] + + opens = positions(OPT_OPEN) + closes = [i + len(self._tokenize(OPT_CLOSE)) - 1 for i in positions(OPT_CLOSE)] + if not opens or len(opens) != len(closes) or any(a >= b for a, b in zip(opens, closes)): + raise ValueError('Invalid OmniJev option markers') + markers = list(zip(opens, closes)) + encoded['decision_meta'] = { + 'n_options': [len(row['omnijev_column_keys'][0])], + 'kinds': [QUESTION_TYPES[row['kinds'][0]]], + 'option_markers': [markers], + 'mm_counts': { + 'image': len(frames), + 'video': 0 + }, + 'question_spans': [(opens[0] - 1, opens[0])], + 'option_spans': [markers] + } + return encoded + + @staticmethod + def _clef_render(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False, separators=(',', ':'), sort_keys=True) + + def _encode_clef(self, inputs) -> Dict[str, Any]: + if getattr(inputs, 'images', None) or getattr(inputs, 'videos', None): + raise NotImplementedError( + 'ClefTemplate supports text-only records this phase (Clef is a text decision model); ' + 'images/videos require the deferred multimodal span-offset port.') + from .heads import QUESTION_TYPES + state, questions = self._clef_fields(inputs) + + schema_ids: List[int] = self._tokenize('\n\nSCHEMA FIELDS:\n') + question_spans: List[Tuple[int, int]] = [] + option_spans: List[List[Tuple[int, int]]] = [] + kinds_int: List[int] = [] + n_options: List[int] = [] + for q_index, (qid, kind, instructions, criteria) in enumerate(questions): + schema_ids += self._tokenize(f'\nFIELD {q_index + 1}\nID: {qid}\nTYPE: {kind}\nINSTRUCTION: ') + question_start = len(schema_ids) + if instructions is None or instructions == '': + instructions = str(qid) + schema_ids += self._tokenize(self._clef_render(instructions)) + question_end = len(schema_ids) + schema_ids += self._tokenize('\nALLOWED OPTIONS:\n') + spans: List[Tuple[int, int]] = [] + for o_index, (option_id, description) in enumerate(self._question_options(kind, criteria)): + schema_ids += self._tokenize(f'OPTION {o_index + 1}: ') + option_start = len(schema_ids) + semantics: Dict[str, Any] = {'option_id': option_id} + if description is not None: + semantics['description'] = description + schema_ids += self._tokenize(self._clef_render(semantics)) + spans.append((option_start, len(schema_ids))) + schema_ids += self._tokenize('\n') + schema_ids += self._tokenize('END FIELD\n') + question_spans.append((question_start, question_end)) + option_spans.append(spans) + kinds_int.append(QUESTION_TYPES[kind]) + n_options.append(len(spans)) + + prefix_ids = self._tokenize(_CLEF_PREFIX_TMPL.format(system=_CLEF_SYSTEM_PROMPT)) + suffix_ids = self._tokenize(_CLEF_SUFFIX) + state_ids = self._tokenize(self._clef_render(state)) + if self.max_length is not None: + fixed_length = len(prefix_ids) + len(schema_ids) + len(suffix_ids) + if fixed_length > self.max_length: + raise ValueError( + f'Clef schema requires {fixed_length} tokens before state; max_length={self.max_length}.') + state_ids = state_ids[:self.max_length - fixed_length] + schema_offset = len(prefix_ids) + len(state_ids) + question_spans = [(s + schema_offset, e + schema_offset) for s, e in question_spans] + option_spans = [[(s + schema_offset, e + schema_offset) for s, e in spans] for spans in option_spans] + input_ids = prefix_ids + state_ids + schema_ids + suffix_ids + if not input_ids or not questions: + raise ValueError('Clef record produced no model input or questions.') + return { + 'input_ids': input_ids, + 'decision_meta': { + 'n_options': n_options, + 'kinds': kinds_int, + 'question_spans': question_spans, + 'option_spans': option_spans, + }, + } + + @staticmethod + def _question_options(kind: str, criteria: Any) -> List[Tuple[str, Any]]: + from .heads import clef_option_ids + ids = clef_option_ids(kind, criteria) + if kind == 'noul': + merged = dict(_CLEF_NOUL_CRITERIA) + merged.update(criteria or {}) + return [(key, merged[key]) for key in ids] + if kind == 'choice': + crit = {str(k): v for k, v in (criteria or {}).items()} + return [(key, crit[key]) for key in ids] + # score + crit = criteria or [] + return [(ids[i], crit[i]) for i in range(len(crit))] diff --git a/src/twinkle/model/decision/heads.py b/src/twinkle/model/decision/heads.py new file mode 100644 index 000000000..a745980ca --- /dev/null +++ b/src/twinkle/model/decision/heads.py @@ -0,0 +1,1105 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# Adapted from ms-swift fbc61cf2b714ade30ed29f28ddd63be61f6ebb50 (Apache-2.0). +"""Data contract shared by Template / Head / Loss / Trainer: + +A batch flattens the questions of all records into `total_q` rows. Because each question +scores a different number of options (noul=2, choice=n, score=6), the per-option logits are +ragged and are stored padded: + + model.forward(...) -> ScoringOutput + logits: FloatTensor [total_q, max_opt] per-option logits, padded + option_mask: BoolTensor [total_q, max_opt] True = a real option slot + kinds: LongTensor [total_q] 0=noul, 1=choice, 2=score (per-kind temp) + counts: LongTensor [num_records] questions per record (regroup total_q at decode) + loss: Optional filled when labels are consumed + ordinal_logits: Optional OmniJev CORAL/CORN cumulative logits + + meta (assembled by DecisionInputProcessor, consumed by ScoringHead.forward): + n_options: List[int] per-question option count (len == total_q) + question_spans: List[Tuple[int, int]] per-question token span (Clef/OmniJev pool) + option_spans: List[List[Tuple[int,int]]] per-question per-option token span + verbalizer_ids: List[List[int]] per-option lm_head column to read (JEV) + bias / temperatures: calibration buffers (JEV/OmniJev), kept on the head + + labels (training), consumed by DecisionLoss and DecisionAccuracy: + target_probs: FloatTensor [total_q, max_opt] padded target distribution + (one-hot for Clef/OmniJev CE, soft for JEV distillation; pad cols = 0) +""" +import json +import math +import os +import string +import torch +import torch.nn as nn +import torch.nn.functional as F +import weakref +from dataclasses import dataclass +from pathlib import Path +from transformers.modeling_outputs import ModelOutput +from typing import Any, Dict, List, Optional, Sequence, Tuple + +# QUESTION_TYPES maps the JEV/Clef encoding: noul -> 0, choice -> 1, score -> 2. +QUESTION_TYPES = {'noul': 0, 'choice': 1, 'score': 2} +KIND_NAMES = ('noul', 'choice', 'score') + + +@dataclass +class ScoringOutput(ModelOutput): + logits: Optional[torch.FloatTensor] = None + option_mask: Optional[torch.BoolTensor] = None + kinds: Optional[torch.LongTensor] = None + counts: Optional[torch.LongTensor] = None + loss: Optional[torch.FloatTensor] = None + ordinal_logits: Optional[torch.FloatTensor] = None + + +def pad_option_logits(per_question: List[torch.Tensor], + pad_value: float = float('-inf'), + width: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]: + total_q = len(per_question) + if width is None: + width = max((int(t.shape[0]) for t in per_question), default=0) + if total_q == 0 or width == 0: + empty = torch.zeros((total_q, width)) + return empty, torch.zeros((total_q, width), dtype=torch.bool) + if any(t.ndim != 1 or t.numel() == 0 or t.numel() > width for t in per_question): + raise ValueError('Invalid option row or padded width would truncate candidates') + device = per_question[0].device + dtype = per_question[0].dtype + logits = torch.full((total_q, width), pad_value, device=device, dtype=dtype) + mask = torch.zeros((total_q, width), device=device, dtype=torch.bool) + for i, t in enumerate(per_question): + n = min(int(t.shape[0]), width) + if n == 0: + continue + logits[i, :n] = t[:n].to(dtype) + mask[i, :n] = True + return logits, mask + + +def masked_softmax(logits: torch.Tensor, option_mask: torch.Tensor, dim: int = -1) -> torch.Tensor: + logits = logits.float().masked_fill(~option_mask, float('-inf')) + probs = torch.softmax(logits, dim=dim) + return torch.nan_to_num(probs, nan=0.0) + + +def masked_argmax(logits: torch.Tensor, option_mask: torch.Tensor, dim: int = -1) -> torch.Tensor: + """Argmax restricted to real option slots (pad slots never win).""" + return logits.float().masked_fill(~option_mask, float('-inf')).argmax(dim=dim) + + +def get_scoring_head(model: nn.Module) -> Optional[nn.Module]: + for name, module in model.named_modules(): + if name == 'scoring_head' or name.endswith('.scoring_head'): + return module + return None + + +def maybe_load_trained_scoring_head(model: nn.Module, adapter_dir: Optional[str]) -> bool: + """After an adapter/checkpoint is applied, overlay a TRAINED head if that dir carries one. + + Two scenarios: + - Adapter checkpoint (LoRA): ``scoring_head.safetensors`` exists → load weights from it. + - Full-finetune checkpoint: ``scoring_head.safetensors`` is absent, ``scoring_head.json`` + present → the head's parameters are inside the model safetensors shards (ignored by + ``from_pretrained`` because ``scoring_head`` was not yet attached). Extract them here. + """ + if not adapter_dir: + return False + head = get_scoring_head(model) + if head is None: + return False + loaded = head.load_pretrained(adapter_dir) + if loaded: + from twinkle import get_logger + get_logger().info(f'decision: loaded trained scoring_head from {adapter_dir}.') + return loaded + + +def clef_option_ids(kind: str, criteria: Any) -> List[str]: + if kind == 'noul': + return ['true', 'false'] + if kind == 'choice': + return [str(key) for key, _ in sorted((str(k), v) for k, v in (criteria or {}).items())] + return [str(index) for index in range(len(criteria or []))] + + +def omnijev_column_keys(kind: str, keys: Sequence[Any]) -> List[str]: + if kind == 'noul': + return ['yes', 'no'] + if kind == 'choice': + return [str(k) for k in keys] + ['abstain'] + return [str(k) for k in keys] + + +class ScoringHead(nn.Module): + """Abstract base for the three decision heads (JEV / Clef / OmniJev). + + A head turns base-model representations into per-option logits over each question's own + option set. `forward` receives BOTH the last hidden states and the lm_head logits so every + head can pick what it needs: + - JEV reads lm_logits at the verbalizer columns of the final token (no nn.Module params + beyond bias/temperature buffers); + - Clef pools hidden_states over question/option spans and runs the JointSchema decoder; + - OmniJev dots option/question MLPs over pooled hidden states and adds an ordinal head. + + Subclasses MUST return a `ScoringOutput` whose `logits`/`option_mask` follow the contract + above, and MUST load their factory head weights with `strict=True`. + """ + + # Set by subclasses that need the base model's output embeddings (Clef lexical prior). + needs_output_embeddings: bool = False + # Set by subclasses that read lm_head logits (JEV verbalizer columns). When False the loader + # passes no lm_logits to the head. Backbone projection behavior depends on the model. + needs_lm_logits: bool = False + keep_fp32: bool = False + + def set_output_embeddings(self, output_embeddings: nn.Module) -> None: + self._embedding_ref = weakref.ref(output_embeddings) + + @property + def _output_embeddings(self): + ref = getattr(self, '_embedding_ref', None) + return ref() if ref is not None else None + + @staticmethod + def _locate(model_dir: str, filename: str, env_var: str, subdir: Optional[str] = None) -> Optional[str]: + """Locate a head artifact file: check an env-var override first, then model_dir/filename + (and optionally model_dir/subdir/filename). Returns the path if found, else None. + Shared by all three head subclasses (JEV passes ``subdir='adapter_vllm'``).""" + env = os.environ.get(env_var) + if env and os.path.exists(env): + return env + candidates = [os.path.join(model_dir, filename)] + if subdir: + candidates.append(os.path.join(model_dir, subdir, filename)) + for candidate in candidates: + if os.path.exists(candidate): + return candidate + return None + + def forward(self, hidden_states: torch.Tensor, lm_logits: Optional[torch.Tensor], meta: Dict[str, + Any]) -> ScoringOutput: + """Compute per-option logits. + + Args: + hidden_states: `[batch, seq_len, hidden]` last-hidden-state of the base model. + lm_logits: `[batch, seq_len, vocab]` lm_head logits (JEV verbalizer columns); + may be None for heads that only use hidden_states. + meta: the per-question locating/calibration dict described in the module contract + (n_options / spans / verbalizer_ids / kinds ...), already moved to device. + + Returns: + ScoringOutput with padded `logits` `[total_q, max_opt]` + `option_mask`. + """ + raise NotImplementedError + + def extra_state_for_save(self) -> Dict[str, Any]: + """Non-parameter meta (verbalizer_ids / bias / temperatures) to persist alongside the + head weights, so a reloaded checkpoint can rebuild the head.""" + return {} + + @classmethod + def from_pretrained(cls, model_dir: str, model=None, config=None) -> 'ScoringHead': + """Construct the head and STRICT-load its factory weights from `model_dir`. + + Each model subclass reads its own artifact (JEV: decision_head.json + calibration.json; + Clef: joint_head.safetensors; OmniJev: head.pt + ord.pt) and asserts a 0-missing / + 0-unexpected strict load. `model`/`config` are passed for heads that size themselves from + the base + """ + raise NotImplementedError + + HEAD_WEIGHTS_NAME = 'scoring_head.safetensors' + HEAD_META_NAME = 'scoring_head.json' + + def head_state_dict(self) -> Dict[str, torch.Tensor]: + """The head's OWN state (params + persistent buffers), EXCLUDING the injected live + `_output_embeddings`. `set_output_embeddings` assigns an nn.Module, which `nn.Module.__setattr__` + auto-registers as a child, so a naive `state_dict()` would drag in the ENTIRE base lm_head + (vocab x hidden, possibly tied to embed_tokens). That reference is re-wired by the loader on + rebuild, so it is never persisted here. + """ + return {k: v.detach().cpu() for k, v in self.state_dict().items() if not k.startswith('_output_embeddings')} + + def save_pretrained(self, save_directory: str, save_weights: bool = True) -> None: + """Persist the head beside the adapter: weights (`head_state_dict`) + the non-parameter + calibration meta (`extra_state_for_save`), so a reloaded checkpoint rebuilds the TRAINED head + without the factory `decision_head.json` / `joint_head.safetensors` (plan decision F). + + When ``save_weights=False`` (full-finetune checkpoint), the head's parameters are already + inside the model ``state_dict`` saved by ``super()._save``; only the meta json is written so + a reloaded checkpoint still carries a human-readable head manifest. + """ + from safetensors.torch import save_file + os.makedirs(save_directory, exist_ok=True) + if save_weights: + save_file(self.head_state_dict(), os.path.join(save_directory, self.HEAD_WEIGHTS_NAME)) + with open(os.path.join(save_directory, self.HEAD_META_NAME), 'w', encoding='utf-8') as f: + json.dump(self.extra_state_for_save(), f, indent=2, ensure_ascii=False) + + def load_pretrained(self, load_directory: str) -> bool: + """Two scenarios: + - Adapter checkpoint (LoRA): ``HEAD_WEIGHTS_NAME`` exists → load weights from it (the head + is NOT inside the LoRA adapter). + - Full-finetune checkpoint: ``HEAD_WEIGHTS_NAME`` is absent → the head's parameters are + already restored from the model ``state_dict`` by ``from_pretrained``; only the meta + json is present, and we return True so the caller knows the head was recognised. + """ + from safetensors.torch import load_file + weights_file = os.path.join(load_directory, self.HEAD_WEIGHTS_NAME) + meta_file = os.path.join(load_directory, self.HEAD_META_NAME) + has_weights = os.path.exists(weights_file) + has_meta = os.path.exists(meta_file) + if not has_weights and not has_meta: + return False + if has_weights: + state = load_file(weights_file) + self._validate_and_load(state, weights_file) + elif has_meta: + state = self._extract_head_state_from_model_dir(load_directory) + if state is not None: + self._validate_and_load(state, 'model safetensors') + return True + + def _validate_and_load(self, state: Dict[str, torch.Tensor], source: str) -> None: + own_keys = {k for k in self.state_dict() if not k.startswith('_output_embeddings')} + own_params = {n for n, _ in self.named_parameters() if not n.startswith('_output_embeddings')} + unexpected = set(state) - own_keys + missing_params = own_params - set(state) + if unexpected or missing_params: + raise ValueError(f'{source} does not match head {type(self).__name__}: ' + f'unexpected={sorted(unexpected)} missing_params={sorted(missing_params)}') + self.load_state_dict(state, strict=False) + + def _extract_head_state_from_model_dir(self, model_dir: str) -> Optional[Dict[str, torch.Tensor]]: + """Used when loading a full-finetune checkpoint: ``from_pretrained`` ignored scoring_head keys + (the submodule was not yet attached); this reads them back from the safetensors files. + Returns ``None`` when the model dir carries no scoring_head weights (e.g. a factory model). + """ + from safetensors.torch import load_file + index_path = os.path.join(model_dir, 'model.safetensors.index.json') + if os.path.exists(index_path): + with open(index_path) as f: + index = json.load(f) + weight_map = index.get('weight_map', {}) + head_keys = {k for k in weight_map if 'scoring_head' in k} + if not head_keys: + return None + shards = {weight_map[k] for k in head_keys} + result: Dict[str, torch.Tensor] = {} + for shard in shards: + data = load_file(os.path.join(model_dir, shard)) + for k, v in data.items(): + if 'scoring_head' in k: + result[k.split('scoring_head.', 1)[1] if 'scoring_head.' in k else k] = v + return result if result else None + single = os.path.join(model_dir, 'model.safetensors') + if os.path.exists(single): + data = load_file(single) + result = {} + for k, v in data.items(): + if 'scoring_head' in k: + result[k.split('scoring_head.', 1)[1] if 'scoring_head.' in k else k] = v + return result if result else None + return None + + +class JevVerbalizerHead(ScoringHead): + + needs_lm_logits = True + keep_fp32 = True + + DEFAULT_RANGES = {'noul': (0, 2), 'choice': (8, 24), 'score': (2, 8)} + DEFAULT_TEMPERATURES = {'noul': 1.0, 'choice': 1.0, 'score': 1.0} + + def __init__(self, bias: Sequence[float], verbalizer_ids: Sequence[int], ranges: Dict[str, Sequence[int]], + temperatures: Dict[str, float]): + super().__init__() + self.bias = nn.Parameter(torch.as_tensor(list(bias), dtype=torch.float32)) + self.register_buffer( + 'verbalizer_ids', torch.as_tensor(list(verbalizer_ids), dtype=torch.long), persistent=False) + self.register_buffer( + 'ranges', + torch.as_tensor([list(ranges[KIND_NAMES[k]]) for k in range(len(KIND_NAMES))], dtype=torch.long), + persistent=False) + self.register_buffer( + 'temperatures', + torch.as_tensor([float(temperatures[KIND_NAMES[k]]) for k in range(len(KIND_NAMES))], dtype=torch.float32), + persistent=False) + + def forward(self, hidden_states: torch.Tensor, lm_logits: Optional[torch.Tensor], meta: Dict[str, + Any]) -> ScoringOutput: + if lm_logits is None: + raise ValueError('JevVerbalizerHead requires lm_logits (needs_lm_logits must be True).') + kinds = [int(k) for k in meta['kinds']] + n_options = [int(n) for n in meta['n_options']] + record_index = [int(r) for r in meta['record_index']] + total_q = len(n_options) + device = lm_logits.device + counts = meta.get('counts') + counts_t = torch.as_tensor([int(c) for c in counts], dtype=torch.long) if counts else None + if total_q == 0: + max_opt = int(meta.get('max_opt') or 0) + empty = torch.zeros((0, max_opt), device=device, dtype=torch.float32) + return ScoringOutput( + logits=empty, + option_mask=torch.zeros((0, max_opt), dtype=torch.bool, device=device), + kinds=torch.zeros((0, ), dtype=torch.long, device=device), + counts=counts_t) + + rows, cols = self._readout_positions(record_index, meta, lm_logits.shape[1], device) + read_logits = lm_logits[rows, cols, :].float() # [total_q, vocab] + log_probs = torch.log_softmax(read_logits, dim=-1) # log P(token), matching vLLM's logprob + + verbalizer_ids = self.verbalizer_ids.to(device) + bias = self.bias.to(device=device, dtype=torch.float32) + ranges = self.ranges.to(device) + temperatures = self.temperatures.to(device=device, dtype=torch.float32) + per_question: List[torch.Tensor] = [] + for i in range(total_q): + k = kinds[i] + n = n_options[i] + lo = int(ranges[k][0]) + hi = int(ranges[k][1]) + if lo + n > verbalizer_ids.shape[0] or n > hi - lo: + raise ValueError( + f'JEV kind {KIND_NAMES[k]} supports at most {hi - lo} options (verbalizer slots), got {n}.') + ids = verbalizer_ids[lo:lo + n] # [n] + lp = log_probs[i].gather(0, ids) # [n] + z = (lp + bias[lo:lo + n]) / temperatures[k] + per_question.append(z) + + max_opt = int(meta.get('max_opt') or max(n_options)) + logits, option_mask = pad_option_logits(per_question, width=max_opt) + return ScoringOutput( + logits=logits, + option_mask=option_mask, + kinds=torch.as_tensor(kinds, dtype=torch.long, device=device), + counts=counts_t) + + @staticmethod + def _readout_positions(record_index: List[int], meta: Dict[str, Any], padded_len: int, + device) -> Tuple[torch.Tensor, torch.Tensor]: + left_padding = bool(meta.get('left_padding', False)) + seq_lens = meta.get('seq_lens') + cols = [] + for ri in record_index: + if left_padding: + cols.append(padded_len - 1) + else: + sl = int(seq_lens[ri]) if seq_lens is not None else padded_len + cols.append(sl - 1) + rows = torch.as_tensor(record_index, dtype=torch.long, device=device) + cols_t = torch.as_tensor(cols, dtype=torch.long, device=device) + return rows, cols_t + + def extra_state_for_save(self) -> Dict[str, Any]: + return { + 'bias': self.bias.detach().float().tolist(), + 'verbalizer_ids': self.verbalizer_ids.tolist(), + 'ranges': { + KIND_NAMES[k]: self.ranges[k].tolist() + for k in range(len(KIND_NAMES)) + }, + 'temperatures': { + KIND_NAMES[k]: float(self.temperatures[k]) + for k in range(len(KIND_NAMES)) + }, + } + + @classmethod + def from_pretrained(cls, model_dir: str, model=None, config=None) -> 'JevVerbalizerHead': + head_file = cls._locate(model_dir, 'decision_head.json', 'JEV_DECISION_HEAD', subdir='adapter_vllm') + if head_file is None: + raise FileNotFoundError( + f'JEV decision_head.json not found under {model_dir!r} (or {model_dir}/adapter_vllm). ' + 'Set JEV_DECISION_HEAD to its path.') + with open(head_file) as f: + head = json.load(f) + for key in ('bias', 'verbalizer_ids'): + if key not in head: + raise KeyError(f'decision_head.json missing required key {key!r}: {head_file}') + ranges = {k: tuple(v) for k, v in head.get('slots', {}).get('ranges', {}).items()} + for kind, default in cls.DEFAULT_RANGES.items(): + ranges.setdefault(kind, default) + + temperatures = dict(cls.DEFAULT_TEMPERATURES) + calib_file = os.environ.get('JEV_DECIDE_CALIBRATION') + if not calib_file or not os.path.exists(calib_file): + calib_file = cls._locate(model_dir, 'calibration.json', 'JEV_DECIDE_CALIBRATION') + if calib_file and os.path.exists(calib_file): + with open(calib_file) as f: + per_kind = json.load(f).get('per_kind', {}) + temperatures.update({k: float(v) for k, v in per_kind.items() if k in temperatures}) + return cls(head['bias'], head['verbalizer_ids'], ranges, temperatures) + + @staticmethod + def _resolve_verbalizer_token_id(tokenizer, text: str) -> int: + + ids = tokenizer.encode(text, add_special_tokens=False) + if len(ids) == 1 and ids[0] != getattr(tokenizer, 'unk_token_id', None): + return ids[0] + raise ValueError(f'Verbalizer {text!r} must encode to one token; got {ids}.') + + @classmethod + def from_tokenizer(cls, tokenizer, config=None) -> 'JevVerbalizerHead': + """Random-init path for training a decision model from scratch on any CausalLM. + + Args: + tokenizer: the base model's tokenizer (any `PreTrainedTokenizerBase`). + config: unused (accepted for signature parity with `from_pretrained`). + + Returns: + A `JevVerbalizerHead` with zero-initialized bias and unit temperatures, ready for + joint training with a LoRA adapter on the base CausalLM. + """ + noul_words = ['false', 'true'] + score_words = ['0', '1', '2', '3', '4', '5'] + choice_words = list(string.ascii_uppercase[:16]) # A..P + + verbalizer_ids = [ + cls._resolve_verbalizer_token_id(tokenizer, w) for w in noul_words + score_words + choice_words + ] + if len(set(verbalizer_ids)) != len(verbalizer_ids): + raise ValueError('Decision verbalizer symbols must resolve to distinct token IDs') + bias = [0.0] * len(verbalizer_ids) + ranges = dict(cls.DEFAULT_RANGES) + temperatures = dict(cls.DEFAULT_TEMPERATURES) + return cls(bias, verbalizer_ids, ranges, temperatures) + + +class EvidenceRoutingLayer(nn.Module): + """One cross-attention block of Clef's `JointSchemaHead`.""" + + def __init__(self, width: int, heads: int, feedforward: int, dropout: float = 0.0) -> None: + super().__init__() + self.query_norm = nn.LayerNorm(width) + self.memory_norm = nn.LayerNorm(width) + self.attention = nn.MultiheadAttention(width, heads, dropout=dropout, batch_first=True) + self.attention_dropout = nn.Dropout(dropout) + self.feedforward_norm = nn.LayerNorm(width) + self.feedforward = nn.Sequential( + nn.Linear(width, feedforward), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(feedforward, width), + nn.Dropout(dropout), + ) + + def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: + normalized_queries = self.query_norm(queries) + routed, _ = self.attention( + normalized_queries, self.memory_norm(memory), self.memory_norm(memory), need_weights=False) + queries = queries + self.attention_dropout(routed) + return queries + self.feedforward(self.feedforward_norm(queries)) + + +class ClefJointSchemaHead(ScoringHead): + """Clef's joint schema head: a real parameterised decoder over span-pooled hidden states.""" + needs_output_embeddings = True + needs_lm_logits = False + keep_fp32 = False + + def __init__(self, + hidden_size: int, + width: int, + routing_layers: int, + layers: int, + heads: int, + feedforward: int, + dropout: float = 0.0) -> None: + super().__init__() + self.hidden_norm = nn.LayerNorm(hidden_size) + self.memory_projection = nn.Linear(hidden_size, width, bias=False) + self.question_projection = nn.Linear(hidden_size, width, bias=False) + self.option_question_projection = nn.Linear(hidden_size, width, bias=False) + self.global_projection = nn.Linear(hidden_size, width, bias=False) + self.option_context_projection = nn.Linear(hidden_size, width, bias=False) + self.option_lexical_projection = nn.Linear(hidden_size, width, bias=False) + self.type_embedding = nn.Embedding(3, width) + self.evidence_layers = nn.ModuleList([ + EvidenceRoutingLayer(width=width, heads=heads, feedforward=feedforward, dropout=dropout) + for _ in range(routing_layers) + ]) + self.option_summary_norm = nn.LayerNorm(width) + self.layers = nn.ModuleList([ + nn.TransformerDecoderLayer( + d_model=width, + nhead=heads, + dim_feedforward=feedforward, + dropout=dropout, + activation='gelu', + batch_first=True, + norm_first=True) for _ in range(layers) + ]) + self.field_norm = nn.LayerNorm(width) + self.option_norm = nn.LayerNorm(width) + self.residual_scorer = nn.Sequential( + nn.Linear(width * 4, width), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(width, 1), + ) + self.prior_logit_scale = nn.Parameter(torch.zeros(())) + self.joint_logit_scale = nn.Parameter(torch.zeros(())) + self.residual_gate = nn.Parameter(torch.zeros(())) + + @staticmethod + def _mean_span(values: torch.Tensor, span: Tuple[int, int]) -> torch.Tensor: + start, end = span + return values[start:end].mean(dim=0) + + def forward(self, hidden_states: torch.Tensor, lm_logits: Optional[torch.Tensor], meta: Dict[str, + Any]) -> ScoringOutput: + input_ids = meta.get('input_ids') + if input_ids is None: + raise ValueError('ClefJointSchemaHead requires meta["input_ids"] ' + '(stashed by ClefTemplate._post_encode before the base drops it).') + if getattr(self, '_output_embeddings', None) is None: + raise RuntimeError('ClefJointSchemaHead needs output embeddings; the loader must call ' + 'set_output_embeddings (needs_output_embeddings=True).') + device = hidden_states.device + n_options = [int(n) for n in meta['n_options']] + kinds = [int(k) for k in meta['kinds']] + total_q = len(n_options) + counts = meta.get('counts') + counts_t = torch.as_tensor([int(c) for c in counts], dtype=torch.long, device=device) if counts else None + if total_q == 0: + max_opt = int(meta.get('max_opt') or 0) + empty = torch.zeros((0, max_opt), device=device, dtype=hidden_states.dtype) + return ScoringOutput( + logits=empty, + option_mask=torch.zeros((0, max_opt), dtype=torch.bool, device=device), + kinds=torch.zeros((0, ), dtype=torch.long, device=device), + counts=counts_t) + + question_spans = meta['question_spans'] + option_spans = meta['option_spans'] + record_index = [int(r) for r in meta['record_index']] + num_records = int(hidden_states.shape[0]) + padded_len = int(hidden_states.shape[1]) + left_padding = bool(meta.get('left_padding', False)) + seq_lens = meta.get('seq_lens') or [padded_len] * num_records + output_embedding_weight = self._output_embeddings.weight + + by_record: List[List[int]] = [[] for _ in range(num_records)] + for qi in range(total_q): + by_record[record_index[qi]].append(qi) + + normalized_hidden = self.hidden_norm(hidden_states.to(self.hidden_norm.weight.dtype)) + per_question: List[torch.Tensor] = [] + per_kind: List[int] = [] + for rec in range(num_records): + qidxs = by_record[rec] + if not qidxs: + continue + seq_len = int(seq_lens[rec]) + pad_offset = (padded_len - seq_len) if left_padding else 0 + sequence_hidden = normalized_hidden[rec, pad_offset:pad_offset + seq_len] + record_logits = self._score_record(sequence_hidden, input_ids[rec], pad_offset, qidxs, question_spans, + option_spans, kinds, output_embedding_weight) + per_question.extend(record_logits) + per_kind.extend(kinds[qi] for qi in qidxs) + + max_opt = int(meta.get('max_opt') or max(n_options)) + logits, option_mask = pad_option_logits(per_question, width=max_opt) + return ScoringOutput( + logits=logits, + option_mask=option_mask, + kinds=torch.as_tensor(per_kind, dtype=torch.long, device=device), + counts=counts_t) + + def _score_record(self, sequence_hidden: torch.Tensor, input_ids_row: torch.Tensor, pad_offset: int, + qidxs: List[int], question_spans: List[Any], option_spans: List[Any], kinds: List[int], + output_embedding_weight: torch.Tensor) -> List[torch.Tensor]: + """Per-record scoring body for the joint schema head.""" + memory = self.memory_projection(sequence_hidden).unsqueeze(0) # [1, seq_len, width] + global_vector = sequence_hidden[-1] # [H] + question_vectors = torch.stack([self._mean_span(sequence_hidden, question_spans[qi]) for qi in qidxs]) + type_ids = torch.as_tensor([kinds[qi] for qi in qidxs], device=sequence_hidden.device, dtype=torch.long) + + option_contexts: List[torch.Tensor] = [] + lexical_options: List[torch.Tensor] = [] + option_counts: List[int] = [] + for qi in qidxs: + spans = option_spans[qi] + context_vectors = torch.stack([self._mean_span(sequence_hidden, span) for span in spans]) + lexical_vectors = [] + for start, end in spans: + token_ids = input_ids_row[pad_offset + int(start):pad_offset + int(end)] + lexical_vectors.append(output_embedding_weight[token_ids].to(sequence_hidden.dtype).mean(dim=0)) + lexical_options.append(torch.stack(lexical_vectors)) + option_contexts.append(context_vectors) + option_counts.append(len(spans)) + + option_queries = [] + for question_index, (context_vectors, lexical) in enumerate(zip(option_contexts, lexical_options)): + option_queries.append( + self.option_context_projection(context_vectors) + self.option_lexical_projection(lexical) + + self.option_question_projection(question_vectors[question_index]).unsqueeze(0)) + routed_options = torch.cat(option_queries, dim=0).unsqueeze(0) + for layer in self.evidence_layers: + routed_options = layer(routed_options, memory) + routed_options = routed_options[0] + split_options = list(torch.split(routed_options, option_counts, dim=0)) + + base_fields = self.question_projection(question_vectors) + option_summaries = [] + for field, options in zip(base_fields, split_options): + routing_weights = torch.softmax(torch.matmul(options, field) / math.sqrt(options.shape[-1]), dim=0) + option_summaries.append(torch.sum(routing_weights.unsqueeze(-1) * options, dim=0)) + fields = ( + base_fields + self.option_summary_norm(torch.stack(option_summaries)) + + self.global_projection(global_vector).unsqueeze(0) + self.type_embedding(type_ids)) + fields = fields.unsqueeze(0) + for layer in self.layers: + fields = layer(fields, memory) + fields = self.field_norm(fields[0]) + + record_logits: List[torch.Tensor] = [] + for field, lexical, routed in zip(fields, lexical_options, split_options): + anchor = F.normalize(question_vectors[len(record_logits)] + global_vector, dim=-1) + lexical_anchor = F.normalize(lexical, dim=-1) + prior_scale = self.prior_logit_scale.clamp(max=math.log(100.0)).exp() + prior = prior_scale * torch.matmul(lexical_anchor, anchor) + options = self.option_norm(routed) + repeated_field = field.unsqueeze(0).expand_as(options) + cosine = F.cosine_similarity(repeated_field, options, dim=-1) + features = torch.cat( + [repeated_field, options, repeated_field * options, + torch.abs(repeated_field - options)], dim=-1) + residual = self.residual_scorer(features).squeeze(-1) + joint_scale = self.joint_logit_scale.clamp(max=math.log(100.0)).exp() + joint = joint_scale * cosine + residual + record_logits.append(prior + torch.sigmoid(self.residual_gate) * joint) + return record_logits + + @classmethod + def from_pretrained(cls, model_dir: str, model=None, config=None) -> 'ClefJointSchemaHead': + from safetensors.torch import load_file + config_file = cls._locate(model_dir, 'joint_head_config.json', 'CLEF_JOINT_HEAD_CONFIG') + weights_file = cls._locate(model_dir, 'joint_head.safetensors', 'CLEF_JOINT_HEAD_WEIGHTS') + if config_file is None: + raise FileNotFoundError(f'Clef joint_head_config.json not found under {model_dir!r}. ' + 'Set CLEF_JOINT_HEAD_CONFIG to its path.') + if weights_file is None: + raise FileNotFoundError(f'Clef joint_head.safetensors not found under {model_dir!r}. ' + 'Set CLEF_JOINT_HEAD_WEIGHTS to its path.') + with open(config_file) as f: + head_config = json.load(f) + head = cls(**head_config) + state = load_file(weights_file) + missing, unexpected = head.load_state_dict(state, strict=False) + if missing or unexpected: + raise RuntimeError(f'Clef joint head strict load failed: missing={missing}, unexpected={unexpected}') + return head + + +class OmniOptionScorer(nn.Module): + """OmniJev's decision head. The parameter + names/shapes match `head.pt` exactly and a strict load has 0 missing / 0 unexpected. + """ + + INIT_SCORE_BIAS = -3.0 + + def __init__(self, d_model: int, d_hidden: int = 1024, n_types: int = 3, norm: str = 'sigmoid') -> None: + super().__init__() + self.opt = nn.Sequential(nn.Linear(d_model, d_hidden), nn.GELU(), nn.Linear(d_hidden, d_hidden)) + self.que = nn.Sequential(nn.Linear(d_model, d_hidden), nn.GELU(), nn.Linear(d_hidden, d_hidden)) + self.score = nn.Linear(d_hidden, 1) + self.log_tau = nn.Parameter(torch.zeros(n_types)) + self.norm = norm + self.abstain = nn.Linear(d_hidden, 1) + self.feat = nn.Linear(6, d_hidden) + self.feat_lin = nn.Linear(6, 1) + + def forward(self, u_opts: torch.Tensor, z_q: torch.Tensor, type_id: int = 0, feats=None) -> torch.Tensor: + lg = self.logits(u_opts, z_q, type_id, feats) + if self.norm == 'softmax' and type_id != 0: + return torch.softmax(lg, dim=0)[:-1] + return torch.sigmoid(lg[:-1]) + + def logits(self, u_opts: torch.Tensor, z_q: torch.Tensor, type_id: int = 0, feats=None) -> torch.Tensor: + """The vector the probabilities come from, abstain last: `[K option logits, abstain]` for a + softmax choice/score, `[logit, 0]` for noul (softmax of [x, 0] == sigmoid(x)). + """ + g = torch.tanh(self.que(z_q)) + h = self.opt(u_opts) * g[None, :] # FiLM-style gating + if feats is not None: + h = h + self.feat(feats) + tau = torch.exp(self.log_tau[type_id]) + logit = self.score(h).squeeze(-1) + if feats is not None: + logit = logit + self.feat_lin(feats).squeeze(-1) + logit = logit / tau + if self.norm == 'softmax' and type_id != 0: + a = self.abstain(g).squeeze(-1) / tau + return torch.cat([logit, a[None]]) + return torch.cat([logit, torch.zeros_like(logit[:1])]) + + +class OmniOrdinalHead(nn.Module): + """OmniJev's ordinal (score) head. Strict-loads""" + + def __init__(self, d_model: int, d_hidden: int = 512) -> None: + super().__init__() + self.z = nn.Sequential(nn.Linear(d_model, d_hidden), nn.GELU(), nn.Linear(d_hidden, 1)) + self.cut = nn.Sequential(nn.Linear(d_model, d_hidden), nn.GELU(), nn.Linear(d_hidden, 1)) + + def forward(self, z_q: torch.Tensor, u_levels: torch.Tensor) -> torch.Tensor: + """z_q [D], u_levels [L, D] (levels in the caller's order) -> probs [L] (sum ~ 1).""" + z = self.z(z_q).squeeze(-1) + raw = self.cut(u_levels).squeeze(-1) + theta = raw[0] + torch.cat([torch.zeros(1, device=raw.device), torch.cumsum(F.softplus(raw[1:]), 0)]) + cdf = torch.sigmoid(theta - z) + cdf = torch.cat([cdf[:-1], torch.ones(1, device=cdf.device)]) + return torch.cat([cdf[:1], cdf[1:] - cdf[:-1]]).clamp(min=1e-8) + + +def omnijev_lm_option_feats(lm_head: nn.Module, h: torch.Tensor, ids: torch.Tensor, opens: Sequence[int], + closes: Sequence[int], q_end: int, open_len: int, close_len: int) -> torch.Tensor: + """Features: [sum log p(option tokens)/10, mean, n_tokens/10, log_softmax over K of the sums, + first-token log p at the question end, validity flag]. + """ + n_feat = 6 + k_opt = min(len(opens), len(closes)) + dev = h.device + feats = torch.zeros(k_opt, n_feat, device=dev, dtype=torch.float32) + if k_opt == 0: + return feats + pred_pos: List[int] = [] + tgt_tok: List[int] = [] + owner: List[int] = [] + firsts: List[int] = [] + for k in range(k_opt): + o, c = int(opens[k]), int(closes[k]) + inner = list(range(o + open_len, c - close_len + 1)) # option text token positions + firsts.append(int(ids[inner[0]]) if inner else -1) + for t in inner: + pred_pos.append(t - 1) + tgt_tok.append(int(ids[t])) + owner.append(k) + sums = torch.zeros(k_opt, device=dev) + cnts = torch.zeros(k_opt, device=dev) + if pred_pos: + pos_t = torch.as_tensor(pred_pos, device=dev, dtype=torch.long) + lg = lm_head(h[pos_t]).float() # [P, vocab] + lp = torch.log_softmax(lg, dim=-1) + tok = torch.as_tensor(tgt_tok, device=dev, dtype=torch.long) + own = torch.as_tensor(owner, device=dev, dtype=torch.long) + picked = lp.gather(1, tok[:, None]).squeeze(1) + sums = sums.index_add(0, own, picked) + cnts = cnts.index_add(0, own, torch.ones_like(picked)) + valid = cnts > 0 + mean = torch.where(valid, sums / cnts.clamp(min=1), torch.zeros_like(sums)) + dist = torch.log_softmax(torch.where(valid, sums, torch.full_like(sums, -1e4)), dim=0) + first = torch.zeros(k_opt, device=dev) + if q_end is not None and q_end >= 0 and any(f >= 0 for f in firsts): + lq = torch.log_softmax(lm_head(h[q_end:q_end + 1]).float(), dim=-1)[0] + for k, f in enumerate(firsts): + if f >= 0: + first[k] = lq[f] + feats[:, 0] = sums / 10.0 + feats[:, 1] = mean + feats[:, 2] = cnts / 10.0 + feats[:, 3] = dist + feats[:, 4] = first + feats[:, 5] = valid.float() + return feats + + +class OmniJevHead(ScoringHead): + """OmniJev System-1 head: `OptionScorer` (choice/noul) + `OrdinalScoreHead` (score) over the hidden + state at each option's `<|/opt|>` marker, gated by the question hidden at the token before the first + `<|opt|>`, plus the 6-dim backbone-opinion features. + + Reads (per question, all positions real-token-relative, `h_row = hidden[rec, pad_offset:+seq_len]`): + * `u_o` = `h_row[close_o]` -- the hidden at each option's OPT_CLOSE last token; + * `z_q` = `h_row[opens[0]-1]` -- the hidden just before the first OPT_OPEN; + * 6-dim feats via `omnijev_lm_option_feats` using the base lm_head (== output embeddings), so + `needs_output_embeddings=True`; the head runs in float32 (`keep_fp32=True`). + + Emitted `logits` are laid out so the shared `masked_softmax` reproduces the expected probabilities: + * noul -> 2 cols `[yes_logit, 0]` (softmax == sigmoid(yes_logit) == P(yes)); + * choice -> K+1 cols `[opt_0..opt_{K-1}, abstain]` (P(abstain) == 1 - sum(option probs)); + * score -> L cols `log(ordinal_probs)` (the cumulative-link head's per-level probabilities). + `meta['n_options']` must carry these column counts (2 / K+1 / L), NOT the option-block count. + + Calibration: the in-head learned `exp(log_tau)` is ALWAYS applied (inside `OptionScorer.logits`). + The post-hoc `head_meta` per-kind temperature and the noul log-odds bias are serving-time only, so + they are folded into the emitted logits ONLY when `not self.training` (a temperature on probs + `p^(1/T)` == dividing the underlying logits by T; the noul bias shifts col0's log-odds first). + """ + needs_output_embeddings = True + needs_lm_logits = False + keep_fp32 = True + + def __init__(self, + hidden_size: int, + head_hidden: int = 1024, + ord_hidden: int = 512, + n_types: int = 3, + norm: str = 'softmax', + lm_feats: bool = True, + ordinal: bool = True, + temperatures: Optional[Dict[str, float]] = None, + noul_bias: float = 0.0) -> None: + super().__init__() + self.head = OmniOptionScorer(hidden_size, head_hidden, n_types, norm) + self.ord = OmniOrdinalHead(hidden_size, ord_hidden) + self.lm_feats = bool(lm_feats) + self.ordinal = bool(ordinal) + temps = temperatures or {} + self.register_buffer( + 'temps', + torch.as_tensor([float(temps.get(KIND_NAMES[k], 1.0)) for k in range(n_types)], dtype=torch.float32), + persistent=False) + self.register_buffer('noul_bias', torch.as_tensor(float(noul_bias), dtype=torch.float32), persistent=False) + self.open_len = 1 + self.close_len = 1 + + def set_marker_lengths(self, open_len: int, close_len: int) -> None: + """Token lengths of `<|opt|>` / `<|/opt|>` (normally 1 each); the loader sets these from the + tokenizer so the option-text inner range in `omnijev_lm_option_feats` is exact.""" + self.open_len = int(open_len) + self.close_len = int(close_len) + + def forward(self, hidden_states: torch.Tensor, lm_logits: Optional[torch.Tensor], meta: Dict[str, + Any]) -> ScoringOutput: + device = hidden_states.device + kinds = [int(k) for k in meta['kinds']] + n_options = [int(n) for n in meta['n_options']] + total_q = len(n_options) + counts = meta.get('counts') + if total_q == 0: + return self._assemble([], kinds, counts, int(meta.get('max_opt') or 0), device) + + lm_head = getattr(self, '_output_embeddings', None) + if self.lm_feats and lm_head is None: + raise RuntimeError('OmniJevHead needs output embeddings for lm_feats; the loader must call ' + 'set_output_embeddings (needs_output_embeddings=True).') + input_ids = meta.get('input_ids') + if self.lm_feats and input_ids is None: + raise ValueError('OmniJevHead requires meta["input_ids"] for lm_feats ' + '(stashed by OmniJevTemplate._post_encode before the base drops it).') + option_markers = meta['option_markers'] # per question: [(open_pos, close_pos), ...] real-token-relative + record_index = [int(r) for r in meta['record_index']] + padded_len = int(hidden_states.shape[1]) + left_padding = bool(meta.get('left_padding', False)) + num_records = int(hidden_states.shape[0]) + seq_lens = meta.get('seq_lens') or [padded_len] * num_records + + per_question: List[torch.Tensor] = [] + for qi in range(total_q): + rec = record_index[qi] + seq_len = int(seq_lens[rec]) + pad_offset = (padded_len - seq_len) if left_padding else 0 + h_row = hidden_states[rec, pad_offset:pad_offset + seq_len] + ids_real = input_ids[rec, pad_offset:pad_offset + seq_len] if input_ids is not None else None + cols = self._score_question(h_row, ids_real, option_markers[qi], kinds[qi], lm_head) + per_question.append(cols) + + if len(per_question) != len(n_options) or any(len(z) != n for z, n in zip(per_question, n_options)): + raise ValueError('OmniJev head columns differ from encoded options') + max_opt = int(meta.get('max_opt') or max(n_options)) + return self._assemble(per_question, kinds, counts, max_opt, device) + + def forward_features(self, per_question_features: Sequence[Tuple[torch.Tensor, torch.Tensor, + Optional[torch.Tensor]]], + meta: Dict[str, Any]) -> ScoringOutput: + """Branch entry point (decision G1). The loader runs the two-stage cached branch forward""" + kinds = [int(k) for k in meta['kinds']] + n_options = [int(n) for n in meta['n_options']] + counts = meta.get('counts') + if not n_options: + return self._assemble([], kinds, counts, int(meta.get('max_opt') or 0), torch.device('cpu')) + device = per_question_features[0][0].device + per_question = [ + self.score_features(u, zq, feats, kinds[qi]) for qi, (u, zq, feats) in enumerate(per_question_features) + ] + if len(per_question) != len(n_options) or any(len(z) != n for z, n in zip(per_question, n_options)): + raise ValueError('OmniJev head columns differ from encoded options') + max_opt = int(meta.get('max_opt') or max(n_options)) + return self._assemble(per_question, kinds, counts, max_opt, device) + + def _assemble(self, per_question: List[torch.Tensor], kinds: List[int], counts: Optional[Any], max_opt: int, + device: torch.device) -> ScoringOutput: + counts_t = torch.as_tensor([int(c) for c in counts], dtype=torch.long, device=device) if counts else None + if not per_question: + return ScoringOutput( + logits=torch.zeros((0, max_opt), device=device, dtype=torch.float32), + option_mask=torch.zeros((0, max_opt), dtype=torch.bool, device=device), + kinds=torch.zeros((0, ), dtype=torch.long, device=device), + counts=counts_t) + logits, option_mask = pad_option_logits(per_question, width=max_opt) + return ScoringOutput( + logits=logits, + option_mask=option_mask, + kinds=torch.as_tensor(kinds, dtype=torch.long, device=device), + counts=counts_t) + + def _score_question(self, h_row: torch.Tensor, ids_real: Optional[torch.Tensor], markers: Sequence[Tuple[int, int]], + kind: int, lm_head: Optional[nn.Module]) -> torch.Tensor: + opens = [int(o) for o, _ in markers] + closes = [int(c) for _, c in markers] + u = h_row[torch.as_tensor(closes, device=h_row.device, dtype=torch.long)].float() # [k, H] + q_end = max(0, opens[0] - 1) if opens else -1 + zq = h_row[q_end].float() if opens else u.mean(0) + feats = None + if self.lm_feats and lm_head is not None and ids_real is not None: + feats = omnijev_lm_option_feats(lm_head, h_row, ids_real, opens, closes, q_end, self.open_len, + self.close_len) + return self.score_features(u, zq, feats, kind) + + def score_features(self, u: torch.Tensor, zq: torch.Tensor, feats: Optional[torch.Tensor], + kind: int) -> torch.Tensor: + if self.ordinal and kind == 2: + probs = self.ord(zq, u) + cols = torch.log(probs.clamp(min=1e-8)) + else: + cols = self.head.logits(u, zq, kind, feats) # noul -> [yes,0]; choice -> [opts..., abstain] + if not self.training: # serving-time post-hoc calibration (temperature + noul log-odds bias) + tau = self.temps[kind].to(cols.device) + if kind == 0: + bias = self.noul_bias.to(cols.device) + cols = torch.cat([(cols[0:1] + bias) / tau, cols[1:] / tau]) + else: + cols = cols / tau + return cols + + def extra_state_for_save(self) -> Dict[str, Any]: + """Persist the calibration + flags so a reloaded checkpoint rebuilds the head without head_meta.json.""" + return { + 'norm': self.head.norm, + 'lm_feats': self.lm_feats, + 'ordinal': self.ordinal, + 'temperatures': { + KIND_NAMES[k]: float(self.temps[k]) + for k in range(len(KIND_NAMES)) + }, + 'noul_bias': float(self.noul_bias), + 'open_len': self.open_len, + 'close_len': self.close_len, + } + + @classmethod + def from_pretrained(cls, model_dir: str, model=None, config=None) -> 'OmniJevHead': + meta_file = cls._locate(model_dir, 'head_meta.json', 'OMNIJEV_HEAD_META') + head_meta: Dict[str, Any] = {} + if meta_file is not None: + with open(meta_file) as f: + head_meta = json.load(f) + norm = head_meta.get('norm', 'sigmoid') + lm_feats = bool(head_meta.get('lm_feats')) + ordinal = bool(head_meta.get('ordinal')) + temps = head_meta.get('temperatures') or {} + biases = head_meta.get('biases') or {} + stray = [k for k in biases if k != 'noul'] + if stray: + raise ValueError(f'head_meta.json has biases for {sorted(stray)}, but only noul has a boundary a ' + 'bias can move (a renormalised softmax over several options is unchanged by a ' + 'constant).') + head_file = cls._locate(model_dir, 'head.pt', 'OMNIJEV_HEAD') + if head_file is None: + raise FileNotFoundError(f'OmniJev head.pt not found under {model_dir!r}. Set OMNIJEV_HEAD to its path.') + head_sd = torch.load(head_file, map_location='cpu') + opt0 = head_sd.get('opt.0.weight') + hidden_size = int(opt0.shape[1]) if opt0 is not None else cls._resolve_hidden_size(model_dir, model, config) + head = cls( + hidden_size, + head_hidden=min(1024, 2 * hidden_size), + ord_hidden=min(512, hidden_size), + norm=norm, + lm_feats=lm_feats, + ordinal=ordinal, + temperatures={ + k: float(v) + for k, v in temps.items() + }, + noul_bias=float(biases.get('noul', 0.0))) + missing, unexpected = head.head.load_state_dict(head_sd, strict=False) + if missing or unexpected: + raise RuntimeError(f'OmniJev head.pt strict load failed: missing={missing}, unexpected={unexpected}') + if ordinal: + ord_file = cls._locate(model_dir, 'ord.pt', 'OMNIJEV_ORD') + if ord_file is None: + raise FileNotFoundError( + f'OmniJev ord.pt not found under {model_dir!r} (ordinal=True). Set OMNIJEV_ORD to its path.') + o_missing, o_unexpected = head.ord.load_state_dict(torch.load(ord_file, map_location='cpu'), strict=False) + if o_missing or o_unexpected: + raise RuntimeError(f'OmniJev ord.pt strict load failed: missing={o_missing}, unexpected={o_unexpected}') + return head + + @staticmethod + def _resolve_hidden_size(model_dir: str, model, config) -> int: + """The backbone text hidden size (== head.pt's `opt.0.weight` in-features). Prefer the live + model/config, else read the checkpoint's `config.json` text_config.""" + for cfg in (config, getattr(model, 'config', None)): + text_cfg = getattr(cfg, 'text_config', None) if cfg is not None else None + hidden = getattr(text_cfg, 'hidden_size', None) or getattr(cfg, 'hidden_size', None) + if hidden: + return int(hidden) + cfg_path = os.path.join(model_dir, 'config.json') + if os.path.exists(cfg_path): + with open(cfg_path) as f: + raw = json.load(f) + hidden = (raw.get('text_config') or {}).get('hidden_size') or raw.get('hidden_size') + if hidden: + return int(hidden) + raise ValueError(f'OmniJevHead cannot resolve the backbone hidden_size from {model_dir!r} / model / config.') + + +class StartLuxHead(ScoringHead): + """Last-position FP32 vocabulary projection restricted to A-Z candidates.""" + + needs_output_embeddings = True + needs_lm_logits = False + keep_fp32 = True + + def __init__(self, letter_ids, temperatures): + super().__init__() + if len(letter_ids) != 26 or len(set(letter_ids)) != 26: + raise ValueError('Expected 26 distinct candidate token IDs') + if any(not math.isfinite(t) or t <= 0 for t in temperatures): + raise ValueError('Decision temperatures must be finite and positive') + self.register_buffer('letter_ids', torch.tensor(letter_ids, dtype=torch.long)) + self.register_buffer('temperatures', torch.tensor(temperatures, dtype=torch.float32)) + + def set_output_embeddings(self, embeddings): + # The backbone owns the vocabulary head; it must stay outside adapter head saves. + self._embedding_ref = weakref.ref(embeddings) + + @classmethod + def from_pretrained(cls, model_dir, model=None, config=None): + data = json.loads((Path(model_dir) / 'decision_config.json').read_text()) + return cls(data['letter_token_ids'], [data['temperature_by_type'][k] for k in QUESTION_TYPES]) + + def extra_state_for_save(self): + return {'letter_token_ids': self.letter_ids.tolist(), 'temperatures': self.temperatures.tolist()} + + def forward(self, hidden_states, lm_logits, meta): + if meta is None: + raise ValueError('StartLux requires decision metadata') + counts = meta['n_options'] + if any(not 2 <= n <= 26 for n in counts): + raise ValueError('Use an explicit wide-choice strategy for more than 26 options') + rows = torch.tensor(meta['record_index'], device=hidden_states.device) + columns = [ + hidden_states.shape[1] - 1 if meta.get('left_padding') else meta['seq_lens'][r] - 1 + for r in meta['record_index'] + ] + columns = torch.tensor(columns, device=hidden_states.device) + embeddings = self._embedding_ref() + if embeddings is None: + raise RuntimeError('Backbone vocabulary head is unavailable') + with torch.autocast(device_type=hidden_states.device.type, enabled=False): + weight = embeddings.weight.index_select(0, self.letter_ids.to(embeddings.weight.device)).float() + scores = F.linear(hidden_states[rows, columns].float(), weight) + per_question = [scores[i, :n] / self.temperatures[int(meta['kinds'][i])] for i, n in enumerate(counts)] + logits, mask = pad_option_logits(per_question, width=meta.get('max_opt')) + return ScoringOutput( + logits=logits, + option_mask=mask, + kinds=torch.tensor(meta['kinds'], device=logits.device), + counts=torch.tensor(meta['counts'], device=logits.device)) diff --git a/src/twinkle/model/decision/model.py b/src/twinkle/model/decision/model.py new file mode 100644 index 000000000..044f13920 --- /dev/null +++ b/src/twinkle/model/decision/model.py @@ -0,0 +1,211 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# Branch execution adapted from ms-swift fbc61cf2 (Apache-2.0). +import torch +from contextlib import contextmanager +from types import MethodType + +from .heads import ClefJointSchemaHead, JevVerbalizerHead, OmniJevHead, StartLuxHead +from .omnijev_branch import (OPT_CLOSE, OPT_OPEN, add_option_tokens, branch_questions, find_rope_owner, is_hybrid, + prefix_inputs, rope_positions, rows_from_ids) + + +def attach_decision_head(model, processor, family='generic', model_dir=None): + """Attach a scoring head before adding PEFT adapters or sharding the backbone.""" + if hasattr(model, 'scoring_head'): + raise ValueError('A scoring head is already attached') + tokenizer = getattr(processor, 'tokenizer', processor) + if family == 'generic': + head = JevVerbalizerHead.from_tokenizer(tokenizer) + else: + classes = { + 'jev': JevVerbalizerHead, + 'clef': ClefJointSchemaHead, + 'omnijev': OmniJevHead, + 'startlux': StartLuxHead + } + if family not in classes or model_dir is None: + raise ValueError('A supported family and checkpoint directory are required') + if family == 'startlux': + import json + from pathlib import Path + from startlux_decision import jevfmt + data = json.loads((Path(model_dir) / 'decision_config.json').read_text()) + if jevfmt.check_tokenizer(tokenizer) != data['letter_token_ids']: + raise ValueError('StartLux verbalizer does not match tokenizer') + if family == 'omnijev': + if not is_hybrid(model): + raise ValueError('OmniJev requires a hybrid Qwen3.5 backbone') + emb = OmniJevHead._locate(model_dir, 'new_tok_emb.pt', 'OMNIJEV_NEW_TOK_EMB') + if emb is None: + raise FileNotFoundError('Missing OmniJev option-token embeddings') + add_option_tokens(model, processor, emb) + head = classes[family].from_pretrained(model_dir, model=model, config=model.config) + device = next(model.parameters()).device + model.scoring_head = head.to(device=device, dtype=torch.float32) + if head.needs_output_embeddings: + head.set_output_embeddings(model.get_output_embeddings()) + origin_forward = model.forward + open_ids, close_ids = None, None + if family == 'omnijev': + open_ids = tokenizer.encode(OPT_OPEN, add_special_tokens=False) + close_ids = tokenizer.encode(OPT_CLOSE, add_special_tokens=False) + head.set_marker_lengths(len(open_ids), len(close_ids)) + + def forward(self, *args, **kwargs): + active_head = self.scoring_head + if family == 'omnijev': + # A PEFT ModulesToSaveWrapper delegates __call__, not custom methods. + actual = active_head + if hasattr(active_head, 'modules_to_save'): + actual = active_head.modules_to_save[active_head.active_adapter] + return _omnijev_branch_forward( + self, + *args, + head=actual, + origin_forward=origin_forward, + backbone_model=self, + tokenizer=tokenizer, + open_ids=open_ids, + close_ids=close_ids, + **kwargs) + meta = kwargs.pop('decision_meta', None) + if meta is None: + raise ValueError('Decision inference requires decision_meta') + kwargs.pop('labels', None) + kwargs.update(output_hidden_states=True, return_dict=True, use_cache=False) + output = origin_forward(*args, **kwargs) + logits = output.logits if family in ('generic', 'jev') else None + with torch.autocast(device_type=output.hidden_states[-1].device.type, enabled=False): + return active_head(output.hidden_states[-1], logits, meta) + + model.forward = MethodType(forward, model) + return head + + +def decision_lora_config(**kwargs): + """Include fp32 head weights in the adapter's normal save/load lifecycle.""" + from peft import LoraConfig + supplied = list(kwargs.pop('modules_to_save', None) or []) + return LoraConfig(modules_to_save=list(dict.fromkeys(supplied + ['scoring_head'])), **kwargs) + + +def activate_decision_head(model): + """Call after PEFT wrapping and before constructing the optimizer.""" + found = False + for name, module in model.named_modules(): + if name.endswith('scoring_head'): + if hasattr(module, 'modules_to_save'): + module = module.modules_to_save[module.active_adapter] + module.float().requires_grad_(True) + if module.needs_output_embeddings: + module.set_output_embeddings(model.get_output_embeddings()) + found = True + if not found: + raise ValueError('No decision head found') + + +@contextmanager +def _grad_ckpt_disabled(backbone_model): + touched = [m for m in backbone_model.modules() if getattr(m, 'gradient_checkpointing', False)] + for m in touched: + m.gradient_checkpointing = False + try: + yield + finally: + for m in touched: + m.gradient_checkpointing = True + + +def _omnijev_branch_forward(self, *args, head, origin_forward, backbone_model, tokenizer, open_ids, close_ids, + **kwargs): + meta = kwargs.pop('decision_meta', None) + kwargs.pop('labels', None) + if meta is None: + raise ValueError('OmniJev branch forward requires decision_meta from OmniJevTemplate._encode.') + input_ids = kwargs.get('input_ids') + if input_ids is None: + raise ValueError('OmniJev branch forward requires input_ids; OmniJevTemplate._post_encode must NOT convert the ' + 'batch to inputs_embeds (the branch re-runs the backbone on the real token ids + pixels).') + attention_mask = kwargs.get('attention_mask') + pixel_values = kwargs.get('pixel_values') + image_grid_thw = kwargs.get('image_grid_thw') + mm_token_type_ids = kwargs.get('mm_token_type_ids') + dev = input_ids.device + B, T = input_ids.shape[0], input_ids.shape[1] + left_padding = bool(meta.get('left_padding', False)) + seq_lens = meta.get('seq_lens') + counts = [int(c) for c in meta['counts']] + option_markers = meta['option_markers'] + mm_counts = meta.get('mm_counts') or [{'image': 0, 'video': 0}] * B + lm_head = backbone_model.get_output_embeddings() + rope_owner = find_rope_owner(backbone_model) + pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0 + + img_off = 0 + patch_off = 0 + q_cursor = 0 + per_question_features = [] + for r in range(B): + if attention_mask is not None: + sl = int(attention_mask[r].sum()) + elif seq_lens is not None: + sl = int(seq_lens[r]) + else: + sl = T + if left_padding: + real_ids = input_ids[r, T - sl:] + real_mm = mm_token_type_ids[r, T - sl:] if mm_token_type_ids is not None else None + else: + real_ids = input_ids[r, :sl] + real_mm = mm_token_type_ids[r, :sl] if mm_token_type_ids is not None else None + enc = {'input_ids': real_ids[None]} + n_img = int(mm_counts[r].get('image', 0)) + if image_grid_thw is not None and n_img > 0: + grid_r = image_grid_thw[img_off:img_off + n_img] + n_patch = int(grid_r.prod(-1).sum()) + if pixel_values is not None: + enc['pixel_values'] = pixel_values[patch_off:patch_off + n_patch] + enc['image_grid_thw'] = grid_r + img_off += n_img + patch_off += n_patch + if real_mm is not None: + enc['mm_token_type_ids'] = real_mm[None] + + nq = counts[r] + real_list = real_ids.tolist() + per_q_spans = [] + for qi in range(nq): + markers = option_markers[q_cursor + qi] + opens = [int(o) for o, _ in markers] + closes = [int(c) for _, c in markers] + per_q_spans.append((opens, closes)) + L_use = max(1, min(opens[0] - 1 for opens, _ in per_q_spans)) + penc = prefix_inputs(enc, L_use) + ppos = rope_positions(rope_owner, penc) + qrows = [rows_from_ids(real_list, L_use, opens, closes, k=len(opens)) for opens, closes in per_q_spans] + with _grad_ckpt_disabled(backbone_model): + outs = branch_questions( + origin_forward, penc, qrows, dev, open_ids, close_ids, lm_head=lm_head, prefix_pos=ppos, pad_id=pad_id) + per_question_features.extend(outs) + q_cursor += nq + + with torch.autocast(device_type=dev.type, enabled=False): + return head.forward_features(per_question_features, meta) + + +def initialize_missing_adapter_head(model, state, adapter_name): + """Seed a released adapter without a saved head from explicitly loaded factory artifacts. + + Only call when the adapter configuration did not declare scoring_head in modules_to_save. + Trained checkpoints declaring a saved head must never silently fall back to factory values. + """ + if any('.scoring_head.' in key for key in state): + raise ValueError('Adapter contains head state but does not declare scoring_head in modules_to_save') + result = dict(state) + for name, module in model.named_modules(): + if name.endswith('scoring_head') and hasattr(module, 'modules_to_save'): + head = module.modules_to_save[adapter_name] + for key, value in head.state_dict().items(): + result[f'{name}.{key}'] = value.detach().cpu() + return result + raise ValueError('Active decision head is unavailable') diff --git a/src/twinkle/model/decision/omnijev_branch.py b/src/twinkle/model/decision/omnijev_branch.py new file mode 100644 index 000000000..b755a167a --- /dev/null +++ b/src/twinkle/model/decision/omnijev_branch.py @@ -0,0 +1,352 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# Adapted from ms-swift fbc61cf2b714ade30ed29f28ddd63be61f6ebb50 (Apache-2.0). +"""OmniJev prefix-cache branching and option-token registration. + +Hybrid linear attention carries recurrent state across tokens, so an attention mask alone +cannot isolate option branches. Each option continues a separate copy of the same prefix +cache. MSO_BRANCH_CHUNK controls chunk size; MSO_BRANCH_CKPT enables recomputation. +Checkpoint-specific output and gradient parity must be checked before deployment. +""" +import copy +import os +import torch +import torch.utils.checkpoint +from typing import Any, Dict, List, Optional, Sequence + +OPT_OPEN = '<|opt|>' +OPT_CLOSE = '<|/opt|>' +N_FEAT = 6 +_PIX_KEYS = ('pixel_values', 'image_grid_thw', 'pixel_values_videos', 'video_grid_thw', 'mm_token_type_ids') + + +def is_hybrid(model) -> bool: + """True when the backbone has linear-attention layers (`config.text_config.layer_types` contains + 'linear'), i.e. a block mask cannot isolate options and branch is required. Qwen3.5-4B is hybrid.""" + cfg = getattr(model, 'config', None) + tc = getattr(cfg, 'text_config', cfg) + lt = getattr(tc, 'layer_types', None) + return bool(lt) and any('linear' in str(t) for t in lt) + + +def find_rope_owner(m, depth: int = 0): + """The submodule that owns `get_rope_index` (Qwen3VLModel); PEFT / DDP bury it a few attrs deep.""" + if depth > 6: + return None + if hasattr(m, 'get_rope_index'): + return m + for attr in ('model', 'base_model', 'module'): + sub = getattr(m, attr, None) + if sub is not None and sub is not m: + r = find_rope_owner(sub, depth + 1) + if r is not None: + return r + return None + + +def rope_positions(rope_owner, enc: Dict[str, Any]) -> Optional[torch.Tensor]: + """mrope `position_ids` `[3, B, T]` of `enc` from the model's own `get_rope_index` (None if the + backbone has no rope owner). The branch prefix needs NO option alignment (each row is scored + from the same prefix state, so order invariance comes from branching, not from shared rope + start positions).""" + if rope_owner is None: + return None + mm = enc.get('mm_token_type_ids') + if mm is None: + # transformers 5.x wants 0=text 1=image 2=video per token; rebuild it from the ids. + cfg = getattr(rope_owner, 'config', None) + img_id = getattr(cfg, 'image_token_id', None) + vid_id = getattr(cfg, 'video_token_id', None) + ids = enc['input_ids'] + mm = torch.zeros_like(ids) + if img_id is not None: + mm[ids == img_id] = 1 + if vid_id is not None: + mm[ids == vid_id] = 2 + kw = dict( + input_ids=enc['input_ids'], + image_grid_thw=enc.get('image_grid_thw'), + video_grid_thw=enc.get('video_grid_thw'), + attention_mask=enc.get('attention_mask')) + try: + pos, _ = rope_owner.get_rope_index(mm_token_type_ids=mm, **kw) # transformers >= 5.x + except TypeError: + pos, _ = rope_owner.get_rope_index(**kw) # older signature + return pos + + +def _rep(v, K): + if torch.is_tensor(v) and v.dim() >= 1: + return v.repeat_interleave(K, dim=0) + if isinstance(v, list): + return [_rep(x, K) for x in v] + if isinstance(v, tuple): + return tuple(_rep(x, K) for x in v) + if isinstance(v, dict): + return {k: _rep(x, K) for k, x in v.items()} + return v + + +def expand_cache(cache, K): + """A copy of `cache` with every tensor repeated K times along the batch axis (KV caches AND recurrent + states). The original stays intact so it can be expanded again for the next chunk.""" + new = copy.copy(cache) + layers = [] + for layer in cache.layers: + l2 = copy.copy(layer) + for k, v in list(vars(layer).items()): + setattr(l2, k, _rep(v, K)) + layers.append(l2) + new.layers = layers + new.update_recurrent_state = lambda *a, **k: None + return new + + +def prefix_inputs(enc: Dict[str, Any], L: int) -> Dict[str, Any]: + """The prefix encoding (`ids[:L]` + the pixel tensors) cut out of a single-question encoding.""" + ids = enc['input_ids'][:, :L] + p = {'input_ids': ids, 'attention_mask': torch.ones_like(ids)} + for k in _PIX_KEYS: + if k in enc: + p[k] = enc[k][:, :L] if k == 'mm_token_type_ids' else enc[k] + return p + + +def rows_from_ids(ids, L: int, opens: Sequence[int], closes: Sequence[int], k: Optional[int] = None) -> List[List[int]]: + """Token ids of ONE single-question sequence (prefix + instruction + option blocks), the prefix length + `L` and the option spans -> K rows, each the shared part `ids[L:opens[0]]` followed by one option block + `ids[o:c+1]`. `L` MUST be `< opens[0]` so every row keeps the token before its option (that is where + `zq` is read).""" + ids = list(ids) + if not opens or not closes: + return [ids[L:]] + head = ids[L:opens[0]] + pairs = [(o, c) for o, c in zip(opens, closes) if c >= o] + if k is not None: + pairs = pairs[:k] + return [head + ids[o:c + 1] for o, c in pairs] + + +def _marks(r: Sequence[int], open_ids: Sequence[int], close_ids: Sequence[int]): + """(index of the first open marker, index of the last token of the last close marker, indices of the + option text tokens) inside one row.""" + ol, cl = len(open_ids), len(close_ids) + r = list(r) + opens = [k for k in range(len(r) - ol + 1) if r[k:k + ol] == list(open_ids)] + closes = [k for k in range(len(r) - cl + 1) if r[k:k + cl] == list(close_ids)] + o_ = opens[0] if opens else 0 + c_ = (closes[-1] + cl - 1) if closes else len(r) - 1 + return o_, c_, list(range(o_ + ol, c_ - cl + 1)) + + +def branch_forward(backbone, + penc: Dict[str, Any], + rows: List[List[int]], + dev, + open_ids: Sequence[int], + close_ids: Sequence[int], + lm_head=None, + prefix_pos: Optional[torch.Tensor] = None, + pos_start: Optional[int] = None, + groups: Optional[List[List[int]]] = None, + chunk: Optional[int] = None, + pad_id: int = 0, + mrope_axes: int = 3) -> Dict[str, Any]: + """`penc`: prefix encoding (`input_ids [1, L]` + pixel tensors); `rows`: token-id lists continuing the + prefix; `groups`: lists of row indices forming one question (the K-way softmax feature is taken inside a + group). `prefix_pos`: `[3, 1, L]` positions of the prefix (`rope_positions`); `pos_start`: rope position + of the first suffix token (default: max prefix position + 1). + -> {'u': [R, D] hidden at each row's closing marker, 'zq': [R, D] hidden just before each row's + opening marker, 'feats': [R, 6] backbone log-prob features (lm_head given) else None} + """ + R = len(rows) + prefix_ids = penc['input_ids'] + L = int(prefix_ids.shape[1]) + kw: Dict[str, Any] = {'input_ids': prefix_ids, 'attention_mask': torch.ones_like(prefix_ids), 'use_cache': True} + kw.update({k: v for k, v in penc.items() if k in _PIX_KEYS}) + if prefix_pos is not None: + kw['position_ids'] = prefix_pos + out = backbone(**kw, logits_to_keep=1) + cache0 = out.past_key_values + if pos_start is None: + pos_start = (int(prefix_pos.max()) + 1) if prefix_pos is not None else L + p0 = int(pos_start) + chunk = chunk or int(os.environ.get('MSO_BRANCH_CHUNK', '64')) + us, zqs, sums_l, cnts_l, first_l = [], [], [], [], [] + for s in range(0, R, chunk): + sub = rows[s:s + chunk] + Rc = len(sub) + n = max(len(r) for r in sub) + ids = torch.full((Rc, n), pad_id, device=dev, dtype=torch.long) + att = torch.zeros((Rc, n), dtype=torch.long, device=dev) + for i, r in enumerate(sub): + ids[i, :len(r)] = torch.tensor(r, device=dev) + att[i, :len(r)] = 1 + full_att = torch.cat([torch.ones((Rc, L), dtype=torch.long, device=dev), att], dim=1) + pos = torch.arange(p0, p0 + n, device=dev)[None, None, :].expand(mrope_axes, Rc, n).contiguous() + cache_pos = torch.arange(L, L + n, device=dev) + + if torch.is_grad_enabled() and os.environ.get('MSO_BRANCH_CKPT', '0') == '1': + # recompute prefix + chunk in backward instead of keeping activations (memory ~ one chunk). + # The prefix pass MUST be inside the checkpoint so the recomputation rebuilds cache0 from + # the same input tensors, keeping the gradient chain to LoRA params intact. If the prefix + # pass were outside (as a closure-captured cache0), checkpoint's recomputation would use + # the stale outer cache0 whose autograd graph is already freed, breaking gradient flow. + def _chunk_h_ckpt(ids_, att_, pos_, cpos_, pids_, patt_, ppos_): + pkw = { + k: v + for k, v in kw.items() if k != 'input_ids' and k != 'attention_mask' and k != 'position_ids' + } + pkw['input_ids'] = pids_ + pkw['attention_mask'] = patt_ + if prefix_pos is not None: + pkw['position_ids'] = ppos_ + pout = backbone(**pkw, logits_to_keep=1) + c0 = pout.past_key_values + o_ = backbone( + input_ids=ids_, + attention_mask=att_, + past_key_values=expand_cache(c0, ids_.shape[0]), + use_cache=True, + output_hidden_states=True, + cache_position=cpos_, + position_ids=pos_, + logits_to_keep=1) + return o_.hidden_states[-1] + + h = torch.utils.checkpoint.checkpoint( + _chunk_h_ckpt, + ids, + full_att, + pos, + cache_pos, + prefix_ids, + torch.ones_like(prefix_ids), + prefix_pos if prefix_pos is not None else torch.zeros(3, 1, 1, device=dev, dtype=torch.long), + use_reentrant=False) + else: + + def _chunk_h(ids_, att_, pos_, cpos_): + o_ = backbone( + input_ids=ids_, + attention_mask=att_, + past_key_values=expand_cache(cache0, ids_.shape[0]), + use_cache=True, + output_hidden_states=True, + cache_position=cpos_, + position_ids=pos_, + logits_to_keep=1) + return o_.hidden_states[-1] + + h = _chunk_h(ids, full_att, pos, cache_pos) # [Rc, n, D] + marks = [_marks(r, open_ids, close_ids) for r in sub] + ar = torch.arange(Rc, device=dev) + u = h[ar, torch.tensor([m[1] for m in marks], device=dev)] + zq = h[ar, torch.tensor([max(0, m[0] - 1) for m in marks], device=dev)] + us.append(u) + zqs.append(zq) + if lm_head is not None: + pred_i, pred_t, tgt_tok, owner, firsts = [], [], [], [], [] + for i, (r, (o_, c_, inner)) in enumerate(zip(sub, marks)): + firsts.append(r[inner[0]] if inner else -1) + for t in inner: + pred_i.append(i) + pred_t.append(t - 1) + tgt_tok.append(r[t]) + owner.append(i) + sums = torch.zeros(Rc, device=dev) + cnts = torch.zeros(Rc, device=dev) + if pred_i: + hp = h[torch.tensor(pred_i, device=dev), torch.tensor(pred_t, device=dev)] + lp = torch.log_softmax(lm_head(hp).float(), dim=-1) + picked = lp.gather(1, torch.tensor(tgt_tok, device=dev)[:, None]).squeeze(1) + own = torch.tensor(owner, device=dev) + sums = sums.index_add(0, own, picked) + cnts = cnts.index_add(0, own, torch.ones_like(picked)) + first = torch.zeros(Rc, device=dev) + fi = [i for i, f in enumerate(firsts) if f >= 0] + if fi: + lq = torch.log_softmax(lm_head(zq[fi]).float(), dim=-1) + first[fi] = lq[torch.arange(len(fi), device=dev), torch.tensor([firsts[i] for i in fi], device=dev)] + sums_l.append(sums) + cnts_l.append(cnts) + first_l.append(first) + u, zq = torch.cat(us), torch.cat(zqs) + feats = None + if lm_head is not None: + sums, cnts, first = torch.cat(sums_l), torch.cat(cnts_l), torch.cat(first_l) + valid = cnts > 0 + feats = torch.zeros(R, N_FEAT, device=dev, dtype=torch.float32) + masked = torch.where(valid, sums, torch.full_like(sums, -1e4)) + dist = torch.zeros(R, device=dev) + for g in (groups or [list(range(R))]): + gi = torch.tensor(g, device=dev) + dist[gi] = torch.log_softmax(masked[gi], dim=0) + feats[:, 0] = sums / 10.0 + feats[:, 1] = torch.where(valid, sums / cnts.clamp(min=1), torch.zeros_like(sums)) + feats[:, 2] = cnts / 10.0 + feats[:, 3] = dist + feats[:, 4] = first + feats[:, 5] = valid.float() + return {'u': u, 'zq': zq, 'feats': feats} + + +def branch_questions(backbone, + penc: Dict[str, Any], + qrows: List[List[List[int]]], + dev, + open_ids: Sequence[int], + close_ids: Sequence[int], + lm_head=None, + prefix_pos: Optional[torch.Tensor] = None, + chunk: Optional[int] = None, + pad_id: int = 0): + """`qrows`: per question, its rows -> per question `(u [K, D] float32, zq [D] float32, feats [K, 6] | + None)`. All questions of the batch ride ONE prefix pass + one chunked batched suffix pass.""" + rows: List[List[int]] = [] + groups: List[List[int]] = [] + for qr in qrows: + groups.append(list(range(len(rows), len(rows) + len(qr)))) + rows.extend(qr) + if not rows: + return [] + br = branch_forward( + backbone, + penc, + rows, + dev, + open_ids, + close_ids, + lm_head=lm_head, + prefix_pos=prefix_pos, + groups=groups, + chunk=chunk, + pad_id=pad_id) + out = [] + for g in groups: + gi = torch.tensor(g, device=dev) + out.append( + (br['u'][gi].float(), br['zq'][g[0]].float(), (br['feats'][gi] if br['feats'] is not None else None))) + return out + + +def add_option_tokens(model, proc, emb_path: Optional[str] = None) -> List[int]: + """Register `<|opt|>` / `<|/opt|>` and give them DETERMINISTIC embeddings. + `resize_token_embeddings` fills new rows randomly; those rows are + frozen in training (only LoRA + head train), so a fresh random draw at eval time would silently change + the very positions the head reads. We init them to the mean embedding and, when `new_tok_emb.pt` exists, + restore it exactly. Returns the two token ids `[open_id, close_id]`.""" + proc.tokenizer.add_special_tokens({'additional_special_tokens': [OPT_OPEN, OPT_CLOSE]}) + n_old = model.get_input_embeddings().weight.shape[0] + model.resize_token_embeddings(len(proc.tokenizer)) + emb = model.get_input_embeddings().weight + ids = proc.tokenizer.convert_tokens_to_ids([OPT_OPEN, OPT_CLOSE]) + with torch.no_grad(): + if emb_path and os.path.exists(emb_path): + saved = torch.load(emb_path, map_location=emb.device) + for tid, row in zip(ids, saved): + emb[tid] = row.to(emb.dtype) + else: + mean = emb[:n_old].float().mean(0) + for k, tid in enumerate(ids): + emb[tid] = (mean + 0.01 * (k + 1)).to(emb.dtype) + return ids diff --git a/src/twinkle/model/decision/omnijev_media.py b/src/twinkle/model/decision/omnijev_media.py new file mode 100644 index 000000000..fb5b8e24a --- /dev/null +++ b/src/twinkle/model/decision/omnijev_media.py @@ -0,0 +1,236 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# Adapted from ms-swift fbc61cf2b714ade30ed29f28ddd63be61f6ebb50 (Apache-2.0). +import hashlib +import io +import os +import subprocess +import tempfile +from concurrent.futures import ThreadPoolExecutor +from PIL import Image, ImageDraw, ImageFont +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union + +from .omnijev_branch import OPT_CLOSE, OPT_OPEN + +# ---- panels (verbatim) ---- +# MSO_NO_PANELS=1 restores the old behaviour (encode only the first image). It exists so the same +# weights can be scored both ways, which is the only way to attribute a change to this fix. +PANELS_OFF = os.environ.get('MSO_NO_PANELS', '').strip().lower() in ('1', 'true', 'yes', 'on') +_BORDER = 6 +_LABEL = 22 + +# ---- video (verbatim) ---- +K, COLS, TILE = 16, 4, 384 + +ImageLike = Union[str, 'Image.Image'] + + +def _grid(n): + return (1, 1) if n <= 1 else (2, 1) if n == 2 else (2, 2) if n <= 4 else (3, (n + 2) // 3) + + +def compose(paths, max_side=448): + ims = [] + for p in paths: + im = _as_pil(p) + im = im.convert('RGB') + if max(im.size) > max_side: + im.thumbnail((max_side, max_side)) + ims.append(im) + cols, rows = _grid(len(ims)) + cw = max(i.width for i in ims) + ch = max(i.height for i in ims) + W = cols * cw + (cols + 1) * _BORDER + H = rows * (ch + _LABEL) + (rows + 1) * _BORDER + canvas = Image.new('RGB', (W, H), (32, 32, 32)) + draw = ImageDraw.Draw(canvas) + for k, im in enumerate(ims): + c, r = k % cols, k // cols + x = _BORDER + c * (cw + _BORDER) + y = _BORDER + r * (ch + _LABEL + _BORDER) + draw.text((x + 2, y + 4), '%d' % (k + 1), fill=(255, 255, 255)) + canvas.paste(im, (x + (cw - im.width) // 2, y + _LABEL + (ch - im.height) // 2)) + return canvas + + +def compose_path(paths, cache_dir, max_side=448): + """same as compose(), written to `cache_dir` and reused on later epochs; returns a path to a single image. A + single path is returned unchanged, so nothing is copied or re-encoded for the common case. + Falls back to the in-memory `compose` when any input is not a statable file path (a PIL image).""" + if not paths: + return None + if len(paths) == 1 or PANELS_OFF: + return paths[0] + if not all(isinstance(p, str) for p in paths): + return compose(paths, max_side) + stamps = [(str(p), os.stat(p).st_size, os.stat(p).st_mtime_ns) for p in paths] + key = hashlib.sha1((repr(stamps) + '|%d' % max_side).encode()).hexdigest()[:20] + out = os.path.join(cache_dir, key + '.jpg') + if os.path.exists(out): + return out + os.makedirs(cache_dir, exist_ok=True) + fd, tmp = tempfile.mkstemp(prefix=key + '.', suffix='.part', dir=cache_dir) + os.close(fd) + try: + compose(paths, max_side).save(tmp, format='JPEG', quality=88) + os.replace(tmp, out) + return out + except Exception: + # Reject incomplete states instead of answering a different question. + if os.path.exists(tmp): + os.unlink(tmp) + raise + + +def font(sz): + for f in ('/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf', '/usr/share/fonts/dejavu/DejaVuSans-Bold.ttf', + '/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf', '/usr/share/fonts/dejavu/DejaVuSans.ttf'): + try: + return ImageFont.truetype(f, sz) + except Exception: + pass + return ImageFont.load_default() + + +FT = font(26) + + +def letterbox(img, t): + img = img.convert('RGB') + w, h = img.size + s = min(TILE / w, TILE / h) + img = img.resize((max(1, int(w * s)), max(1, int(h * s))), Image.BILINEAR) + tile = Image.new('RGB', (TILE, TILE), (0, 0, 0)) + tile.paste(img, ((TILE - img.size[0]) // 2, (TILE - img.size[1]) // 2)) + d = ImageDraw.Draw(tile) + label = 't=%ds' % t + d.rectangle([0, 0, 16 + 15 * len(label), 36], fill=(0, 0, 0)) + d.text((6, 4), label, fill=(255, 255, 255), font=FT) + return tile + + +def probe_duration(path): + out = subprocess.run(['ffprobe', '-v', 'error', '-show_entries', 'format=duration', '-of', 'csv=p=0', path], + capture_output=True, + text=True, + timeout=60).stdout.strip() + try: + return float(out) + except ValueError: + raise ValueError('could not read the video (ffprobe gave no duration)') + + +def grab(path, t): + out = subprocess.run([ + 'ffmpeg', '-v', 'error', '-ss', + '%.3f' % t, '-i', path, '-frames:v', '1', '-f', 'image2pipe', '-vcodec', 'mjpeg', '-q:v', '3', 'pipe:1' + ], + capture_output=True, + timeout=120).stdout + return Image.open(io.BytesIO(out)).convert('RGB') if out else None + + +def video_mosaic(path, out_jpg): + """16 uniformly spaced frames, letterboxed to 384 px, timestamp stamped, one 4x4 mosaic (video.py).""" + dur = probe_duration(path) + if dur < 1.0: + raise ValueError('video shorter than one second') + ts = [(i + 0.5) / K * dur for i in range(K)] + with ThreadPoolExecutor(max_workers=4) as ex: + imgs = list(ex.map(lambda t: grab(path, min(t, max(0.0, dur - 0.05))), ts)) + last = None + tiles = [] + for t, im in zip(ts, imgs): + if im is None: + im = last if last is not None else Image.new('RGB', (TILE, TILE), (0, 0, 0)) + last = im + tiles.append(letterbox(im, int(round(t)))) + mosaic = Image.new('RGB', (COLS * TILE, (K // COLS) * TILE)) + for k, tl in enumerate(tiles): + r, c = divmod(k, COLS) + mosaic.paste(tl, (c * TILE, r * TILE)) + mosaic.save(out_jpg, quality=85) + return dur, [int(round(t)) for t in ts] + + +def video_state(path, out_jpg=None): + """OFFLINE tool: a raw clip -> `{"images": [mosaic path], "video": {...}}`, the state to hand + `OmniJevTemplate`.""" + out_jpg = out_jpg or (os.path.splitext(path)[0] + '.mosaic.jpg') + dur, ts = video_mosaic(path, out_jpg) + return { + 'images': [out_jpg], + 'video': { + 'n_frames': K, + 'cols': COLS, + 'tile': TILE, + 'timestamps': ts, + 'duration': round(dur, 1) + } + } + + +def video_content(img, video, instr) -> Tuple[List['Image.Image'], List[Dict[str, Any]]]: + """One image, or a 4x4 mosaic split back into its frames, each preceded by its timestamp (adapted + to accept a PIL image OR a path). Returns (list of PIL + images in prompt order, chat content list of {'type':'image'|'text', ...}).""" + im = _as_pil(img).convert('RGB') + if not video: + return [im], [{'type': 'image'}, {'type': 'text', 'text': instr}] + n_frames, cols, tile = video['n_frames'], video['cols'], video['tile'] + frames: List[Image.Image] = [] + content: List[Dict[str, Any]] = [{ + 'type': + 'text', + 'text': + 'A video of {:.0f} seconds; {} frames sampled uniformly in order.'.format(video.get('duration', 0), n_frames) + }] + for k in range(n_frames): + r, c = divmod(k, cols) + frames.append(im.crop((c * tile, r * tile, (c + 1) * tile, (r + 1) * tile))) + content.append({'type': 'text', 'text': 't={}s:'.format(video['timestamps'][k])}) + content.append({'type': 'image'}) + content.append({'type': 'text', 'text': instr}) + return frames, content + + +# Media adapters for the DecisionEncoder. +def _as_pil(img) -> 'Image.Image': + """Resolve PIL images, paths and dataset image dictionaries with Twinkle's loader.""" + if isinstance(img, Image.Image): + return img + if isinstance(img, dict) and 'bytes' in img: + img = img['bytes'] or img['path'] + from twinkle.template.vision_utils import load_image + return load_image(img) + + +def compose_state(images: Sequence[ImageLike], + video: Optional[Dict[str, Any]], + cache_dir: Optional[str] = None, + max_side: int = 448) -> ImageLike: + """Media choice: several stills and NO video -> one numbered + panel (`compose_path`); otherwise `images[0]` (the pre-composed video mosaic, or the single still). + Returns the ONE image (path or PIL) that `video_content` then reads.""" + if not images: + raise ValueError('OmniJev state carries no images; `system_one` encodes exactly one image per record.') + if len(images) > 1 and not video: + cache_dir = cache_dir or os.environ.get('MSO_PANELS') or os.path.join(tempfile.gettempdir(), 'mso_panels') + return compose_path(list(images), cache_dir, max_side) + return images[0] + + +def content_to_text(content: Sequence[Dict[str, Any]]) -> str: + """Render chat content with one image placeholder per frame.""" + parts: List[str] = [] + for c in content: + if c.get('type') == 'image': + parts.append('') + else: + parts.append(c.get('text', '')) + return ''.join(parts) + + +def option_body(opt_texts: Sequence[str]) -> str: + """The assistant body appended after the chat prompt: every option block wrapped in the markers + (`"".join("{}{}{}".format(OPT_OPEN, t, OPT_CLOSE) ...)`).""" + return ''.join(f'{OPT_OPEN}{t}{OPT_CLOSE}' for t in opt_texts) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 5540f058f..16f88ad83 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -1348,6 +1348,17 @@ def _get_adapter_state_dict_for_save(self, adapter_name: str) -> dict: adapter_state = self.strategy.get_adapter_state_dict(self.model, adapter_name) adapter_suffix = f'.{adapter_name}.' modules_to_save_infix = f'.modules_to_save.{adapter_name}.' + # Strategies collect parameters without materializing the full backbone state. + # PEFT modules_to_save also require persistent buffers (e.g. calibration state). + model = self.strategy.unwrap_model(self.model) + for module_name, module in model.named_modules(): + if modules_to_save_infix not in f'{module_name}.': + continue + for buffer_name, buffer in module.named_buffers(recurse=False): + if buffer_name in module._non_persistent_buffers_set: + continue + key = f'{module_name}.{buffer_name}' + adapter_state[key] = Torch.to_local_tensor(buffer).detach().cpu() processed_state_dict = {} for key, value in adapter_state.items(): if modules_to_save_infix in key: diff --git a/src/twinkle/preprocessor/decision.py b/src/twinkle/preprocessor/decision.py new file mode 100644 index 000000000..0878e9573 --- /dev/null +++ b/src/twinkle/preprocessor/decision.py @@ -0,0 +1,953 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# Adapted from ms-swift fbc61cf2b714ade30ed29f28ddd63be61f6ebb50 (Apache-2.0). +"""Preprocessor base for the `decision` task_type (typed-decision System-1 scoring models). + +A typed-decision record carries a `state` plus one or more questions; each question has a +`kind in {noul, choice, score}` and its own option set. `ScoringPreprocessor` normalizes a raw +dataset row into Twinkle decision fields consumed by DecisionEncoder and DecisionInputProcessor. +The encoder renders the prompt and records question spans; target_probs is supervision only. + +Output row contract: + messages: [{'role': 'user', 'content': state}] + questions: List[str] per-question text + kinds: List[str] per-question 'noul' | 'choice' | 'score' + options: List[List[str]] per-question option texts (noul -> ['false','true'], score -> ['0'..'5']) + target_probs: List[List[float]] per-question gold distribution over its OWN options (training only) + +Subclasses override `parse_record` to turn a model-specific raw row into `(state, questions)`, +where each question is a dict `{'kind', 'question', 'options', 'gold'}` and `gold` is either an +option index (int) or a full distribution (sequence of floats, for distillation targets). +""" +import json +import math +import random as _random +import re +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Sequence, Tuple + +from .base import Preprocessor + + +@dataclass +class DecisionAugmentConfig: + """Training-time schema augmentation for decision datasets. + + Only the operations safe for all three decision models (JEV / Clef / OmniJev) are + included. All probabilities are per-question; state_as_text is per-record. Disabled by + default (all zeros / False) so augmentation must be explicitly opted into. + """ + + drop_instructions: float = 0.0 + subsample_options: int = 0 + state_as_text: float = 0.0 + shuffle_fields: bool = False + + +def _humanize(label: str) -> str: + return re.sub(r'[_\-/]+', ' ', str(label)).strip() + + +def _state_to_text(state: Any) -> str: + if not isinstance(state, dict): + if isinstance(state, str): + return state + return json.dumps(state, ensure_ascii=False) + if len(state) == 0: + return '' + if len(state) == 1: + return str(next(iter(state.values()))) + return '\n'.join(f'{_humanize(key).capitalize()}: {value}' for key, value in state.items()) + + +def _subsample_options(question: Dict[str, Any], gold: Any, rng: _random.Random, + max_options: int) -> Tuple[Dict[str, Any], Any]: + """Subsample choice options to ``max_options``, always keeping the highest-probability option. + + Only applies to ``choice`` questions whose options exceed ``max_options``. For an int gold + (one-hot), the gold index is always kept and ``max_options - 1`` distractors are sampled. + For a distribution gold (soft labels), the argmax option is always kept and the distribution + is truncated and renormalized over the kept options. ``noul`` / ``score`` questions have + fixed option sets and are never subsampled. + """ + options = question.get('options') + if not options or len(options) <= max_options: + return question, gold + kind = question.get('kind') + if kind in ('noul', 'score'): + return question, gold + + column_keys = question.get('column_keys') + has_abstain = column_keys is not None and len(column_keys) == len(options) + 1 + target_width = len(options) + int(has_abstain) + if isinstance(gold, int) and 0 <= gold < target_width: + peak = gold + elif isinstance(gold, (list, tuple)) and len(gold) == target_width: + peak = max(range(len(gold)), key=lambda i: gold[i]) + else: + return question, gold + + required = {peak} if peak < len(options) else set() + others = [i for i in range(len(options)) if i not in required] + kept = sorted(required | set(rng.sample(others, min(max_options - len(required), len(others))))) + target_indices = kept + ([len(options)] if has_abstain else []) + if isinstance(gold, int): + new_gold: Any = target_indices.index(gold) + else: + new_gold = [float(gold[i]) for i in target_indices] + total = sum(new_gold) + new_gold = [p / total for p in new_gold] + + new_q = {**question, 'options': [options[i] for i in kept]} + if isinstance(question.get('criteria'), dict): + new_q['criteria'] = {options[i]: question['criteria'][options[i]] for i in kept} + if column_keys is not None: + new_q['column_keys'] = [column_keys[i] for i in target_indices] + return new_q, new_gold + + +def augment_decision_row( + state: Any, + questions: List[Dict[str, Any]], + config: Optional[DecisionAugmentConfig], + rng: Optional[_random.Random] = None, +) -> Tuple[Any, List[Dict[str, Any]]]: + """Apply schema augmentation to a decision record's ``(state, questions)``. + + Works on the pre-``_build_row`` representation -- after ``parse_record`` and before + ``_build_row`` -- so it is agnostic to the downstream template/head. Returns a (possibly + modified) ``(state, questions)`` tuple. + + When *config* is None or all-disabled, returns the inputs unchanged (zero-cost passthrough). + """ + if config is None: + return state, questions + if (config.drop_instructions == 0.0 and config.subsample_options <= 0 and config.state_as_text == 0.0 + and not config.shuffle_fields): + return state, questions + rng = rng or _random.Random(42) + + # drop_instructions: set the question text to empty (template will still render the layout, + # model must infer intent from options only). + if config.drop_instructions > 0: + for q in questions: + if rng.random() < config.drop_instructions: + q['question'] = '' + + # subsample_options: keep gold + N distractors for choice questions. + if config.subsample_options > 0: + new_questions = [] + for q in questions: + q, q['gold'] = _subsample_options(q, q.get('gold'), rng, config.subsample_options) + new_questions.append(q) + questions = new_questions + + # state_as_text: randomly flatten dict state to plain text so the model sees different + # renderings (JSON vs human-readable lines). Per-record probability. + if config.state_as_text > 0 and isinstance(state, dict) and state: + if rng.random() < config.state_as_text: + state = _state_to_text(state) + + # shuffle_fields: randomize field order in dict state AND question order. The state-side + # shuffle makes the model robust to permutation of schema fields (important for Clef which + # renders SCHEMA FIELDS in the state's key order); the question-side shuffle only matters + # for multi-question records (also Clef). + if config.shuffle_fields: + if isinstance(state, dict) and len(state) > 1: + keys = list(state.keys()) + rng.shuffle(keys) + state = {k: state[k] for k in keys} + if len(questions) > 1: + rng.shuffle(questions) + + return state, questions + + +class ScoringPreprocessor(Preprocessor): + """Base: assemble the decision row contract; subclasses supply `parse_record`. + + Supports optional training-time schema augmentation (``augment_config``) applied between + ``parse_record`` and ``_build_row`` — after the raw row is parsed into ``(state, questions)`` + and before the decision row is built. This keeps augmentation model-agnostic: subclasses + only override ``parse_record`` and ``_build_row``, and the base hooks augmentation in + between automatically. + """ + + NOUL_OPTIONS = ['false', 'true'] + SCORE_OPTIONS = ['0', '1', '2', '3', '4', '5'] + + # Canonical aliases for question kinds. + KIND_ALIASES = { + 'noul': 'noul', + 'yes_no': 'noul', + 'yn': 'noul', + 'bool': 'noul', + 'choice': 'choice', + 'multiple_choice': 'choice', + 'mcq': 'choice', + 'select': 'choice', + 'score': 'score', + 'rating': 'score', + 'ordinal': 'score', + 'likert': 'score', + } + + def __init__(self, + *, + augment_config: Optional[DecisionAugmentConfig] = None, + augment_seed: int = 42, + **kwargs) -> None: + if kwargs: + raise TypeError(f'Unknown processor arguments: {sorted(kwargs)}') + self.augment_config = augment_config + self._augment_rng = _random.Random(augment_seed) + + def __call__(self, rows): + kept, dropped = [], [] + for row in self.map_col_to_row(rows): + result = self.preprocess(row) + if result is None: + raise ValueError('Decision record has no questions') + kept.extend(result if isinstance(result, list) else [result]) + return kept, dropped + + def _augment(self, state: Any, questions: List[Dict[str, Any]]) -> Tuple[Any, List[Dict[str, Any]]]: + return augment_decision_row(state, questions, self.augment_config, self._augment_rng) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + state, questions = self.parse_record(row) + if not questions: + return None + state, questions = self._augment(state, questions) + return self._build_row(row, state, questions) + + def _build_row(self, row: Dict[str, Any], state: Any, questions: List[Dict[str, Any]]) -> Dict[str, Any]: + """Assemble ONE decision row (the contract in the module docstring) from `(state, questions)`. + + Split out of `preprocess` so a subclass can fan a record out into several rows (JEV renders one + question per prompt) while reusing this row-building logic unchanged. + """ + kinds = [q['kind'] for q in questions] + options = [self._resolve_options(q) for q in questions] + out: Dict[str, Any] = { + 'messages': [{ + 'role': 'user', + 'content': self._state_to_text(state) + }], + 'questions': [q['question'] for q in questions], + 'kinds': kinds, + 'options': options, + } + for key in ('images', 'videos', 'audios'): + if row.get(key) is not None: + out[key] = row[key] + if self._has_gold(questions): + out['target_probs'] = [self.to_target_dist(q['gold'], len(opt)) for q, opt in zip(questions, options)] + return out + + @classmethod + def _normalize_kind(cls, kind: Any) -> str: + key = str(kind).strip().lower() + if key not in cls.KIND_ALIASES: + raise ValueError(f'unknown question kind {kind!r}; expected one of noul/choice/score (or an alias).') + return cls.KIND_ALIASES[key] + + @staticmethod + def _question_items(questions_raw: Any) -> List[Tuple[Any, Any]]: + if isinstance(questions_raw, dict): + return list(questions_raw.items()) + if isinstance(questions_raw, (list, tuple)): + return list(enumerate(questions_raw)) + return [] + + @staticmethod + def _state_to_text(state: Any) -> str: + """Coerce the ``state`` value to a plain string for the message ``content``. + """ + if isinstance(state, str): + return state + if state is None: + return '' + if isinstance(state, dict): + parts = [v for v in state.values() if isinstance(v, str) and v.strip()] + return '\n'.join(parts) + return str(state) + + def parse_record(self, row: Dict[str, Any]) -> Tuple[Any, List[Dict[str, Any]]]: + """Return `(state, questions)`; each question is `{'kind','question','options','gold'}`. + + `options` may be omitted for noul/score (defaults are filled by `_resolve_options`); `gold` + may be None at inference. Must be implemented by each model subclass. + """ + raise NotImplementedError + + def _resolve_options(self, question: Dict[str, Any]) -> List[str]: + options = question.get('options') + if options: + return list(options) + default = self.default_options(question['kind']) + if default is None: + raise ValueError(f"A 'choice' question must provide its options: {question}") + return default + + @classmethod + def default_options(cls, kind: str) -> Optional[List[str]]: + if kind == 'noul': + return list(cls.NOUL_OPTIONS) + if kind == 'score': + return list(cls.SCORE_OPTIONS) + return None + + @staticmethod + def _has_gold(questions: Sequence[Dict[str, Any]]) -> bool: + return any(q.get('gold') is not None for q in questions) + + @staticmethod + def to_target_dist(gold: Any, n_options: int) -> List[float]: + """Gold -> a distribution over the question's own `n_options` options. + + An int gold is the correct option index (one-hot); a sequence gold is already a full + distribution (JEV distillation), and must match `n_options`. + """ + if isinstance(gold, (list, tuple)): + probs = [float(x) for x in gold] + if len(probs) != n_options: + raise ValueError(f'target distribution length {len(probs)} != n_options {n_options}') + return probs + idx = int(gold) + if not 0 <= idx < n_options: + raise ValueError(f'gold index {idx} out of range for n_options {n_options}') + dist = [0.0] * n_options + dist[idx] = 1.0 + return dist + + @staticmethod + def _first(source: Dict[str, Any], keys: Sequence[str], default: Any = None) -> Any: + """Return the first non-None value among `keys` (a permissive column lookup shared by the + model subclasses, whose datasets use differing column names). + """ + for key in keys: + value = source.get(key) + if value is not None: + return value + return default + + @staticmethod + def _as_option_dict(option: Any) -> Dict[str, Any]: + """Normalize ONE unified option element to a dict.""" + if isinstance(option, str): + return {'text': option} + if isinstance(option, dict): + return option + raise ValueError(f'an option must be a string or an object, got {option!r}') + + @classmethod + def _option_text(cls, option: Any) -> str: + """The display text of a unified option: its `text`, else its `key`.""" + d = cls._as_option_dict(option) + text = d.get('text') + return str(text if text is not None else d.get('key', '')) + + @classmethod + def _option_key(cls, option: Any, index: int) -> str: + """The answer key of a unified option: its `key`, else its `text`, else `option_{index}`.""" + d = cls._as_option_dict(option) + key = d.get('key') or d.get('text') + return str(key) if key is not None else f'option_{index}' + + +class JevPreprocessor(ScoringPreprocessor): + """JEV distill-corpus preprocessor: fan a record out into ONE row per question. + + JEV's System-1 pass scores a single question per forward, so a record + carrying several questions becomes several rows (`RowPreprocessor.batched_preprocess` accepts a + list return), each rendering one bare prompt and carrying `counts=[1]` after the + collator (one CompletionRequest per question). + + `parse_record` maps the corpus columns onto the shared question dict `{kind, question, options, + gold}`. `gold` is the teacher's soft distribution over the question's OWN options (JEV + distillation) or an option index; `to_target_dist` validates its length against the option count. + + NOTE(schema): `SargeDev/jev-distill-corpus-v3`'s exact column names could not be verified offline + (HF page/API unreachable in this environment), so the field lookup below is intentionally + permissive (it tries the natural names plus common aliases). It MUST be re-checked + against the real dataset before a training run; a `columns={...}` remap can be passed to + `__init__` to rename the corpus columns onto the canonical names used here. + """ + + KIND_KEYS = ('kind', 'type', 'question_type') + QUESTION_KEYS = ('question', 'prompt', 'query') + OPTION_KEYS = ('options', 'opts', 'choices') + STATE_KEYS = ('state', 'context', 'passage', 'text') + GOLD_KEYS = ('target_probs', 'target', 'probs', 'dist', 'distribution', 'gold', 'label', 'answer') + # Optional: a record that already packs a list of question dicts under one column. + QUESTIONS_KEYS = ('questions', ) + + def preprocess(self, row: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]: + state, questions = self.parse_record(row) + if not questions: + return None + state, questions = self._augment(state, questions) + # One row per question (JEV renders a single question per prompt). + return [self._build_row(row, state, [q]) for q in questions] + + def _resolve_options(self, question: Dict[str, Any]) -> List[str]: + """Force the canonical option set for noul/score; only `choice` uses caller-supplied options. + + JEV's System-1 head scores a FIXED per-kind verbalizer: noul always reads the two tokens for + ['false','true'] and score always reads the six digit tokens ['0'..'5'], regardless of any + option text a record carries. The canonical names force + `['false','true']` / `[str(i) for i in range(6)]` and ignore caller options for those kinds + (only `choice` consumes them). + """ + canonical = self.default_options(question['kind']) + if canonical is not None: # noul / score + return canonical + return super()._resolve_options(question) # choice: caller options are required + + def parse_record(self, row: Dict[str, Any]) -> Tuple[Any, List[Dict[str, Any]]]: + state = self._first(row, self.STATE_KEYS, default='') + return state, self._parse_questions(row) + + def _parse_questions(self, row: Dict[str, Any]) -> List[Dict[str, Any]]: + packed = self._first(row, self.QUESTIONS_KEYS, default=None) + if isinstance(packed, (list, tuple)) and packed and isinstance(packed[0], dict): + return [self._normalize_question(q) for q in packed] + kind = self._first(row, self.KIND_KEYS, default=None) + question = self._first(row, self.QUESTION_KEYS, default=None) + if kind is None or question is None: + return [] + return [self._normalize_question(row)] + + def _normalize_question(self, source: Dict[str, Any]) -> Dict[str, Any]: + kind = self._first(source, self.KIND_KEYS, default=None) + question = self._first(source, self.QUESTION_KEYS, default=None) + if kind is None or question is None: + raise ValueError(f'JEV question needs a kind and a question text; got {source!r}') + out: Dict[str, Any] = {'kind': self._normalize_kind(kind), 'question': question} + options = self._first(source, self.OPTION_KEYS, default=None) + if options is not None: + # Unified schema: an option may be a plain string or an object; JEV renders only the + # display text (no abstain/region), so collapse each element to its text. + out['options'] = [self._option_text(o) for o in options] + gold = self._first(source, self.GOLD_KEYS, default=None) + if gold is not None: + out['gold'] = gold + return out + + +class ClefPreprocessor(ScoringPreprocessor): + """Clef joint-schema preprocessor: ONE record -> ONE row carrying ALL its questions. + + Clef decodes a record's questions jointly (Clef is the only + multi-question-per-forward model); the base`preprocess` builds a single row whose + `kinds` / `questions` / `options` / `target_probs` are parallel lists over the + record's questions, in the original `questions` order. + + A Clef record is `{id?, state, questions: {qid: {type, instructions, criteria}}, }` where + `criteria` is a dict {option_id: description} for `choice`, a list of level descriptions for + `score`, and optional for `noul`. `state` may be ANY JSON value (rendered as-is by the + template), so the raw value is carried in a `state` key and `question_ids` / `criteria` are added + for `DecisionEncoder` to rebuild the byte-exact SCHEMA FIELDS layout. + """ + + STATE_KEYS = ('state', 'context', 'input') + QUESTIONS_KEYS = ('questions', 'fields', 'schema') + KIND_KEYS = ('type', 'kind', 'question_type') + INSTRUCTION_KEYS = ('instructions', 'instruction', 'question', 'prompt', 'text') + CRITERIA_KEYS = ('criteria', 'options', 'choices', 'levels') + GOLD_KEYS = ('answer', 'gold', 'label', 'target', 'value') + ANSWERS_KEYS = ('answers', 'labels', 'golds', 'targets') + ID_KEYS = ('id', 'record_id') + + def parse_record(self, row: Dict[str, Any]) -> Tuple[Any, List[Dict[str, Any]]]: + state = self._first(row, self.STATE_KEYS, default='') + answers = self._first(row, self.ANSWERS_KEYS, default=None) + questions_raw = self._first(row, self.QUESTIONS_KEYS, default=None) + items = self._question_items(questions_raw) + questions: List[Dict[str, Any]] = [] + for qid, q in items: + if not isinstance(q, dict): + raise ValueError(f'Clef question {qid!r} must be a dict of type/instructions/criteria; got {q!r}') + kind = self._normalize_kind(self._first(q, self.KIND_KEYS, default=None)) + criteria = self._normalize_criteria(kind, self._first(q, self.CRITERIA_KEYS, default=None)) + instructions = self._first(q, self.INSTRUCTION_KEYS, default=None) + from twinkle.model.decision.heads import clef_option_ids + option_ids = clef_option_ids(kind, criteria) + gold = self._resolve_gold(q, answers, qid, option_ids) + questions.append({ + 'kind': kind, + 'question': instructions if instructions is not None else str(qid), + 'options': option_ids, + 'question_id': str(qid), + 'criteria': criteria, + }) + if gold is not None: + questions[-1]['gold'] = gold + return state, questions + + def _build_row(self, row: Dict[str, Any], state: Any, questions: List[Dict[str, Any]]) -> Dict[str, Any]: + out = super()._build_row(row, state, questions) + # `state` may be a non-str JSON value; keep the message content a string and carry the raw + # value separately so DecisionEncoder can render it byte-exactly (json for non-str). + out['messages'] = [{'role': 'user', 'content': self._state_to_text(state)}] + if isinstance(state, str): + out['state'] = state + elif state is None: + out['state'] = '' + else: + out['state'] = json.dumps(state, ensure_ascii=False, separators=(',', ':'), sort_keys=True) + out['question_ids'] = [q['question_id'] for q in questions] + out['criteria'] = [ + json.dumps(c, ensure_ascii=False, separators=(',', ':'), sort_keys=True) + if c is not None and not isinstance(c, str) else c for c in [q.get('criteria') for q in questions] + ] + return out + + @staticmethod + def _state_to_text(state: Any) -> str: + if isinstance(state, str): + return state + if state is None: + return '' + return json.dumps(state, ensure_ascii=False, separators=(',', ':'), sort_keys=True) + + @classmethod + def _normalize_criteria(cls, kind: str, raw: Any) -> Any: + """Coerce the unified `options` list into Clef's native `criteria` shape. + + Clef renders byte-exact against the SCHEMA FIELDS layout, which needs a + `{option_id: description}` map for `choice` (columns sorted by id) and a list of level + descriptions for `score`. + """ + if raw is None or isinstance(raw, dict): + return raw # native criteria map (choice) or omitted (noul) + if not isinstance(raw, (list, tuple)): + return raw + if kind == 'score': + return [cls._option_text(o) for o in raw] # level descriptions + # choice (noul ignores criteria): {option_id: description}, ids from key/text/index. + return {cls._option_key(o, i): cls._option_text(o) for i, o in enumerate(raw)} + + def _resolve_gold(self, question: Dict[str, Any], answers: Any, qid: Any, option_ids: List[str]) -> Optional[Any]: + """Resolve the correct answer to an option INDEX (or a full distribution), or None at inference. + + Accepts, at question level then in a record-level `answers` map: a distribution (list/tuple of + floats), a bool (noul), an option id string, or an int index. + """ + gold = self._first(question, self.GOLD_KEYS, default=None) + if gold is None and isinstance(answers, dict): + gold = answers.get(str(qid), answers.get(qid)) + if gold is None: + return None + if isinstance(gold, (list, tuple)): + return [float(x) for x in gold] + if isinstance(gold, bool): + gold = 'true' if gold else 'false' + if isinstance(gold, int): + return gold + text = str(gold).strip() + for index, option_id in enumerate(option_ids): + if text == option_id or text.lower() == option_id.lower(): + return index + if text.lstrip('-').isdigit(): + return int(text) + raise ValueError(f'Clef gold {gold!r} for question {qid!r} matches no option id {option_ids}.') + + +def _omnijev_render_option(o: Dict[str, Any]) -> str: + """How one option block is written into the prompt. An abstain option renders as + 'none of the above'; a region option as its integer box; otherwise the option text + truncated to 200 chars. + """ + if o.get('abstain'): + return 'none of the above' + if 'region' in o: + box = o['region']['box'] + b = [int(round(float(v))) for v in box] + return 'region [{},{},{},{}]'.format(*b) + return str(o.get('text', ''))[:200] + + +def _omnijev_options(q: Dict[str, Any]) -> Tuple[List[Any], List[Dict[str, Any]]]: + """Compute ``(keys, opts)`` for a question: the answer keys and option-block dicts. + + `keys` are the answer keys (noul -> ['yes'], score -> the level labels, choice -> the option + keys); `opts` are the option-BLOCK dicts the template renders with `_omnijev_render_option` + (noul -> ONE block, score -> L blocks, choice -> K blocks, abstain excluded -- it is the head's + extra learned column, not a rendered block). Accepts both the Jev `criteria`-map shape and the + OmniJev `options`-list shape (whose entries may be region options). + """ + kind = q['kind'] + if kind == 'noul': + opts = [{'region': q['region']}] if q.get('region') else [{'text': ''}] + return ['yes'], opts + if kind == 'score': + lv = q.get('levels') or list((q.get('criteria') or {}).keys()) + if not lv and q.get('options') is not None: + # Unified schema: levels given as a plain `options` list (strings or {text/key}). + lv = [o if isinstance(o, str) else (o.get('text') or o.get('key')) for o in q['options']] + return list(lv), [{'text': str(x)} for x in lv] + if q.get('options') is not None: + keys: List[Any] = [] + opts = [] + for i, o in enumerate(q['options']): + if isinstance(o, str): # unified sugar: a plain string is {'text': s} + o = {'text': o} + if o.get('abstain'): + continue + keys.append(o.get('key') or o.get('text') or f'option_{i}') + opts.append(o) + return keys, opts + crit = q.get('criteria') or {} + keys = list(crit.keys()) + return keys, [{'text': k if crit[k] in (None, '') else f'{k}: {crit[k]}'} for k in keys] + + +class OmniJevPreprocessor(ScoringPreprocessor): + """OmniJev preprocessor: fan a record out into ONE row per question (like `JevPreprocessor`). + + OmniJev scores a single question per forward (the K option + rows of decision G1's branch are expanded at forward time, not here), so a record carrying + several questions becomes several rows. Each row carries the option-BLOCK structure + (`_omnijev_options`) plus the width-matched `column_keys`, and the gold as a distribution over + the COLUMN width (noul 2 / choice K+1 incl. the abstain column / score L) -- NOT the block count, + because for noul (1 block -> 2 cols) and choice (K blocks -> K+1 cols) they differ. + + The state is multimodal (`{'images': [...], 'video': {...}}` in the request); the raw + `images` / `videos` are passed through and the panel/mosaic composition (decision H) is applied + by `DecisionEncoder`, not here. The per-question `instructions` become the message text. + """ + + STATE_KEYS = ('state', 'context', 'input') + QUESTIONS_KEYS = ('questions', 'fields') + KIND_KEYS = ('type', 'kind', 'question_type') + INSTRUCTION_KEYS = ('instructions', 'instruction', 'question', 'prompt', 'text') + LEVELS_KEYS = ('levels', ) + CRITERIA_KEYS = ('criteria', ) + OPTIONS_KEYS = ('options', ) + REGION_KEYS = ('region', ) + GOLD_KEYS = ('answer', 'gold', 'label', 'target', 'value') + ANSWERS_KEYS = ('answers', 'labels', 'golds', 'targets') + IMAGES_KEYS = ('images', 'image') + VIDEOS_KEYS = ('videos', 'video') + + def preprocess(self, row: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]: + state, questions = self.parse_record(row) + if not questions: + return None + state, questions = self._augment(state, questions) + return [self._build_row(row, state, [q]) for q in questions] + + def parse_record(self, row: Dict[str, Any]) -> Tuple[Any, List[Dict[str, Any]]]: + state = self._first(row, self.STATE_KEYS, default=None) + answers = self._first(row, self.ANSWERS_KEYS, default=None) + questions_raw = self._first(row, self.QUESTIONS_KEYS, default=None) + questions: List[Dict[str, Any]] = [] + for qid, q in self._question_items(questions_raw): + if not isinstance(q, dict): + raise ValueError(f'OmniJev question {qid!r} must be a dict of type/instructions/...; got {q!r}') + kind = self._normalize_kind(self._first(q, self.KIND_KEYS, default=None)) + instruction = self._first(q, self.INSTRUCTION_KEYS, default=None) + instruction = str(qid) if instruction is None else instruction + src: Dict[str, Any] = {'kind': kind} + for keys, dst in ((self.LEVELS_KEYS, 'levels'), (self.CRITERIA_KEYS, 'criteria'), + (self.OPTIONS_KEYS, 'options'), (self.REGION_KEYS, 'region')): + value = self._first(q, keys, default=None) + if value is not None: + src[dst] = value + keys, opts = _omnijev_options(src) + from twinkle.model.decision.heads import omnijev_column_keys + column_keys = omnijev_column_keys(kind, keys) + gold = self._resolve_gold(q, answers, qid, kind, column_keys) + item: Dict[str, Any] = { + 'kind': kind, + 'question': instruction, + 'options': [_omnijev_render_option(o) for o in opts], + 'column_keys': column_keys, + } + if gold is not None: + item['gold'] = gold + questions.append(item) + return state, questions + + def _build_row(self, row: Dict[str, Any], state: Any, questions: List[Dict[str, Any]]) -> Dict[str, Any]: + q = questions[0] + column_keys = q['column_keys'] + out: Dict[str, Any] = { + 'messages': [{ + 'role': 'user', + 'content': self._content_text(state, q['question']) + }], + 'questions': [q['question']], + 'kinds': [q['kind']], + 'options': [q['options']], + 'omnijev_column_keys': [column_keys], + } + images, video = self._resolve_media(row, state) + if images is not None: + out['images'] = images + if video is not None: + # The `video` DESCRIPTOR dict, NOT a raw video file: route it to extra_kwargs so + # DecisionEncoder's `video_content` can crop the mosaic back into frames. Putting it in + # `videos` would treat it as a clip list and prepend a spurious `