28 lines
1.2 KiB
Python
28 lines
1.2 KiB
Python
"""One 8-token real text, bounded memory/semantic capacity smoke check."""
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
import resource
|
|
import time
|
|
import numpy as np
|
|
from .model import StreamingLM
|
|
from .run import tokenize,write_json
|
|
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser();p.add_argument('--root',required=True,type=Path);a=p.parse_args()
|
|
samples,ids=tokenize(a.root)
|
|
start=time.time()
|
|
lm=StreamingLM(a.root/'weights',json.loads((a.root/'model_config.json').read_text()),seconds=1200)
|
|
def progress(layer):print(json.dumps({'preflight_layer':layer,'seconds':time.time()-start}),flush=True)
|
|
h,_=lm.forward(ids[:1,:8],progress=progress)
|
|
nll=lm.nll(h,ids[:1,:8])
|
|
if not np.isfinite(nll).all():raise FloatingPointError('Nonfinite NLL')
|
|
write_json(a.root/'preflight.json',{'passed':True,'sample_id':samples[0]['id'],'tokens':8,
|
|
'elapsed_seconds':time.time()-start,'peak_rss_kib':resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
|
|
'mean_nll':float(nll.mean()),'token_nll':nll.tolist(),
|
|
'limitation':'Real-weight capacity check, not an independent correctness oracle or quality benchmark'})
|
|
|
|
|
|
if __name__=='__main__':main()
|