Source code for ts3.models.transductive.poyo_plus.poyo_plus_extractor

"""POYO+'s unit-embedding extractor.

A POYO+ unit embedding is its row of ``unit_emb``, the input token the model learned for
that unit. Rows exist only for units the checkpoint was trained on, and the pretrain
vocabulary is built from the pretrain sessions alone, so the eval units are simply absent
from the pretrain checkpoint. That is the whole reason this model is transductive and not a
config away from being inductive.
"""

import numpy as np
import torch

from core.dataset import BenchmarkRegime, WholeSessionSpikeDataset
from ts3.models.transductive.base import TransductiveExtractor




class POYOPlusExtractor(TransductiveExtractor):
    """POYO+'s ``unit_emb`` rows, optionally length-matched to the pretrain part.

    ``normalize="scale_only"`` multiplies the eval part by one scalar so its mean vector
    length matches the pretrain part's. Every eval vector keeps its direction and the eval
    cloud keeps its shape; only its radius moves.
    """

    def __init__(self, *args, normalize: str = "none", **kwargs):
        super().__init__(*args, **kwargs)
        assert normalize in ("none", "scale_only"), f"unknown normalize: {normalize}"
        self.normalize = normalize
        self._train_norm: float | None = None

    @property
    def name(self) -> str:
        return super().name + ("" if self.normalize == "none" else f"_{self.normalize}")



    def encode(self, regime: BenchmarkRegime) -> tuple[torch.Tensor, np.ndarray]:
        embs, uids = super().encode(regime)

        if regime == "pretrain":
            self._train_norm = float(embs.float().norm(dim=1).mean())
            return embs, uids

        if self.normalize == "scale_only":
            assert self._train_norm is not None, (
                "scale_only reads the pretrain part's norm, so encode('pretrain') runs first"
            )
            norm = float(embs.float().norm(dim=1).mean())
            if norm == 0.0:
                self.logger.warning("eval embeddings are all zero, not rescaling")
            else:
                embs = embs * (self._train_norm / norm)
                self.logger.info(
                    f"scale_only: eval mean ||x|| {norm:.3f} -> {self._train_norm:.3f}"
                )

        return embs, uids




    def read(
        self, state_dict: dict, dataset: WholeSessionSpikeDataset, uids: set[str]
    ) -> tuple[torch.Tensor, np.ndarray]:
        weights = state_dict["unit_emb.weight"]  # [V, D]
        vocab = state_dict["unit_emb.vocab"]  # OrderedDict[uid, idx]
        assert len(vocab) == len(weights), (
            f"vocab size {len(vocab)} != embedding rows {len(weights)}"
        )

        # one ordering, used for both the lookup and the uids, so they cannot disagree
        out_uids = sorted(uids)
        missing = [u for u in out_uids if u not in vocab]
        assert not missing, f"{len(missing)} uids not found in vocab, e.g. {missing[:5]}"

        idx = torch.tensor([vocab[u] for u in out_uids], dtype=torch.long)
        return weights[idx].contiguous(), np.asarray(out_uids).astype(str)