reproduction/scripts/check_training.py

"""Diagnostic learning and exact optimizer/RNG resume proof; not held-out quality."""
import sys,json,time,random,hashlib,gc
from pathlib import Path
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
sys.path.insert(0,str(Path(__file__).resolve().parent))
import numpy as np
import torch
from transformers import AutoTokenizer,AutoModelForTokenClassification,DataCollatorForTokenClassification
from pii.core import LABELS,BIO,decode_window,byte_positions
from train import prepare

out=Path('work/training-diagnostic');out.mkdir(parents=True,exist_ok=True)
random.seed(42);np.random.seed(42);torch.manual_seed(42)
tok=AutoTokenizer.from_pretrained('artifacts/backbone',local_files_only=True)
# Shortest complete documents providing every target label keeps diagnostic bounded.
rows=[json.loads(x) for x in open('data/prepared/train.jsonl')]
rows.sort(key=lambda r:(len(r['text']),r['id']))
selected=[];missing=set(LABELS)
for r in rows:
    labels={s['label'] for s in r['spans']}
    if labels & missing:selected.append(r);missing-=labels
    if not missing:break
chunks,stats=prepare(selected,tok,512,42)
collate=DataCollatorForTokenClassification(tok,return_tensors='pt')
batch={k:v.to('mps') for k,v in collate(chunks).items() if k in ('input_ids','attention_mask','labels')}
model=AutoModelForTokenClassification.from_pretrained('artifacts/backbone',num_labels=len(BIO),id2label=dict(enumerate(BIO)),label2id={v:k for k,v in enumerate(BIO)},attn_implementation='eager',local_files_only=True).to('mps')
opt=torch.optim.AdamW(model.parameters(),lr=2e-5,weight_decay=.01)
offsets=[tok(r['text'],return_offsets_mapping=True)['offset_mapping'] for r in selected]
gold=[]
for r in selected:
    bp=byte_positions(r['text']);gold.append({(s['label'],bp[s['start']],bp[s['end']]) for s in r['spans']})
def assess():
    model.eval()
    with torch.no_grad():o=model(**batch)
    found=0;pred=0
    for r,off,logits,g in zip(selected,offsets,o.logits.cpu().numpy(),gold):
        p={(f['label'],f['start'],f['end']) for f in decode_window(logits[:len(off)],off,r['text'],512)}
        found+=len(p&g);pred+=len(p)
    return {'loss':o.loss.item(),'exact_entity_true_positives':found,'gold_entities':sum(map(len,gold)),'predicted_entities':pred,'recall':found/sum(map(len,gold))}
def step():
    model.train();opt.zero_grad(set_to_none=True);loss=model(**batch).loss
    assert torch.isfinite(loss);loss.backward();torch.nn.utils.clip_grad_norm_(model.parameters(),1.,error_if_nonfinite=True);opt.step();return loss.item()
def fingerprint():
    h=hashlib.sha256()
    for n,p in model.named_parameters():h.update(n.encode());h.update(p.detach().cpu().numpy().tobytes())
    return h.hexdigest()
start=time.perf_counter();before=assess();history=[]
while True:
    loss=step();a=assess();history.append(a)
    if a['loss']<before['loss'] and a['recall']>before['recall']:break
    if time.perf_counter()-start>180:raise RuntimeError('Diagnostic 3-minute coordination budget exhausted; no learning claim')
model.save_pretrained(out/'model')
state={'optimizer':opt.state_dict(),'random':random.getstate(),'numpy':np.random.get_state(),'torch':torch.get_rng_state(),'mps':torch.mps.get_rng_state()}
torch.save(state,out/'state.pt');del state
uninterrupted_loss=step();uninterrupted_hash=fingerprint();uninterrupted=assess();torch.save({k:v.detach().cpu() for k,v in model.state_dict().items()},out/'uninterrupted.pt')
del model,opt;gc.collect();torch.mps.empty_cache()
model=AutoModelForTokenClassification.from_pretrained(out/'model',local_files_only=True,attn_implementation='eager').to('mps')
opt=torch.optim.AdamW(model.parameters(),lr=2e-5,weight_decay=.01)
state=torch.load(out/'state.pt',map_location='cpu',weights_only=False);opt.load_state_dict(state['optimizer'])
random.setstate(state['random']);np.random.set_state(state['numpy']);torch.set_rng_state(state['torch']);torch.mps.set_rng_state(state['mps']);del state
resumed_loss=step();resumed_hash=fingerprint();resumed=assess()
reference=torch.load(out/'uninterrupted.pt',weights_only=True)
max_parameter_difference=max((v.detach().cpu()-reference[k]).abs().max().item() for k,v in model.state_dict().items())
del reference
report={'purpose':'Training functionality diagnostic only; same examples fitted and scored, no benchmark','device':'mps','seed':42,'recipe':{'lr':2e-5,'weight_decay':.01,'gradient_clip':1.},'selected_ids':[r['id'] for r in selected],'labels_covered':LABELS,'chunk_stats':stats,'steps_to_first_learning_evidence':len(history),'stopping_rule':'First evaluation with both lower token loss and greater exact-entity recall than random-head baseline','before':before,'after':history[-1],'resume':{'checkpoint':'work/training-diagnostic','uninterrupted_next_update_loss':uninterrupted_loss,'resumed_next_update_loss':resumed_loss,'uninterrupted_parameters_sha256':uninterrupted_hash,'resumed_parameters_sha256':resumed_hash,'all_model_parameters_bitwise_equal':uninterrupted_hash==resumed_hash,'max_parameter_abs_difference':max_parameter_difference,'resume_status':'Same next-update loss and evaluation; MPS parameter differences reported without a bitwise guarantee','uninterrupted_metrics':uninterrupted,'resumed_metrics':resumed},'elapsed_seconds':time.perf_counter()-start}
Path('artifacts/training-check.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)
assert uninterrupted_loss==resumed_loss and uninterrupted==resumed,'Resumed loss or outputs differ'