evaluation/esm_embed_util.py

"""Pad/truncate per-token ESM representations to match trained feature dimension."""
from __future__ import annotations

import numpy as np


def pca_input_dim(pca) -> int:
    """Number of features PCA was fit on (sklearn 0.24+ uses n_features_in_)."""
    if hasattr(pca, 'n_features_in_') and pca.n_features_in_ is not None:
        return int(pca.n_features_in_)
    if hasattr(pca, 'components_'):
        return int(pca.components_.shape[1])
    raise ValueError('Cannot determine PCA input dimension from saved artifact')


def infer_per_token_flat_dim(scaler, pca) -> int:
    """Feature length of flattened L x embed_dim before inference preprocessing."""
    if pca is not None:
        return pca_input_dim(pca)
    return int(scaler.mean_.shape[0])


def infer_target_token_count(scaler, pca, embed_dim: int) -> int:
    if embed_dim <= 0:
        raise ValueError(f'Invalid embed_dim={embed_dim}')
    flat_dim = infer_per_token_flat_dim(scaler, pca)
    if flat_dim % embed_dim != 0:
        raise ValueError(
            f'Cannot infer token count: flat_dim={flat_dim} not divisible by embed_dim={embed_dim}'
        )
    return flat_dim // embed_dim


def align_per_token_representation(
    representation: np.ndarray, target_tokens: int
) -> tuple[np.ndarray, bool]:
    """
    representation: (n_tokens, embed_dim)
    Returns flattened vector of length target_tokens * embed_dim, and whether truncation occurred.
    """
    rep = np.asarray(representation, dtype=np.float32)
    if rep.ndim != 2:
        raise ValueError(f'Expected 2D per-token representation, got shape {rep.shape}')
    n_tokens, embed_dim = rep.shape
    truncated = n_tokens > target_tokens
    if n_tokens > target_tokens:
        rep = rep[:target_tokens, :]
    elif n_tokens < target_tokens:
        pad = np.zeros((target_tokens - n_tokens, embed_dim), dtype=np.float32)
        rep = np.vstack([rep, pad])
    return rep.reshape(-1), truncated


def preprocess_per_token_flat(Zs: np.ndarray, scaler, pca) -> np.ndarray:
    """
    Apply the same scaler/PCA order as training.

    Early PCA (use_pca_early): flatten -> PCA -> StandardScaler -> model.
    Late PCA: flatten -> StandardScaler -> PCA -> model.
    """
    n_scaler = int(scaler.mean_.shape[0])
    n_raw = int(Zs.shape[1])

    if pca is None:
        if n_scaler != n_raw:
            raise ValueError(
                f'Per-token scaler expects {n_scaler} dims but embeddings are {n_raw}. '
                'Missing PCA from training (use_pca_early).'
            )
        return scaler.transform(Zs)

    n_pca_in = pca_input_dim(pca)
    if n_raw == n_pca_in:
        if n_scaler < n_pca_in:
            # Early PCA: reduce high-D flat vector, then scale.
            return scaler.transform(pca.transform(Zs))
        # Late PCA: scale full flat vector, then reduce.
        return pca.transform(scaler.transform(Zs))
    if n_raw == n_scaler:
        return scaler.transform(Zs)
    raise ValueError(
        f'Per-token feature mismatch: flat={n_raw}, scaler={n_scaler}, pca_input={n_pca_in}. '
        'Check alignment token count and checkpoint artifacts.'
    )