Buckets:
| """Train frozen-feature heads. Primary fitting never opens human attribute labels.""" | |
| import argparse,copy,json,time,os,hashlib,datetime | |
| from pathlib import Path | |
| import numpy as np,pandas as pd,torch,yaml | |
| from models.scalar_bt import ScalarBT | |
| from models.latent_linear import LatentLinear | |
| from models.latent_mlp import LatentMLP | |
| from models.latent_ortho import LatentOrtho | |
| from models.multi_attribute import MultiAttribute | |
| from losses.preference import bradley_terry | |
| from losses.intervention import intervention_losses | |
| from analysis.tracking import save_run,rebuild | |
| from representations.cache import validate_feature_table,file_hash,atomic_json | |
| IMPORTED_SOURCE={str(p):p.read_bytes() for p in [Path(__file__),*Path('models').glob('*.py'),*Path('losses').glob('*.py')]} | |
| TRAIN_SOURCE_HASHES={str(Path(k).relative_to(Path.cwd())) if Path(k).is_absolute() else k:hashlib.sha256(v).hexdigest() for k,v in IMPORTED_SOURCE.items()} | |
| def fit_specification(args,seed,cfg,manifest,layer,feature_identity): | |
| requested={'method':args.method,'backbone':args.backbone,'layer':layer,'pooling':args.pooling,'arch':args.arch,'latent_dim':args.latent_dim,'n':args.n,'filter':args.filter,'seed':seed,'shuffle_labels':args.shuffle_labels,'flip_intervention_signs':args.flip_intervention_signs,'criteria':getattr(args,'criteria',[0,1,2,3])} | |
| spec={'config':cfg,'requested':requested,'features':feature_identity,'preferences_sha256':file_hash('data/processed/preferences.parquet'),'training_source_hashes':TRAIN_SOURCE_HASHES} | |
| if args.method=='attribute':spec['attributes_sha256']=file_hash('data/processed/attributes.parquet') | |
| if args.method=='ilr' and args.n: | |
| spec['intervention_identity']=validate_feature_table(f'outputs/features/{args.backbone}_interventions','data/processed/interventions.parquet',f'data/processed/tokenized_{args.backbone}_interventions.parquet') | |
| return spec | |
| def resume_checked(path,spec): | |
| required=['config.json','metrics.json','model.pt','predictions.npz','learning_curve.csv','selected_interventions.json','source_manifest.json'] | |
| assert all((path/x).is_file() for x in required),'Incomplete trained artifact inventory' | |
| old=json.loads((path/'config.json').read_text()) | |
| if (path/'fit_spec.json').exists(): | |
| assert json.loads((path/'fit_spec.json').read_text())==spec,'Fitted specification changed; use a fresh run-tag or original source/config' | |
| else: | |
| # One-time compatibility audit for existing baseline fits predating this | |
| # guard. Preserve their actual snapshots and disclose the later binding. | |
| assert old['method']!='ilr','Legacy ILR needs explicit provenance review' | |
| assert all(old.get(k)==v for k,v in spec['config'].items() if k not in ['latent_dim','backbone']),'Legacy hyperparameters changed' | |
| r=spec['requested'];expected={'method':r['method'],'backbone_tag':r['backbone'],'representation_layer':r['layer'],'pooling_method':r['pooling'],'latent_architecture':r['arch'],'seed':r['seed'],'latent_dim':1 if r['method']=='scalar' else 4 if r['method']=='attribute' else r['latent_dim'],'control':'shuffle' if r['shuffle_labels'] else 'flipped' if r['flip_intervention_signs'] else 'standard'} | |
| assert all(old.get(k)==v for k,v in expected.items()),'Legacy fitted design changed' | |
| assert old['feature_manifest']==spec['features']['manifest'] | |
| if old.get('feature_identity'):assert old['feature_identity']==spec['features'] | |
| sources=json.loads((path/'source_manifest.json').read_text()) | |
| assert all(sources.get(k)==v for k,v in spec['training_source_hashes'].items() if k.startswith(('models/','losses/'))),'Legacy model/loss source changed' | |
| pred=np.load(path/'predictions.npz');ex=pd.read_parquet('data/processed/responses.parquet').sort_values('feature_index') | |
| assert np.array_equal(pred['response_ids'],ex.response_id) | |
| atomic_json(path/'legacy_fit_binding.json',{'bound_at':datetime.datetime.now(datetime.timezone.utc).isoformat(),'note':'Baseline fitted before immutable fit-spec guard. Existing config/model/loss hashes, original feature manifest, response IDs and prior full token/source audit checked. Original trainer snapshot retained; later binding is not retrospective proof of absent historical preference-file changes. Engine observed no dataset or fitting-objective changes during these baseline runs.','original_trainer_sha256':sources.get('scripts/train_heads.py')}) | |
| atomic_json(path/'fit_spec.json',spec) | |
| marker=json.loads((path/'TRAINED.json').read_text()) | |
| for name,digest in marker.get('fit_artifact_hashes',{}).items():assert file_hash(path/name)==digest,f'Changed fitted artifact {name}' | |
| def load_features(backbone,layer,pooling): | |
| root=Path(f'outputs/features/{backbone}'); manifest=json.loads((root/'manifest.json').read_text()); assert (root/'COMPLETE.json').exists() | |
| layer=layer or manifest['layers'][2] | |
| x=np.load(root/f'layer{layer}_{pooling}.npy',mmap_mode='r'); return x,manifest,layer | |
| def fit(args,seed): | |
| assert os.environ.get('CUDA_VISIBLE_DEVICES')=='GPU-2a9c799b-ab9e-c677-db0e-379f480cb00f', 'Physical GPU1 UUID binding required' | |
| assert torch.cuda.device_count()==1 | |
| cfg=yaml.safe_load(Path('configs/latent_bt.yaml').read_text()) | |
| if args.method=='ilr':cfg.update(yaml.safe_load(Path('configs/ilr.yaml').read_text())) | |
| criteria=getattr(args,'criteria',[0,1,2,3]);q=len(criteria) | |
| assert len(set(criteria))==q and set(criteria)<=set(range(4)) and q>=2 | |
| assert args.latent_dim>=q or args.method=='scalar' | |
| assert args.method!='attribute' or criteria==[0,1,2,3] | |
| x,manifest,layer=load_features(args.backbone,args.layer,args.pooling) | |
| feature_identity=validate_feature_table(f'outputs/features/{args.backbone}','data/processed/responses.parquet',f'data/processed/tokenized_{args.backbone}.parquet') | |
| eid=f'hs2_{args.backbone}_l{layer}_{args.pooling}_{args.method}_{args.arch}_k{args.latent_dim}_n{args.n}_{args.filter}_s{seed}' | |
| control='shuffle' if args.shuffle_labels else 'flipped' if args.flip_intervention_signs else 'standard' | |
| if control!='standard':eid+='_'+control | |
| if criteria!=[0,1,2,3]:eid+='_criteria'+''.join(map(str,criteria)) | |
| path=Path('outputs/real')/eid | |
| if getattr(args,'run_tag',None):path=Path(str(path)+'_'+args.run_tag);eid=path.name | |
| fit_spec=fit_specification(args,seed,cfg,manifest,layer,feature_identity) | |
| if (path/'TRAINED.json').exists():resume_checked(path,fit_spec);return | |
| path.mkdir(parents=True,exist_ok=True); torch.manual_seed(seed); torch.set_num_threads(4); device=cfg['device']; begin=time.monotonic() | |
| ex=pd.read_parquet('data/processed/responses.parquet').sort_values('feature_index').reset_index(drop=True); pref=pd.read_parquet('data/processed/preferences.parquet'); pref=pref[~pref.tie].reset_index(drop=True) | |
| assert len(x)==len(ex) | |
| tri=np.flatnonzero(ex.split=='train'); h=np.asarray(x,dtype='float32'); mean=h[tri].mean(0); scale=np.maximum(h[tri].std(0),.1); h=(h-mean)/scale | |
| assert np.isfinite(h).all(),'Nonfinite original feature values' | |
| ht=torch.as_tensor(h,device=device) | |
| pt={s:pref[pref.split==s] for s in ['train','dev','test']} | |
| tensor=lambda a:torch.as_tensor(a,device=device) | |
| tr={k:tensor(pt['train'][k].to_numpy()) for k in ['a_index','b_index','label']};dv={k:tensor(pt['dev'][k].to_numpy()) for k in ['a_index','b_index','label']} | |
| if args.shuffle_labels:tr['label']=tr['label'][torch.randperm(len(tr['label']),device=device)] | |
| if args.method=='scalar':model=ScalarBT(h.shape[1]) | |
| elif args.method=='attribute':model=MultiAttribute(h.shape[1]) | |
| elif args.arch=='mlp':model=LatentMLP(h.shape[1],args.latent_dim) | |
| elif args.method=='ortho':model=LatentOrtho(h.shape[1],args.latent_dim) | |
| else:model=LatentLinear(h.shape[1],args.latent_dim) | |
| model=model.to(device); optimizer=torch.optim.AdamW(model.parameters(),lr=cfg['learning_rate'],weight_decay=cfg['weight_decay']) | |
| int_data=None; selected=[]; actual_counts={}; intervention_identity=None;selected_base_counts={};selected_sign_counts={} | |
| if args.method=='ilr' and args.n: | |
| ir=pd.read_parquet('data/processed/interventions.parquet'); ir=ir[(ir.int_split=='train') & ir[args.filter]] | |
| for k in criteria: | |
| g=np.random.default_rng(seed+459+100003*k) | |
| candidates=ir[ir.criterion_index==k].sort_values('intervention_id'); order=g.permutation(len(candidates)); chosen=candidates.iloc[order[:args.n]] | |
| actual_counts[str(k)]=len(chosen) | |
| if len(chosen)<args.n:raise ValueError(f'Insufficient {args.filter} interventions for {k}: {len(chosen)} < {args.n}') | |
| selected.extend(chosen.intervention_id.tolist()) | |
| selected_base_counts[str(k)]=int(chosen.base_response_id.nunique());selected_sign_counts[str(k)]={str(s):int((chosen.direction==s).sum()) for s in [-1,1]} | |
| ir=ir.set_index('intervention_id').loc[selected].reset_index() | |
| ipath=Path(f'outputs/features/{args.backbone}_interventions'); assert (ipath/'COMPLETE.json').exists() | |
| intervention_identity=validate_feature_table(ipath,'data/processed/interventions.parquet',f'data/processed/tokenized_{args.backbone}_interventions.parquet') | |
| for key in ['model','model_revision','tokenizer_revision','layers','d','pooling_sha256']: | |
| assert manifest[key]==intervention_identity['manifest'][key],f'Original/modified feature-space mismatch: {key}' | |
| ih=np.asarray(np.load(ipath/f'layer{layer}_{args.pooling}.npy',mmap_mode='r'),dtype='float32'); ih=(ih-mean)/scale | |
| assert np.isfinite(ih).all(),'Nonfinite intervention feature values' | |
| int_data={'original':ht[tensor(ir.base_feature_index.to_numpy())], 'modified':tensor(ih[ir.feature_index.to_numpy()]),'k':tensor(ir.criterion_index.map({k:i for i,k in enumerate(criteria)}).to_numpy()),'s':tensor(ir.direction.to_numpy()).float()} | |
| if args.flip_intervention_signs:int_data['s']=-int_data['s'] | |
| attr=None | |
| if args.method=='attribute': | |
| # Explicit dense-supervision upper-bound exception. No other branch reads this table. | |
| from data.helpsteer import CRITERIA | |
| a=ex[['response_id','split']].merge(pd.read_parquet('data/processed/attributes.parquet'),on='response_id',how='left'); arr=a[CRITERIA].to_numpy('float32') | |
| am=arr[tri].mean(0); ast=arr[tri].std(0); arr=(arr-am)/ast | |
| attr={s:(tensor(np.flatnonzero(ex.split==s)),tensor(arr[ex.split==s])) for s in ['train','dev']} | |
| best=float('inf'); best_step=0; best_state=None; curves=[] | |
| # Linear BT uses feature differences; nonlinear heads must evaluate responses separately. | |
| def logits(a,b): | |
| return model(ht[a])-model(ht[b]) if args.arch=='mlp' else model(ht[a]-ht[b]) | |
| for step in range(cfg['steps']): | |
| model.train();optimizer.zero_grad(set_to_none=True) | |
| pref_loss=bradley_terry(logits(tr['a_index'],tr['b_index']),tr['label'].float()); loss=pref_loss | |
| target=leak=res=torch.tensor(0.,device=device) | |
| if attr: | |
| ids,values=attr['train']; aloss=torch.nn.functional.mse_loss(model.encode(ht[ids]),values) | |
| # Dense attributes + preference readout. Preferences also count as supervision. | |
| loss=loss+aloss | |
| if int_data: | |
| dz=model.encode(int_data['modified'])-model.encode(int_data['original']);target,leak,res=intervention_losses(dz,int_data['k'],int_data['s'],criteria=q,residual=args.latent_dim>q,tau=cfg['tau']) | |
| loss=loss+cfg['lambda_target']*target+cfg['lambda_leak']*leak+cfg['lambda_residual']*res | |
| if args.method=='ortho':loss=loss+cfg['lambda_ortho']*model.penalty() | |
| loss.backward();torch.nn.utils.clip_grad_norm_(model.parameters(),10.);optimizer.step() | |
| if step%cfg['eval_every']==0 or step==cfg['steps']-1: | |
| model.eval() | |
| with torch.inference_mode(): | |
| dl=bradley_terry(logits(dv['a_index'],dv['b_index']),dv['label'].float()).item(); selection=dl | |
| if attr: | |
| ids,values=attr['dev'];selection+=torch.nn.functional.mse_loss(model.encode(ht[ids]),values).item() | |
| curves.append({'step':step,'loss':loss.item(),'train_preference_loss':pref_loss.item(),'dev_preference_loss':dl,'selection_loss':selection,'target_loss':target.item(),'leakage_loss':leak.item(),'residual_loss':res.item()}) | |
| if selection<best-1e-5:best=selection;best_step=step;best_state={k:v.detach().cpu().clone() for k,v in model.state_dict().items()} | |
| if step>=cfg['min_steps'] and step-best_step>=cfg['patience']:break | |
| model.load_state_dict(best_state);model.eval() | |
| with torch.inference_mode(): | |
| reward=np.concatenate([model(batch).cpu().numpy() for batch in ht.split(1024)]) | |
| z=np.concatenate([model.encode(batch).cpu().numpy() for batch in ht.split(1024)]) if args.method!='scalar' else np.empty((len(ex),0)) | |
| te=pt['test'];pred=reward[te.a_index]-reward[te.b_index];acc=float(np.mean(pred*(2*te.label.to_numpy()-1)>0)) | |
| np.savez_compressed(path/'predictions.npz',reward=reward,z=z,response_ids=ex.response_id.to_numpy(str),mean=mean,scale=scale,test_pair_ids=te.pair_id.to_numpy(str),test_logits=pred,test_labels=te.label.to_numpy()) | |
| torch.save(best_state,path/'model.pt');pd.DataFrame(curves).to_csv(path/'learning_curve.csv',index=False);(path/'selected_interventions.json').write_text(json.dumps(selected,indent=2)) | |
| config=dict(cfg,experiment_id=eid,dataset='HelpSteer2-Preference',split='test',backbone=manifest['model'],model_revision=manifest['model_revision'],representation_layer=layer,pooling_method=args.pooling,latent_architecture=args.arch,seed=seed,latent_dim=z.shape[1] if z.shape[1] else 1,criteria_count=4,residual_enabled=args.latent_dim>4 and args.method not in ['scalar','attribute'],n_interventions=args.n if args.method=='ilr' else 0,intervention_filter=args.filter if args.method=='ilr' else 'none',lambda_target=cfg.get('lambda_target',0),lambda_leak=cfg.get('lambda_leak',0),lambda_residual=cfg.get('lambda_residual',0),method=args.method,backbone_tag=args.backbone,actual_intervention_counts=actual_counts,feature_manifest=manifest) | |
| config['control']=control | |
| config.update(criteria_count=q,criteria_indices=criteria,residual_enabled=args.latent_dim>q and args.method not in ['scalar','attribute'],selected_base_counts=selected_base_counts,selected_sign_counts=selected_sign_counts,intervention_budget_unit='edits_per_criterion') | |
| config.update(feature_identity=feature_identity,intervention_identity=intervention_identity) | |
| stats={'preference_accuracy':acc,'best_step':best_step,'best_dev_selection_loss':best,'runtime_seconds':time.monotonic()-begin,'evaluation_status':'preference_only_pending_semantic_evaluation'} | |
| if hasattr(model.encoder,'weight'): | |
| c=torch.nn.functional.normalize(model.encoder.weight.detach(),dim=1) | |
| stats['direction_orthogonality_error']=float((c@c.T-torch.eye(len(c),device=device)).square().sum().sqrt()) | |
| save_run(path,config,stats) | |
| for name,content in IMPORTED_SOURCE.items(): | |
| rel=Path(name).relative_to(Path.cwd()) if Path(name).is_absolute() else Path(name);dest=path/'imported_training_source'/rel;dest.parent.mkdir(parents=True,exist_ok=True);dest.write_bytes(content) | |
| atomic_json(path/'fit_spec.json',fit_spec) | |
| atomic_json(path/'TRAINED.json',{'best_step':best_step,'actual_steps':step+1,'selection':'dev_preference_loss' if not attr else 'dev_preference_plus_attribute_mse','fit_artifact_hashes':{name:file_hash(path/name) for name in ['config.json','model.pt','predictions.npz','learning_curve.csv','selected_interventions.json','fit_spec.json']}});print(eid,acc,best_step,flush=True) | |
| if __name__=='__main__': | |
| p=argparse.ArgumentParser();p.add_argument('--method',choices=['scalar','latent','ortho','attribute','ilr'],default='latent');p.add_argument('--backbone',default='32b');p.add_argument('--layer',type=int);p.add_argument('--pooling',choices=['mean','last'],default='mean');p.add_argument('--arch',choices=['linear','mlp'],default='linear');p.add_argument('--latent-dim',type=int,default=5);p.add_argument('--n',type=int,default=0);p.add_argument('--filter',choices=['raw','direction_filtered','strict'],default='strict');p.add_argument('--seeds',type=int,nargs='+',default=[0,1,2,3,4]);p.add_argument('--shuffle-labels',action='store_true');p.add_argument('--flip-intervention-signs',action='store_true');p.add_argument('--run-tag');p.add_argument('--criteria',type=int,nargs='+',default=[0,1,2,3]);args=p.parse_args() | |
| # Optional explicit identity for an intentionally changed fitted specification. | |
| for seed in args.seeds:fit(args,seed) | |
| rebuild() | |
Xet Storage Details
- Size:
- 16.7 kB
- Xet hash:
- 9c5952362edfb534d2f616aef0cdb845b7811d516a9c705f6a61af683684e6f1
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.