File size: 5,716 Bytes
3fd1a35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | #!/usr/bin/env python3
"""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))
# Only read the remote artifact. Stream it without a second local disk copy.
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()
|