| """Native MiniCPM answer readout on the fixed development set, without our pointer head. |
| |
| Purpose: separate what the pinned foundation already knows from what the trained |
| pointer/LoRA adds. The state, instructions and option texts are rendered by the same |
| vendored kev.api.to_record path the pointer uses, so only the readout differs. |
| |
| Two readouts are measured on the identical prompts-per-question: |
| letter restricted next-token distribution over single-token option letters (<=26 options) |
| likelihood conditional log probability of each rendered option text, summed and length normalised |
| |
| No training, no new API calls, no release holdout. The pinned chat template is required: |
| the asset set shipped without chat_template.jinja, so the file of the pinned revision is |
| supplied separately and hash checked here. |
| """ |
| import argparse |
| import collections |
| import gzip |
| import json |
| import re |
| import sys |
| import time |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT / 'vendor/kev-pinned')) |
| from decision_data import api_payload, file_digest, gold_key, question_keys, read_jsonl |
| from decision_metrics import metric |
| from kev.api import SystemOneRequest, to_record |
| from train_minicpm_pilot import atomic |
|
|
| CONTROL = re.compile(r'<\|([A-Za-z0-9_]+)\|>') |
| LETTERS = [chr(ord('A') + i) for i in range(26)] |
| SYSTEM = ('You are a precise decision engine. Read the state and answer the question ' |
| 'using only the options that are given. Do not explain.') |
|
|
|
|
| def safe(text): |
| """Caller text must never be able to forge chat control tokens.""" |
| return CONTROL.sub(r'<¦\1¦>', text) |
|
|
|
|
| def letter_prompt(state, instr, options): |
| lines = [LETTERS[i] + ') ' + o for i, o in enumerate(options)] |
| return ('State:\n' + state + '\n\nQuestion: ' + instr + '\n\nOptions:\n' + '\n'.join(lines) |
| + '\n\nReply with only the letter of the correct option.') |
|
|
|
|
| def text_prompt(state, instr, options): |
| lines = ['- ' + o for o in options] |
| return ('State:\n' + state + '\n\nQuestion: ' + instr + '\n\nOptions:\n' + '\n'.join(lines) |
| + '\n\nReply with only the text of the correct option.') |
|
|
|
|
| def build(tok, user): |
| messages = [{'role': 'system', 'content': SYSTEM}, {'role': 'user', 'content': safe(user)}] |
| return tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False) |
|
|
|
|
| def questions_of(record): |
| """Flatten one record into independent native questions, keeping the pointer's option order.""" |
| wire = api_payload(record) |
| rec, meta = to_record(SystemOneRequest.model_validate(wire)) |
| out = [] |
| for qid, packed, m in zip(record['questions'], rec['questions'], meta): |
| q = record['questions'][qid] |
| out.append({'qid': qid, 'type': q['type'], 'state': rec['state'], 'instr': packed['instr'], |
| 'options': packed['options'], 'keys': question_keys(q), |
| 'label': question_keys(q).index(gold_key(q))}) |
| return out |
|
|
|
|
| def anchor_readout_for(options, mode, max_options): |
| """Which frozen readout should teach a question with this many options.""" |
| if mode == 'letter': |
| return 'letter' |
| if mode == 'likelihood': |
| return 'likelihood' |
| return 'likelihood' if options <= max_options else 'letter' |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--tag', default='minicpm5_native_readout_v1') |
| ap.add_argument('--assets', default='/workspace/minicpm5-2b-assets') |
| ap.add_argument('--assets-report', default='minicpm5_assets_campaign3', |
| help='Campaign report folder holding the hash verification of those assets') |
| ap.add_argument('--chat-template', default=str(ROOT / 'assets/chat_template.jinja')) |
| ap.add_argument('--chat-template-sha256', default='cc945752db555d60949b16989df4ccfeb52a313d6b4b5c5229dd786e2e9fcf1c') |
| ap.add_argument('--data', default='data/minicpm_pilot_v1/dev.jsonl.gz', |
| help='One jsonl(.gz) file, or a directory with a hash manifest of files') |
| ap.add_argument('--output-name', default='native_dev.json.gz') |
| ap.add_argument('--max-tokens', type=int, default=3072) |
| ap.add_argument('--batch', type=int, default=2) |
| ap.add_argument('--limit', type=int, default=0) |
| ap.add_argument('--sources', default='', |
| help='Comma separated source prefixes to keep, for anchoring only what must not drift') |
| ap.add_argument('--sample', type=int, default=0, |
| help='Random subset of records; anchoring is a regulariser and does not need full coverage') |
| ap.add_argument('--sample-seed', type=int, default=17) |
| ap.add_argument('--precision', choices=['fp32', 'bf16'], default='fp32', |
| help='bf16 halves the cost of an anchoring pass; measured numbers stay FP32') |
| ap.add_argument('--reserve-seconds', type=float, default=180) |
| ap.add_argument('--anchor-out', default='', |
| help='Write {targets: {record id: {question id: {option key: p}}}} for the anchoring loss') |
| ap.add_argument('--letter-only', action='store_true', |
| help='Skip the option-likelihood readouts; one forward pass per question') |
| ap.add_argument('--anchor-readout', choices=['letter', 'likelihood', 'auto'], default='letter', |
| help='Which frozen voice the anchor should carry. On knowledge questions the letter ' |
| 'readout is the weaker one: mmlu 52.5 percent by letter against 62.5 percent by ' |
| 'option likelihood, and anchoring to the weaker teacher is how v19 ended 16 points ' |
| 'below its own backbone there. "auto" takes the likelihood where the option count ' |
| 'is small and the cheap letter elsewhere.') |
| ap.add_argument('--anchor-max-options', type=int, default=6, |
| help='Option count up to which "auto" pays for the likelihood readout') |
| ap.add_argument('--adapter-checkpoint', default='', |
| help='Optional trained LoRA checkpoint directory, applied to the backbone without the pointer head') |
| args = ap.parse_args() |
| if not args.tag.replace('_', '').isalnum(): |
| raise ValueError('Invalid run') |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| template = Path(args.chat_template) |
| if file_digest(template) != args.chat_template_sha256: |
| raise ValueError('Chat template does not match the pinned revision digest') |
| assets = Path(args.assets) |
| if not args.assets_report.replace('_', '').isalnum(): |
| raise ValueError('Invalid assets report') |
| verified = json.loads((ROOT / 'reports/campaign' / args.assets_report / 'result.json').read_text()) |
| assert verified['status'] == 'HASH_VERIFIED' |
| deadline = json.loads(Path('/workspace/open-jev-watchdog.json').read_text())['deadline'] |
| assert time.time() + args.reserve_seconds < deadline |
| assert torch.cuda.device_count() == 1 and '4090' in torch.cuda.get_device_name(0) |
| torch.set_num_threads(4) |
| torch.manual_seed(17) |
| torch.backends.cuda.matmul.allow_tf32 = False |
| torch.backends.cudnn.allow_tf32 = False |
|
|
| tok = AutoTokenizer.from_pretrained(str(assets), local_files_only=True, trust_remote_code=False) |
| tok.chat_template = template.read_text() |
| rendered = build(tok, 'probe') |
| tail = '<|im_start|>assistant\n<think>\n\n</think>\n\n' |
| if not rendered.endswith(tail): |
| raise ValueError('Unexpected generation prompt boundary: ' + repr(rendered[-80:])) |
| |
| letter_ids, letter_variant = {}, None |
| for variant in ['', ' ']: |
| ids = {l: tok(variant + l, add_special_tokens=False).input_ids for l in LETTERS} |
| if all(len(v) == 1 for v in ids.values()) and len({v[0] for v in ids.values()}) == len(LETTERS): |
| letter_ids = {l: v[0] for l, v in ids.items()} |
| letter_variant = 'bare' if variant == '' else 'leading_space' |
| break |
| if not letter_ids: |
| raise ValueError('No single-token letter labelling available') |
|
|
| model = AutoModelForCausalLM.from_pretrained(str(assets), local_files_only=True, trust_remote_code=False, |
| dtype=torch.bfloat16 if args.precision == 'bf16' else torch.float32, |
| attn_implementation='sdpa') |
| adapter = {} |
| if args.adapter_checkpoint: |
| |
| |
| from peft import LoraConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict |
| from safetensors.torch import load_file |
| checkpoint = Path(args.adapter_checkpoint) |
| config = json.loads((checkpoint / 'adapter_config.json').read_text()) |
| if config['peft_type'] != 'LORA' or config['task_type'] != 'FEATURE_EXTRACTION': |
| raise ValueError('Unexpected adapter configuration') |
| model.model = get_peft_model(model.model, LoraConfig( |
| task_type='FEATURE_EXTRACTION', r=config['r'], lora_alpha=config['lora_alpha'], |
| lora_dropout=config['lora_dropout'], target_modules=config['target_modules'], |
| trainable_token_indices=config['trainable_token_indices'])) |
| saved = load_file(str(checkpoint / 'adapter.safetensors')) |
| set_peft_model_state_dict(model.model, dict(saved)) |
| actual = get_peft_model_state_dict(model.model, save_embedding_layers=False) |
| if set(saved) != set(actual) or not all(torch.equal(v.cpu(), saved[k]) for k, v in actual.items()): |
| raise ValueError('Adapter state did not round-trip into the native backbone') |
| adapter = {'checkpoint': str(checkpoint), 'tensor_keys': len(saved), |
| 'sha256': file_digest(checkpoint / 'adapter.safetensors'), 'pointer_head_loaded': False} |
| model.to('cuda') |
| model.config.use_cache = False |
| model.eval() |
|
|
| source = ROOT / args.data |
| if source.is_dir(): |
| manifest = json.loads((source / 'manifest.json').read_text()) |
| records = [] |
| if 'files' in manifest: |
| for item in manifest['files']: |
| path = source / item['name'] |
| if file_digest(path) != item['sha256']: |
| raise ValueError('Evaluation data digest mismatch: ' + item['name']) |
| records.extend(read_jsonl(path)) |
| else: |
| |
| for split in ['train']: |
| path = source / (split + '.jsonl.gz') |
| if file_digest(path) != manifest['artifacts'][split]['sha256']: |
| raise ValueError('Training data digest mismatch: ' + split) |
| records.extend(read_jsonl(path)) |
| else: |
| records = read_jsonl(source) |
| if args.sources: |
| wanted = tuple(x for x in args.sources.split(',') if x) |
| records = [r for r in records if r['_meta']['source'].startswith(wanted)] |
| if not records: |
| raise ValueError('No record matches the requested sources') |
| if args.sample and args.sample < len(records): |
| import random as _random |
| records = _random.Random(args.sample_seed).sample(records, args.sample) |
| records.sort(key=lambda r: r['_meta']['id']) |
| if args.limit: |
| records = records[:args.limit] |
| report = ROOT / 'reports/campaign' / args.tag |
| report.mkdir(parents=True, exist_ok=True) |
| raw_dir = ROOT / 'artifacts/evaluation' / args.tag |
| raw_dir.mkdir(parents=True, exist_ok=True) |
| atomic(report / 'progress.json', {'status': 'LOADED', 'done': 0, 'total': len(records), 'epoch': time.time()}) |
|
|
| def logits_last(prompt): |
| ids = tok(prompt, add_special_tokens=False, return_tensors='pt').input_ids.to('cuda') |
| if ids.shape[1] > args.max_tokens: |
| return None, ids.shape[1] |
| with torch.no_grad(): |
| out = model(input_ids=ids, use_cache=False, logits_to_keep=1) |
| return out.logits[0, -1].float(), ids.shape[1] |
|
|
| def logits_last_batch(prompts): |
| """Final-position logits for several prompts at once. |
| |
| One prompt per forward pass costs about two hours over a full training mixture, which is why the |
| anchoring targets could not be produced in the lease that needed them. Prompts are left padded so the |
| last position is the real final token for every row, and explicit position ids keep the padding from |
| shifting the sequence. |
| """ |
| encoded = [tok(prompt, add_special_tokens=False).input_ids for prompt in prompts] |
| widest = max(len(ids) for ids in encoded) |
| if widest > args.max_tokens: |
| return None, widest |
| pad = tok.pad_token_id if tok.pad_token_id is not None else 0 |
| ids = torch.full((len(encoded), widest), pad, dtype=torch.long) |
| mask = torch.zeros((len(encoded), widest), dtype=torch.long) |
| for row, sequence in enumerate(encoded): |
| ids[row, widest - len(sequence):] = torch.tensor(sequence) |
| mask[row, widest - len(sequence):] = 1 |
| positions = (mask.cumsum(-1) - 1).clamp_min(0) |
| ids, mask, positions = ids.to('cuda'), mask.to('cuda'), positions.to('cuda') |
| with torch.no_grad(): |
| out = model(input_ids=ids, attention_mask=mask, position_ids=positions, |
| use_cache=False, logits_to_keep=1) |
| return out.logits[:, -1].float(), [len(sequence) for sequence in encoded] |
|
|
| def option_logprobs(prompt, options): |
| """Score every option under the same prefix. |
| |
| Options of equal token length are batched without padding, and only the logits of |
| the scored positions are materialised: the full vocabulary over a padded batch does |
| not fit next to FP32 weights on one 24GB card. |
| """ |
| prefix = tok(prompt, add_special_tokens=False).input_ids |
| pieces = [tok(o, add_special_tokens=False).input_ids for o in options] |
| widest = len(prefix) + max(len(p) for p in pieces) |
| if widest > args.max_tokens: |
| return None, None, widest |
| totals = [None] * len(pieces) |
| by_length = {} |
| for i, piece in enumerate(pieces): |
| by_length.setdefault(len(piece), []).append(i) |
| prefix_tensor = torch.tensor(prefix) |
| for length, indices in by_length.items(): |
| for start in range(0, len(indices), args.batch): |
| chunk = indices[start:start + args.batch] |
| ids = torch.stack([torch.cat([prefix_tensor, torch.tensor(pieces[i])]) for i in chunk]).to('cuda') |
| with torch.no_grad(): |
| logits = model(input_ids=ids, use_cache=False, logits_to_keep=length + 1).logits.float() |
| logprobs = torch.log_softmax(logits, dim=-1) |
| for row, i in enumerate(chunk): |
| piece = pieces[i] |
| totals[i] = float(sum(logprobs[row, j, token] for j, token in enumerate(piece))) |
| means = [t / len(p) for t, p in zip(totals, pieces)] |
| return totals, means, len(prefix) |
|
|
| def distribution(scores): |
| top = max(scores) |
| weights = [pow(2.718281828459045, s - top) for s in scores] |
| total = sum(weights) |
| return [w / total for w in weights] |
|
|
| anchor_targets = {} if args.anchor_out else None |
| anchor_used = collections.Counter() |
| rows, skipped, started = [], [], time.time() |
| precomputed = None |
| if args.letter_only and args.batch > 1: |
| |
| |
| precomputed = {} |
| jobs = [] |
| for record in records: |
| items = questions_of(record) |
| precomputed[record['_meta']['id']] = items |
| for item in items: |
| if len(item['options']) <= len(LETTERS): |
| jobs.append((record['_meta']['id'], item, |
| build(tok, letter_prompt(item['state'], item['instr'], item['options'])))) |
| jobs.sort(key=lambda job: len(job[2])) |
| atomic(report / 'progress.json', {'status': 'BATCHED_READOUT', 'done': 0, 'total': len(jobs), |
| 'records': len(records), 'epoch': time.time()}) |
| for start in range(0, len(jobs), args.batch): |
| if time.time() + args.reserve_seconds >= deadline: |
| skipped.append({'reason': 'DEADLINE_RESERVE', 'remaining_questions': len(jobs) - start}) |
| break |
| chunk = jobs[start:start + args.batch] |
| batched, widths = logits_last_batch([job[2] for job in chunk]) |
| if batched is None: |
| skipped.append({'reason': 'LETTER_PROMPT_TOO_LONG', 'tokens': widths}) |
| continue |
| for (_, item, _), logits, width in zip(chunk, batched, widths): |
| selected = [letter_ids[LETTERS[i]] for i in range(len(item['options']))] |
| p = torch.softmax(logits[selected], dim=-1).cpu().tolist() |
| item['letter_result'] = {'p': p, 'prediction': max(range(len(p)), key=p.__getitem__), |
| 'prompt_tokens': width, |
| 'mass_on_letters': float(torch.softmax(logits, dim=-1)[selected].sum())} |
| if (start // args.batch) % 50 == 0: |
| atomic(report / 'progress.json', {'status': 'BATCHED_READOUT', 'done': start + len(chunk), |
| 'total': len(jobs), 'elapsed_seconds': time.time() - started, |
| 'epoch': time.time()}) |
| for index, record in enumerate(records): |
| if time.time() + args.reserve_seconds >= deadline: |
| skipped.append({'record': record['_meta']['id'], 'reason': 'DEADLINE_RESERVE'}) |
| break |
| meta = record['_meta'] |
| items = precomputed.get(meta['id']) if precomputed is not None else None |
| if items is None: |
| items = questions_of(record) |
| for item in items: |
| |
| base = {'id': meta['id'], 'group_id': meta['group_id'], 'qid': item['qid'], |
| 'suite': meta.get('benchmark_suite', meta.get('benchmark_split', 'train')), |
| 'variant': meta.get('variant', 'clean'), |
| 'source': meta['source'], 'type': item['type'], 'label': item['label'], |
| 'options': len(item['options'])} |
| row = dict(base) |
| if 'letter_result' in item: |
| row['letter'] = item['letter_result'] |
| elif len(item['options']) <= len(LETTERS): |
| logits, width = logits_last(build(tok, letter_prompt(item['state'], item['instr'], item['options']))) |
| if logits is None: |
| skipped.append({'id': meta['id'], 'qid': item['qid'], 'reason': 'LETTER_PROMPT_TOO_LONG', 'tokens': width}) |
| else: |
| selected = [letter_ids[LETTERS[i]] for i in range(len(item['options']))] |
| p = torch.softmax(logits[selected], dim=-1).cpu().tolist() |
| row['letter'] = {'p': p, 'prediction': max(range(len(p)), key=p.__getitem__), |
| 'prompt_tokens': width, |
| 'mass_on_letters': float(torch.softmax(logits, dim=-1)[selected].sum())} |
| wanted = anchor_readout_for(len(item['options']), args.anchor_readout, args.anchor_max_options) |
| if args.letter_only or (anchor_targets is not None and wanted == 'letter' |
| and args.anchor_readout != 'letter'): |
| rows.append(row) |
| if 'letter' in row and anchor_targets is not None: |
| anchor_targets.setdefault(meta['id'], {})[item['qid']] = dict( |
| zip(item['keys'], row['letter']['p'])) |
| anchor_used['letter'] += 1 |
| continue |
| sums, means, width = option_logprobs(build(tok, text_prompt(item['state'], item['instr'], item['options'])), |
| item['options']) |
| if sums is None: |
| skipped.append({'id': meta['id'], 'qid': item['qid'], 'reason': 'TEXT_PROMPT_TOO_LONG', 'tokens': width}) |
| else: |
| psum, pmean = distribution(sums), distribution(means) |
| row['likelihood_sum'] = {'p': psum, 'prediction': max(range(len(psum)), key=psum.__getitem__)} |
| row['likelihood_mean'] = {'p': pmean, 'prediction': max(range(len(pmean)), key=pmean.__getitem__)} |
| row['prefix_tokens'] = width |
| if anchor_targets is not None and wanted == 'likelihood': |
| anchor_targets.setdefault(meta['id'], {})[item['qid']] = dict(zip(item['keys'], psum)) |
| anchor_used['likelihood'] += 1 |
| if anchor_targets is not None and wanted == 'likelihood' and 'likelihood_sum' not in row \ |
| and 'letter' in row: |
| |
| anchor_targets.setdefault(meta['id'], {})[item['qid']] = dict(zip(item['keys'], row['letter']['p'])) |
| anchor_used['letter_fallback'] += 1 |
| rows.append(row) |
| if (index + 1) % 20 == 0: |
| atomic(report / 'progress.json', {'status': 'RUNNING', 'done': index + 1, 'total': len(records), |
| 'questions': len(rows), 'elapsed_seconds': time.time() - started, |
| 'epoch': time.time()}) |
|
|
| def view(name): |
| out = [] |
| for r in rows: |
| if name in r and r['variant'] == 'clean': |
| out.append({'p': r[name]['p'], 'prediction': r[name]['prediction'], 'label': r['label'], |
| 'type': r['type'], 'suite': r['suite'], 'source': r['source']}) |
| return out |
|
|
| readouts = {} |
| for name in ['letter', 'likelihood_sum', 'likelihood_mean']: |
| selected = view(name) |
| readouts[name] = {'overall': metric(selected), |
| 'suites': {s: metric([r for r in selected if r['suite'] == s]) |
| for s in sorted({r['suite'] for r in selected})}, |
| 'types': {t: metric([r for r in selected if r['type'] == t]) |
| for t in sorted({r['type'] for r in selected})}} |
| common = [r for r in rows if r['variant'] == 'clean' and 'letter' in r and 'likelihood_sum' in r] |
| for name in ['letter', 'likelihood_sum', 'likelihood_mean']: |
| readouts[name]['common_subset'] = metric([{'p': r[name]['p'], 'prediction': r[name]['prediction'], |
| 'label': r['label'], 'type': r['type']} for r in common]) |
|
|
| if anchor_targets is not None: |
| anchor_path = ROOT / args.anchor_out |
| anchor_path.parent.mkdir(parents=True, exist_ok=True) |
| anchored_questions = sum(len(v) for v in anchor_targets.values()) |
| anchor_path.write_text(json.dumps({ |
| 'targets': anchor_targets, 'readout': args.anchor_readout, |
| 'readout_counts': dict(anchor_used), 'anchor_max_options': args.anchor_max_options, |
| 'model': 'openbmb/MiniCPM5-2B', |
| 'pointer_head': False, 'adapter': bool(args.adapter_checkpoint), |
| 'chat_template_sha256': args.chat_template_sha256, 'data': args.data, |
| 'records': len(anchor_targets), 'questions': anchored_questions, |
| 'skipped_questions_over_letter_limit': len([s for s in skipped if 'LETTER' in s.get('reason', '')]), |
| 'note': 'Frozen zero-shot distributions of the base model, for KL anchoring during training. ' |
| 'Questions with more options than single-token letters are absent by design.', |
| 'epoch': time.time()}, separators=(',', ':'))) |
| print(json.dumps({'anchor_records': len(anchor_targets), 'anchor_questions': anchored_questions, |
| 'path': str(anchor_path)}), flush=True) |
| raw = raw_dir / args.output_name |
| raw.write_bytes(gzip.compress(json.dumps(rows, separators=(',', ':')).encode(), mtime=0)) |
| result = {'status': 'MEASURED' if not skipped else 'MEASURED_WITH_SKIPS', |
| 'model': 'openbmb/MiniCPM5-2B', 'readout': 'native_answer_tokens', 'pointer_head': False, |
| 'new_training': False, 'release_holdout_opened': False, |
| 'data': args.data, 'records': len(records), 'questions': len(rows), |
| 'common_subset_questions': len(common), 'letter_variant': letter_variant, |
| 'chat_template_sha256': args.chat_template_sha256, 'assets_manifest_sha256': verified['manifest_sha256'], |
| 'sources_filter': args.sources, 'sampled_records': args.sample, 'sample_seed': args.sample_seed, |
| 'generation_prompt_tail': tail, 'precision': args.precision.upper(), 'adapter': adapter or None, |
| 'readouts': readouts, 'skipped': skipped, |
| 'prediction_sha256': file_digest(raw), 'elapsed_seconds': time.time() - started, 'epoch': time.time()} |
| atomic(report / 'result.json', result) |
| atomic(report / 'progress.json', {'status': result['status'], 'done': len(records), 'total': len(records), |
| 'questions': len(rows), 'epoch': time.time()}) |
| print(json.dumps({'status': result['status'], 'questions': result['questions'], |
| 'letter': readouts['letter']['overall'], 'likelihood_mean': readouts['likelihood_mean']['overall']}, |
| indent=2), flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
