Blame for source/multitask_evaluate.py

"""Matched Kumo Large probes on each joint checkpoint and the exact author holdout."""
from pathlib import Path
import sys
sys.path.insert(0,str(Path(__file__).resolve().parent))
from multitask_model import *
import random
import bigsoldb_evaluate as reference

EVAL=OUT/'evaluation'


def embeddings(epoch):
    path=EVAL/'features'/f'epoch_{epoch:02d}.npz'
    checkpoint=OUT/'checkpoints'/f'epoch_{epoch:02d}'
    hashes={n:sha(checkpoint/n) for n in ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth']}
    geometry=pd.read_parquet(DATA/'bigsoldb_fixed_geometries.parquet')
    if path.exists():
        saved=np.load(path);assert saved['smiles'].tolist()==geometry.smiles.tolist()
        assert json.loads(str(saved['checkpoint_sha256']))==hashes
        return saved['X'],hashes
    model=MultiTask(checkpoint);model.eval().requires_grad_(False)
    cache=RawCache('bigsoldb');X=np.empty((len(geometry),640),dtype=np.float32)
    with torch.inference_mode():
        for ids in common.batch_indices(np.arange(len(geometry)),cache.lengths,max_size=24,atom_cap=2400):
            _,_,z=model.encode(cache.sample(ids));X[ids]=z.float().cpu().numpy()
    assert np.isfinite(X).all()
    np.savez_compressed(path,X=X,smiles=geometry.smiles.to_numpy(dtype=str),checkpoint_sha256=json.dumps(hashes))
    del model;gc.collect();torch.cuda.empty_cache()
    return X,hashes


def build_report(run):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    results=pd.read_csv(EVAL/'results.csv')
    old=pd.read_csv(WORK/'bigsoldb_evaluation/results.csv')
    starting=old.loc[old.representation=='cosmo_epoch10'].iloc[0]
    rows=[dict(model='PCQMep2 COSMOep10 + Kumo L (starting checkpoint)',origin='Prior matched evaluation',epoch=0,
        **{k:float(starting[k]) for k in reference.METRICS})]
    rows += [dict(model=f'Joint epoch {int(r.epoch)} + Kumo L',origin='This continuation',epoch=int(r.epoch),
        **{k:getattr(r,k) for k in reference.METRICS}) for r in results.itertuples()]
    monroe=old.loc[old.representation=='monroe'].iloc[0]
    rows.append(dict(model='Monroe + Kumo L',origin='Prior matched evaluation',epoch=None,**{k:float(monroe[k]) for k in reference.METRICS}))
    rows += [dict(model=name,origin='Solvaformer published Table 1',epoch=None,
        **dict(zip(reference.METRICS,values))) for name,values in reference.PUBLISHED.items()]
    comparison=pd.DataFrame(rows);comparison.to_csv(EVAL/'analysis/comparison.csv',index=False)
    fig,axes=plt.subplots(1,3,figsize=(12,3.7),layout='constrained')
    curve=pd.concat([pd.DataFrame([dict(epoch=0,**{k:float(starting[k]) for k in reference.METRICS})]),results],ignore_index=True)
    for ax,metric in zip(axes,['mae','rmse','r2']):
        ax.plot(curve.epoch,curve[metric],'-o',label='REM3DI + Kumo L',color='#0072B2')
        ax.axhline(float(monroe[metric]),color='#CC79A7',ls='--',label='Monroe + Kumo L')
        ax.axhline(dict(zip(reference.METRICS,reference.PUBLISHED['Solvaformer']))[metric],color='#666666',ls=':',label='Published Solvaformer')
        ax.set_xlabel('Joint pretraining epoch');ax.set_ylabel(metric.upper());ax.set_xticks(curve.epoch);ax.grid(alpha=.2)
    axes[0].legend(fontsize=7);fig.savefig(EVAL/'analysis/epoch_curve.png',dpi=180)
    fig.savefig(EVAL/'analysis/epoch_curve.svg');plt.close(fig)
    manifest=json.loads((DATA/'manifest.json').read_text())
    training=[]
    for p in sorted((OUT/'checkpoints').glob('epoch_[0-9][0-9]/metrics.json')):
        r=json.loads(p.read_text());training.append(dict(epoch=r['epoch'],**r['validation'],**{'train_'+k:v for k,v in r['train'].items()}))
    pd.DataFrame(training).to_csv(EVAL/'analysis/training_metrics.csv',index=False)
    lines=['# PCQMep2 → COSMOep10 → CombiSolv-QM multi-task continuation','',
        'All encoder, learned atomic/geometric preprocessor, and COSMO-head weights were initialized from the exact PCQM epoch 2 → COSMO epoch 10 checkpoint used in the preceding BigSolDB 2.0 evaluation. Byte-level hashes were checked against that evaluation protocol. The omitted embedding-denoising decoder alone was restored from PCQM epoch 2.','',
        f"Completed {CONFIG['epochs']} full joint epochs. Every retained COSMO training molecule and every retained CombiSolv-QM training pair occurs in every epoch; the shorter stream cycles to match the longer one. Per-epoch coverage and repeats are saved with each checkpoint. Frozen MACE inputs are cached in float32; both learned preprocessors, the REM3DI encoder and all three task heads train. The new solvation head is 1281 → 256 → 1 (SiLU), over the ordered solute640/solvent640/T÷298 concatenation.",'',
        'CombiSolv-QM predicts signed solvation free energy at 298 K, in kcal/mol; it is not a direct solubility target. Its labels use training-only mean/std standardization and no logarithm. The existing COSMO log1p/ln transforms and training-set standardization are preserved exactly so its pretrained output head keeps the same meaning. The COSMO task averages sigma-bin MSE, area MSE and volume MSE with equal task-component weight.','',
        'REM3DI denoising uses Gaussian noise σ=0.3 on preprocessed atomic embeddings, the original decoder, original atom mask, original geometric/pair detach paths and the original non-detached clean target. The historical atom-denoising loss is divided by its 1088 channels to express it per feature; the unnormalized original value is also logged. No coordinate-noise objective is substituted.','',
        'The three task log variances start at −0.5 and are learned with Monroe’s exact objective `sum(L_t / (2*exp(s_t)) + s_t/2)`. They are not softmax-normalized or constrained to sum to one. The separate decorrelation penalty is `8 * mean(CovOffDiag(COSMO embeddings), CovOffDiag(unique CombiSolv embeddings))`, where each penalty is the squared off-diagonal sample covariance sum divided by `D*(D-1)`. This uses Monroe’s released checkpoint coefficient of 8. Value and gradient comparisons against the pinned author source pass exactly.','',
        '## Data and exclusions','',
        f"The unchanged raw CombiSolv table contains {manifest['raw_combisolv_rows']:,} rows. Prepared memberships: `{json.dumps(manifest['combisolv_splits'],sort_keys=True)}`. COSMO memberships: `{json.dumps(manifest['cosmo_splits'],sort_keys=True)}`. Raw rows remain available in the dataset; exclusions affect new training membership only.",'',
        'The published BigSolDB held-out solute connectivity groups are excluded from all new training objectives and from both CombiSolv input roles. The original COSMO ESOL/FreeSolv quarantine is preserved; CombiSolv is protected for the requested BigSolDB holdout and is not claimed to preserve ESOL/FreeSolv quarantine. CombiSolv validation holds out 2% of eligible solute connectivity groups across objectives; the original COSMO validation groups are also excluded from CombiSolv training. Supplied source InChI defines CombiSolv molecular identity; the original SMILES/InChI mapping is retained. These exclusions cannot undo earlier exposure already present in the starting COSMO10/PCQM2 checkpoint, so this is a continuation comparison, not a claim of fully decontaminated foundation pretraining.','',
        '## BigSolDB 2.0 frozen probes','',
        'The exact previously audited author test partition is reused: 9,686 measurements from 126 held-out solutes, with all 75,731 non-test measurements as Kumo context. Solute/solvent conformers are byte-identical to those used for the starting checkpoint. Kumo Large uses four estimators, seed 42, float16 inference, query chunks 128, and all 1281 concatenated features including temperature in Kelvin. Its transforms fit on training rows only. Each checkpoint is reported; test scores do not select a pretraining checkpoint or alter the protocol.','',
        '| Model | MAE | MSE | RMSE | R² | Pearson | Spearman |','|---|---:|---:|---:|---:|---:|---:|']
    for r in comparison.itertuples():lines.append('| '+r.model+' | '+' | '.join(f'{getattr(r,k):.4f}' for k in reference.METRICS)+' |')
    lines += ['', 'Lower MAE/MSE/RMSE and higher R²/correlations are better. Error units are log10 mol/L (squared for MSE). Published models use different pretraining and model classes. The original paper prose says 9,250 test rows; its committed prediction file contains 9,686 and reproduces all six published MPNN+Charges metrics. Its private neural train/validation subdivision is unavailable; Kumo uses the audited non-test complement. No CV standard deviation or significance claim is attached to this single released holdout.','',
        '![Epoch-wise frozen Kumo probes](epoch_curve.png)','',
        f'[CombiSolv source](https://zenodo.org/records/5970538) · [Prepared dataset](https://huggingface.co/datasets/{DATA_REPO}) · [Checkpoints](https://huggingface.co/{MODEL_REPO}) · [Probe W&B]({run.url}) · [Training W&B]({json.loads((OUT/"wandb.json").read_text())["url"]})','',
        f'[Monroe UW source](https://github.com/blazejba/Monroe/blob/{MONROE_REV}/monroe/mtl/UW.py) · [Monroe covariance source](https://github.com/blazejba/Monroe/blob/{MONROE_REV}/monroe/train/loss.py) · [Solvaformer Table 1]({reference.PAPER_URL})','']
    report='\n'.join(lines);(EVAL/'analysis/REPORT.md').write_text(report)
    (OUT/'README.md').write_text('---\ntags:\n- chemistry\n- multitask\n---\n\n'+report.replace('](epoch_curve.png)','](evaluation/analysis/epoch_curve.png)'))
    return comparison


def publish(run):
    files=[]
    for p in OUT.rglob('*'):
        if not p.is_file():continue
        rel=p.relative_to(OUT)
        if rel.parts[0] in ['cache','data','checkpoints','wandb'] or '__pycache__' in rel.parts:continue
        if p.suffix in ['.json','.csv','.npz','.png','.svg','.md']:files.append(p)
    sources=[Path(__file__).with_name(n) for n in ['multitask_data.py','multitask_model.py','multitask_train.py','multitask_evaluate.py','multitask_driver.py',
        'eval_common.py','train_cosmo.py','bigsoldb_evaluate.py','openadmet_evaluate.py']]
    ops=[CommitOperationAdd(path_in_repo=str(p.relative_to(OUT)),path_or_fileobj=str(p)) for p in files]
    ops += [CommitOperationAdd(path_in_repo='source/'+p.name,path_or_fileobj=str(p)) for p in sources]
    commit=HfApi().create_commit(repo_id=MODEL_REPO,operations=ops,commit_message='Publish complete joint-training provenance and every matched BigSolDB Kumo probe')
    entries={s.rfilename:s for s in HfApi().model_info(MODEL_REPO,revision=commit.oid,files_metadata=True).siblings}
    for p in files:
        entry=entries[str(p.relative_to(OUT))]
        if entry.lfs:assert entry.lfs.sha256==sha(p)
    import wandb
    artifact=wandb.Artifact('rem3di-cosmo10-combisolv-multitask-report',type='evaluation-report')
    for p in files:artifact.add_file(str(p),name=str(p.relative_to(OUT)))
    for p in sources:artifact.add_file(str(p),name='source/'+p.name)
    run.log_artifact(artifact,aliases=['complete','latest']).wait(timeout=60)
    return dict(repository=MODEL_REPO,revision=commit.oid,verified=True,report=f'https://huggingface.co/{MODEL_REPO}/blob/{commit.oid}/evaluation/analysis/REPORT.md')


def main():
    import sdm,wandb
    EVAL.mkdir(exist_ok=True)
    for name in ['features','predictions','analysis','data']:(EVAL/name).mkdir(exist_ok=True)
    lock=(OUT/'probing.lock').open('a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
    assert json.loads((OUT/'training_complete.json').read_text())['epochs']==CONFIG['epochs']
    train=pd.read_csv(WORK/'bigsoldb_evaluation/data/train.csv');test=pd.read_csv(WORK/'bigsoldb_evaluation/data/test.csv')
    assert len(train)==75731 and len(test)==9686
    for name in ['train.csv','test.csv','split_validation.json']:
        source=WORK/'bigsoldb_evaluation/data'/name;target=EVAL/'data'/name
        if not target.exists():shutil.copyfile(source,target)
        assert sha(source)==sha(target)
    indices=np.load(WORK/'bigsoldb_evaluation/features/pair_indices.npz')
    assert np.array_equal(indices['test_clean_rows'],test['_clean_row']) and np.array_equal(indices['train_clean_rows'],train['_clean_row'])
    solute=indices['solute'];solvent=indices['solvent'];temperature=indices['temperature_K']
    y=pd.concat([train,test])['LogS(mol/L)'].to_numpy(dtype=np.float64)
    shutil.copyfile(WORK/'bigsoldb_evaluation/features/pair_indices.npz',EVAL/'features/pair_indices.npz')
    write(EVAL/'protocol.json',dict(upstream_revision=reference.UPSTREAM_REVISION,
        preceding_evaluation_repository='Queimo/rem3di-pcqm-epoch2-cosmo',
        preceding_evaluation_revision='9266a41fbe3b36afe0615e0c3ba232ca74bdeef7',
        train_rows=len(train),test_rows=len(test),split_files_sha256={n:sha(EVAL/'data'/n) for n in ['train.csv','test.csv','split_validation.json']},
        pair_indices_sha256=sha(EVAL/'features/pair_indices.npz'),
        shared_original_geometries_sha256=sha(WORK/'bigsoldb_evaluation/features/shared_graphs.parquet'),
        prepared_geometries_sha256=sha(DATA/'bigsoldb_fixed_geometries.parquet'),
        kumo=dict(size='large',checkpoint_repository='nvidia/Kumo-Tabular',checkpoint_revision='v1.0.1',
            num_estimators=4,seed=42,precision='float16 autocast',full_context=75731,columns=1281,query_chunk=128),
        versions=json.loads((WORK/'bigsoldb_evaluation/protocol.json').read_text())['versions']))
    wb_path=OUT/'probe_wandb.json';wb=json.loads(wb_path.read_text()) if wb_path.exists() else dict(run_id=wandb.util.generate_id());write(wb_path,wb)
    run=wandb.init(entity='queimo',project='rem3di-reproduction',id=wb['run_id'],resume='allow',mode='online',
        name='pcqm2-cosmo10-combisolv-every-epoch-bigsoldb-kumo-large',group='rem3di-cosmo-combisolv',job_type='frozen-probing',
        config=CONFIG,dir=str(OUT),settings=wandb.Settings(init_timeout=90))
    write(wb_path,dict(run_id=run.id,url=run.url))
    records=pd.read_csv(EVAL/'results.csv').to_dict('records') if (EVAL/'results.csv').exists() else []
    reference.status=status
    try:
        for epoch in range(1,CONFIG['epochs']+1):
            if any(int(r['epoch'])==epoch for r in records):continue
            status('embedding_for_probe',epoch=epoch,wandb_url=run.url)
            z,hashes=embeddings(epoch)
            X=np.concatenate([z[solute],z[solvent],temperature[:,None]],axis=1).astype(np.float32)
            assert X.shape==(len(y),1281) and np.isfinite(X).all()
            torch.manual_seed(42);torch.cuda.manual_seed_all(42);np.random.seed(42);random.seed(42)
            model=sdm.models.KumoTabular(task='regression',size='large',device='cuda')
            started=time.perf_counter();status('fitting_kumo',epoch=epoch,wandb_url=run.url)
            shapes=reference.fit_context(model,X,y,len(train),1)
            pred=reference.predict(model,X[len(train):]);model.clear();del model,X,z;gc.collect();torch.cuda.empty_cache()
            result=dict(epoch=epoch,method='frozen_concat_kumo_large',n_train=len(train),n_test=len(test),raw_columns=1281,
                context_blocks=1,num_estimators=4,seed=42,seconds=time.perf_counter()-started,**reference.metrics(y[len(train):],pred))
            np.savez_compressed(EVAL/'predictions'/f'epoch_{epoch:02d}.npz',y_true=y[len(train):],y_pred=pred,
                test_clean_rows=test['_clean_row'].to_numpy(),test_raw_rows=test['_raw_row'].to_numpy(),checkpoint_sha256=json.dumps(hashes))
            write(EVAL/'features'/f'epoch_{epoch:02d}_context.json',dict(processed_context_shapes=shapes,n_train=len(train),full_context=True,
                all_raw_columns=1281,seed=42,num_estimators=4,checkpoint_sha256=hashes))
            records.append(result);pd.DataFrame(records).sort_values('epoch').to_csv(EVAL/'results.csv',index=False)
            run.log(dict(epoch=epoch,**{'bigsoldb/'+k:v for k,v in result.items() if isinstance(v,(int,float))}))
            write(EVAL/'progress.json',dict(completed=len(records),total=CONFIG['epochs'],latest=result))
        assert sorted(int(r['epoch']) for r in records)==list(range(1,CONFIG['epochs']+1))
        for r in records:
            z=np.load(EVAL/'predictions'/f'epoch_{int(r["epoch"]):02d}.npz')
            assert np.array_equal(z['y_true'],test['LogS(mol/L)']) and np.array_equal(z['test_clean_rows'],test['_clean_row'])
            for key,value in reference.metrics(z['y_true'],z['y_pred']).items():assert np.isclose(value,r[key],rtol=1e-8,atol=1e-9)
        comparison=build_report(run)
        run.log({'bigsoldb/comparison':wandb.Table(dataframe=comparison),'bigsoldb/epoch_curve':wandb.Image(str(EVAL/'analysis/epoch_curve.png'))})
        write(EVAL/'analysis/validation.json',dict(passed=True,epochs=CONFIG['epochs'],all_metrics_recomputed=True,
            exact_author_test_rows=True,all_training_rows_in_each_context=True,all1281_features=True))
        hub=publish(run);write(OUT/'hub_verified.json',hub)
        run.summary.update(dict(complete=True,epochs=CONFIG['epochs'],hub_revision=hub['revision']));run.finish()
        write(OUT/'complete.json',dict(completed=True,epochs=CONFIG['epochs'],probes=CONFIG['epochs'],hub=hub,completed_unix=time.time()))
        status('complete',epochs=CONFIG['epochs'],probes=len(records),report=hub['report'])
    except Exception:
        status('failed_evaluation',error=traceback.format_exc());run.finish(exit_code=1);raise


if __name__=='__main__':
    try:main()
    except Exception:status('failed_evaluation',error=traceback.format_exc());raise