| """Compare teacher-forced full-vocabulary distributions, never top-k approximations.""" |
| import argparse |
| import json |
| import math |
| from pathlib import Path |
| import numpy as np |
|
|
|
|
| def log_probs(logits): |
| values = np.asarray(logits, dtype=np.float64) |
| values -= values.max(axis=-1, keepdims=True) |
| return values - np.log(np.exp(values).sum(axis=-1, keepdims=True)) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--reference', type=Path, required=True) |
| parser.add_argument('--candidate', type=Path, required=True) |
| parser.add_argument('--output', type=Path, required=True) |
| parser.add_argument('--assistant-masks', type=Path, help='Verified assistant-target masks; omit for whole-sequence diagnostics') |
| args = parser.parse_args() |
| mask_info = json.loads(args.assistant_masks.read_text()) if args.assistant_masks else None |
| for directory in (args.reference, args.candidate): |
| metadata = json.loads((directory / 'metadata.json').read_text()) |
| if metadata.get('status') != 'complete': |
| raise ValueError(f'Incomplete inference artifacts: {directory}') |
| if metadata.get('generation_max_new_tokens', 0): |
| raise ValueError('Free-generation targets are not heldout teacher-forced data') |
| reference = {r['id']: r for r in map(json.loads, (args.reference / 'records.jsonl').read_text().splitlines())} |
| candidates = list(map(json.loads, (args.candidate / 'records.jsonl').read_text().splitlines())) |
| if not candidates or len({r['id'] for r in candidates}) != len(candidates): |
| raise ValueError('Missing or duplicate candidate records') |
| results = [] |
| token_kl_by_split = {} |
| for candidate in candidates: |
| ref = reference[candidate['id']] |
| if ref['token_ids'] != candidate['token_ids'] or ref['split'] != candidate['split']: |
| raise ValueError(f'Misaligned inputs: {candidate["id"]}') |
| a = np.load(args.reference / ref['logits_file'], mmap_mode='r') |
| b = np.load(args.candidate / candidate['logits_file'], mmap_mode='r') |
| targets = np.asarray(ref['token_ids'][1:], dtype=np.int64) |
| if a.shape != b.shape or a.ndim != 2 or a.shape[0] != len(targets): |
| raise ValueError(f'Misaligned logits: {candidate["id"]}: {a.shape}, {b.shape}') |
| target_mask = np.ones(len(targets), dtype=bool) |
| domain = None |
| if mask_info: |
| mask_record = mask_info['records'][candidate['id']] |
| if mask_record['token_ids'] != ref['token_ids'] or mask_record['split'] != ref['split']: |
| raise ValueError('Assistant mask belongs to different input tokens or split') |
| target_mask = np.asarray(mask_record['target_mask'], dtype=bool) |
| domain = mask_record['domain'] |
| if target_mask.shape != targets.shape or not target_mask.any(): |
| raise ValueError('Missing or misaligned assistant targets') |
| kl, ref_loss, quant_loss, agreements = [], [], [], [] |
| for start in range(0, len(targets), 16): |
| end = min(start + 16, len(targets)) |
| keep = target_mask[start:end] |
| if not keep.any(): |
| continue |
| lp, lq = log_probs(a[start:end][keep]), log_probs(b[start:end][keep]) |
| if not np.isfinite(lp).all() or not np.isfinite(lq).all(): |
| raise ValueError(f'Nonfinite logits: {candidate["id"]}') |
| kl.extend(np.sum(np.exp(lp) * (lp - lq), axis=-1).tolist()) |
| row = np.arange(int(keep.sum())) |
| selected_targets = targets[start:end][keep] |
| ref_loss.extend((-lp[row, selected_targets]).tolist()) |
| quant_loss.extend((-lq[row, selected_targets]).tolist()) |
| agreements.extend((lp.argmax(axis=-1) == lq.argmax(axis=-1)).tolist()) |
| results.append({'id': candidate['id'], 'split': candidate['split'], 'domain': domain, 'tokens': len(kl), 'kl_sum': sum(kl), 'kl_mean': float(np.mean(kl)), 'kl_p95': float(np.percentile(kl, 95)), 'reference_nll_sum': sum(ref_loss), 'candidate_nll_sum': sum(quant_loss), 'top1_agreements': sum(agreements)}) |
| token_kl_by_split.setdefault(candidate['split'], []).extend(kl) |
| if domain: |
| token_kl_by_split.setdefault(candidate['split'] + '/' + domain, []).extend(kl) |
| results[-1].update({'kl_p99_9': float(np.percentile(kl, 99.9)), 'kl_max': max(kl)}) |
| summaries = {} |
| for split in sorted(token_kl_by_split): |
| rows = [r for r in results if r['split'] == split or (r['domain'] and r['split'] + '/' + r['domain'] == split)] |
| n = sum(r['tokens'] for r in rows) |
| ref_nll = sum(r['reference_nll_sum'] for r in rows) / n |
| cand_nll = sum(r['candidate_nll_sum'] for r in rows) / n |
| summaries[split] = {'records': len(rows), 'tokens': n, 'kl_mean': sum(r['kl_sum'] for r in rows) / n, 'reference_nll': ref_nll, 'candidate_nll': cand_nll, 'nll_delta': cand_nll - ref_nll, 'reference_perplexity': math.exp(ref_nll), 'candidate_perplexity': math.exp(cand_nll), 'top1_agreement': sum(r['top1_agreements'] for r in rows) / n} |
| values = token_kl_by_split[split] |
| summaries[split].update({'kl_median': float(np.median(values)), 'kl_p95': float(np.percentile(values, 95)), 'kl_p99': float(np.percentile(values, 99)), 'kl_p99_9': float(np.percentile(values, 99.9)), 'kl_max': max(values)}) |
| result = {'metric': 'KL(reference || candidate), natural logarithms, all vocabulary entries, teacher-forced positions', 'reference': str(args.reference), 'candidate': str(args.candidate), 'splits': summaries, 'records': results} |
| result['position_scope'] = 'assistant outputs, tool calls and EOS; supplied headers excluded' if mask_info else 'all sequence targets including user/tool prompts' |
| result['assistant_masks'] = str(args.assistant_masks) if args.assistant_masks else None |
| args.output.write_text(json.dumps(result, indent=2) + '\n') |
| print(json.dumps(summaries, indent=2)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
