"""Inference-only development diagnostics; original labels are reference labels. Residue shuffling is an out-of-distribution intervention, not a label-preserving biological transformation. Original calibration temperatures remain fixed. """ from __future__ import annotations import argparse from collections import Counter import gc import hashlib import json from pathlib import Path import random import numpy as np import torch import laya_formal_experiment as core from laya_direct_bpe import DirectBPE, digest from laya_direct_experiment import make_items ROOT=core.ROOT SEEDS=(20260922,20260923,20260924) def perturbed(row,variant): row=dict(row) if variant=='sequence_removed':row['sequence']='' elif variant.startswith('residue_shuffle_'): raw=row['sequence'];characters=list(raw) seed=int.from_bytes(hashlib.sha256(f"diagnostic-v1:{variant}:{row['id']}".encode()).digest()[:8],'big') random.Random(seed).shuffle(characters) row['sequence']=''.join(characters) if Counter(raw)!=Counter(row['sequence']):raise ValueError('Shuffle changed residue composition') return row def permutation(row): order=list(range(len(row['choices']))) seed=int.from_bytes(hashlib.sha256(f"candidate-diagnostic-v1:{row['id']}".encode()).digest()[:8],'big') random.Random(seed).shuffle(order) if order==list(range(len(order))):order=order[1:]+order[:1] return order def canonicalize(records,orders): result=[] for record in records: order=orders[record['id']] logits=[0.]*len(order) for position,canonical in enumerate(order):logits[canonical]=record['logits'][position] item=dict(record,logits=logits,label=order[record['label']]) values=np.asarray(logits,dtype=float);exponent=np.exp(values-values.max()) item['probs']=(exponent/exponent.sum()).tolist() result.append(item) return result def compare(reference,changed): if [(r['id'],r['label'],r['task']) for r in reference]!=[(r['id'],r['label'],r['task']) for r in changed]: raise ValueError('Diagnostic membership/labels changed') out={} for task in sorted({r['task'] for r in reference}): pairs=[(a,b) for a,b in zip(reference,changed) if a['task']==task] p=np.asarray([a['probs'] for a,b in pairs]);q=np.asarray([b['probs'] for a,b in pairs]) if not np.isfinite(p).all() or not np.isfinite(q).all():raise FloatingPointError('Non-finite diagnostic') middle=(p+q)/2 js=.5*np.sum(p*np.log(np.maximum(p,1e-300)/np.maximum(middle,1e-300))+ q*np.log(np.maximum(q,1e-300)/np.maximum(middle,1e-300)),axis=1) out[task]={'n':len(pairs),'prediction_agreement':float((p.argmax(1)==q.argmax(1)).mean()), 'mean_probability_l1':float(np.abs(p-q).sum(1).mean()), 'mean_max_probability_difference':float(np.abs(p-q).max(1).mean()), 'mean_js_divergence':float(js.mean())} return out def run_one(run_dir,out,batch_size=16): if (out/'summary.json').exists():return json.loads((out/'summary.json').read_text()) saved=json.loads((run_dir/'summary.json').read_text()) if not saved['formal'] or saved['test_access'] or not saved['checkpoint_reload_logits_match']: raise ValueError('Source checkpoint failed protocol gate') rep=DirectBPE(run_dir/'checkpoint/representation') eligible=core.load_eligible_ids(saved['eligible_id_filter']) rows=[r for r in core.load_split(ROOT/'artifacts/laya_formal_data','both','selection_dev') if r['id'] in eligible] _,build_model,builder=core.import_laya('vendor/laya') cfg=json.loads((run_dir/'checkpoint/rl_agent_config.json').read_text()) model=core.fresh_reload(run_dir/'checkpoint',cfg,build_model,torch.device('cuda')) tokenizer=rep.expanded if saved['condition']=='full_bpe' else rep.base base_items=make_items(rows,rep,saved['condition'],builder,saved['seed']) baseline=core.evaluate(model,base_items,tokenizer,torch.device('cuda'),batch_size) previous=[json.loads(x) for x in (run_dir/'selection_dev_predictions.jsonl').read_text().splitlines()] if len(previous)!=len(baseline) or any(a['id']!=b['id'] or max(abs(x-y) for x,y in zip(a['logits'],b['logits']))>1e-5 for a,b in zip(previous,baseline)): raise ValueError('Unperturbed predictions do not reproduce source result') temps={t:v['temperature'] for t,v in saved['calibration_temperature'].items()} result={'condition':saved['condition'],'seed':saved['seed'],'split':'selection_dev','test_access':False, 'no_training':True,'reference_labels_not_assumed_valid_under_sequence_intervention':True, 'source_summary_sha256':digest(run_dir/'summary.json'), 'baseline_logits_reproduced':True,'temperature_refitted':False, 'original':{'raw':core.metric_from_records(baseline),'calibrated':core.metric_from_records(baseline,temps)}, 'variants':{}} out.mkdir(parents=True,exist_ok=True) for variant in ('sequence_removed','residue_shuffle_0','residue_shuffle_1','residue_shuffle_2','candidate_permutation'): if variant=='candidate_permutation': items=[];orders={} for row in rows: order=permutation(row);orders[row['id']]=order ids,markers=rep.build(row,saved['condition'],builder,order) items.append({'ids':ids,'markers':markers,'label':order.index(row['label']), 'task':row['task'],'id':row['id']}) predictions=canonicalize(core.evaluate(model,items,tokenizer,torch.device('cuda'),batch_size),orders) unchanged=None else: changed=[perturbed(row,variant) for row in rows] unchanged=sum(a['sequence']==b['sequence'] for a,b in zip(rows,changed)) items=make_items(changed,rep,saved['condition'],builder,saved['seed']) predictions=core.evaluate(model,items,tokenizer,torch.device('cuda'),batch_size) result['variants'][variant]={'reference_label_metrics':{'raw':core.metric_from_records(predictions), 'calibrated':core.metric_from_records(predictions,temps)}, 'comparison_to_original':compare(baseline,predictions), 'unchanged_sequence_count':unchanged} (out/f'{variant}_predictions.jsonl').write_text(''.join(json.dumps(r)+'\n' for r in predictions)) print(json.dumps({'condition':saved['condition'],'seed':saved['seed'],'variant':variant, 'comparison':result['variants'][variant]['comparison_to_original']}),flush=True) (out/'summary.json').write_text(json.dumps(result,indent=2)+'\n') del model gc.collect();torch.cuda.empty_cache() return result def main(): p=argparse.ArgumentParser(description=__doc__) p.add_argument('--output-dir',type=Path,default=ROOT/'artifacts/laya_controls/sequence_diagnostics') p.add_argument('--seeds',type=int,nargs='+',default=list(SEEDS)) p.add_argument('--conditions',nargs='+',choices=['raw','full_bpe'],default=['full_bpe','raw']) args=p.parse_args() for seed in args.seeds: for condition in args.conditions: run_one(ROOT/f'artifacts/laya_direct_legacy/{condition}_seed{seed}', args.output_dir/f'{condition}_seed{seed}') print('All requested development diagnostics completed.',flush=True) if __name__=='__main__':main()