"""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\n\n\n\n' if not rendered.endswith(tail): raise ValueError('Unexpected generation prompt boundary: ' + repr(rendered[-80:])) # The answer starts exactly at the next token; letter ids must be single tokens. 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: # Separates two effects: what the LoRA did to the backbone, and what the pointer readout adds. # The pointer head is deliberately not loaded here. 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: # A frozen training mixture: its manifest records artifacts rather than a file list. 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: # Batching across records is what makes a whole training mixture affordable: one prompt per forward # pass costs hours, eight at a time costs minutes. Prompts are sorted by length so padding is small. 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: # Training records carry no benchmark suite; that field exists only on evaluation data. 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: # A prompt too long for the likelihood readout still deserves the cheap anchor. 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()