| |
| """Full-vocabulary distances at fixed prefixes; no target-answer likelihood score. |
| |
| Requires numpy only. These metrics remain prefix dependent, not task accuracy. |
| All computations use float64; raw logits are centered before RMSE comparison. |
| """ |
| import argparse |
| import json |
| import shlex |
| import subprocess |
| from pathlib import Path |
| import numpy as np |
|
|
|
|
| def candidate_chunks(path, count, remote): |
| expected = count*157184*4 |
| if not remote: |
| if path.stat().st_size != expected: raise ValueError('candidate shape mismatch') |
| data = np.memmap(path, dtype='<f4', mode='r', shape=(count,157184)) |
| for start in range(0,count,4): yield data[start:start+4] |
| return |
| quoted = shlex.quote(str(path)) |
| |
| command = 'stat -c %s -- '+quoted+' && cat -- '+quoted |
| with subprocess.Popen(['ssh','-o','BatchMode=yes',remote,command],stdout=subprocess.PIPE) as proc: |
| try: |
| if int(proc.stdout.readline()) != expected: raise ValueError('remote candidate shape mismatch') |
| for start in range(0,count,4): |
| rows = min(4,count-start) |
| raw = proc.stdout.read(rows*157184*4) |
| if len(raw) != rows*157184*4: raise ValueError('truncated remote logits') |
| yield np.frombuffer(raw,dtype='<f4').reshape(rows,157184) |
| if proc.stdout.read(1) or proc.wait() != 0: raise ValueError('remote logit read failed') |
| finally: |
| if proc.poll() is None: proc.kill() |
|
|
|
|
| def metrics(reference, candidate): |
| r, c = np.asarray(reference, dtype=np.float64), np.asarray(candidate, dtype=np.float64) |
| if not np.isfinite(r).all() or not np.isfinite(c).all(): raise ValueError('nonfinite logits') |
| def distribution(x): |
| shifted = x-x.max(axis=-1, keepdims=True) |
| lp = shifted-np.log(np.exp(shifted).sum(axis=-1, keepdims=True)) |
| return lp, np.exp(lp) |
| lp, p = distribution(r) |
| lq, q = distribution(c) |
| lm = np.logaddexp(lp, lq)-np.log(2.) |
| diff = c-r |
| diff -= diff.mean(axis=-1, keepdims=True) |
| rtop, ctop = r.argmax(axis=-1), c.argmax(axis=-1) |
| top2 = np.partition(r, -2, axis=-1)[:, -2:] |
| gap = top2.max(axis=-1)-top2.min(axis=-1) |
| return dict(kl=(p*(lp-lq)).sum(axis=-1), |
| js=(p*(lp-lm)+q*(lq-lm)).sum(axis=-1)/2, |
| tv=np.abs(p-q).sum(axis=-1)/2, |
| centered_logit_rmse=np.sqrt(np.square(diff).mean(axis=-1)), |
| top1_equal=(rtop == ctop).astype(np.float64), |
| reference_top_gap=gap, |
| reference_preference_regret=lp[np.arange(len(r)), rtop]-lp[np.arange(len(r)), ctop]) |
|
|
|
|
| def summarize(data): |
| if not data or not len(data['kl']): return None |
| out = dict(tokens=len(data['kl'])) |
| for key, value in data.items(): |
| x = np.asarray(value) |
| out[key] = dict(mean=float(x.mean()), p50=float(np.median(x)), |
| p95=float(np.quantile(x, .95)), max=float(x.max())) |
| confident = np.asarray(data['reference_top_gap']) >= 1. |
| out['reference_gap_ge_1_nat'] = dict(tokens=int(confident.sum()), |
| top1_agreement=float(np.asarray(data['top1_equal'])[confident].mean()) if confident.any() else None) |
| return out |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument('--suite', type=Path, required=True) |
| p.add_argument('--reference', type=Path, required=True) |
| p.add_argument('--candidate', type=Path, required=True) |
| p.add_argument('--candidate-ssh', help='Read remote logits over SSH; compute on this host') |
| p.add_argument('--suffix', default='rknn') |
| p.add_argument('--output', type=Path, required=True) |
| a = p.parse_args() |
| suite = json.loads(a.suite.read_text()) |
| pooled = {'think':{}, 'nothink':{}} |
| result = dict(reference=str(a.reference), candidate=str(a.candidate), candidate_ssh=a.candidate_ssh, suffix=a.suffix, |
| scope='Fixed teacher prefixes, full vocabulary, natural logarithms; not free generation or accuracy', cases=[]) |
| for case in suite['cases']: |
| count = len(case['teacher_ids']) |
| file = a.reference/(case['id']+'.bf16.f32') |
| if file.stat().st_size != count*157184*4: raise ValueError('shape mismatch: '+str(file)) |
| reference = np.memmap(file, dtype='<f4', mode='r', shape=(count,157184)) |
| candidate_path = a.candidate/(case['id']+'.'+a.suffix+'.f32') |
| values = {} |
| for i, candidate in enumerate(candidate_chunks(candidate_path, count, a.candidate_ssh)): |
| start = i*4 |
| for key, x in metrics(reference[start:start+4], candidate).items(): |
| values.setdefault(key, []).extend(x.tolist()) |
| mode = 'think' if case['thinking'] else 'nothink' |
| for key, x in values.items(): pooled[mode].setdefault(key, []).extend(x) |
| def part(lo, hi): return summarize({key:x[lo:hi] for key,x in values.items()}) |
| row = dict(id=case['id'], mode=mode, whole=summarize(values), |
| reasoning=part(0, case['answer_start']), answer=part(case['answer_start'], count), |
| windows={f'{i}-{min(i+256,count)}':part(i,min(i+256,count)) for i in range(0,count,256)}) |
| result['cases'].append(row) |
| print(json.dumps(dict(case=case['id'], tv=row['whole']['tv'], kl=row['whole']['kl'])), flush=True) |
| result['aggregate'] = {mode:summarize(values) for mode,values in pooled.items()} |
| a.output.parent.mkdir(parents=True, exist_ok=True) |
| a.output.write_text(json.dumps(result, ensure_ascii=False, indent=2)+'\n') |
| print(json.dumps(result['aggregate']), flush=True) |
|
|
|
|
| if __name__ == '__main__': main() |
|
|
