40 lines
1.9 KiB
Python
40 lines
1.9 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)
|
|
snapshots=[]
|
|
result={'passed':False,'sample_id':samples[0]['id'],'tokens':8}
|
|
def progress(layer):
|
|
status=dict(line.split(':',1) for line in Path('/proc/self/status').read_text().splitlines() if ':' in line)
|
|
snapshots.append({'layer':layer,'seconds':time.time()-start,'rss_kib':int(status['VmRSS'].split()[0]),
|
|
'process_swap_kib':int(status.get('VmSwap','0 kB').split()[0]),
|
|
'available_kib':int(next(line.split()[1] for line in Path('/proc/meminfo').read_text().splitlines() if line.startswith('MemAvailable:')))})
|
|
print(json.dumps(snapshots[-1]),flush=True)
|
|
try:
|
|
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')
|
|
result.update(passed=True,mean_nll=float(nll.mean()),token_nll=nll.tolist())
|
|
except Exception as exc:
|
|
result.update(error_type=type(exc).__name__,error=str(exc))
|
|
raise
|
|
finally:
|
|
result.update(elapsed_seconds=time.time()-start,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
|
|
resource_snapshots=snapshots,limitation='Real-weight capacity check, not an independent correctness oracle or quality benchmark')
|
|
write_json(a.root/'preflight.json',result)
|
|
|
|
|
|
if __name__=='__main__':main()
|