#!/usr/bin/env python3 from __future__ import annotations import argparse import hashlib import json from pathlib import Path from typing import Any from datetime import datetime, timezone import numpy as np import torch from scipy.spatial.distance import pdist from scipy.stats import spearmanr from torch.utils.data import DataLoader from pino.heads import PIMTHeads from pino.pimt_model import ( DEFAULT_EMBEDDING_DIM, FragranceTrajectoryDataset, PhysicsInformedMixtureTransformer, ) from pino.train import pad_trajectory_collate from pino.upload_data import DEFAULT_TRAIN_RATIO, create_molecule_disjoint_split_from_records PYRFUME_SOURCE = "pyrfume-cas" SYNTHETIC_SOURCES = {"openpom-predicted", "volatility-fallback"} DEFAULT_GENRES = ["citrus_cologne", "fougere", "floral_woody", "amber_oriental", "wildcard"] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Evaluate a PINO PIMT checkpoint.") parser.add_argument("--checkpoint", default="models/pimt_v6.pt", help="Local checkpoint path.") parser.add_argument("--data", default="data/empirical_dataset_v8.jsonl", help="Evaluation dataset JSONL.") parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--train-ratio", type=float, default=DEFAULT_TRAIN_RATIO, help="Molecule-disjoint train molecule ratio used to rebuild the 15% holdout validation set.") parser.add_argument("--structural-source", choices=["morgan", "openpom_256", "pom_alltags"], default="morgan", help="Structural embedding block matching the checkpoint's training arm.") parser.add_argument("--output", default="metrics_v6_current.json", help="Metrics JSON output path.") parser.add_argument( "--hf-repo", default=None, help="Optional HF model repo to download --hf-filename from before evaluation.", ) parser.add_argument("--hf-filename", default=None, help="Checkpoint filename in --hf-repo.") parser.add_argument("--hf-token", default=None, help="HF token override; defaults to cached auth/env.") return parser.parse_args() def resolve_checkpoint(args: argparse.Namespace) -> Path: checkpoint = Path(args.checkpoint) if checkpoint.exists(): return checkpoint if not args.hf_repo: raise FileNotFoundError(f"Checkpoint not found: {checkpoint}") from huggingface_hub import hf_hub_download filename = args.hf_filename or checkpoint.name downloaded = hf_hub_download( repo_id=args.hf_repo, filename=filename, token=args.hf_token, ) return Path(downloaded) def read_records(path: str | Path) -> list[dict[str, Any]]: with Path(path).open("r", encoding="utf-8") as f: return [json.loads(line) for line in f if line.strip()] def sha256_file(path: str | Path) -> str: digest = hashlib.sha256() with Path(path).open("rb") as f: for chunk in iter(lambda: f.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def make_validation_dataset(records: list[dict[str, Any]], *, train_ratio: float, seed: int, structural_source: str = "morgan", objective_dim: int = 138) -> FragranceTrajectoryDataset: split = create_molecule_disjoint_split_from_records(records, train_ratio=train_ratio, seed=seed) genre_map = {genre: idx for idx, genre in enumerate(DEFAULT_GENRES)} return FragranceTrajectoryDataset( records=split["validation"], use_embedding_fallback=True, structural_source=structural_source, genre_map=genre_map, objective_dim=objective_dim, ) def load_vocab(path: str | Path = "data/pyrfume_vocabulary.json") -> list[str]: with Path(path).open("r", encoding="utf-8") as f: raw = json.load(f) if isinstance(raw, dict) and isinstance(raw.get("vocabulary"), list): return [str(v) for v in raw["vocabulary"]] if isinstance(raw, dict): return [str(raw[k]) for k in sorted(raw)] if isinstance(raw, list): return [str(v) for v in raw] return [] def infer_model_dims(state_dict: dict[str, torch.Tensor]) -> dict[str, int]: input_weight = state_dict["input_proj.weight"] hidden_dim, embedding_dim = input_weight.shape state_dim = state_dict["gating.physics_scaler"].shape[0] num_layers = max( int(k.split(".")[2]) for k in state_dict if k.startswith("encoder.layers.") and k.endswith(".self_attn.in_proj_weight") ) + 1 in_proj_rows = state_dict["encoder.layers.0.self_attn.in_proj_weight"].shape[0] num_heads = 8 if hidden_dim % 8 == 0 and hidden_dim >= 512 else 4 if in_proj_rows != hidden_dim * 3: raise ValueError("Unexpected transformer attention weight shape") return { "hidden_dim": hidden_dim, "embedding_dim": embedding_dim, "state_dim": state_dim, "num_layers": num_layers, "num_heads": num_heads, } def evaluate(args: argparse.Namespace) -> dict[str, Any]: torch.manual_seed(args.seed) np.random.seed(args.seed) checkpoint_path = resolve_checkpoint(args) ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) model_state = ckpt["model_state_dict"] dims = infer_model_dims(model_state) model = PhysicsInformedMixtureTransformer( embedding_dim=dims["embedding_dim"], state_dim=dims["state_dim"], hidden_dim=dims["hidden_dim"], num_heads=dims["num_heads"], num_layers=dims["num_layers"], ) model.load_state_dict(model_state) model.eval() objective_dim = 575 if args.structural_source == "pom_alltags" else 138 heads = PIMTHeads(hidden_dim=dims["hidden_dim"], objective_dim=objective_dim) heads_missing, heads_unexpected = heads.load_state_dict(ckpt["heads_state_dict"], strict=False) heads.eval() records = read_records(args.data) engine_source = "openpom_256" if args.structural_source == "pom_alltags" else args.structural_source val_dataset = make_validation_dataset(records, train_ratio=args.train_ratio, seed=args.seed, structural_source=engine_source, objective_dim=objective_dim) val_loader = DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, collate_fn=pad_trajectory_collate, ) all_preds: list[np.ndarray] = [] all_targets: list[np.ndarray] = [] with torch.no_grad(): for batch in val_loader: latent = model( batch["tokens"], batch["physics"], src_key_padding_mask=batch["src_key_padding_mask"], ) output = heads(latent, batch["physics"], batch["src_key_padding_mask"]) all_preds.append(output["objective"].numpy()) all_targets.append(batch["target_obj"].numpy()) preds = np.concatenate(all_preds, axis=0) targets = np.concatenate(all_targets, axis=0) tier_names = ["top", "middle", "base"] tier_metrics: dict[str, Any] = {} for tier_idx, tier_name in enumerate(tier_names): pred_tier = preds[:, tier_idx, :] target_tier = targets[:, tier_idx, :] sims = 1.0 - pdist(pred_tier, metric="cosine") rho_list = [] for idx in range(len(pred_tier)): rho, _ = spearmanr(pred_tier[idx], target_tier[idx]) if rho is not None and not np.isnan(rho): rho_list.append(float(rho)) tier_metrics[tier_name] = { "mse": float(((pred_tier - target_tier) ** 2).mean()), "mae": float(np.abs(pred_tier - target_tier).mean()), "prediction_variance": float(pred_tier.var()), "target_variance": float(target_tier.var()), "mean_pairwise_cosine_similarity": float(sims.mean()) if len(sims) else None, "mean_spearman_rho": float(np.mean(rho_list)) if rho_list else None, } flat_rho, flat_p = spearmanr(preds.reshape(-1), targets.reshape(-1)) profile_pred = preds.mean(axis=1) profile_target = targets.mean(axis=1) profile_rho, profile_p = spearmanr(profile_pred.reshape(-1), profile_target.reshape(-1)) provenance_metrics: dict[str, Any] = {} if val_dataset.records and any("pyramid_target_provenance" in r for r in val_dataset.records): provenance = np.array([ r.get("pyramid_target_provenance", [["none"] * 138 for _ in range(3)]) for r in val_dataset.records ], dtype=object) active = targets > 0 provenance_slices = { "pyrfume_cas_only": active & (provenance == PYRFUME_SOURCE), "openpom_or_volatility_only": active & np.isin(provenance, list(SYNTHETIC_SOURCES)), "full_reconstructed": active, } for name, mask in provenance_slices.items(): n = int(mask.sum()) if n >= 3: labelled_cells = mask.any(axis=0) sliced_preds = preds[:, labelled_cells] sliced_targets = np.where(mask, targets, 0.0)[:, labelled_cells] rho, p_value = spearmanr(sliced_preds.reshape(-1), sliced_targets.reshape(-1)) provenance_metrics[name] = { "n_positive_values": n, "n_tier_descriptor_cells": int(labelled_cells.sum()), "spearman_rho": float(rho), "spearman_p_value": float(p_value), } else: provenance_metrics[name] = { "n_positive_values": n, "n_tier_descriptor_cells": 0, "spearman_rho": None, "spearman_p_value": None, } synth_n = provenance_metrics["openpom_or_volatility_only"]["n_positive_values"] full_n = provenance_metrics["full_reconstructed"]["n_positive_values"] provenance_metrics["interpretation"] = ( "openpom-predicted + volatility-fallback are the majority of active labels; full-set rho is product/internal model-agreement, not a rigorous chemistry result." if full_n and synth_n / full_n > 0.5 else "full-set rho is still split by provenance; only pyrfume-cas-only is suitable as the closest trustworthy profile subset." ) vocab = load_vocab() top_descriptors: dict[str, Any] = {} for tier_idx, tier_name in enumerate(tier_names): pred_mean = preds[:, tier_idx, :].mean(axis=0) target_mean = targets[:, tier_idx, :].mean(axis=0) pred_top = np.argsort(pred_mean)[-5:][::-1] target_top = np.argsort(target_mean)[-5:][::-1] top_descriptors[tier_name] = { "predicted": [ {"descriptor": vocab[i] if i < len(vocab) else f"d{i}", "score": float(pred_mean[i])} for i in pred_top ], "target": [ {"descriptor": vocab[i] if i < len(vocab) else f"d{i}", "score": float(target_mean[i])} for i in target_top ], } return { "generated_at_utc": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"), "checkpoint": { "path": str(checkpoint_path), "sha256": sha256_file(checkpoint_path), "epoch": ckpt.get("epoch"), "val_loss": ckpt.get("val_loss"), "missing_head_keys_on_load": list(heads_missing), "unexpected_head_keys_on_load": list(heads_unexpected), **dims, }, "data": { "path": args.data, "sha256": sha256_file(args.data), "n_records": len(records), "n_validation_samples": len(val_dataset), }, "prediction": { "profile_mse": float(((profile_pred - profile_target) ** 2).mean()), "profile_mae": float(np.abs(profile_pred - profile_target).mean()), "profile_prediction_variance": float(profile_pred.var()), "tier_prediction_variance": float(preds.var()), "tier_metrics": tier_metrics, }, "spearman": { "flat_spearman_rho": float(flat_rho), "flat_spearman_p_value": float(flat_p), "profile_spearman_rho": float(profile_rho), "profile_spearman_p_value": float(profile_p), }, "provenance_spearman": provenance_metrics, "top_descriptors": top_descriptors, } def main() -> None: args = parse_args() metrics = evaluate(args) output = Path(args.output) output.write_text(json.dumps(metrics, indent=2), encoding="utf-8") print(f"Checkpoint: {metrics['checkpoint']['path']}") print(f"Validation samples: {metrics['data']['n_validation_samples']}") print(f"Profile MSE: {metrics['prediction']['profile_mse']:.6f}") print(f"Profile Spearman rho: {metrics['spearman']['profile_spearman_rho']:.4f}") if metrics.get("provenance_spearman"): for name, vals in metrics["provenance_spearman"].items(): if isinstance(vals, dict) and "spearman_rho" in vals: rho = vals["spearman_rho"] rho_s = f"{rho:.4f}" if rho is not None else "n/a" print(f"{name}: rho={rho_s} n={vals['n_positive_values']}") for tier, vals in metrics["prediction"]["tier_metrics"].items(): print( f"{tier.title():<7} MSE={vals['mse']:.4f} " f"MAE={vals['mae']:.4f} Var(pred)={vals['prediction_variance']:.6f} " f"rho={vals['mean_spearman_rho']:.4f}" ) print(f"Wrote {output}") if __name__ == "__main__": main()