from __future__ import annotations import argparse import json import logging from pathlib import Path from typing import Any import numpy as np import torch import torch.nn as nn from scipy.spatial.distance import cdist from scipy.stats import spearmanr from torch.utils.data import DataLoader from pino.heads import PIMTHeads from pino.pimt_model import FragranceTrajectoryDataset, PhysicsInformedMixtureTransformer from pino.train import pad_trajectory_collate logger = logging.getLogger("pino.evaluate") @torch.no_grad() def extract_embeddings( model: nn.Module, heads: nn.Module, loader: DataLoader, device: torch.device, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Return predicted/target descriptor profiles plus genre labels.""" model.eval() heads.eval() predictions: list[np.ndarray] = [] targets: list[np.ndarray] = [] genre_labels: list[np.ndarray] = [] tier_predictions: list[np.ndarray] = [] for batch in loader: tokens = batch["tokens"].to(device) physics = batch["physics"].to(device) src_key_padding_mask = batch["src_key_padding_mask"].to(device) target_obj = batch["target_obj"].to(device) latent = model(tokens, physics, src_key_padding_mask) pred = heads(latent, physics, src_key_padding_mask) tier_pred = pred["objective"].cpu().numpy() tier_target = target_obj.cpu().numpy() tier_predictions.append(tier_pred) # Collapse the top/middle/base pyramid into one perceptual descriptor # profile for retrieval/distance metrics. predictions.append(tier_pred.mean(axis=1)) targets.append(tier_target.mean(axis=1)) genre_labels.append(batch["genre_label"].cpu().numpy()) return ( np.vstack(predictions), np.vstack(targets), np.concatenate(genre_labels), np.concatenate(tier_predictions, axis=0), ) def spearman_distance_correlation(embeddings: np.ndarray, targets: np.ndarray) -> dict[str, float]: """ Compute Spearman correlation between latent-space distances and target-space distances. Returns the mean and per-sample Spearman rho between latent and descriptor distances. """ latent_dist = cdist(embeddings, embeddings, metric="euclidean") target_dist = cdist(targets, targets, metric="euclidean") # Use upper triangular indices to avoid self-pairs and duplicates. n = len(embeddings) idx = np.triu_indices(n, k=1) lat_vec = latent_dist[idx] tgt_vec = target_dist[idx] rho, p_value = spearmanr(lat_vec, tgt_vec) return { "spearman_rho": float(rho), "spearman_p_value": float(p_value), "n_pairs": int(len(lat_vec)), } def mean_average_precision(embeddings: np.ndarray, labels: np.ndarray, k: int | None = None) -> dict[str, float]: """ Compute retrieval mAP in the latent space using genre labels as relevance. For each sample, positives are same-genre samples; negatives are all others. """ k = k or len(embeddings) - 1 distances = cdist(embeddings, embeddings, metric="euclidean") np.fill_diagonal(distances, np.inf) aps: list[float] = [] for i in range(len(embeddings)): sorted_idx = np.argsort(distances[i])[:k] relevant = (labels[sorted_idx] == labels[i]).astype(np.int32) if relevant.sum() == 0: continue # Cumulative precision at each relevant rank. precisions = [] num_relevant = 0 for rank, is_rel in enumerate(relevant, start=1): if is_rel: num_relevant += 1 precisions.append(num_relevant / rank) aps.append(np.mean(precisions) if precisions else 0.0) return { "mAP": float(np.mean(aps)) if aps else 0.0, "AP_std": float(np.std(aps)) if aps else 0.0, "n_queries": len(aps), } def _infer_checkpoint_config(checkpoint: dict[str, Any]) -> dict[str, int]: """Infer architecture knobs saved by the current trainer.""" model_state = checkpoint["model_state_dict"] input_proj = model_state["input_proj.weight"] hidden_dim = int(input_proj.shape[0]) embedding_dim = int(input_proj.shape[1]) state_dim = int(model_state["gating.physics_scaler"].shape[0]) layer_prefixes = { key.split(".")[2] for key in model_state if key.startswith("encoder.layers.") and key.endswith("self_attn.in_proj_weight") } num_layers = len(layer_prefixes) or 4 # num_heads is not represented in the state dict. Current HF training uses # four heads; fall back to the largest common divisor up to 8 for older runs. num_heads = 4 if hidden_dim % 4 == 0 else next(h for h in (8, 2, 1) if hidden_dim % h == 0) return { "hidden_dim": hidden_dim, "embedding_dim": embedding_dim, "state_dim": state_dim, "num_layers": num_layers, "num_heads": num_heads, } def evaluate(checkpoint_path: str, data_path: str, batch_size: int = 16, hidden_dim: int | None = None) -> dict[str, Any]: """Load a checkpoint and run perceptual evaluation metrics.""" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") checkpoint = torch.load(checkpoint_path, map_location=device) cfg = _infer_checkpoint_config(checkpoint) if hidden_dim is not None and hidden_dim != cfg["hidden_dim"]: logger.warning("Ignoring --hidden-dim=%d; checkpoint requires hidden_dim=%d", hidden_dim, cfg["hidden_dim"]) all_genres = ["citrus_cologne", "fougere", "floral_woody", "amber_oriental", "wildcard"] genre_map = {genre: idx for idx, genre in enumerate(all_genres)} dataset = FragranceTrajectoryDataset(data_path=data_path, use_embedding_fallback=True, genre_map=genre_map) loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=pad_trajectory_collate) model = PhysicsInformedMixtureTransformer( embedding_dim=cfg["embedding_dim"], state_dim=cfg["state_dim"], hidden_dim=cfg["hidden_dim"], num_heads=cfg["num_heads"], num_layers=cfg["num_layers"], ).to(device) heads = PIMTHeads(hidden_dim=cfg["hidden_dim"], objective_dim=138).to(device) model.load_state_dict(checkpoint["model_state_dict"]) heads.load_state_dict(checkpoint["heads_state_dict"]) predictions, targets, genre_labels, tier_predictions = extract_embeddings(model, heads, loader, device) spearman_metrics = spearman_distance_correlation(predictions, targets) map_metrics = mean_average_precision(predictions, genre_labels) mse = float(np.mean((predictions - targets) ** 2)) mae = float(np.mean(np.abs(predictions - targets))) metrics = { "spearman": spearman_metrics, "retrieval_mAP": map_metrics, "prediction": { "profile_mse": mse, "profile_mae": mae, "prediction_variance": float(np.var(predictions)), "tier_prediction_variance": float(np.var(tier_predictions)), }, "n_samples": len(predictions), "embedding_dim": predictions.shape[1], "checkpoint": { "epoch": checkpoint.get("epoch"), "val_loss": checkpoint.get("val_loss"), **cfg, }, } logger.info("=" * 50) logger.info("PERCEPTUAL EVALUATION RESULTS") logger.info("=" * 50) logger.info("Spearman distance correlation: rho=%.4f | p=%.4e | n_pairs=%d", spearman_metrics["spearman_rho"], spearman_metrics["spearman_p_value"], spearman_metrics["n_pairs"]) logger.info("Retrieval mAP@all: %.4f +/- %.4f | n_queries=%d", map_metrics["mAP"], map_metrics["AP_std"], map_metrics["n_queries"]) logger.info("Profile prediction: mse=%.4f | mae=%.4f | variance=%.6f", mse, mae, metrics["prediction"]["prediction_variance"]) logger.info("=" * 50) return metrics if __name__ == "__main__": parser = argparse.ArgumentParser(description="Evaluate PINO perceptual metrics") parser.add_argument("--checkpoint", default="models/pimt_v1.pt", help="Path to trained checkpoint") parser.add_argument("--data", default="data/synthetic_dataset_v2.jsonl", help="Dataset JSON-L") parser.add_argument("--batch-size", type=int, default=16, help="Evaluation batch size") parser.add_argument("--hidden-dim", type=int, default=None, help="Deprecated; architecture is inferred from checkpoint") parser.add_argument("--output", default="metrics.json", help="Path to write JSON metrics") parser.add_argument("--log-level", default="INFO") args = parser.parse_args() logging.basicConfig( level=getattr(logging, args.log_level.upper(), logging.INFO), format="%(asctime)s %(levelname)s %(name)s: %(message)s", ) metrics = evaluate(args.checkpoint, args.data, args.batch_size, args.hidden_dim) Path(args.output).write_text(json.dumps(metrics, indent=2)) print(json.dumps(metrics, indent=2))