scripts/evaluate_heads.py

import argparse,json,itertools,shutil,datetime
from pathlib import Path
import numpy as np,pandas as pd,torch
from eval.metrics import corr,matching,axis_metrics,cca,cka,selectivity,composition
from data.helpsteer import CRITERIA
from intervention.natural_pairs import natural_pairs
from analysis.tracking import rebuild,save_evaluation_provenance
from representations.cache import validate_feature_table,file_hash,atomic_json

def savez_atomic(path,**values):
    path=Path(path);tmp=path.with_suffix('.tmp.npz');np.savez_compressed(tmp,**values);tmp.replace(path)

def validate_predictions(pred,ex,pref):
    assert np.array_equal(ex.response_id.to_numpy(),pred['response_ids']),'Prediction response identity mismatch'
    assert len(pred['reward'])==len(ex) and len(pred['z'])==len(ex)
    assert np.isfinite(pred['reward']).all() and np.isfinite(pred['z']).all(),'Nonfinite predictions'
    te=pref[(pref.split=='test')&~pref.tie]
    assert np.array_equal(te.pair_id.to_numpy(),pred['test_pair_ids']),'Prediction preference pair identity mismatch'
    assert np.array_equal(te.label.to_numpy(),pred['test_labels']),'Prediction preference labels changed'
    logits=pred['reward'][te.a_index]-pred['reward'][te.b_index]
    assert np.array_equal(logits,pred['test_logits']),'Saved preference logits inconsistent with rewards'

def compose_axes(z,a,pairs,axes):
    if axes==[0,1,2,3]:return composition(z,a,pairs)
    from eval.metrics import rubric_weights
    from scipy.stats import kendalltau,spearmanr
    w=rubric_weights()[:,axes];z=(z-z.mean(0))/(z.std(0)+1e-10);a=(a-a.mean(0))/(a.std(0)+1e-10)
    sm=z@w.T;sh=a@w.T;i,j=np.asarray(pairs).T;dh=sh[i]-sh[j];dm=sm[i]-sm[j];mask=abs(dh)>1e-8
    acc=np.array([np.mean((dm[:,r]*dh[:,r]>0)[mask[:,r]]) if mask[:,r].any() else np.nan for r in range(len(w))])
    tau=np.array([kendalltau(sm[:,r],sh[:,r]).statistic for r in range(len(w))]);rho=np.array([spearmanr(sm[:,r],sh[:,r]).statistic for r in range(len(w))])
    return {'composition_accuracy':float(np.nanmean(acc)),'composition_kendall_tau':float(np.nanmean(tau)),'composition_spearman':float(np.nanmean(rho)),'composition_valid_rubrics':int(np.isfinite(acc).sum()),'composition_reduced_projection':True},{'weights':w,'accuracy':acc,'tau':tau,'rho':rho,'pair_correct':(dm*dh>0).astype('float32'),'pair_valid':mask,'pairs':pairs}

def evaluate(path):
    save_evaluation_provenance(path)
    cfg=json.loads((path/'config.json').read_text()); metrics=json.loads((path/'metrics.json').read_text())
    spec=json.loads((path/'fit_spec.json').read_text())
    identity=validate_feature_table(f'outputs/features/{cfg["backbone_tag"]}','data/processed/responses.parquet',f'data/processed/tokenized_{cfg["backbone_tag"]}.parquet')
    assert identity==spec['features'],'Original features changed since fit'
    assert file_hash('data/processed/preferences.parquet')==spec['preferences_sha256'],'Preference table changed since fit'
    for name,digest in json.loads((path/'TRAINED.json').read_text()).get('fit_artifact_hashes',{}).items():
        assert file_hash(path/name)==digest,f'Fitted artifact changed: {name}'
    ex=pd.read_parquet('data/processed/responses.parquet').sort_values('feature_index').reset_index(drop=True)
    validate_predictions(np.load(path/'predictions.npz'),ex,pd.read_parquet('data/processed/preferences.parquet'))
    if (path/'EVALUATED.json').exists():(path/'EVALUATED.json').replace(path/'EVALUATED.previous.json')
    population={'human_test':'pending','natural':'pending','heldout_generated':'unavailable','stability':'pending'}
    heldout_keys=['target_effect','criterion_leakage','leakage_median','leakage_ratio_of_means','residual_movement']
    for key in heldout_keys:metrics[key]=None;metrics['natural_'+key]=None
    metrics['natural_n']=0;metrics['heldout_intervention_n']=0
    if (path/'heldout_interventions.npz').exists():(path/'heldout_interventions.npz').replace(path/'heldout_interventions.previous.npz')
    source={}
    for name in ['scripts/evaluate_heads.py','eval/metrics.py','intervention/natural_pairs.py','representations/cache.py']:
        source[name]=file_hash(name);dest=path/'evaluation_source'/name;dest.parent.mkdir(parents=True,exist_ok=True);shutil.copyfile(name,dest)
    (path/'evaluation_source_manifest.json').write_text(json.dumps(source,indent=2))
    pref_pred=np.load(path/'predictions.npz')
    metrics['preference_accuracy']=float(np.mean(pref_pred['test_logits']*(2*pref_pred['test_labels']-1)>0))
    if cfg['method']=='scalar':
        metrics['evaluation_status']='preference_evaluated_no_latent_axes'
        atomic_json(path/'metrics.json',metrics)
        atomic_json(path/'EVALUATION_POPULATIONS.json',{'human_test':'not_applicable','natural':'not_applicable','heldout_generated':'not_applicable','stability':'not_applicable'})
        return
    pred=np.load(path/'predictions.npz');z=pred['z'];ex=pd.read_parquet('data/processed/responses.parquet').sort_values('feature_index').reset_index(drop=True)
    a=ex[['response_id','split']].merge(pd.read_parquet('data/processed/attributes.parquet'),on='response_id',how='left')[CRITERIA].to_numpy()
    assert np.array_equal(ex.response_id.to_numpy(),pred['response_ids'])
    axes=cfg.get('criteria_indices',[0,1,2,3]);q=len(axes)
    tr=ex.split=='train';dev=ex.split=='dev';test=ex.split=='test';zt=z[test];at=a[test][:,axes]
    singular=np.linalg.svd(z[dev]-z[dev].mean(0),compute_uv=False)
    variance=singular**2/(np.sum(singular**2)+1e-12)
    metrics['dev_coordinate_variance_fractions']=variance.tolist()
    metrics['dev_coordinate_effective_rank']=float(np.exp(-np.sum(variance*np.log(variance+1e-12))))
    named=cfg['method'] in ['ilr','attribute'] and (cfg['method']=='attribute' or cfg['n_interventions']>0)
    stats,mat=axis_metrics(zt,at,named=named,rank=True);metrics.update(stats)
    # Test Hungarian assignment is descriptive only; all deployed matching uses DEV.
    if named:order=np.arange(q);sign=np.ones(q)
    else:order,sign=matching(z[dev],a[dev][:,axes],rank=True)
    remaining=[i for i in range(z.shape[1]) if i not in order];all_order=np.r_[order,remaining].astype(int);all_sign=np.r_[sign,np.ones(len(remaining))]
    aligned=z[:,all_order]*all_sign; norm=aligned[tr].std(0)+1e-10
    indices=np.flatnonzero(test);tex=ex[test].reset_index(drop=True);pairs=[];pids=[]
    for pid,ids in tex.groupby('prompt_id').indices.items():
        for i,j in itertools.combinations(ids,2):pairs.append((i,j));pids.append(pid)
    pairs=np.array(pairs)
    cm,cd=compose_axes(aligned[test,:q],at,pairs,axes);metrics.update(cm)
    if not named:
        metrics['composition_mapping']='Oracle Dev Axis Matching (Hungarian + signs)'
        rawcm,_=compose_axes(zt[:,:q],at,pairs,axes);metrics.update({f'unmatched_{k}':v for k,v in rawcm.items()})
    else:metrics['composition_mapping']='intervention-defined' if cfg['method']=='ilr' else 'dense-supervised'
    natural=natural_pairs(tex,a[test],min_target=1,max_off=.5)
    natural=natural[np.isin(natural[:,2],axes)]
    if len(natural):natural[:,2]=np.array([{k:i for i,k in enumerate(axes)}[int(k)] for k in natural[:,2]])
    ns={};nd={}
    if len(natural):
        ai,bi,k,s=natural.T;dz=(aligned[indices[bi.astype(int)]]-aligned[indices[ai.astype(int)]])/norm
        ns,nd=selectivity(dz,k.astype(int),s,criteria=q);metrics.update({f'natural_{key}':v for key,v in ns.items()});metrics['natural_n']=len(natural)
    population.update(human_test='complete',natural='complete' if len(natural) else 'empty')
    savez_atomic(path/'evaluation.npz',correlation=corr(zt,at,rank=True),matched_correlation=mat,oracle_dev_order=order,oracle_dev_sign=sign,z_aligned=aligned[test],attributes=at,test_prompt_ids=tex.prompt_id.to_numpy(str),natural_pairs=natural,criteria_indices=np.array(axes),**{f'natural_{k}':v for k,v in nd.items()},**{f'composition_{k}':v for k,v in cd.items()},composition_prompt_ids=np.array(pids,dtype=str))
    # Held-out generated interventions use TRAIN-origin prompts excluded from intervention training.
    ipath=Path(f'outputs/features/{cfg["backbone_tag"]}_interventions')
    if (ipath/'COMPLETE.json').exists():
        identity=validate_feature_table(ipath,'data/processed/interventions.parquet',f'data/processed/tokenized_{cfg["backbone_tag"]}_interventions.parquet')
        if cfg.get('intervention_identity') is not None:assert identity==cfg['intervention_identity'],'Intervention cache differs from fitted run'
        for key in ['model','model_revision','tokenizer_revision','layers','d','pooling_sha256']:
            assert cfg['feature_manifest'][key]==identity['manifest'][key],f'Original/modified feature-space mismatch: {key}'
        ir=pd.read_parquet('data/processed/interventions.parquet');ir=ir[(ir.int_split=='test')&ir.strict&ir.criterion_index.isin(axes)]
        population['heldout_generated']='empty' if not len(ir) else 'complete'
        if len(ir):
            from models.latent_linear import LatentLinear
            from models.latent_mlp import LatentMLP
            from models.multi_attribute import MultiAttribute
            m=MultiAttribute(len(pred['mean'])) if cfg['method']=='attribute' else (LatentMLP if cfg['latent_architecture']=='mlp' else LatentLinear)(len(pred['mean']),cfg['latent_dim'])
            m.load_state_dict(torch.load(path/'model.pt',map_location='cpu',weights_only=True));m.eval()
            ih=np.asarray(np.load(ipath/f'layer{cfg["representation_layer"]}_{cfg["pooling_method"]}.npy',mmap_mode='r')[ir.feature_index],dtype='float32');ih=(ih-pred['mean'])/pred['scale']
            assert np.isfinite(ih).all()
            with torch.inference_mode():iz=np.concatenate([m.encode(x).numpy() for x in torch.tensor(ih).split(256)])
            delta=(iz[:,all_order]*all_sign-aligned[ir.base_feature_index])/norm
            mapped=ir.criterion_index.map({k:i for i,k in enumerate(axes)}).to_numpy()
            im,idd=selectivity(delta,mapped,ir.direction.to_numpy(),criteria=q);metrics.update(im)
            metrics['heldout_intervention_n']=len(ir)
            savez_atomic(path/'heldout_interventions.npz',intervention_ids=ir.intervention_id.to_numpy(str),prompt_ids=ir.prompt_id.to_numpy(str),criterion=ir.criterion_index.to_numpy(),coordinate_criterion=mapped,direction=ir.direction.to_numpy(),delta_z=delta,**idd)
    metrics['evaluation_status']='semantic_evaluated';atomic_json(path/'metrics.json',metrics)
    atomic_json(path/'EVALUATION_POPULATIONS.json',population)
    print(path.name,{k:metrics.get(k) for k in ['preference_accuracy','axis_purity','composition_accuracy','natural_n','target_effect']},flush=True)

def stability(paths):
    groups={}
    for p in paths:
        r=json.loads((p/'metrics.json').read_text())
        if r['method']=='scalar':continue
        key=tuple(r[k] for k in ['backbone','representation_layer','pooling_method','method','latent_architecture','latent_dim','n_interventions','intervention_filter'])+(r.get('control','standard'),tuple(r.get('criteria_indices',[0,1,2,3])))
        groups.setdefault(key,[]).append((p,r))
    ex=pd.read_parquet('data/processed/responses.parquet').sort_values('feature_index');test=ex.split.to_numpy()=='test'
    for items in groups.values():
        for p,r in items:
            x=np.load(p/'predictions.npz')['z'][test];ax=[];sub=[];cks=[];semantic=[];semantic_sub=[];residual=[];pair_records=[]
            xa=np.load(p/'evaluation.npz')['z_aligned'];dim=r['criteria_count']
            for q,s in items:
                if q==p:continue
                y=np.load(q/'predictions.npz')['z'][test];c=corr(x,y)
                if r['method'] in ['ilr','attribute'] and (r['method']=='attribute' or r['n_interventions']):ax.append(np.diag(c)[:r['criteria_count']].mean())
                else:
                    order,_=matching(x,y);ax.append(abs(c[order,np.arange(y.shape[1])]).mean())
                sub.append(cca(x,y));cks.append(cka(x,y))
                ya=np.load(q/'evaluation.npz')['z_aligned'];semantic.append(float(np.diag(corr(xa[:,:dim],ya[:,:dim])).mean()));semantic_sub.append(cca(xa[:,:dim],ya[:,:dim]))
                if xa.shape[1]>dim:residual.append(float(abs(corr(xa[:,dim:],ya[:,dim:])).mean()))
                pair_records.append({'other_seed':s['seed'],'full_axis_stability':float(ax[-1]),'full_subspace_overlap':float(sub[-1]),'dev_named_semantic_axis_stability':semantic[-1],'dev_named_semantic_subspace_overlap':semantic_sub[-1]})
            r.update(axis_stability=float(np.mean(ax)) if ax else None,subspace_stability=float(np.mean(sub)) if sub else None,subspace_linear_cka=float(np.mean(cks)) if cks else None)
            r.update(dev_named_semantic_axis_stability=float(np.mean(semantic)) if semantic else None,dev_named_semantic_subspace_stability=float(np.mean(semantic_sub)) if semantic_sub else None,residual_axis_stability=float(np.mean(residual)) if residual else None)
            atomic_json(p/'seed_pair_stability.json',pair_records)
            atomic_json(p/'metrics.json',r)

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('--pattern',default='*');args=p.parse_args()
    torch.set_num_threads(2);paths=[p for p in Path('outputs/real').glob(args.pattern) if (p/'TRAINED.json').exists()]
    for path in paths:evaluate(path)
    stability(paths)
    for path in paths:
        save_evaluation_provenance(path)
        populations=json.loads((path/'EVALUATION_POPULATIONS.json').read_text());populations['stability']='complete' if populations['stability']!='not_applicable' else 'not_applicable'
        atomic_json(path/'EVALUATION_POPULATIONS.json',populations)
        inputs={str(p):file_hash(p) for p in [Path('data/processed/responses.parquet'),Path('data/processed/attributes.parquet'),Path('data/processed/preferences.parquet'),path/'predictions.npz',path/'model.pt']}
        if populations['heldout_generated'] in ['complete','empty']:inputs['data/processed/interventions.parquet']=file_hash('data/processed/interventions.parquet')
        atomic_json(path/'EVALUATED.json',{'completed_at':datetime.datetime.now(datetime.timezone.utc).isoformat(),'populations':populations,'input_hashes':inputs,'evaluation_source':json.loads((path/'evaluation_source_manifest.json').read_text()),'artifact_hashes':{p.name:file_hash(p) for p in [path/'metrics.json',path/'evaluation.npz',path/'heldout_interventions.npz'] if p.exists()}})
    rebuild()