API Reference

This page provides automated API documentation extracted directly from the Python source code using mkdocstrings.

---

Core Embedder

embpy.embedder.BioEmbedder

Central class for generating biological embeddings.

Manages different embedding models and provides a unified interface for embedding genes (via DNA or protein sequences) and small molecules.

Attributes:

Name	Type	Description
device	device	The computing device (CPU or GPU) used for model inference.
model_cache	dict[str, BaseModelWrapper]	A cache to store loaded model instances.
gene_resolver	GeneResolver	An instance to handle gene identifier resolution and sequence fetching.

Source code in src/embpy/embedder.py

class BioEmbedder:
    """
    Central class for generating biological embeddings.

    Manages different embedding models and provides a unified interface
    for embedding genes (via DNA or protein sequences) and small molecules.

    Attributes
    ----------
    device : torch.device
        The computing device (CPU or GPU) used for model inference.
    model_cache : dict[str, BaseModelWrapper]
        A cache to store loaded model instances.
    gene_resolver : GeneResolver
        An instance to handle gene identifier resolution and sequence fetching.
    """

    def __init__(
        self,
        device: str | torch.device | None = "auto",  # type: ignore[name-defined]
        organism: str = "human",
        resolver_backend: Literal["api", "local"] = "api",
        mart_file: str | None = None,
        chromosome_folder: str | None = None,
    ):
        """
        Initializes the BioEmbedder.

        Args:
            device: 'auto', 'cuda', 'mps', or 'cpu', or torch.device.
            organism: Default organism for sequence resolution and annotation
                (e.g. 'human', 'mouse', 'zebrafish'). Any species supported
                by Ensembl can be used.
            resolver_backend: 'api' to use online APIs, 'local' to use local FASTAs.
            mart_file: path to Mart CSV (required if resolver_backend='local').
            chromosome_folder: path to folder with chr*.fa files (required if 'local').
        """
        self.organism = organism

        # Device setup
        if isinstance(device, str):
            if device == "auto":
                self.device = get_device()
            else:
                self.device = torch.device(device)  # type: ignore[attr-defined]
        elif isinstance(device, torch.device):  # type: ignore[attr-defined]
            self.device = device
        else:
            raise ConfigError("Invalid device; use 'auto','cpu','cuda','mps', or torch.device.")

        # GeneResolver setup
        self.resolver_backend = resolver_backend
        if resolver_backend == "local":
            if not mart_file or not chromosome_folder:
                raise ConfigError("mart_file and chromosome_folder must be provided for local resolver.")
            self.gene_resolver = GeneResolver(
                mart_file=mart_file,
                chromosome_folder=chromosome_folder,
            )
        else:
            self.gene_resolver = GeneResolver(species=organism)

        # Protein resolver
        self.protein_resolver = ProteinResolver(organism=organism)

        # Text resolver for description-based embeddings
        self.text_resolver = TextResolver(organism=organism)

        # Model cache and discovery.
        #
        # `model_cache` holds DNA/protein/molecule/text/morphology wrappers
        # that subclass `BaseModelWrapper`. Single-cell foundation-model
        # wrappers (scGPT, Geneformer, UCE, Tahoe, ...) subclass
        # `SingleCellWrapper` instead, so they live in a separate cache.
        # Both caches are keyed so that repeated calls to embed_cells /
        # embed_perturbation / etc. reuse an already-loaded torch model
        # instead of re-instantiating it on every call -- critical for
        # chunked inference over large datasets.
        self.model_cache: dict[str, BaseModelWrapper] = {}
        self._singlecell_cache: dict[tuple[str, str], Any] = {}
        self._available_models = self._discover_models()

        # Layer 2: per-call resolution report. Populated by
        # ``embed_genes_batch`` and exposed for callers (e.g.
        # ``BioEmbedderProvider`` and ``embed_perturbations.py``)
        # that want a per-identifier audit trail without parsing logs.
        # Each call to ``embed_genes_batch`` REPLACES this attribute,
        # so callers should grab it immediately after the call returns.
        self.last_report: ResolutionReport | None = None

        logging.info(
            "BioEmbedder initialized: device=%s, organism=%s, backend=%s",
            self.device,
            self.organism,
            self.resolver_backend,
        )

    def _discover_models(self) -> dict[str, tuple[type[BaseModelWrapper], str]]:
        """Filters the MODEL_REGISTRY based on available wrapper classes."""
        # Deferred (see the import note at the top of this module): building
        # the registry imports every wrapper module, so it happens here -- at
        # BioEmbedder construction -- rather than at `import embpy.embedder`.
        from .embedder_registry.flat import MODEL_REGISTRY
        from .models.base import BaseModelWrapper

        available = {}
        for name, (wrapper_class, model_path) in MODEL_REGISTRY.items():
            if wrapper_class is None:
                logging.debug(
                    f"Skipping model '{name}': Wrapper class not imported (optional dependency likely missing)."
                )
                continue
            if model_path is None:
                logging.warning(f"Skipping model '{name}': No model path/identifier defined in registry.")
                continue
            # Check if the wrapper class itself is valid (was imported successfully)
            if not issubclass(wrapper_class, BaseModelWrapper):
                logging.error(f"Internal Error: Registered item for '{name}' is not a valid BaseModelWrapper subclass.")
                continue

            available[name] = (wrapper_class, model_path)
            logging.debug(f"Registered model '{name}' using wrapper {wrapper_class.__name__} for path '{model_path}'")

        if not available:
            logging.warning("No models were successfully registered. Check imports and MODEL_REGISTRY definition.")

        return available

    @staticmethod
    def _wrapper_supports_layer_selection(wrapper: Any) -> bool:
        """Whether ``wrapper`` (class or instance) can embed from a chosen layer.

        The wrapper must declare ``target_layer`` explicitly on ``embed`` or
        ``embed_batch``. A bare ``**kwargs`` does not count: it swallows the
        argument and returns default-layer vectors, which is exactly the silent
        wrong answer this check exists to prevent.
        """
        import inspect

        for method_name in ("embed_batch", "embed"):
            method = getattr(wrapper, method_name, None)
            if method is None:
                continue
            try:
                params = inspect.signature(method).parameters
            except (TypeError, ValueError):  # pragma: no cover - builtins/C funcs
                continue
            if "target_layer" in params:
                return True
        return False

    def _assert_layer_selection_supported(self, model: str, layer: int) -> None:
        """Fail loudly when a layer was requested but the model cannot honour it."""
        entry = self._available_models.get(model)
        wrapper: Any = entry[0] if entry else self.model_cache.get(model)

        if wrapper is not None and self._wrapper_supports_layer_selection(wrapper):
            return

        if entry is None and model not in self.model_cache:
            hint = (
                f"'{model}' resolves to a static lookup table or a non-wrapper "
                "embedding source, which has no layers to choose from."
            )
        else:
            name = wrapper.__name__ if isinstance(wrapper, type) else type(wrapper).__name__
            hint = f"{name} does not implement layer selection (no 'target_layer' argument)."

        raise ValueError(
            f"layer={layer} was requested but model '{model}' cannot provide it. {hint} "
            "Re-run without 'layer' to use the model's default output, or pick a model "
            "that supports it (protein, DNA and text transformer wrappers do)."
        )

    def _get_model(self, model_name: str, *, load: bool = True) -> BaseModelWrapper:
        """Loads a model or retrieves it from the cache using the registry or direct HF loading for text models.

        With ``load=False`` the wrapper is constructed but its weights are not
        fetched, so class-level attributes (``has_attention``, ``model_type``)
        can be read without paying for a download. Such an instance is NOT
        cached: callers that later want a working model must go through the
        normal path and get one that is actually loaded.
        """
        if model_name in self.model_cache:
            return self.model_cache[model_name]

        from .embedder_registry.flat import HUMAN_ONLY_MODELS, MOUSE_ONLY_MODELS

        is_human = self.organism.lower() in ("human", "homo_sapiens")
        is_mouse = self.organism.lower() in ("mouse", "mus_musculus")
        if not is_human and model_name in HUMAN_ONLY_MODELS:
            logging.warning(
                "Model '%s' was trained on human data only. Results for organism='%s' may not be meaningful.",
                model_name,
                self.organism,
            )
        if not is_mouse and model_name in MOUSE_ONLY_MODELS:
            logging.warning(
                "Model '%s' was trained on mouse data only. Results for organism='%s' may not be meaningful.",
                model_name,
                self.organism,
            )

        if model_name in self._available_models:
            logging.info(f"Loading registered model '{model_name}' onto device '{self.device}'...")
            WrapperClass, model_path_or_name = self._available_models[model_name]
            try:
                extra_kwargs: dict[str, Any] = {}
                if model_name == "boltz2_pairwise":
                    extra_kwargs["output_type"] = "pairwise"
                elif model_name == "boltz2_both":
                    extra_kwargs["output_type"] = "both"
                from .models.api_models import APIEmbeddingWrapper

                if WrapperClass is APIEmbeddingWrapper:
                    provider_map = {
                        "openai_small": "openai",
                        "openai_large": "openai",
                        "cohere_v3": "cohere",
                        "cohere_multilingual": "cohere",
                        "voyage_3": "voyage",
                        "voyage_3_lite": "voyage",
                        "google_embed": "google",
                    }
                    extra_kwargs["provider"] = provider_map.get(model_name, "openai")
                model_instance = WrapperClass(model_path_or_name=model_path_or_name, **extra_kwargs)
                if not load:
                    return model_instance
                model_instance.load(self.device)
                self.model_cache[model_name] = model_instance
                logging.info(f"Model '{model_name}' loaded successfully.")
                return model_instance
            except ImportError as e:
                # Missing optional dependency -- forward as typed DependencyError
                # so the operator (and SLURM exit code) can distinguish
                # "needs `pixi install` of a different env" from "the model
                # is genuinely broken on HF".
                #
                # Prefer the embpy extra: it is the install line that actually
                # works (`boltz` alone is not enough for Boltz-2, which needs
                # `boltz[cuda]`). Fall back to the module named anywhere in the
                # exception chain, and only then to a placeholder.
                from .embedder_registry.extras import WRAPPER_EXTRAS

                extra = WRAPPER_EXTRAS.get(getattr(WrapperClass, "__name__", ""))
                pkg = f"embpy[{extra}]" if extra else (
                    _missing_package_from_exception(e) or "unknown"
                )
                logging.error(f"Failed to load model '{model_name}': missing dependency '{pkg}'.")
                # Keep whatever the wrapper said. It frequently knows a constraint the
                # extra name cannot express -- Evo is unresolvable on macOS, evo2 needs
                # Python < 3.13 -- and dropping it sent users to an install line that
                # could not work on their machine.
                raise DependencyError(
                    package=pkg, feature=f"model '{model_name}'", detail=str(e) or None
                ) from e
            except Exception as e:
                logging.error(
                    f"Failed to load model '{model_name}' using wrapper {WrapperClass.__name__} and path '{model_path_or_name}': {e}"
                )
                raise ModelLoadError(
                    model_name=model_name,
                    cause=f"{type(e).__name__}: {e}",
                ) from e

        # If not in registry, try to load as a text model from Hugging Face
        else:
            logging.info(f"Model '{model_name}' not in registry. Attempting to load as text model from Hugging Face...")
            try:
                from .models.text_models import TextLLMWrapper

                model_instance = TextLLMWrapper(model_path_or_name=model_name)
                if not load:
                    return model_instance
                model_instance.load(self.device)
                self.model_cache[model_name] = model_instance
                logging.info(f"Successfully loaded text model '{model_name}' from Hugging Face.")
                return model_instance
            except Exception as e:
                available_model_names = self.list_available_models()
                message = (
                    f"Model '{model_name}' could not be loaded from Hugging Face and is not in the predefined registry."
                )
                if available_model_names:
                    message += f" Available predefined models: {available_model_names}"
                message += f" Original error: {str(e)}"
                raise ModelNotFoundError(message) from e

    @staticmethod
    def _detect_vocab_type(
        var_names: Sequence[str], sample: int = 200
    ) -> Literal["symbol", "ensembl_id", "mixed", "unknown"]:
        r"""Guess the gene-identifier convention of an AnnData's ``var_names``.

        Heuristic: sample up to ``sample`` names and check what fraction
        look like Ensembl gene IDs (``ENSG`` / ``ENSMUSG`` / ``ENSRNOG`` /
        ``ENSG0``... or any ``ENS[A-Z]+G\\d+`` prefix, optionally with a
        trailing version suffix).

        Returns
        -------
        str
            ``"ensembl_id"`` if >= 80% match the Ensembl pattern,
            ``"symbol"`` if <= 20% match (rest look like gene symbols),
            ``"mixed"`` if in between, ``"unknown"`` if empty input.
        """
        import re

        if len(var_names) == 0:
            return "unknown"
        head = list(var_names[: min(sample, len(var_names))])
        pattern = re.compile(r"^ENS[A-Z]*G\d+(\.\d+)?$")
        n_ens = sum(1 for v in head if pattern.match(str(v)))
        frac = n_ens / len(head)
        if frac >= 0.8:
            return "ensembl_id"
        if frac <= 0.2:
            return "symbol"
        return "mixed"

    def _ensure_singlecell_vocabulary(
        self,
        adata,
        model_key: str,
        organism: str = "human",
    ):
        """Auto-convert ``adata.var_names`` for a single-cell foundation model.

        Does nothing if the model's ``vocab_type`` is ``"any"`` or
        ``"either"`` (the wrapper handles it internally), or if the
        detected format already matches. Otherwise uses
        :class:`GeneResolver` to remap symbols <-> Ensembl IDs via
        pyensembl (fast, local) with MyGene.info and Ensembl REST
        fallbacks for misses.

        Parameters
        ----------
        adata
            AnnData to inspect / convert. If conversion is needed the
            returned AnnData is a sliced copy with rekeyed ``var_names``;
            the original is left untouched.
        model_key
            Key into the single-cell model registry.
        organism
            Species name passed to ``GeneResolver`` (default ``"human"``).

        Returns
        -------
        tuple[AnnData, dict]
            The (possibly rewritten) AnnData and a small report dict
            with the detected/target formats and number of mapped genes.
        """
        from .models.singlecell_models import _SC_MODEL_REGISTRY

        card = _SC_MODEL_REGISTRY.get(model_key)
        if card is None or card.vocab_type in ("any", "either"):
            return adata, {"action": "none", "reason": f"vocab_type={getattr(card, 'vocab_type', '?')}"}

        target = card.vocab_type  # "symbol" or "ensembl_id"
        current = self._detect_vocab_type(adata.var_names)
        if current == target or current == "unknown":
            return adata, {"action": "none", "reason": f"already {current}"}
        if current == "mixed":
            logging.warning(
                "adata.var_names for model '%s' look mixed (both symbols and "
                "Ensembl IDs). Leaving untouched; expect poor vocabulary match.",
                model_key,
            )
            return adata, {"action": "none", "reason": "mixed var_names"}

        # Need to convert from `current` to `target`.
        logging.info(
            "Auto-converting adata.var_names for model '%s': %s -> %s (%d genes)",
            model_key,
            current,
            target,
            adata.n_vars,
        )

        if current == "ensembl_id" and target == "symbol":
            mapping = self.gene_resolver.ensembl_to_symbols_batch(list(adata.var_names), organism=organism)
        elif current == "symbol" and target == "ensembl_id":
            mapping = self.gene_resolver.symbols_to_ensembl_batch(list(adata.var_names), organism=organism)
        else:
            return adata, {"action": "none", "reason": f"unhandled {current}->{target}"}

        new_names = [mapping.get(v) for v in adata.var_names]
        keep_mask = np.array([n is not None and n != "" for n in new_names], dtype=bool)
        n_mapped = int(keep_mask.sum())
        if n_mapped == 0:
            logging.error(
                "Vocabulary conversion %s -> %s produced 0 mapped genes for "
                "model '%s'. Returning original adata; embedding will likely fail.",
                current,
                target,
                model_key,
            )
            return adata, {"action": "failed", "detected": current, "target": target, "n_mapped": 0}

        out = adata[:, keep_mask].copy()
        new_names_arr = np.array([n for n, keep in zip(new_names, keep_mask, strict=False) if keep])
        # Drop duplicates (multiple Ensembl IDs can map to the same symbol
        # and vice versa). Keep first occurrence.
        seen: set[str] = set()
        unique_mask = np.zeros(len(new_names_arr), dtype=bool)
        for i, n in enumerate(new_names_arr):
            if n not in seen:
                seen.add(n)
                unique_mask[i] = True
        if (~unique_mask).any():
            logging.info(
                "Dropped %d duplicate names during %s -> %s conversion.",
                int((~unique_mask).sum()),
                current,
                target,
            )
            out = out[:, unique_mask].copy()
            new_names_arr = new_names_arr[unique_mask]

        # Preserve the original identifiers as a var column so they
        # remain recoverable (e.g. for round-trip mapping after embedding).
        orig_col = "original_ensembl_id" if current == "ensembl_id" else "original_gene_symbol"
        orig_names = np.asarray(adata.var_names)[keep_mask]
        out.var[orig_col] = orig_names[unique_mask] if (~unique_mask).any() else orig_names
        out.var_names = new_names_arr
        logging.info(
            "Vocabulary conversion OK: %d -> %d genes mapped (%s -> %s).",
            adata.n_vars,
            out.n_vars,
            current,
            target,
        )
        return out, {
            "action": "converted",
            "detected": current,
            "target": target,
            "n_mapped": n_mapped,
            "n_out": out.n_vars,
        }

    def _get_or_load_singlecell_wrapper(
        self,
        model_key: str,
        batch_size: int,
        device_str: str,
        wrapper_kwargs: dict[str, Any] | None = None,
    ):
        """Return a cached single-cell foundation-model wrapper.

        The cache is keyed by ``(model_key, device_str)`` plus any
        constructor kwargs; ``batch_size`` is applied on every call so it
        can change between chunks without forcing a reload.

        This is what makes chunked inference over large datasets cheap:
        a caller looping ``for chunk in chunks: embedder.embed_cells(chunk)``
        pays the model-instantiation cost once, not per chunk.

        ``wrapper_kwargs`` reaches the wrapper constructor, which is how
        checkpoint-only models are usable at all: ``StateEmbeddingWrapper``
        and ``StackWrapper`` take ``checkpoint`` there and raise from
        ``load()`` without it. The kwargs are part of the cache key, so
        asking for two different checkpoints of the same model does not
        silently return the first one.
        """
        from .models.singlecell_models import get_singlecell_wrapper

        wrapper_kwargs = wrapper_kwargs or {}
        kwargs_key = tuple(sorted((k, str(v)) for k, v in wrapper_kwargs.items()))
        cache_key = (model_key, device_str, kwargs_key) if kwargs_key else (
            model_key, device_str,
        )
        cached = self._singlecell_cache.get(cache_key)
        if cached is not None:
            # Cheap attribute update -- do not re-run `.load()`.
            try:
                cached.batch_size = batch_size
            except Exception:  # noqa: BLE001
                pass
            logging.debug("Reusing cached single-cell wrapper for '%s'", model_key)
            return cached

        logging.info("Loading single-cell wrapper for '%s' on %s ...", model_key, device_str)
        wrapper = get_singlecell_wrapper(
            model_key, batch_size=batch_size, **wrapper_kwargs
        )
        wrapper.load(device_str)
        self._singlecell_cache[cache_key] = wrapper
        return wrapper

    def get_model(self, model: str, *, load: bool = True) -> Any:
        """Return the model wrapper behind a model key, for direct introspection.

        ``embed()`` covers the common path, but the wrappers expose more than a
        pooled vector -- per-layer hidden states, per-layer attention, and pooled
        embeddings for every layer at once. Reaching those used to require the
        private ``_get_model``; this is the supported way.

        Parameters
        ----------
        model
            A key from :meth:`list_available_models`.
        load
            Load the weights onto this embedder's device if they are not loaded
            yet. Pass ``False`` to inspect class-level attributes (such as
            ``has_attention``) without paying for a download.

        Returns
        -------
        BaseModelWrapper
            The wrapper, loaded unless ``load=False``. Wrappers are cached, so
            repeated calls return the same instance.

        Examples
        --------
        >>> embedder = BioEmbedder(device="cpu")               # doctest: +SKIP
        >>> wrapper = embedder.get_model("esm2_8M")            # doctest: +SKIP
        >>> ids = wrapper.tokenizer("MTEYKLVVVG", return_tensors="pt")["input_ids"]
        >>> attn = wrapper.extract_attention(ids, layers=[-1])  # doctest: +SKIP

        See Also
        --------
        embpy.models.base.BaseModelWrapper.extract_attention
        embpy.models.base.BaseModelWrapper.extract_hidden_states
        embpy.models.base.BaseModelWrapper.embed_all_layers
        """
        inst = self._get_model(model, load=load)
        if load and getattr(inst, "model", None) is None:
            inst.load(self.device)
        return inst

    def clear_model_cache(self, *, which: Literal["all", "singlecell", "other"] = "all") -> None:
        """Drop cached model wrappers and free associated GPU memory.

        Parameters
        ----------
        which
            Which cache(s) to clear. ``"all"`` drops both the single-cell
            foundation-model cache and the DNA/protein/molecule/text/
            morphology cache. ``"singlecell"`` clears only the single-cell
            cache (useful when switching from foundation-model inference
            to structure/morphology work). ``"other"`` clears only the
            non-single-cell cache.

        After dropping references, runs ``gc.collect()`` +
        ``torch.cuda.empty_cache()`` so the allocator actually releases
        the GPU memory back to the system.
        """
        import gc

        if which in ("all", "singlecell"):
            n = len(self._singlecell_cache)
            self._singlecell_cache.clear()
            logging.info("Cleared %d cached single-cell wrapper(s).", n)
        if which in ("all", "other"):
            n = len(self.model_cache)
            self.model_cache.clear()
            logging.info("Cleared %d cached model wrapper(s).", n)

        gc.collect()
        try:
            import torch

            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                torch.cuda.ipc_collect()
        except Exception:  # noqa: BLE001
            pass

    # ------------------------------------------------------------------
    # Standardized output API (canonical EmbeddingResult + exporters)
    # ------------------------------------------------------------------

    def embed(
        self,
        identifiers: Any | None = None,
        *,
        entity_type: str | Sequence[str] | None = None,
        model: str | Sequence[str],
        id_type: str | Mapping[str, str] | None = None,
        organism: str | None = None,
        pooling_strategy: str = "mean",
        layer: int | None = None,
        output: Literal["anndata", "table", "payload"] | None = None,
        target: Any = None,
        input_path: str | os.PathLike[str] | None = None,
        output_path: str | os.PathLike[str] | None = None,
        attach_to: Literal["auto", "obs", "var", "uns"] = "auto",
        harmonize_dim: int | None = None,
        path: str | os.PathLike[str] | None = None,
        fmt: Literal["csv", "npz", "zarr"] = "npz",
        missing: Literal["error", "nan"] = "error",
        id_column: str | None = None,
        identifier_column: str | None = None,
        anndata_axis: Literal["obs", "var"] | None = None,
        obs_column: str | None = None,
        var_column: str | None = None,
        is_perturbation: bool = False,
        embedding_source: Literal["auto", "model", "hf", "huggingface", "static", "precomputed"] = "auto",
        hf_repo_id: str = DEFAULT_STATIC_EMBEDDING_REPO,
        hf_cache_dir: str | os.PathLike[str] | None = None,
        hf_token: str | None = None,
        hf_id_key: str | None = None,
        hf_id_type: str | None = None,
        whole_genome: bool = False,
        biotype: str = "protein_coding",
        show_progress: bool = False,
        key: str | None = None,
        metadata_mode: Literal["minimal", "full"] = "full",
        include_matrix: bool = True,
        random_state: int = 0,
        **embed_kwargs: Any,
    ):
        """Embed entities and return the standardized embpy output.

        This is the package-wide output path. It accepts direct in-memory
        inputs (sequences, NumPy arrays, pandas Series/DataFrames, AnnData)
        and CSV/TSV/Parquet paths, normalizes them before inference,
        canonicalizes identifiers by entity type, then exports either
        AnnData or a table. Generated embeddings are never stored in
        AnnData ``.X``; standalone AnnData uses sparse placeholder ``.X``
        and stores matrices in ``.obsm``, ``.varm`` or ``.uns``.
        ``output="payload"`` returns the entity-aligned embedding payload
        directly, including canonical ids, model/provenance metadata and
        the requested embedding matrix.

        Gene embeddings are feature embeddings by default and therefore
        route to ``.varm`` for AnnData output. Set ``is_perturbation=True``
        when gene identifiers are perturbation/action labels; those
        embeddings route to ``.obsm`` and the flag is recorded in
        provenance. When the labels come from ``obs_column`` (or
        ``var_column``), rows are aligned on that column of the target
        rather than on its ``obs_names``/``var_names``, so every cell
        receives the vector of its label.

        ``embedding_source="auto"`` routes known static lookup models
        such as ``genept`` and ``gene2vec`` to a Hugging Face dataset
        repository and routes all other models through normal inference.
        The default static repository is ``theislab/Embpy_Data`` and
        files are resolved via ``embeddings/{model}.npz`` / ``.parquet``
        / ``.npy`` using :class:`embpy.pp.HFHandler`. Static lookup
        models return the same standardized outputs as model inference,
        but only for genes present in the downloaded table.

        Defaults are explicit: if ``output`` is omitted, an output path
        chooses compact NPZ/CSV/Zarr file output, otherwise AnnData is
        returned. Multiple models or multiple entity types return the
        same output family with deterministic keys that include entity
        type and model name.

        ``layer`` selects which transformer layer the embedding is pooled
        from, instead of the default final layer. Intermediate layers often
        transfer better to downstream tasks than the last one, which is
        specialised for the pretraining objective. Negative values count
        from the end, so ``layer=-1`` reproduces the default and ``layer=-4``
        takes the fourth-from-last::

            embedder.embed(adata, entity_type="protein", model="esm2_650M", layer=-4)

        Not every architecture supports this. Models that cannot select a
        layer (fingerprints, static lookup tables, and wrappers that have not
        implemented it) raise a :class:`ValueError` rather than silently
        returning default-layer embeddings. Use
        :meth:`~embpy.models.base.BaseModelWrapper.extract_hidden_states` for
        finer-grained access to every layer at once.
        """
        from .io.exporters import route_output
        from .io.normalize import normalize_embedding_input

        # ``layer`` is the documented, discoverable spelling; ``target_layer``
        # is what the wrappers accept. Normalise here so both work and the
        # downstream support check sees a single canonical key.
        if layer is not None:
            embed_kwargs.setdefault("target_layer", layer)

        # Validate here, at the single public entry point, so static lookup
        # tables and other sources that bypass _embed_to_result are covered too.
        requested_layer = next(
            (
                embed_kwargs[k]
                for k in ("target_layer", "layer", "embedding_layer", "layer_name")
                if embed_kwargs.get(k) is not None
            ),
            None,
        )
        if requested_layer is not None:
            for _m in [model] if isinstance(model, str) else list(model):
                self._assert_layer_selection_supported(str(_m), requested_layer)

        org = organism or self.organism
        out_path = output_path if output_path is not None else path
        data = input_path if input_path is not None else identifiers

        if data is None and target is not None and not whole_genome:
            data = target
        if data is None and not whole_genome:
            raise ValueError(
                "input normalization: no identifiers or input_path were provided. "
                "Pass identifiers, an input file path, an AnnData object, or "
                "whole_genome=True."
            )

        if target is None and data is not None and self._is_anndata_like(data):
            target = data

        entity_types = self._resolve_entity_types(entity_type, data, whole_genome)
        if is_perturbation and any(et != "gene" for et in entity_types):
            raise ValueError(
                "input normalization: is_perturbation=True is currently only "
                "supported for entity_type='gene'. Use entity_type='perturbation' "
                "for morphology/action-image embeddings, or attach_to='obs' for "
                "other row-aligned entity embeddings."
            )
        models = self._as_string_list(model, arg_name="model")
        if "cell" in entity_types:
            if entity_types != ["cell"]:
                raise ValueError(
                    "input normalization: entity_type='cell' cannot be mixed "
                    "with other entity types in one BioEmbedder.embed(...) call. "
                    "Run cell embeddings and feature/entity embeddings as "
                    "separate calls so their AnnData alignment stays explicit."
                )
            return self._embed_cells_standardized(
                data,
                models=models,
                output=output or "anndata",
                target=target,
                path=out_path,
                fmt=fmt,
                missing=missing,
                key=key,
                metadata_mode=metadata_mode,
                include_matrix=include_matrix,
                **embed_kwargs,
            )

        normalized_by_entity = {}
        for et in entity_types:
            if whole_genome:
                normalized_by_entity[et] = self._whole_genome_input(
                    entity_type=et,
                    organism=org,
                    biotype=biotype,
                )
                continue
            entity_data = data[et] if isinstance(data, Mapping) and et in data else data
            try:
                normalized_by_entity[et] = normalize_embedding_input(
                    entity_data,
                    entity_type=et,
                    id_column=id_column,
                    identifier_column=identifier_column,
                    anndata_axis=anndata_axis,
                    obs_column=obs_column,
                    var_column=var_column,
                )
            except Exception as exc:
                if str(exc).startswith(("input loading:", "input normalization:")):
                    raise
                raise ValueError(
                    f"input normalization: failed for entity_type={et!r}: {type(exc).__name__}: {exc}"
                ) from exc

        tasks = [(et, m) for et in entity_types for m in models]
        iterator = tasks
        if show_progress:
            from tqdm.auto import tqdm

            iterator = tqdm(tasks, desc="embpy embed", unit="result")

        results = []
        for et, model_name in iterator:
            result = self._embed_to_result(
                normalized_by_entity[et],
                entity_type=et,
                model=model_name,
                id_type=self._id_type_for_entity(id_type, et),
                organism=org,
                pooling_strategy=pooling_strategy,
                is_perturbation=bool(is_perturbation and et == "gene"),
                embedding_source=embedding_source,
                hf_repo_id=hf_repo_id,
                hf_cache_dir=hf_cache_dir,
                hf_token=hf_token,
                hf_id_key=hf_id_key,
                hf_id_type=hf_id_type,
                show_progress=show_progress,
                **embed_kwargs,
            )
            results.append(result)

        resolved_output = output or ("table" if out_path is not None else "anndata")
        return route_output(
            results,
            output=resolved_output,
            target=target,
            attach_to=attach_to,
            harmonize_dim=harmonize_dim,
            path=out_path,
            fmt=fmt,
            missing=missing,
            key=key,
            metadata_mode=metadata_mode,
            include_matrix=include_matrix,
            random_state=random_state,
            obs_column=obs_column,
            var_column=var_column,
        )

    def _embed_cells_standardized(
        self,
        data: Any,
        *,
        models: Sequence[str],
        output: Literal["anndata", "table", "payload"],
        target: Any = None,
        path: str | os.PathLike[str] | None = None,
        fmt: Literal["csv", "npz", "zarr"] = "npz",
        missing: Literal["error", "nan"] = "error",
        key: str | None = None,
        metadata_mode: Literal["minimal", "full"] = "full",
        include_matrix: bool = True,
        **cell_kwargs: Any,
    ):
        """Route single-cell embeddings through ``BioEmbedder.embed``.

        ``embed_cells`` remains the implementation backend, but this method
        converts each produced ``.obsm`` matrix into the same standardized
        ``EmbeddingResult``/exporter path used by the rest of
        ``BioEmbedder.embed``. That keeps tutorials and public workflows on
        one entry point while preserving existing single-cell behavior.
        """
        from .io.exporters import to_anndata_many, to_payloads, to_tables
        from .io.result import EmbeddingProvenance, EmbeddingResult

        if output not in ("anndata", "table", "payload"):
            raise ValueError(f"output must be 'anndata', 'table', or 'payload', got {output!r}.")
        if not self._is_anndata_like(data):
            raise ValueError("input normalization: entity_type='cell' expects an AnnData object as identifiers/target.")
        if key is not None and len(models) > 1:
            raise ValueError(
                "output routing: a single key cannot name multiple cell "
                "embedding results. Omit key or request one model."
            )

        obsm_prefix = str(cell_kwargs.get("obsm_prefix", "X_"))
        adata_out = self.embed_cells(
            data,
            models=list(models),
            **cell_kwargs,
        )
        ids = tuple(str(x) for x in adata_out.obs_names)

        results = []
        keys = []
        cell_meta = adata_out.uns.get("embpy_cell_embeddings", {})
        for model_name in models:
            default_key = f"{obsm_prefix}{model_name}"
            out_key = key or default_key
            if default_key not in adata_out.obsm:
                details = cell_meta.get(model_name, {}) if isinstance(cell_meta, dict) else {}
                message = details.get("error") if isinstance(details, dict) else None
                raise RuntimeError(
                    f"embedding generation: cell model {model_name!r} did not "
                    f"produce adata.obsm[{default_key!r}]." + (f" Backend error: {message}" if message else "")
                )
            matrix = np.asarray(adata_out.obsm[default_key], dtype=np.float32)
            if matrix.ndim != 2 or matrix.shape[0] != adata_out.n_obs:
                raise ValueError(
                    f"embedding generation: adata.obsm[{default_key!r}] must "
                    f"have shape (n_obs, n_dims), got {matrix.shape!r} for "
                    f"n_obs={adata_out.n_obs}."
                )
            extra: dict[str, Any] = {
                "entity_type": "cell",
                "organism": self.organism,
                "input_kind": "anndata",
                "input_source": None,
                "input_id_column": "obs_names",
                "input_id_type": "obs_names",
                "n_requested_inputs": int(adata_out.n_obs),
                "n_successfully_embedded_entities": int(adata_out.n_obs),
                "n_embedding_failures": 0,
                "n_unresolved_identifiers": 0,
                "duplicate_canonical_ids_dropped": 0,
                "n_dropped_or_unresolved_entities": 0,
                "canonical_id_scheme": "obs_names",
            }
            if isinstance(cell_meta, dict) and isinstance(cell_meta.get(model_name), dict):
                extra["cell_embedding"] = dict(cell_meta[model_name])
            prov = EmbeddingProvenance.create(
                model=str(model_name),
                pooling=None,
                layer=None,
                extra=extra,
            )
            results.append(
                EmbeddingResult(
                    matrix=matrix,
                    entity_ids=ids,
                    entity_type="cell",
                    id_scheme="obs_names",
                    provenance=prov,
                )
            )
            keys.append(out_key)

        if output == "payload":
            return to_payloads(
                results,
                keys=keys,
                metadata_mode=metadata_mode,
                include_matrix=include_matrix,
            )
        if output == "table":
            return to_tables(results, path=path, fmt=fmt)
        if path is not None:
            logging.warning(
                "output='anndata' ignores path/output_path; call `.write_h5ad(...)` on the returned AnnData."
            )
        return to_anndata_many(
            results,
            target=adata_out,
            attach_to="obs",
            keys=keys,
            missing=missing,
        )

    def _embed_to_result(
        self,
        identifiers: Any,
        *,
        entity_type: str,
        model: str,
        id_type: str | None,
        organism: str,
        pooling_strategy: str,
        is_perturbation: bool = False,
        embedding_source: str = "model",
        hf_repo_id: str = DEFAULT_STATIC_EMBEDDING_REPO,
        hf_cache_dir: str | os.PathLike[str] | None = None,
        hf_token: str | None = None,
        hf_id_key: str | None = None,
        hf_id_type: str | None = None,
        show_progress: bool = False,
        **embed_kwargs: Any,
    ):
        """Dispatch to a batch embedder + canonicalize -> EmbeddingResult.

        The numeric embedding comes from the existing ``embed_*_batch``
        methods; canonicalization reuses the existing resolvers (the sole
        owners of id-mapping logic). Rows whose embedding or id-resolution
        fails are dropped; duplicate canonical ids collapse to the first.
        """
        from .io._canon import SCHEME, build_aliases, canonicalize, drop_and_dedup
        from .io.normalize import NormalizedInput, normalize_embedding_input
        from .io.result import EmbeddingProvenance, EmbeddingResult

        norm = (
            identifiers
            if isinstance(identifiers, NormalizedInput)
            else normalize_embedding_input(identifiers, entity_type=entity_type)
        )
        ids = list(norm.identifiers)
        alias_cols = {k: list(v) for k, v in norm.alias_columns.items()}
        requested_n = len(ids)
        source = self._resolve_embedding_source(
            embedding_source,
            entity_type=entity_type,
            model=model,
        )
        if source == "hf":
            return self._embed_static_hf_to_result(
                norm,
                entity_type=entity_type,
                model=model,
                id_type=id_type,
                organism=organism,
                pooling_strategy=pooling_strategy,
                is_perturbation=is_perturbation,
                repo_id=hf_repo_id,
                cache_dir=hf_cache_dir,
                token=hf_token,
                id_key=hf_id_key,
                table_id_type=hf_id_type,
            )

        extra: dict[str, Any] = {
            "entity_type": entity_type,
            "organism": organism,
            "input_kind": norm.input_kind,
            "input_source": norm.source,
            "input_id_column": norm.id_column,
            "input_id_type": id_type,
            "is_perturbation": bool(is_perturbation),
            "n_requested_inputs": requested_n,
        }
        # First key that is actually present -- not ``or``-chained, because
        # layer 0 is a legitimate (and falsy) selection that must not be
        # mistaken for "no layer requested".
        layer = next(
            (
                embed_kwargs[k]
                for k in ("target_layer", "layer", "embedding_layer", "layer_name")
                if embed_kwargs.get(k) is not None
            ),
            None,
        )
        if layer is not None:
            self._assert_layer_selection_supported(model, layer)
        if "region" in embed_kwargs:
            extra["region"] = embed_kwargs["region"]
        if "isoform" in embed_kwargs:
            extra["isoform_mode"] = embed_kwargs["isoform"]
        whole_genome_sequences = norm.metadata.get("prefetched_sequences")
        output_entity_type = entity_type

        if entity_type == "gene":
            it = "ensembl_id" if norm.metadata.get("whole_genome") else (id_type or "symbol")
            vecs = self.embed_genes_batch(
                model=model,
                identifiers=ids,
                id_type=it,
                organism=organism,
                pooling_strategy=pooling_strategy,
                **embed_kwargs,
                fetch_all_dna=bool(norm.metadata.get("whole_genome")),
                biotype=str(norm.metadata.get("biotype", "protein_coding")),
                prefetched_sequences=whole_genome_sequences,
            )
            raw, matrix = self._aligned_matrix(ids, vecs)
            canon, keep = canonicalize(
                raw,
                "gene",
                organism,
                id_type=it,
                gene_resolver=self.gene_resolver,
            )
            scheme = SCHEME["gene"]
        elif entity_type == "molecule":
            it = id_type or "smiles"
            smiles, name_alias = self._molecule_inputs_to_smiles(ids, it, extra)
            vecs = self.embed_molecules_batch(
                smiles,
                model,
                pooling_strategy=pooling_strategy,
                **embed_kwargs,
            )
            raw, matrix = self._aligned_matrix(smiles, vecs)
            alias_cols = self._molecule_alias_columns(raw, name_alias)
            canon, keep = canonicalize(raw, "molecule", organism, id_type=it)
            scheme = SCHEME["molecule"]
        elif entity_type == "protein":
            it = id_type or "symbol"
            d = self.embed_proteins_batch(
                ids,
                model,
                id_type=it,
                organism=organism,
                pooling_strategy=pooling_strategy,
                **embed_kwargs,
            )
            raw, matrix, protein_aliases, output_entity_type, scheme = self._dict_matrix(d)
            if output_entity_type == "protein_isoform":
                extra.setdefault("isoform_mode", "all")
                canon = [str(x) for x in raw]
                keep = np.array([True] * len(canon), dtype=bool)
                aliases = protein_aliases
            else:
                canon, keep = canonicalize(
                    raw,
                    "protein",
                    organism,
                    id_type=it,
                    protein_resolver=self.protein_resolver,
                )
                aliases = protein_aliases
                scheme = SCHEME["protein"]
        elif entity_type == "perturbation":
            morph_kwargs = dict(embed_kwargs)
            morphology_dataset = morph_kwargs.pop(
                "morphology_dataset",
                morph_kwargs.pop("dataset", "hpa"),
            )
            morphology_source = morph_kwargs.pop(
                "morphology_source",
                morph_kwargs.pop("source", "subcell"),
            )
            morphology_local_dir = morph_kwargs.pop(
                "morphology_local_dir",
                morph_kwargs.pop("local_dir", None),
            )
            aggregate = morph_kwargs.pop(
                "aggregation",
                morph_kwargs.pop("aggregate", "mean"),
            )
            perturbation_type = morph_kwargs.pop("perturbation_type", "genetic")
            max_images = morph_kwargs.pop("max_images", 5)
            plate_type = morph_kwargs.pop("plate_type", None)
            jump_profiles_dir = morph_kwargs.pop("jump_profiles_dir", None)
            skip_failures = bool(morph_kwargs.pop("skip_failures", True))
            n_workers_value = morph_kwargs.pop("n_workers", None)
            if n_workers_value is None:
                n_workers_value = morph_kwargs.pop("morphology_workers", 8)
            n_workers = int(n_workers_value)
            verbose = bool(morph_kwargs.pop("verbose", True))
            matrix, successful_labels = self.embed_perturbation_morphology_batch(
                ids,
                perturbation_type=perturbation_type,
                dataset=morphology_dataset,
                source=morphology_source,
                model=model,
                pooling_strategy=pooling_strategy,
                local_dir=morphology_local_dir,
                aggregate=aggregate,
                max_images=max_images,
                plate_type=plate_type,
                jump_profiles_dir=jump_profiles_dir,
                verbose=verbose,
                skip_failures=skip_failures,
                n_workers=n_workers,
                **morph_kwargs,
            )
            raw = [str(x) for x in successful_labels]
            alias_cols = {}
            matrix = np.asarray(matrix, dtype=np.float32)
            if matrix.ndim != 2:
                raise ValueError(f"perturbation morphology embeddings must be 2D, got {matrix.shape!r}.")
            if matrix.shape[0] == 0:
                raise ValueError(
                    "embedding generation: no perturbation morphology embeddings were produced. "
                    "Check the morphology dataset/source, local cache, max_images, and perturbation labels."
                )
            canon = raw
            keep = np.ones((len(canon),), dtype=bool)
            scheme = "perturbation_label"
            aliases = {str(label): {"perturbation_label": str(label)} for label in successful_labels}
            extra.update(
                {
                    "morphology_dataset": morphology_dataset,
                    "morphology_source": morphology_source,
                    "morphology_local_dir": morphology_local_dir,
                    "aggregation": aggregate,
                    "perturbation_type": perturbation_type,
                    "max_images": max_images,
                    "plate_type": plate_type,
                    "jump_profiles_dir": jump_profiles_dir,
                    "skip_failures": skip_failures,
                    "n_workers": n_workers,
                }
            )
        elif entity_type in ("sequence", "text"):
            raw, matrix = self._embed_raw_inputs(
                ids,
                model=model,
                pooling_strategy=pooling_strategy,
                show_progress=show_progress,
                **embed_kwargs,
            )
            canon, keep = canonicalize(
                raw,
                entity_type,
                organism,
                id_type=entity_type,
            )
            scheme = SCHEME[entity_type]
            aliases = {}
        else:
            raise ValueError(
                "embedding generation: entity_type must be one of "
                "'gene', 'molecule', 'protein', 'sequence', 'text', or 'perturbation', "
                f"got {entity_type!r}."
            )

        n_embedding_failures = max(0, requested_n - len(raw))
        n_canonical_unresolved = int((~keep).sum())
        kept_canonical = [str(c) for c, ok in zip(canon, keep, strict=False) if ok and c]
        n_duplicate_canonical = len(kept_canonical) - len(set(kept_canonical))
        canon, matrix, alias_cols, kept_raw = drop_and_dedup(
            canon,
            matrix,
            keep,
            raw,
            alias_cols,
        )
        built_aliases = build_aliases(output_entity_type, canon, kept_raw, alias_cols)
        aliases = self._merge_aliases(built_aliases, aliases if "aliases" in locals() else {})
        extra.update(
            {
                "canonical_id_scheme": scheme,
                "n_successfully_embedded_entities": len(canon),
                "n_embedding_failures": n_embedding_failures,
                "n_unresolved_identifiers": n_canonical_unresolved,
                "duplicate_canonical_ids_dropped": n_duplicate_canonical,
                "n_dropped_or_unresolved_entities": (
                    n_embedding_failures + n_canonical_unresolved + n_duplicate_canonical
                ),
            }
        )
        prov = EmbeddingProvenance.create(
            model=model,
            pooling=pooling_strategy,
            layer=layer,
            extra=extra,
        )
        return EmbeddingResult(
            matrix=matrix,
            entity_ids=tuple(canon),
            entity_type=output_entity_type,
            id_scheme=scheme,
            provenance=prov,
            aliases=aliases or None,
        )

    @staticmethod
    def _normalize_embedding_source(value: str) -> Literal["auto", "model", "hf"]:
        source = str(value).lower().replace("-", "_")
        if source in ("auto", "default"):
            return "auto"
        if source in ("model", "inference", "local_model"):
            return "model"
        if source in ("hf", "huggingface", "huggingface_hub", "static", "precomputed"):
            return "hf"
        raise ValueError(
            "embedding generation: embedding_source must be one of "
            f"'auto', 'model' or 'hf'/'static', got {value!r}."
        )

    @staticmethod
    def _resolve_embedding_source(source: str, *, entity_type: str, model: str) -> Literal["model", "hf"]:
        normalized = BioEmbedder._normalize_embedding_source(source)
        if normalized != "auto":
            return normalized
        if entity_type == "gene" and model in DEFAULT_STATIC_EMBEDDING_MODELS:
            return "hf"
        return "model"

    def _embed_static_hf_to_result(
        self,
        identifiers: Any,
        *,
        entity_type: str,
        model: str,
        id_type: str | None,
        organism: str,
        pooling_strategy: str,
        is_perturbation: bool,
        repo_id: str,
        cache_dir: str | os.PathLike[str] | None,
        token: str | None,
        id_key: str | None,
        table_id_type: str | None,
    ):
        """Return an EmbeddingResult by looking up rows in a static HF table."""
        from .io._canon import SCHEME, build_aliases, canonicalize
        from .io.normalize import NormalizedInput, normalize_embedding_input
        from .io.result import EmbeddingProvenance, EmbeddingResult

        if entity_type != "gene":
            raise ValueError(
                "static embedding lookup currently supports entity_type='gene' only; "
                f"got {entity_type!r}."
            )

        norm = (
            identifiers
            if isinstance(identifiers, NormalizedInput)
            else normalize_embedding_input(identifiers, entity_type=entity_type)
        )
        requested_ids = [str(x) for x in norm.identifiers]
        table = self._load_static_embedding_table(
            model,
            repo_id=repo_id,
            cache_dir=cache_dir,
            token=token,
            id_key=id_key,
            id_type=table_id_type,
        )
        table_ids = [str(x) for x in table["ids"]]
        table_matrix = np.asarray(table["matrix"], dtype=np.float32)
        table_id_type = str(table["id_type"])

        table_canon, table_keep = canonicalize(
            table_ids,
            entity_type,
            organism,
            id_type=table_id_type,
            gene_resolver=self.gene_resolver,
        )
        row_by_canonical: dict[str, int] = {}
        table_alias_by_canonical: dict[str, str] = {}
        for i, (cid, ok) in enumerate(zip(table_canon, table_keep, strict=False)):
            if not ok or not cid:
                continue
            canonical_id = str(cid)
            row_by_canonical.setdefault(canonical_id, i)
            if table_ids[i] != canonical_id:
                table_alias_by_canonical.setdefault(canonical_id, table_ids[i])

        request_id_type = id_type or "symbol"
        req_canon, req_keep = canonicalize(
            requested_ids,
            entity_type,
            organism,
            id_type=request_id_type,
            gene_resolver=self.gene_resolver,
        )

        out_ids: list[str] = []
        out_rows: list[np.ndarray] = []
        kept_raw: list[str] = []
        unresolved_ids: list[str] = []
        lookup_missing_ids: list[str] = []
        seen: set[str] = set()
        unresolved = 0
        lookup_misses = 0
        duplicate_canonical = 0

        for raw_id, cid, ok in zip(requested_ids, req_canon, req_keep, strict=False):
            if not ok or not cid:
                unresolved += 1
                unresolved_ids.append(raw_id)
                continue
            canonical_id = str(cid)
            row = row_by_canonical.get(canonical_id)
            if row is None:
                lookup_misses += 1
                lookup_missing_ids.append(raw_id)
                continue
            if canonical_id in seen:
                duplicate_canonical += 1
                continue
            seen.add(canonical_id)
            out_ids.append(canonical_id)
            out_rows.append(np.asarray(table_matrix[row], dtype=np.float32))
            kept_raw.append(raw_id)

        if not out_rows:
            raise ValueError(
                "static embedding lookup: no requested identifiers were found in "
                f"model={model!r} from repo={repo_id!r}. "
                f"Requested {len(requested_ids)} ids; unresolved={unresolved}, "
                f"lookup_misses={lookup_misses}."
            )

        if unresolved_ids or lookup_missing_ids:
            missing_ids = list(dict.fromkeys(unresolved_ids + lookup_missing_ids))
            preview = missing_ids[:10]
            warnings.warn(
                "Static embedding model "
                f"{model!r} from {repo_id!r} is available for "
                f"{len(out_ids)}/{len(requested_ids)} requested gene(s). "
                f"Unavailable or unresolved genes were dropped; first missing: {preview}.",
                UserWarning,
                stacklevel=2,
            )

        matrix = np.vstack(out_rows).astype(np.float32, copy=False)
        aliases = build_aliases(entity_type, out_ids, kept_raw, {})
        for cid, table_id in table_alias_by_canonical.items():
            if cid in out_ids and table_id != cid:
                alias_key = "gene_symbol" if table_id_type == "symbol" else "source_id"
                aliases.setdefault(cid, {}).setdefault(alias_key, table_id)

        extra: dict[str, Any] = {
            "entity_type": entity_type,
            "organism": organism,
            "input_kind": norm.input_kind,
            "input_source": norm.source,
            "input_id_column": norm.id_column,
            "input_id_type": id_type,
            "is_perturbation": bool(is_perturbation),
            "embedding_source": "hf",
            "hf_repo_id": repo_id,
            "hf_model_key": model,
            "hf_id_key": table.get("id_key"),
            "hf_id_type": table_id_type,
            "canonical_id_scheme": SCHEME[entity_type],
            "n_requested_inputs": len(requested_ids),
            "n_static_available_entities": int(table_matrix.shape[0]),
            "n_successfully_embedded_entities": len(out_ids),
            "n_embedding_failures": lookup_misses,
            "n_unresolved_identifiers": unresolved,
            "duplicate_canonical_ids_dropped": duplicate_canonical,
            "n_dropped_or_unresolved_entities": unresolved + lookup_misses + duplicate_canonical,
        }
        # A static table is a precomputed lookup: no forward pass, no tokens,
        # so nothing was pooled. Recording the caller's `pooling_strategy`
        # here -- which defaults to "mean" whether or not they asked for it --
        # stamped a pooling that never happened, and `_default_key` then put it
        # in the key, so `genept` produced `X_emb__gene__genept__pool_mean`.
        # `None` is the same convention the cell path already uses.
        prov = EmbeddingProvenance.create(
            model=model,
            pooling=None,
            extra=extra,
        )
        return EmbeddingResult(
            matrix=matrix,
            entity_ids=tuple(out_ids),
            entity_type=entity_type,
            id_scheme=SCHEME[entity_type],
            provenance=prov,
            aliases=aliases or None,
        )

    def _load_static_embedding_table(
        self,
        model: str,
        *,
        repo_id: str,
        cache_dir: str | os.PathLike[str] | None,
        token: str | None,
        id_key: str | None,
        id_type: str | None,
    ) -> dict[str, Any]:
        from .pp.hf_handler import HFHandler

        obj = HFHandler(repo_id, token=token).download_embedding(model, cache_dir=cache_dir)
        return self._coerce_static_embedding_table(
            obj,
            model=model,
            id_key=id_key,
            id_type=id_type,
        )

    def _coerce_static_embedding_table(
        self,
        obj: Any,
        *,
        model: str,
        id_key: str | None,
        id_type: str | None,
    ) -> dict[str, Any]:
        """Normalize HFHandler output into ids + float32 matrix."""
        import pandas as pd

        if isinstance(obj, pd.DataFrame):
            return self._coerce_static_embedding_frame(
                obj,
                model=model,
                id_key=id_key,
                id_type=id_type,
            )
        if not isinstance(obj, Mapping):
            raise TypeError(
                f"static embedding lookup: expected a DataFrame or dict for {model!r}, "
                f"got {type(obj).__name__}."
            )

        matrix_key = self._first_present_key(obj, ("embeddings", "matrix", "X", "arr_0"))
        if matrix_key is None:
            raise KeyError(
                "static embedding lookup: embedding file must contain one of "
                "'embeddings', 'matrix', 'X', or 'arr_0'."
            )
        matrix = np.asarray(obj[matrix_key], dtype=np.float32)
        if matrix.ndim != 2:
            raise ValueError(f"static embedding lookup: matrix must be 2D, got {matrix.shape!r}.")

        resolved_id_key = id_key or self._first_present_key(
            obj,
            (
                "entity_ids",
                "ids",
                "ensembl_gene_ids",
                "ensembl_gene_id",
                "ensembl_ids",
                "ensembl_id",
                "symbols",
                "gene_symbols",
                "gene_symbol",
                "genes",
                "index",
            ),
        )
        if resolved_id_key is None:
            raise KeyError(
                "static embedding lookup: embedding file must contain row identifiers. "
                "Use one of 'symbols', 'gene_symbols', 'ensembl_gene_ids', "
                "'entity_ids', 'ids', or pass hf_id_key=..."
            )
        ids_obj = obj[resolved_id_key]

        if isinstance(ids_obj, pd.DataFrame):
            ids = [] if ids_obj.empty else ids_obj.iloc[:, 0].astype(str).tolist()
        else:
            ids = [str(x) for x in np.asarray(ids_obj).reshape(-1).tolist()]

        if matrix.shape[0] != len(ids):
            raise ValueError(
                f"static embedding lookup: {model!r} has {matrix.shape[0]} embedding rows "
                f"but {len(ids)} ids in {resolved_id_key!r}."
            )
        return {
            "ids": ids,
            "matrix": matrix,
            "id_key": resolved_id_key,
            "id_type": id_type or self._infer_static_gene_id_type(resolved_id_key, ids),
        }

    def _coerce_static_embedding_frame(
        self,
        frame: Any,
        *,
        model: str,
        id_key: str | None,
        id_type: str | None,
    ) -> dict[str, Any]:
        import pandas as pd

        df = pd.DataFrame(frame)
        resolved_id_key = id_key or self._first_present_key(
            df,
            (
                "entity_ids",
                "ids",
                "ensembl_gene_ids",
                "ensembl_gene_id",
                "ensembl_ids",
                "ensembl_id",
                "symbols",
                "gene_symbols",
                "gene_symbol",
                "genes",
                "id",
            ),
        )
        if resolved_id_key is not None:
            ids = df[resolved_id_key].astype(str).tolist()
            value_df = df.drop(columns=[resolved_id_key])
        elif not isinstance(df.index, pd.RangeIndex):
            ids = df.index.astype(str).tolist()
            resolved_id_key = str(df.index.name or "index")
            value_df = df
        else:
            raise KeyError(
                "static embedding lookup: parquet/static table must have an id column "
                "or a non-default index. Pass hf_id_key=... if needed."
            )

        if "embedding" in value_df.columns:
            matrix = np.vstack(value_df["embedding"].map(np.asarray).to_numpy()).astype(np.float32)
        else:
            dim_cols = [c for c in value_df.columns if str(c).startswith("dim_")]
            numeric = value_df[dim_cols] if dim_cols else value_df.select_dtypes(include=[np.number])
            matrix = numeric.to_numpy(dtype=np.float32)
        if matrix.ndim != 2 or matrix.shape[1] == 0:
            raise ValueError(f"static embedding lookup: no numeric embedding columns found for {model!r}.")
        if matrix.shape[0] != len(ids):
            raise ValueError(
                f"static embedding lookup: {model!r} has {matrix.shape[0]} embedding rows "
                f"but {len(ids)} ids."
            )
        return {
            "ids": ids,
            "matrix": matrix,
            "id_key": resolved_id_key,
            "id_type": id_type or self._infer_static_gene_id_type(resolved_id_key, ids),
        }

    @staticmethod
    def _first_present_key(container: Any, keys: Sequence[str]) -> str | None:
        for key in keys:
            if key in container:
                return key
        return None

    @staticmethod
    def _infer_static_gene_id_type(id_key: str, ids: Sequence[str]) -> str:
        key = str(id_key).lower()
        if "ensembl" in key or key in {"ensembl_id", "ensembl_gene_id"}:
            return "ensembl_id"
        if "symbol" in key or key in {"gene", "genes"}:
            return "symbol"
        sample = [str(x).upper() for x in ids[:50]]
        if sample and all(x.startswith(("ENSG", "ENSMUSG", "ENS")) for x in sample):
            return "ensembl_id"
        return "symbol"

    @staticmethod
    def _as_string_list(value: str | Sequence[str], *, arg_name: str) -> list[str]:
        if isinstance(value, str):
            return [value]
        values = [str(x) for x in value]
        if not values:
            raise ValueError(f"input normalization: {arg_name} cannot be empty.")
        return values

    def _resolve_entity_types(
        self,
        entity_type: str | Sequence[str] | None,
        data: Any,
        whole_genome: bool,
    ) -> list[str]:
        if entity_type is None:
            if whole_genome:
                return ["gene"]
            if isinstance(data, Mapping):
                return [str(k) for k in data.keys()]
            raise ValueError(
                "input normalization: entity_type is required unless identifiers "
                "is a mapping keyed by entity type or whole_genome=True."
            )
        values = self._as_string_list(entity_type, arg_name="entity_type")
        allowed = {"gene", "molecule", "protein", "sequence", "text", "perturbation", "cell"}
        bad = [x for x in values if x not in allowed]
        if bad:
            raise ValueError(
                f"input normalization: unsupported entity_type(s) {bad}; expected one of {sorted(allowed)}."
            )
        if isinstance(data, Mapping):
            missing = [x for x in values if x not in data]
            if missing and not whole_genome:
                raise ValueError(
                    "input normalization: identifiers mapping is missing "
                    f"entity_type key(s) {missing}. Available keys: {list(data.keys())}."
                )
        return values

    @staticmethod
    def _id_type_for_entity(
        id_type: str | Mapping[str, str] | None,
        entity_type: str,
    ) -> str | None:
        if isinstance(id_type, Mapping):
            return id_type.get(entity_type)
        return id_type

    @staticmethod
    def _is_anndata_like(value: Any) -> bool:
        try:
            from anndata import AnnData
        except ImportError:  # pragma: no cover
            return False
        return isinstance(value, AnnData)

    def _whole_genome_input(
        self,
        *,
        entity_type: str,
        organism: str,
        biotype: str,
    ):
        from .io.normalize import NormalizedInput

        if entity_type != "gene":
            raise ValueError(
                "input loading: whole_genome=True is currently supported for "
                f"entity_type='gene' only, got {entity_type!r}."
            )
        try:
            sequences = self.gene_resolver.get_gene_sequences(biotype=biotype)
        except Exception as exc:
            raise RuntimeError(
                "input loading: whole-genome gene request failed while reading "
                f"the genome annotation resource for organism={organism!r}, "
                f"biotype={biotype!r}: {type(exc).__name__}: {exc}. "
                "Install/index the resolver backend (for example pyensembl or "
                "a local Mart/genome resource) or pass explicit identifiers."
            ) from exc
        if not sequences:
            raise RuntimeError(
                "input loading: whole-genome gene request could not load any "
                f"Ensembl gene IDs for organism={organism!r}, biotype={biotype!r}. "
                "The resolver backend or genome annotation resource is missing "
                "or empty; install/index the resource or pass explicit identifiers."
            )
        return NormalizedInput(
            identifiers=tuple(str(x) for x in sequences.keys()),
            input_kind="whole_genome",
            source=f"whole_genome:{organism}:{biotype}",
            id_column="ensembl_gene_id",
            metadata={
                "whole_genome": True,
                "organism": organism,
                "biotype": biotype,
                "prefetched_sequences": sequences,
            },
        )

    @staticmethod
    def _merge_aliases(
        left: dict[str, dict[str, str]],
        right: dict[str, dict[str, str]],
    ) -> dict[str, dict[str, str]]:
        out = {k: dict(v) for k, v in left.items()}
        for key, mapping in right.items():
            out.setdefault(key, {}).update(mapping)
        return out

    @staticmethod
    def _molecule_alias_columns(
        raw_smiles: list[str],
        name_alias: dict[str, str],
    ) -> dict[str, list[str]]:
        names = [name_alias.get(s) for s in raw_smiles]
        return {"name": list(names)} if any(names) else {}

    def _embed_raw_inputs(
        self,
        identifiers: list[str],
        *,
        model: str,
        pooling_strategy: str,
        show_progress: bool = False,
        **kwargs: Any,
    ) -> tuple[list[str], np.ndarray]:
        inst = self._get_model(model)
        raw_iter = identifiers
        rows: list[np.ndarray] = []
        kept: list[str] = []

        if hasattr(inst, "embed_batch"):
            try:
                batch = inst.embed_batch(
                    inputs=identifiers,
                    pooling_strategy=pooling_strategy,
                    **kwargs,
                )
            except TypeError:
                batch = inst.embed_batch(
                    input=identifiers,
                    pooling_strategy=pooling_strategy,
                    **kwargs,
                )
            return self._aligned_matrix(identifiers, list(batch))

        if show_progress:
            from tqdm.auto import tqdm

            raw_iter = tqdm(identifiers, desc=f"{model} inputs", unit="input")
        for item in raw_iter:
            try:
                rows.append(
                    np.asarray(
                        inst.embed(input=item, pooling_strategy=pooling_strategy, **kwargs),
                        dtype=np.float32,
                    ).ravel()
                )
                kept.append(item)
            except Exception as exc:  # noqa: BLE001
                logging.warning("Failed to embed input %r with %s: %s", item, model, exc)
        if not rows:
            raise ValueError("embedding generation: no embeddings were produced.")
        return kept, np.stack(rows, axis=0)

    @staticmethod
    def _aligned_matrix(
        ids: list[str],
        vecs: list[np.ndarray | None],
    ) -> tuple[list[str], np.ndarray]:
        """Drop None embeddings; return (kept_ids, stacked matrix)."""
        kept_ids = [i for i, v in zip(ids, vecs, strict=False) if v is not None]
        kept = [np.asarray(v, dtype=np.float32).ravel() for v in vecs if v is not None]
        if not kept:
            raise ValueError("embedding generation: no embeddings were produced (all inputs failed).")
        return kept_ids, np.stack(kept, axis=0)

    @staticmethod
    def _dict_matrix(
        d: dict,
    ) -> tuple[list[str], np.ndarray, dict[str, dict[str, str]], str, str]:
        """Flatten a {id: vector} batch-protein result into (ids, matrix)."""
        ids: list[str] = []
        rows: list[np.ndarray] = []
        aliases: dict[str, dict[str, str]] = {}
        has_isoforms = any(isinstance(v, dict) for v in d.values())

        for k, v in d.items():
            if isinstance(v, dict):
                for iso_id, iso_vec in v.items():
                    ids.append(str(iso_id))
                    rows.append(np.asarray(iso_vec, dtype=np.float32).ravel())
                    aliases.setdefault(str(iso_id), {})["parent_input_id"] = str(k)
                continue
            ids.append(str(k))
            rows.append(np.asarray(v, dtype=np.float32).ravel())
        if not rows:
            raise ValueError("embedding generation: no protein embeddings were produced.")
        if has_isoforms:
            return ids, np.stack(rows, axis=0), aliases, "protein_isoform", "uniprot_isoform"
        return ids, np.stack(rows, axis=0), aliases, "protein", "uniprot"

    def _molecule_inputs_to_smiles(
        self,
        ids: list[str],
        id_type: str,
        extra: dict[str, Any],
    ) -> tuple[list[str], dict[str, str]]:
        """Return SMILES to embed + a {smiles: name} alias map.

        For ``id_type="smiles"`` the inputs are passed through. For
        ``id_type="name"`` each name is resolved to SMILES via the drug
        resolver chain, recording which API answered (provenance).
        """
        if id_type in ("smiles", "canonical_smiles"):
            return ids, {}
        if id_type != "name":
            raise ValueError(
                f"canonicalization: molecule id_type must be 'smiles', 'canonical_smiles', or 'name', got {id_type!r}."
            )
        from .resources.molecule.resolver import DrugResolver

        resolver = DrugResolver()
        smiles: list[str] = []
        name_alias: dict[str, str] = {}
        sources: dict[str, int] = {}
        for name in ids:
            res = resolver.name_to_smiles_resolved(name)
            if res.smiles is None:
                logging.warning("Dropping drug name %r: could not resolve to SMILES.", name)
                continue
            smiles.append(res.smiles)
            name_alias[res.smiles] = name
            sources[res.source] = sources.get(res.source, 0) + 1
        if sources:
            extra["name_resolution_sources"] = ", ".join(f"{k}:{v}" for k, v in sorted(sources.items()))
        return smiles, name_alias

    def embed_gene(
        self,
        identifier: str,
        model: str,
        id_type: Literal["symbol", "ensembl_id", "uniprot_id", "sequence"] = "symbol",
        organism: str = "human",
        pooling_strategy: str = "mean",
        region: Literal["full", "exons", "introns"] = "full",
        protein_isoform: Literal["canonical", "all"] = "canonical",
        gene_description_format: str = "Gene: {identifier}. Type: {id_type}. Organism: {organism}.",
        **kwargs: Any,
    ) -> np.ndarray:
        """
        Generates an embedding for a single gene using the specified model.

        Determines whether to use sequence (DNA/protein) or text description
        based on the selected model's type.

        Args:
            identifier (str): The gene identifier (e.g., "TP53", "ENSG00000141510")
                or a raw sequence when ``id_type="sequence"``.
            model (str): The user-facing name of the model to use (e.g., "enformer_human_rough",
                        "esm2_650M", "minilm_l6_v2"). Must be a key in MODEL_REGISTRY.
            id_type (Literal): Type of the identifier provided. Use ``"sequence"`` to
                pass a raw DNA/protein/text string directly. Defaults to "symbol".
            organism (str): Organism name (e.g., "human", "mouse"). Used for sequence/description lookup.
                            Defaults to "human".
            pooling_strategy (str): Pooling strategy if the model outputs per-token/residue embeddings.
                                    Defaults to "mean". Check model's `available_pooling_strategies`.
            region (Literal): Gene region to embed. ``"full"`` uses the complete genomic
                sequence (default), ``"exons"`` concatenates only exonic regions, and
                ``"introns"`` concatenates only intronic regions. Only applies to DNA
                models when *id_type* is ``"symbol"`` or ``"ensembl_id"``.
            protein_isoform (Literal): For protein models: ``"canonical"`` (default)
                embeds only the canonical UniProt sequence; ``"all"`` is not valid
                here -- use :meth:`embed_protein_isoforms` for all isoforms.
            gene_description_format (str): Format string used to generate input for text models.
                                        Defaults to "Gene: {identifier}. Type: {id_type}. Organism: {organism}.".
            **kwargs: Additional arguments passed to the specific model's embed method
                    (e.g., `target_layer` for transformer models).

        Returns
        -------
            np.ndarray: The computed gene embedding.

        Raises
        ------
            ModelNotFoundError: If the requested model name is not registered or available.
            IdentifierError: If the gene identifier cannot be resolved or required sequence/description fetched.
            ValueError: If the model type is ambiguous or incompatible inputs are generated.
            RuntimeError: If model loading or inference fails.
        """
        inst = self._get_model(model)
        mtype = inst.model_type

        if id_type != "sequence" and mtype == "protein":
            from embpy.resources.gene_resolver import detect_identifier_type

            detected = detect_identifier_type(identifier)
            if detected == "protein_sequence":
                logging.info(
                    "Auto-detected raw protein sequence (id_type was %r); "
                    "embedding directly. Pass id_type='sequence' explicitly "
                    "to suppress this message.",
                    id_type,
                )
                id_type = "sequence"

        if id_type == "sequence":
            input_data = identifier
        elif mtype == "dna":
            if id_type not in ("symbol", "ensembl_id"):
                raise IdentifierError(
                    f"DNA models require id_type 'symbol', 'ensembl_id', or 'sequence', got '{id_type}'."
                )
            dna_id_type: Literal["symbol", "ensembl_id"] = id_type  # type: ignore[assignment]
            if region in ("exons", "introns"):
                seq = self.gene_resolver.get_gene_region_sequence(
                    identifier,
                    id_type=dna_id_type,
                    organism=organism,
                    region=region,
                )
            elif self.resolver_backend == "local":
                seq = self.gene_resolver.get_local_dna_sequence(identifier, dna_id_type)
            else:
                seq = self.gene_resolver.get_dna_sequence(identifier, dna_id_type, organism)
            if not seq:
                raise IdentifierError(f"DNA not found for {id_type}='{identifier}' (region={region})")
            input_data = seq
        elif mtype == "protein":
            prot = self.protein_resolver.get_canonical_sequence(identifier, id_type, organism)
            if not prot:
                raise IdentifierError(f"Protein not found for {id_type}='{identifier}'")
            input_data = prot
        elif mtype == "text":
            desc = self.gene_resolver.get_gene_description(
                identifier, id_type, organism, format_string=gene_description_format
            )
            if not desc:
                raise IdentifierError(f"Description not found for {id_type}='{identifier}'")
            input_data = desc
        elif mtype == "ppi":
            input_data = identifier
        else:
            raise ValueError(f"Unsupported model type '{mtype}' for embedding.")
        # Embed
        try:
            emb = inst.embed(input=input_data, pooling_strategy=pooling_strategy, **kwargs)
            return emb
        except Exception as e:
            logging.error(f"Embedding error for {identifier} with model {model}: {e}")
            raise RuntimeError(f"Embedding failed for {identifier}") from e

    def embed_protein(
        self,
        identifier: str,
        model: str,
        id_type: Literal["symbol", "ensembl_id", "uniprot_id", "sequence"] = "symbol",
        organism: str = "human",
        pooling_strategy: str = "mean",
        isoform: Literal["canonical", "all"] = "canonical",
        **kwargs: Any,
    ) -> np.ndarray | dict[str, np.ndarray]:
        """Embed a protein sequence using a protein language model.

        When ``isoform="canonical"``, returns a single embedding of the
        canonical UniProt sequence.  When ``isoform="all"``, returns a
        dict mapping isoform accession IDs to their embeddings.

        Parameters
        ----------
        identifier
            Gene symbol, Ensembl ID, UniProt accession, or raw amino
            acid sequence (when ``id_type="sequence"``).
        model
            Protein model name (e.g. ``"esm2_650M"``, ``"prot_t5_xl"``).
        id_type
            Type of identifier.
        organism
            Organism name (default ``"human"``).
        pooling_strategy
            Pooling strategy (``"mean"``, ``"max"``, ``"cls"``).
        isoform
            ``"canonical"`` for the canonical sequence only, ``"all"``
            to embed every UniProt isoform.
        **kwargs
            Forwarded to the model's ``embed`` method.

        Returns
        -------
        np.ndarray
            When ``isoform="canonical"``.
        dict[str, np.ndarray]
            When ``isoform="all"`` -- maps isoform accession to embedding.
        """
        inst = self._get_model(model)
        if inst.model_type != "protein":
            raise ValueError(f"Model '{model}' is not a protein model (type={inst.model_type}).")

        if id_type == "sequence":
            emb = inst.embed(input=identifier, pooling_strategy=pooling_strategy, **kwargs)
            return emb

        if isoform == "canonical":
            seq = self.protein_resolver.get_canonical_sequence(identifier, id_type, organism)
            if not seq:
                raise IdentifierError(f"Canonical protein not found for {id_type}='{identifier}'")
            return inst.embed(input=seq, pooling_strategy=pooling_strategy, **kwargs)

        # isoform == "all"
        isoforms = self.protein_resolver.get_isoforms(
            identifier,
            id_type,
            organism,
            include_canonical=True,
        )
        if not isoforms:
            raise IdentifierError(f"No isoforms found for {id_type}='{identifier}'")

        results: dict[str, np.ndarray] = {}
        for iso_id, seq in isoforms.items():
            try:
                results[iso_id] = inst.embed(
                    input=seq,
                    pooling_strategy=pooling_strategy,
                    **kwargs,
                )
            except Exception as e:  # noqa: BLE001
                logging.warning(f"Failed to embed isoform {iso_id}: {e}")
        return results

    def embed_proteins_batch(
        self,
        identifiers: list[str],
        model: str,
        id_type: Literal["symbol", "ensembl_id", "uniprot_id", "sequence"] = "symbol",
        organism: str = "human",
        pooling_strategy: str = "mean",
        isoform: Literal["canonical", "all"] = "canonical",
        **kwargs: Any,
    ) -> dict[str, np.ndarray] | dict[str, dict[str, np.ndarray]]:
        """Embed proteins for a batch of gene identifiers.

        Parameters
        ----------
        identifiers
            List of gene identifiers or raw amino-acid sequences
            (when ``id_type="sequence"``).
        model
            Protein model name.
        id_type
            Type of identifiers.  Use ``"sequence"`` to embed raw
            amino-acid strings directly without resolver lookup.
        organism
            Organism name.
        pooling_strategy
            Pooling strategy.
        isoform
            ``"canonical"`` or ``"all"``.  Ignored when
            ``id_type="sequence"``.
        **kwargs
            Forwarded to the model.

        Returns
        -------
        dict[str, np.ndarray]
            When ``isoform="canonical"`` or ``id_type="sequence"``
            -- maps identifier (or sequence) to embedding.
        dict[str, dict[str, np.ndarray]]
            When ``isoform="all"`` -- maps identifier to
            ``{isoform_accession: embedding}``.
        """
        inst = self._get_model(model)
        if inst.model_type != "protein":
            raise ValueError(f"Model '{model}' is not a protein model.")

        if id_type == "sequence":
            results: dict[str, np.ndarray] = {}
            for seq in identifiers:
                try:
                    results[seq] = inst.embed(
                        input=seq,
                        pooling_strategy=pooling_strategy,
                        **kwargs,
                    )
                except Exception as e:  # noqa: BLE001
                    logging.warning(f"Failed to embed protein sequence: {e}")
            return results

        if isoform == "canonical":
            seqs = self.protein_resolver.get_canonical_sequences_batch(
                identifiers,
                id_type,
                organism,
            )
            results = {}
            for ident, seq in seqs.items():
                try:
                    results[ident] = inst.embed(
                        input=seq,
                        pooling_strategy=pooling_strategy,
                        **kwargs,
                    )
                except Exception as e:  # noqa: BLE001
                    logging.warning(f"Failed to embed protein for {ident}: {e}")
            return results

        # isoform == "all"
        all_isoforms = self.protein_resolver.get_isoforms_batch(
            identifiers,
            id_type,
            organism,
            include_canonical=True,
        )
        results_iso: dict[str, dict[str, np.ndarray]] = {}
        for ident, isoforms_map in all_isoforms.items():
            results_iso[ident] = {}
            for iso_id, seq in isoforms_map.items():
                try:
                    results_iso[ident][iso_id] = inst.embed(
                        input=seq,
                        pooling_strategy=pooling_strategy,
                        **kwargs,
                    )
                except Exception as e:  # noqa: BLE001
                    logging.warning(f"Failed to embed isoform {iso_id} for {ident}: {e}")
        return results_iso

    def embed_genes_batch(
        self,
        model: str,
        identifiers: Sequence[str] | None = None,
        id_type: Literal["symbol", "ensembl_id", "uniprot_id", "sequence"] = "symbol",
        organism: str = "human",
        pooling_strategy: str = "mean",
        region: Literal["full", "exons", "introns"] = "full",
        gene_description_format: str | None = None,
        fetch_all_dna: bool = False,
        biotype: str = "protein_coding",
        prefetched_sequences: dict[str, str] | None = None,
        **kwargs: Any,
    ) -> list[np.ndarray | None]:
        """
        Generates embeddings for a batch of genes.

        If `identifiers` is None and `fetch_all_dna` is True, it will automatically
        fetch ALL genes of the specified `biotype` and embed them.

        Set ``region`` to ``"exons"`` or ``"introns"`` to embed only those
        gene regions (DNA models only).
        """
        inst = self._get_model(model)
        mtype = inst.model_type

        # Dictionary to hold pre-fetched sequences (if any)
        prefetched_data: dict[str, str] = dict(prefetched_sequences or {})
        if prefetched_data and identifiers is None:
            identifiers = list(prefetched_data.keys())
            id_type = "ensembl_id"
            logging.info("Using %d prefetched whole-genome sequences.", len(prefetched_data))

        # --- LOGIC CHANGE: Discovery Mode ---
        # If identifiers is None OR we explicitly want to pre-fetch
        if (fetch_all_dna or identifiers is None) and not prefetched_data:
            if self.resolver_backend == "api":
                logging.info(f"Fetching list of all '{biotype}' genes via API...")
                # This returns {ensembl_id: dna_sequence} or None
                fetched = self.gene_resolver.get_gene_sequences(biotype=biotype)
                if fetched is not None:
                    prefetched_data = fetched

                if identifiers is None:
                    # USE DISCOVERED GENES
                    if prefetched_data:
                        identifiers = list(prefetched_data.keys())
                        logging.info(f"Auto-discovered {len(identifiers)} genes.")
                        # Force ID type to Ensembl ID as that's what get_gene_sequences returns
                        id_type = "ensembl_id"
                    else:
                        logging.error(f"No genes found for biotype '{biotype}'.")
                        return []
            else:
                if identifiers is None:
                    logging.error("Cannot auto-discover genes with 'local' backend. Provide identifiers.")
                    return []
                logging.warning("fetch_all_dna is ignored in local mode.")

        if not identifiers:
            self.last_report = ResolutionReport(model_name=model, organism=organism)
            return []

        # Layer 2: fresh ResolutionReport for this batch. Every input
        # gets exactly one ``ResolutionRecord`` regardless of whether
        # it resolves, fails the resolver, fails the embedder, or hits
        # the OOM bisector. Replaces ``self.last_report`` so callers
        # can grab it after the call returns.
        report = ResolutionReport(model_name=model, organism=organism)
        self.last_report = report

        input_data_list: list[str | None] = []
        # Track per-input resolver state so we can attach reason codes
        # AFTER the embed_batch result is known (a resolver hit that
        # later fails to embed should be reported separately from a
        # resolver miss).
        per_input_state: list[dict[str, Any]] = []
        logging.info(f"Batch: {len(identifiers)} items for '{model}' ({mtype})...")

        # Narrow id_type for DNA resolver methods (only accept symbol/ensembl_id)
        dna_id: Literal["symbol", "ensembl_id"] | None = (
            id_type if id_type in ("symbol", "ensembl_id") else None  # type: ignore[assignment]
        )

        for ident in identifiers:
            data = None
            # Per-input audit state for Layer 2 reporting.
            state: dict[str, Any] = {
                "source": None,
                "reason": None,
                "attempted": [],
                "latency_ms": 0.0,
            }
            try:
                with time_block(
                    "resolver_input",
                    model=model,
                    identifier=ident,
                    id_type=id_type,
                    mtype=mtype,
                ) as ev:
                    if id_type == "sequence":
                        data = ident
                        state["source"] = "passthrough"
                    elif mtype == "dna":
                        state["attempted"].append("gene_resolver_dna")
                        if region in ("exons", "introns") and dna_id is not None:
                            data = self.gene_resolver.get_gene_region_sequence(
                                ident,
                                id_type=dna_id,
                                organism=organism,
                                region=region,
                            )
                            state["source"] = f"gene_resolver_region:{region}"
                        else:
                            if prefetched_data:
                                if id_type == "ensembl_id":
                                    data = prefetched_data.get(ident)
                                    if data is not None:
                                        state["source"] = "prefetched"

                            if data is None:
                                if dna_id is None:
                                    logging.warning(
                                        "DNA models require id_type 'symbol', 'ensembl_id', or 'sequence', "
                                        "got '%s'; skipping %s.",
                                        id_type,
                                        ident,
                                    )
                                    state["reason"] = f"resolver:unsupported_id_type:{id_type}"
                                elif self.resolver_backend == "local":
                                    data = self.gene_resolver.get_local_dna_sequence(ident, dna_id)
                                    state["source"] = "gene_resolver_local"
                                else:
                                    data = self.gene_resolver.get_dna_sequence(ident, dna_id, organism)
                                    state["source"] = "gene_resolver_api"

                    elif mtype == "protein":
                        state["attempted"].append("protein_resolver")
                        data = self.protein_resolver.get_canonical_sequence(ident, id_type, organism)
                        state["source"] = "protein_resolver"

                    elif mtype == "text":
                        # Default template combines the input identifier
                        # with MyGene.info fields. ``_SafeFormatDict`` in
                        # the resolver tolerates missing fields, so genes
                        # with no ``summary`` on MyGene still produce a
                        # thin but non-None description (e.g. just "Gene
                        # TP53 (human). TP53: tumor protein p53.").
                        state["attempted"].append("mygene")
                        fmt = gene_description_format or ("Gene {identifier} ({organism}). {symbol}: {name}. {summary}")
                        data = self.gene_resolver.get_gene_description(
                            ident,
                            id_type,
                            organism,
                            format_string=fmt,
                        )
                        state["source"] = "mygene_description"

                    if data is None:
                        logging.warning(f"No data for {ident}; skipping.")
                        ev["status"] = "no_hit"
                        if state["reason"] is None:
                            state["reason"] = f"{mtype}_resolver:no_hit"
            except Exception as e:  # noqa: BLE001
                # We deliberately swallow resolver-side exceptions per
                # input (the resolver wraps its own errors in None
                # returns; anything still escaping here is unexpected
                # but should not nuke a 2000-gene batch). Layer 2
                # records the failure so callers can audit.
                logging.warning(f"Error fetching {ident}: {e}")
                state["reason"] = f"resolver_exception:{type(e).__name__}"
                log_event(
                    "resolver_input_exception",
                    level="warn",
                    model=model,
                    identifier=ident,
                    error=type(e).__name__,
                    error_msg=str(e)[:200],
                )
            input_data_list.append(data)
            per_input_state.append(state)

        valid_inputs = [d for d in input_data_list if d is not None]
        valid_indices = [i for i, d in enumerate(input_data_list) if d is not None]

        if not valid_inputs:
            # Every input failed the resolver -- record each one as
            # UNRESOLVED with the per-input reason captured above and
            # return early. This is the failure mode that used to be
            # silently swallowed under "BioEmbedder returned no
            # embeddings"; now the sidecar carries the actionable
            # detail.
            for ident, state in zip(identifiers, per_input_state, strict=False):
                report.record(
                    ident,
                    "unresolved",
                    source=state["source"],
                    reason=state["reason"] or "resolver:no_data",
                    latency_ms=float(state["latency_ms"]),
                    attempted_sources=list(state["attempted"]),
                )
            return [None] * len(identifiers)

        # Layer 3: wrap the model forward pass in an OOM-bisection
        # helper. If the wrapper already does its own batching (most
        # of them do, after the protein / DNA fixes), the outer call
        # succeeds and the wrapper is responsible for chunking. If
        # the wrapper raises an OOM-shaped exception, the bisector
        # halves the input list and retries, isolating the specific
        # input that genuinely cannot fit. The single-item case that
        # still OOMs propagates so it can be classified as
        # ContextOverflowError / ModelOOMError below.
        def _do_embed(batch: Sequence[str]) -> list[Any]:
            return list(
                inst.embed_batch(
                    inputs=list(batch),
                    pooling_strategy=pooling_strategy,
                    **kwargs,
                )
            )

        try:
            with time_block(
                "model_forward",
                model=model,
                mtype=mtype,
                n_inputs=len(valid_inputs),
                device=str(self.device),
            ):
                batch_results = embed_batch_with_oom_recovery(
                    _do_embed,
                    valid_inputs,
                    on_split=lambda left, right: log_event(
                        "oom_bisect",
                        level="warn",
                        model=model,
                        left=left,
                        right=right,
                    ),
                )
        except EmbpyError:
            # Already a typed embpy exception (e.g. ModelOOMError raised by
            # a wrapper that classifies OOMs itself). Let it propagate
            # unchanged -- wrapping would only obscure the original fields.
            traceback.print_exc()
            raise
        except Exception as e:
            # Classify the failure so the upstream entry point can give a
            # category-specific exit code without parsing stderr. The
            # historical behaviour ("swallow + return [None]*N") masked
            # CUDA OOMs, context overflows, missing CUDA kernels, and
            # missing pip packages behind a single misleading "no
            # embeddings" error several layers up. We refuse to do that.
            logging.error(
                "Batch embed failed for model='%s' (n_inputs=%d): %s",
                model,
                len(valid_inputs),
                e,
            )
            traceback.print_exc()
            typed = _classify_embedder_exception(
                e,
                model_name=model,
                n_inputs=len(valid_inputs),
                device=str(self.device),
            )
            raise typed from e

        results: list[np.ndarray | None] = [None] * len(identifiers)
        for idx, emb in zip(valid_indices, batch_results, strict=False):
            results[idx] = emb

        # Layer 2: now that we know which inputs ended up with an
        # embedding (and which were resolver-misses), record one
        # ResolutionRecord per input. Embedder-side failures (OOM,
        # context overflow) raise typed errors above and never reach
        # this point, so anything in ``results[i] is None`` here is a
        # resolver miss with state already captured.
        valid_index_set = set(valid_indices)
        for i, (ident, state) in enumerate(zip(identifiers, per_input_state, strict=False)):
            if i in valid_index_set and results[i] is not None:
                report.record(
                    ident,
                    "resolved",
                    source=state["source"],
                    latency_ms=float(state["latency_ms"]),
                    attempted_sources=list(state["attempted"]),
                )
            else:
                report.record(
                    ident,
                    "unresolved",
                    source=state["source"],
                    reason=state["reason"] or f"{mtype}_resolver:no_data",
                    latency_ms=float(state["latency_ms"]),
                    attempted_sources=list(state["attempted"]),
                )

        return results

    def embed_molecule(
        self,
        identifier: str,  # Expecting SMILES string
        model: str,  # e.g. "chemberta_zinc_v1"
        pooling_strategy: str = "mean",
        **kwargs: Any,
    ) -> np.ndarray:
        """
        Generate an embedding for a single small molecule.

        Validates that the input is a proper SMILES string before embedding.

        Args:
            identifier (str):
                A SMILES string representing the molecule (e.g., "CCO" for ethanol).
            model (str):
                Key of the molecule embedding model to use (e.g., "chemberta_zinc_v1").
            pooling_strategy (str, optional):
                Pooling strategy to aggregate token‐level embeddings.
                Defaults to "mean".
            **kwargs:
                Additional keyword arguments forwarded to the model’s `embed` method.

        Returns
        -------
            np.ndarray:
                A 1D NumPy array representing the molecule embedding.

        Raises
        ------
            ModelNotFoundError:
                If the specified model key is not registered.
            ValueError:
                If the chosen model is not of type "molecule", or if the SMILES string is invalid.
            RuntimeError:
                If the embedding process itself fails.
        """
        inst = self._get_model(model)
        if inst.model_type != "molecule":
            raise ValueError(f"Model '{model}' is not a molecule embedder.")

        smiles = identifier
        from rdkit import Chem
        # Validate SMILES
        if Chem.MolFromSmiles(smiles) is None:
            raise ValueError(f"Invalid SMILES string: '{smiles}'")

        logging.debug(f"Embedding molecule SMILES: {smiles} with model '{model}'")
        try:
            emb = inst.embed(input=smiles, pooling_strategy=pooling_strategy, **kwargs)
            logging.debug(f"Molecule embedding shape: {emb.shape}")
            return emb
        except Exception as e:
            logging.error(f"Embedding failed for SMILES '{smiles}': {e}")
            raise RuntimeError(f"Embedding failed for SMILES '{smiles}'") from e

    def embed_molecules_batch(
        self,
        identifiers: Sequence[str],
        model: str,
        pooling_strategy: str = "mean",
        **kwargs: Any,
    ) -> list[np.ndarray | None]:
        """
        Generate embeddings for a batch of small molecules.

        Invalid SMILES strings are skipped and return `None` in their positions.

        Args:
            identifiers (Sequence[str]):
                A list of SMILES strings to embed.
            model (str):
                Key of the molecule embedding model to use.
            pooling_strategy (str, optional):
                Pooling strategy to apply. Defaults to "mean".
            **kwargs:
                Additional keyword arguments forwarded to `embed_batch`.

        Returns
        -------
            List[Optional[np.ndarray]]:
                A list of embeddings (NumPy arrays) for valid SMILES, with `None`
                for any invalid or failed entries. Order aligns with input list.

        Raises
        ------
            ModelNotFoundError:
                If the specified model key is not registered.
            ValueError:
                If the chosen model is not of type "molecule".
        """
        inst = self._get_model(model)
        if inst.model_type != "molecule":
            raise ValueError(f"Model '{model}' is not a molecule embedder.")

        from rdkit import Chem
        valid_inputs: list[str] = []
        valid_indices: list[int] = []
        for idx, smi in enumerate(identifiers):
            if Chem.MolFromSmiles(smi) is not None:
                valid_inputs.append(smi)
                valid_indices.append(idx)
            else:
                logging.warning(f"Skipping invalid SMILES: '{smi}'")

        # Prepare output list
        results: list[np.ndarray | None] = [None] * len(identifiers)
        if not valid_inputs:
            logging.warning("No valid SMILES provided; returning all None.")
            return results

        logging.info(f"Embedding {len(valid_inputs)} valid SMILES with model '{model}'")
        try:
            # Wrappers disagree on the batch parameter name: most declare
            # ``inputs`` (base.py), while ChemBERTa declares ``input``. Try the
            # common spelling first and fall back, mirroring the generic embed
            # path -- otherwise every molecule model except ChemBERTa raises a
            # TypeError that is silently swallowed as "all inputs failed".
            try:
                batch_embs = inst.embed_batch(inputs=valid_inputs, pooling_strategy=pooling_strategy, **kwargs)
            except TypeError:
                batch_embs = inst.embed_batch(input=valid_inputs, pooling_strategy=pooling_strategy, **kwargs)
            for out_idx, emb in zip(valid_indices, batch_embs, strict=False):
                results[out_idx] = emb
        except Exception as e:  # noqa: BLE001
            logging.error(f"Batch embedding failed: {e}")
            # Leave all as None

        return results

    def embed_text(
        self,
        text: str,
        model: str,  # Can be registry name OR any HF model identifier
        pooling_strategy: str = "mean",
        **kwargs: Any,
    ) -> np.ndarray:
        """
        Generates an embedding for an arbitrary text string using a text model.

        Args:
            text (str): The input text string.
            model (str): The name of the text embedding model. Can be either:
                        - A predefined model name (e.g., "minilm_l6_v2", "bert_base_uncased")
                        - Any Hugging Face model identifier (e.g., "sentence-transformers/all-MiniLM-L6-v2")
            pooling_strategy (str): Pooling strategy. Defaults to "mean".
            **kwargs: Additional arguments for the text model's embed method.

        Returns
        -------
            np.ndarray: The computed text embedding.

        Raises
        ------
            ModelNotFoundError: If the model cannot be loaded from HF or found in registry.
            ValueError: If the loaded model is not a text model.
            RuntimeError: If model loading or inference fails.
        """
        model_instance = self._get_model(model)
        if model_instance.model_type != "text":
            raise ValueError(f"Model '{model}' is not a text embedder. Use embed_gene or embed_molecule.")

        logging.debug(f"Embedding text: '{text[:100]}...' using model '{model}'")
        try:
            embedding = model_instance.embed(input=text, pooling_strategy=pooling_strategy, **kwargs)
            logging.debug(f"Text embedding generated with shape: {embedding.shape}")
            return embedding
        except Exception as e:
            logging.error(f"Error during text embedding generation for model '{model}': {e}")
            raise RuntimeError(f"Text embedding failed for input '{text[:50]}...'.") from e

    def embed_text_api(
        self,
        text: str,
        model: str = "text-embedding-3-small",
        provider: str = "openai",
        api_key: str | None = None,
        base_url: str | None = None,
        **kwargs: Any,
    ) -> np.ndarray:
        """Embed text using an API-based model with explicit provider config.

        Convenience method for one-off API calls where you want to
        specify the provider, API key, and/or base URL directly
        without pre-registering the model.

        Parameters
        ----------
        text
            Text string to embed.
        model
            Model name as expected by the API.
        provider
            ``"openai"``, ``"cohere"``, ``"voyage"``, ``"google"``,
            or ``"generic"`` (OpenAI-compatible endpoint).
        api_key
            API key (or set the provider's env var).
        base_url
            Custom API endpoint URL.
        **kwargs
            Forwarded to the API wrapper.

        Returns
        -------
        np.ndarray
            Embedding vector.
        """
        from .models.api_models import APIEmbeddingWrapper

        cache_key = f"_api_{provider}_{model}"
        if cache_key not in self.model_cache:
            wrapper = APIEmbeddingWrapper(
                model_path_or_name=model,
                provider=provider,
                api_key=api_key,
                base_url=base_url,
            )
            wrapper.load(self.device)
            self.model_cache[cache_key] = wrapper

        return self.model_cache[cache_key].embed(input=text, **kwargs)

    def embed_texts_batch(
        self,
        texts: Sequence[str],
        model: str,
        pooling_strategy: str = "mean",
        batch_size: int | None = None,
        **kwargs: Any,
    ) -> list[np.ndarray | None]:  # Return None for errors
        """
        Generates embeddings for a batch of arbitrary text strings.

        Args:
            texts (Sequence[str]): A list or tuple of text strings.
            model (str): The name of the text embedding model.
            pooling_strategy (str): Pooling strategy. Defaults to "mean".
            batch_size (Optional[int]): Maximum batch size for processing. If None, processes all texts at once.
                                       Use smaller values to avoid OOM errors with large datasets.
            **kwargs: Additional arguments for the model's embed_batch method.

        Returns
        -------
            list[Optional[np.ndarray]]: List of embeddings, with None for texts that failed.
        """
        model_instance = self._get_model(model)
        if model_instance.model_type != "text":
            raise ValueError(f"Model '{model}' is not a text embedder.")

        valid_inputs = list(texts)
        logging.info(f"Embedding batch of {len(valid_inputs)} texts using model '{model}'...")

        try:
            batch_results = model_instance.embed_batch(
                inputs=valid_inputs,
                pooling_strategy=pooling_strategy,
                batch_size=batch_size,
                **kwargs,
            )
            if len(batch_results) != len(valid_inputs):
                logging.error(
                    f"Batch embedding returned {len(batch_results)} results for {len(valid_inputs)} inputs. Mismatch!"
                )
                return [None] * len(texts)
            logging.info("Batch text embedding successful.")
            results: list[np.ndarray | None] = list(batch_results)

        except (ValueError, KeyError) as e:
            logging.error(f"Error during batch text embedding generation for model '{model}': {e}")
            results = [None] * len(texts)

        return results

    def embed_description(
        self,
        identifier: str,
        model: str = "minilm_l6_v2",
        entity_type: Literal["gene", "protein", "molecule", "auto"] = "auto",
        sources: list[str] | str = "all",
        pooling_strategy: str = "mean",
        **kwargs: Any,
    ) -> np.ndarray:
        """Fetch text description(s) for a biological entity and embed them.

        Resolves the identifier to rich text descriptions from public
        knowledge sources (MyGene, NCBI, Ensembl, UniProt, Wikipedia,
        PubChem), combines them, and embeds the result with a text model.

        Parameters
        ----------
        identifier
            Gene symbol, Ensembl ID, UniProt accession, drug name,
            or SMILES string.
        model
            Text embedding model (default ``"minilm_l6_v2"``).
        entity_type
            ``"gene"``, ``"protein"``, ``"molecule"``, or ``"auto"``
            to auto-detect.
        sources
            Knowledge sources to query. ``"all"`` uses the defaults
            for the detected entity type.
        pooling_strategy
            Pooling strategy for the text model.
        **kwargs
            Forwarded to the text model's ``embed`` method.

        Returns
        -------
        np.ndarray
            The text embedding vector.

        Examples
        --------
        >>> emb = embedder.embed_description("TP53", model="minilm_l6_v2")
        >>> emb.shape
        (384,)
        >>> emb = embedder.embed_description("aspirin", entity_type="molecule")
        """
        description = self.text_resolver.get_combined_description(
            identifier,
            entity_type=entity_type,
            sources=sources,
        )
        logging.info(
            "Description for '%s' (%d chars): %s...",
            identifier,
            len(description),
            description[:100],
        )
        return self.embed_text(
            text=description,
            model=model,
            pooling_strategy=pooling_strategy,
            **kwargs,
        )

    def embed_descriptions_batch(
        self,
        identifiers: Sequence[str],
        model: str = "minilm_l6_v2",
        entity_type: Literal["gene", "protein", "molecule", "auto"] = "auto",
        sources: list[str] | str = "all",
        pooling_strategy: str = "mean",
        batch_size: int | None = None,
        **kwargs: Any,
    ) -> list[np.ndarray | None]:
        """Fetch and embed text descriptions for a batch of entities.

        Parameters
        ----------
        identifiers
            List of biological entity identifiers.
        model
            Text embedding model.
        entity_type
            Entity type (applied to all identifiers).
        sources
            Knowledge sources to query.
        pooling_strategy
            Pooling strategy for the text model.
        batch_size
            Maximum batch size for embedding inference.
        **kwargs
            Forwarded to the text model.

        Returns
        -------
        List of embedding arrays, with ``None`` for failed lookups.
        """
        texts: list[str] = []
        valid_indices: list[int] = []

        for i, ident in enumerate(identifiers):
            try:
                desc = self.text_resolver.get_combined_description(
                    ident,
                    entity_type=entity_type,
                    sources=sources,
                )
                if desc:
                    texts.append(desc)
                    valid_indices.append(i)
                else:
                    logging.warning("No description found for '%s'", ident)
            except Exception as e:  # noqa: BLE001
                logging.warning("Failed to fetch description for '%s': %s", ident, e)

        if not texts:
            return [None] * len(identifiers)

        batch_embs = self.embed_texts_batch(
            texts=texts,
            model=model,
            pooling_strategy=pooling_strategy,
            batch_size=batch_size,
            **kwargs,
        )

        results: list[np.ndarray | None] = [None] * len(identifiers)
        for idx, emb in zip(valid_indices, batch_embs, strict=False):
            results[idx] = emb

        return results

    @staticmethod
    def _singlecell_preprocessing_warnings(
        adata: Any,
        *,
        models: Sequence[str],
        resolved_preprocessing: str,
        pca_use_hvg: bool,
    ) -> dict[str, list[str]]:
        """Return actionable warnings about model/input preprocessing fit."""
        from .models.singlecell_models import singlecell_info

        warnings_by_model: dict[str, list[str]] = {}
        for model_key in models:
            card = singlecell_info(model_key)
            warnings: list[str] = []
            if card.input_layer == "log_normalized" and "log_normalized" not in adata.layers:
                warnings.append(
                    "Model consumes .layers['log_normalized'], but that layer "
                    "is absent. Run with preprocessing='auto' or "
                    "preprocessing='standard', or provide that layer yourself."
                )
            if card.input_layer == "counts" and "counts" not in adata.layers:
                warnings.append(
                    "Model is configured to consume raw counts from "
                    ".layers['counts'], but that layer is absent. The wrapper "
                    "may fall back to .X; use preprocessing='auto' or "
                    "preprocessing='raw' to materialize the counts layer."
                )
            if card.uses_hvg and pca_use_hvg and "highly_variable" not in adata.var.columns:
                warnings.append(
                    "Model can restrict to highly-variable genes, but "
                    "adata.var['highly_variable'] is absent. Use "
                    "preprocessing='auto'/'standard' with select_hvg=True, "
                    "or pass pca_use_hvg=False to use all genes."
                )
            if warnings:
                warnings_by_model[model_key] = warnings
                for message in warnings:
                    logging.warning(
                        "Single-cell preprocessing warning for %s (resolved=%s): %s",
                        model_key,
                        resolved_preprocessing,
                        message,
                    )
        return warnings_by_model

    def embed_cells(
        self,
        adata,  # anndata.AnnData
        models: list[str] | str = "scgpt",
        preprocessing: Literal["auto", "raw", "standard", "none"] = "auto",
        *,
        # Preprocessing params (forwarded to preprocess_counts)
        target_sum: float | None = 1e4,
        n_top_genes: int = 2000,
        select_hvg: bool = True,
        hvg_flavor: Literal["auto", "seurat", "seurat_v3", "cell_ranger"] = "auto",
        log_transform: bool = True,
        scale: bool = False,
        max_value: float | None = 10.0,
        min_genes: int = 200,
        min_cells: int = 3,
        max_pct_mito: float | None = None,
        # PCA-specific
        n_pca_components: int = 50,
        pca_use_hvg: bool = True,
        # scVI-specific
        n_latent: int = 30,
        n_layers_scvi: int = 2,
        n_hidden_scvi: int = 128,
        max_epochs: int = 200,
        early_stopping: bool = True,
        batch_key: str | None = None,
        labels_key: str | None = None,
        protein_expression_obsm_key: str | None = None,
        # General
        batch_size: int = 32,
        obsm_prefix: str = "X_",
        copy: bool = True,
        backend: Literal["cpu", "gpu"] = "cpu",
        vocab_conversion: Literal["auto", "off"] = "auto",
        model_kwargs: dict[str, dict[str, Any]] | None = None,
    ):
        """Embed single cells from an AnnData object.

        Orchestrates preprocessing and multi-model embedding in a single
        call.  Each requested model's embeddings are stored in
        ``adata.obsm`` under ``"{obsm_prefix}{model_name}"``.

        Parameters
        ----------
        adata : anndata.AnnData
            AnnData with raw counts in ``.X``.
        models : str or list[str]
            Model key(s) from the single-cell model registry.  Accepts
            any key from :func:`~embpy.models.singlecell_models.list_singlecell_models`,
            e.g. ``"scgpt"``, ``"geneformer_v2_12L"``, ``"pca"``,
            ``"scvi"``, ``"scanvi"``, ``"totalvi"``.
        preprocessing : {"auto", "raw", "standard", "none"}
            ``"auto"`` chooses a model-aware default from the single-cell
            registry: raw-count models get QC/counts-layer preparation, while
            models that consume processed expression (for example PCA) get the
            standard pipeline.
            ``"standard"`` runs log-normalize + HVG.
            ``"raw"`` applies only QC filtering.
            ``"none"`` skips preprocessing entirely.
        target_sum
            Target total counts for normalization (standard pipeline).
        n_top_genes
            Number of highly variable genes (standard pipeline).
        select_hvg
            Whether to compute highly variable genes in the standard pipeline.
        hvg_flavor
            Highly-variable-gene method. ``"auto"`` chooses a scanpy flavor
            compatible with ``log_transform``.
        log_transform
            Whether to log1p-transform (standard pipeline).
        scale
            Whether to scale to unit variance (standard pipeline).
        max_value
            Max value after scaling.
        min_genes
            Minimum genes per cell for QC.
        min_cells
            Minimum cells per gene for QC.
        max_pct_mito
            Maximum mitochondrial fraction (``None`` to skip).
        n_pca_components
            Number of PCA components (when ``"pca"`` is in models).
        pca_use_hvg
            Whether PCA restricts to HVGs.
        n_latent
            scvi-tools latent dimensionality.
        n_layers_scvi
            Number of hidden layers in scvi-tools encoder/decoder.
        n_hidden_scvi
            Number of nodes per hidden layer in scvi-tools.
        max_epochs
            Maximum training epochs for scvi-tools.
        early_stopping
            Whether to use early stopping for scvi-tools.
        batch_key
            Batch column in ``adata.obs`` for scvi-tools batch correction.
        labels_key
            Label column in ``adata.obs`` for scANVI.
        protein_expression_obsm_key
            Key in ``adata.obsm`` with protein counts for totalVI.
        batch_size
            Batch size for foundation model inference.
        obsm_prefix
            Prefix for ``.obsm`` keys (default ``"X_"``).
        copy
            If ``True``, operate on a copy of adata.
        model_kwargs
            Per-model constructor arguments, keyed by model key, e.g.
            ``{"state": {"checkpoint": "/path/to/se600m.ckpt"}}``. Needed by
            models whose weights are not fetched for you: ``state`` and
            ``stack`` take a ``checkpoint`` (``state`` also accepts
            ``model_folder``), and raise from ``load()`` without one, so
            before this existed they could not be reached through this
            method at all. The kwargs are part of the wrapper cache key,
            so two checkpoints of one model do not collide.

        Returns
        -------
        anndata.AnnData
            The AnnData with:

            - ``.X`` = original raw counts
            - ``.layers["counts"]`` = raw counts copy
            - ``.layers["log_normalized"]`` = processed (standard pipeline)
            - ``.obsm["{prefix}{model}"]`` = embeddings per model
            - ``.uns["embpy_cell_embeddings"]`` = metadata dict
        """
        from .models.singlecell_models import (
            PCAEmbedding,
            ScVIToolsWrapper,
            list_singlecell_models,
            resolve_singlecell_preprocessing,
            singlecell_info,
        )
        from .pp.sc_preprocessing import preprocess_counts

        if isinstance(models, str):
            models = [models]

        available = list_singlecell_models()
        for m in models:
            if m not in available:
                raise ValueError(f"Unknown single-cell model '{m}'. Available: {available}")

        resolved_preprocessing, preprocessing_plan = resolve_singlecell_preprocessing(
            models,
            preprocessing,
        )

        if copy:
            adata = adata.copy()

        # ---- Preprocessing -----------------------------------------------
        preprocessing_options = {
            "target_sum": target_sum,
            "n_top_genes": n_top_genes,
            "select_hvg": select_hvg,
            "hvg_flavor": hvg_flavor,
            "log_transform": log_transform,
            "scale": scale,
            "max_value": max_value,
            "min_genes": min_genes,
            "min_cells": min_cells,
            "max_pct_mito": max_pct_mito,
            "backend": backend,
        }
        if resolved_preprocessing != "none":
            adata = preprocess_counts(
                adata,
                pipeline=resolved_preprocessing,
                min_genes=min_genes,
                min_cells=min_cells,
                max_pct_mito=max_pct_mito,
                target_sum=target_sum,
                n_top_genes=n_top_genes,
                select_hvg=select_hvg,
                hvg_flavor=hvg_flavor,
                log_transform=log_transform,
                scale=scale,
                max_value=max_value,
                copy=False,
                backend=backend,
            )
        else:
            logging.warning(
                "Skipping single-cell preprocessing because preprocessing='none'. "
                "Models that expect .layers['log_normalized'], .layers['counts'], "
                "or adata.var['highly_variable'] may fall back to .X or run "
                "without HVG restriction."
            )

        preprocessing_warnings = self._singlecell_preprocessing_warnings(
            adata,
            models=models,
            resolved_preprocessing=resolved_preprocessing,
            pca_use_hvg=pca_use_hvg,
        )
        preprocessing_metadata: dict[str, Any] = {
            **preprocessing_plan,
            "options": preprocessing_options,
            "warnings": preprocessing_warnings,
        }

        # ---- Embed with each model ---------------------------------------
        device_str = str(self.device)
        metadata: dict[str, dict[str, Any]] = {
            "__preprocessing__": preprocessing_metadata,
        }

        for model_key in models:
            obsm_key = f"{obsm_prefix}{model_key}"
            logging.info("Embedding cells with '%s' ...", model_key)

            try:
                if model_key == "pca":
                    wrapper = PCAEmbedding(
                        n_components=n_pca_components,
                        use_hvg=pca_use_hvg,
                        layer="log_normalized",
                        scale=True,
                        backend=backend,
                    )
                    wrapper.load(device_str)
                    embs = wrapper.embed_cells(adata)
                    # Park the fitted wrapper so decode_cells() can reuse
                    # the sklearn/cuml PCA + scaler for inverse_transform.
                    self._singlecell_cache[(model_key, device_str)] = wrapper

                elif model_key in ("scvi", "scanvi", "totalvi"):
                    model_cls_name = model_key.upper()
                    if model_cls_name == "SCVI":
                        model_cls_name = "SCVI"
                    elif model_cls_name == "SCANVI":
                        model_cls_name = "SCANVI"
                    elif model_cls_name == "TOTALVI":
                        model_cls_name = "TOTALVI"

                    wrapper = ScVIToolsWrapper(
                        model_class=model_cls_name,
                        n_latent=n_latent,
                        n_layers=n_layers_scvi,
                        n_hidden=n_hidden_scvi,
                        max_epochs=max_epochs,
                        early_stopping=early_stopping,
                        batch_key=batch_key,
                        labels_key=labels_key,
                        protein_expression_obsm_key=protein_expression_obsm_key,
                        layer="counts",
                        batch_size=batch_size,
                    )
                    wrapper.load(device_str)
                    embs = wrapper.embed_cells(adata)
                    # Park the trained scvi-tools model so decode_cells()
                    # can route arbitrary latents through its generative
                    # module without retraining.
                    self._singlecell_cache[(model_key, device_str)] = wrapper

                else:
                    # Foundation models (scGPT, Geneformer, UCE, Tahoe, ...)
                    # are cached by (model_key, device) so repeated calls
                    # to embed_cells -- e.g. chunked inference over a large
                    # AnnData -- reuse the already-loaded torch model
                    # instead of re-instantiating it on every call.
                    wrapper = self._get_or_load_singlecell_wrapper(
                        model_key, batch_size, device_str,
                        (model_kwargs or {}).get(model_key),
                    )
                    # Auto-adapt var_names to the model's vocabulary.
                    # This is the difference between a silent zero-match
                    # failure (e.g. passing Ensembl IDs to scGPT) and a
                    # working embedding. Set vocab_conversion="off" to
                    # opt out of the automatic GeneResolver roundtrip.
                    if vocab_conversion == "auto":
                        adata_for_model, report = self._ensure_singlecell_vocabulary(
                            adata, model_key, organism=self.organism
                        )
                        metadata[f"{model_key}__vocab"] = report
                    else:
                        adata_for_model = adata
                    embs = wrapper.embed_cells(adata_for_model)

                adata.obsm[obsm_key] = embs
                card = singlecell_info(model_key)
                metadata[model_key] = {
                    "obsm_key": obsm_key,
                    "embedding_dim": embs.shape[1],
                    "n_cells": embs.shape[0],
                    "wrapper_class": type(wrapper).__name__,
                    "preprocessing": preprocessing_metadata,
                    "model_requirements": {
                        "default_preprocessing": card.default_preprocessing,
                        "input_layer": card.input_layer,
                        "uses_hvg": card.uses_hvg,
                        "vocab_type": card.vocab_type,
                    },
                }
                logging.info(
                    "  -> stored in .obsm['%s'], shape %s",
                    obsm_key,
                    embs.shape,
                )

            except Exception as e:  # noqa: BLE001
                logging.error("Failed to embed with '%s': %s", model_key, e)
                metadata[model_key] = {
                    "error": str(e),
                    "preprocessing": preprocessing_metadata,
                }

        adata.uns["embpy_cell_embeddings"] = metadata
        logging.info(
            "embed_cells complete: %d models, %d cells",
            len(models),
            adata.n_obs,
        )
        return adata

    def decode_cells(
        self,
        adata=None,
        *,
        latent=None,
        model: str = "scvi",
        obsm_key: str | None = None,
        gene_names: Sequence[str] | None = None,
        write_layer: str | None = None,
        **decode_kwargs: Any,
    ):
        """Decode cell embeddings back to gene-expression space.

        Companion to :meth:`embed_cells` that routes a latent matrix
        through the *same* fitted encoder-decoder wrapper that produced
        it. This is the decoder hook used by flow-matching / cellflow
        style perturbation pipelines: encode basal + perturbed cells to
        a shared latent, learn a transport map there, then decode the
        predicted latent back to counts with this method.

        Requires a prior :meth:`embed_cells` call with the same
        ``model`` key so the trained wrapper is cached for this device.
        For STATE / Stack, instantiate the wrapper manually (checkpoint +
        gene-list paths cannot be supplied through :meth:`embed_cells`).

        Parameters
        ----------
        adata : anndata.AnnData, optional
            AnnData that holds the latent to decode. Used both to pick
            the latent from ``obsm[obsm_key]`` and to recover
            ``var_names`` for gene-parametric decoders (STATE).
        latent : np.ndarray, optional
            Alternative to ``adata`` + ``obsm_key``: pass the latent
            matrix directly.
        model : str
            Encoder-decoder registry key: ``"pca"``, ``"scvi"``,
            ``"scanvi"``, ``"totalvi"``, or ``"state"``.
        obsm_key : str, optional
            Override key used to pull the latent from ``adata.obsm``.
            Defaults to ``f"X_{model}"``.
        gene_names : sequence of str, optional
            Target genes to decode to (STATE only). Defaults to
            ``adata.var_names`` when ``adata`` is supplied.
        write_layer : str, optional
            If set and ``adata`` is supplied, the decoded expression is
            written to ``adata.layers[write_layer]``.
        **decode_kwargs
            Forwarded to the wrapper's ``decode_cells`` (e.g.
            ``library_size``, ``batch_index`` for scVI, ``read_depth``
            for STATE).

        Returns
        -------
        np.ndarray of shape ``(n_cells, n_genes)``
            Decoded expression (NB mean / log-probs / linear recon,
            depending on the model; see the wrapper docstring).
        """
        from .models.singlecell_models import singlecell_info

        card = singlecell_info(model)
        if not card.supports_decode:
            raise ValueError(
                f"Model '{model}' does not expose a decoder "
                f"(wrapper_class={card.wrapper_class_name}). "
                "Supported: pca, scvi, scanvi, totalvi, state."
            )

        if latent is None:
            if adata is None:
                raise ValueError("decode_cells needs either `latent` or `adata` (+ optional `obsm_key`).")
            key = obsm_key or f"X_{model}"
            if key not in adata.obsm:
                raise KeyError(
                    f"{key!r} not present in adata.obsm. Run embed_cells "
                    f"with models=[{model!r}] first, or pass `obsm_key`."
                )
            latent_arr = np.asarray(adata.obsm[key])
        else:
            latent_arr = np.asarray(latent)

        if gene_names is None and adata is not None:
            gene_names = list(adata.var_names)

        device_str = str(self.device)
        cache_key = (model, device_str)
        wrapper = self._singlecell_cache.get(cache_key)
        if wrapper is None:
            raise RuntimeError(
                f"No fitted '{model}' wrapper cached for device "
                f"{device_str!r}. Call `embed_cells(adata, models=["
                f"'{model}'])` first (for STATE/Stack, instantiate the "
                "wrapper manually with the checkpoint path and pass it to "
                "`wrapper.embed_cells` / `wrapper.decode_cells` directly)."
            )

        decoded = wrapper.decode_cells(
            latent_arr,
            gene_names=gene_names,
            adata=adata,
            **decode_kwargs,
        )

        if write_layer is not None and adata is not None:
            if decoded.shape[1] != adata.n_vars:
                raise ValueError(
                    f"Decoded matrix has {decoded.shape[1]} genes but adata has "
                    f"{adata.n_vars} var_names; cannot write to adata.layers."
                )
            adata.layers[write_layer] = decoded

        return decoded

    def generate_cells(
        self,
        base_adata,
        test_adata,
        *,
        model: str = "stack",
        wrapper=None,
        **generate_kwargs: Any,
    ):
        """Run in-context generation to synthesise new cell profiles.

        Currently only supported by the Stack wrapper. Unlike
        :meth:`decode_cells`, this does NOT consume a latent matrix: it
        conditions on a donor-specific ``base_adata`` context and
        synthesises gene-expression for every cell in ``test_adata``.

        Parameters
        ----------
        base_adata : anndata.AnnData
            Donor-conditioned base context.
        test_adata : anndata.AnnData
            Cells to predict for.
        model : str
            Registry key (currently only ``"stack"``).
        wrapper : SingleCellWrapper, optional
            Already-constructed wrapper. Required for Stack because its
            checkpoint and gene-list paths cannot be expressed through
            ``embed_cells``.
        **generate_kwargs
            Forwarded to the wrapper's ``generate_cells`` (split_column,
            split_values, prompt_ratio, ...).
        """
        from .models.singlecell_models import singlecell_info

        card = singlecell_info(model)
        if not card.supports_generation:
            raise ValueError(f"Model '{model}' does not expose an in-context generation head. Supported: stack.")

        if wrapper is None:
            device_str = str(self.device)
            wrapper = self._singlecell_cache.get((model, device_str))
        if wrapper is None:
            raise RuntimeError(
                f"No '{model}' wrapper available. For Stack, instantiate "
                "it manually with `StackWrapper(checkpoint=..., "
                "genelist=...)`, call `.load(device)`, and pass it via "
                "the `wrapper=` kwarg."
            )

        return wrapper.generate_cells(base_adata, test_adata, **generate_kwargs)

    def embed_adata(
        self,
        adata,  # anndata.AnnData
        *,
        # --- Cell-level embeddings (from expression) ---
        cell_models: list[str] | str | None = None,
        preprocessing: Literal["auto", "raw", "standard", "none"] = "auto",
        target_sum: float | None = 1e4,
        n_top_genes: int = 2000,
        select_hvg: bool = True,
        hvg_flavor: Literal["auto", "seurat", "seurat_v3", "cell_ranger"] = "auto",
        log_transform: bool = True,
        scale: bool = False,
        max_value: float | None = 10.0,
        min_genes: int = 200,
        min_cells: int = 3,
        max_pct_mito: float | None = None,
        n_pca_components: int = 50,
        pca_use_hvg: bool = True,
        n_latent: int = 30,
        n_layers_scvi: int = 2,
        n_hidden_scvi: int = 128,
        max_epochs: int = 200,
        early_stopping: bool = True,
        batch_key: str | None = None,
        labels_key: str | None = None,
        protein_expression_obsm_key: str | None = None,
        # --- Perturbation-level embeddings (gene / protein / molecule) ---
        perturbation_models: list[str] | str | None = None,
        perturbation_column: str | None = None,
        perturbation_type: Literal[
            "auto",
            "symbol",
            "ensembl_id",
            "uniprot_id",
            "smiles",
        ] = "auto",
        perturbation_organism: str = "human",
        pooling_strategy: str = "mean",
        # --- General ---
        batch_size: int = 32,
        obsm_prefix: str = "X_",
        copy: bool = True,
        backend: Literal["cpu", "gpu"] = "cpu",
    ):
        """Unified embedding of an AnnData -- cells and/or perturbations.

        Combines cell-level embeddings (from expression via single-cell
        foundation models, PCA, or scVI) with perturbation-level
        embeddings (from gene/protein/molecule identifiers in ``.obs``)
        in a single call.

        Parameters
        ----------
        adata : anndata.AnnData
            Input AnnData with raw counts in ``.X`` and (optionally)
            perturbation annotations in ``.obs``.

        cell_models : str or list[str], optional
            Single-cell model key(s) for expression-based cell
            embeddings (e.g. ``"scgpt"``, ``"pca"``, ``"scvi"``).
            ``None`` skips cell embedding.
        preprocessing
            Preprocessing pipeline for cell models.
        target_sum, n_top_genes, log_transform, scale, max_value
            Normalization and scaling options forwarded to
            :func:`~embpy.pp.preprocess_counts`.
        select_hvg, hvg_flavor
            Highly-variable-gene options forwarded to
            :func:`~embpy.pp.preprocess_counts`.
        min_genes, min_cells, max_pct_mito
            QC filtering options forwarded to
            :func:`~embpy.pp.preprocess_counts`.
        n_pca_components, pca_use_hvg
            PCA-specific parameters.
        n_latent, n_layers_scvi, n_hidden_scvi, max_epochs,
        early_stopping, batch_key, labels_key,
        protein_expression_obsm_key
            scvi-tools parameters.

        perturbation_models : str or list[str], optional
            Sequence/molecule model key(s) for perturbation embeddings
            (e.g. ``"esm2_650M"``, ``"enformer_human_rough"``,
            ``"chemberta2MTR"``).  ``None`` skips perturbation embedding.
        perturbation_column : str, optional
            Column in ``adata.obs`` containing perturbation identifiers.
            Required when *perturbation_models* is set.
        perturbation_type
            Type of identifiers in the perturbation column.
            ``"auto"`` auto-detects per identifier.
        perturbation_organism
            Organism for gene/protein resolution.
        pooling_strategy
            Pooling for sequence/molecule models.

        batch_size
            Batch size for model inference.
        obsm_prefix
            Prefix for ``.obsm`` keys (default ``"X_"``).
        copy
            If ``True``, operate on a copy of adata.

        Returns
        -------
        anndata.AnnData
            The enriched AnnData with:

            - ``.X`` = original raw counts
            - ``.layers["counts"]`` = raw counts copy
            - ``.layers["log_normalized"]`` = processed expression
              (standard pipeline)
            - ``.obsm["{prefix}{cell_model}"]`` = cell embeddings
            - ``.uns["perturbations"]["{prefix}{pert_model}"]``
              = perturbation embeddings, one row per unique perturbation
            - ``.uns["embpy_embeddings"]`` = metadata dict

        Examples
        --------
        >>> result = embedder.embed_adata(
        ...     adata,
        ...     cell_models=["pca", "scvi", "scgpt"],
        ...     perturbation_models=["esm2_650M", "chemberta2MTR"],
        ...     perturbation_column="perturbation",
        ...     perturbation_type="auto",
        ... )
        >>> result.obsm["X_scgpt"].shape  # cell embeddings
        (5000, 512)
        >>> result.uns["perturbations"]["X_esm2_650M"]["matrix"].shape
        (n_perturbations, 1280)
        """
        from .io.exporters import to_anndata
        from .io.result import EmbeddingProvenance, EmbeddingResult
        from .resources.gene_resolver import detect_identifier_type

        if copy:
            adata = adata.copy()

        metadata: dict[str, dict[str, Any]] = {}

        # ==================================================================
        # Cell-level embeddings (expression-based)
        # ==================================================================
        if cell_models is not None:
            adata = self.embed_cells(
                adata,
                models=cell_models,
                preprocessing=preprocessing,
                target_sum=target_sum,
                n_top_genes=n_top_genes,
                select_hvg=select_hvg,
                hvg_flavor=hvg_flavor,
                log_transform=log_transform,
                scale=scale,
                max_value=max_value,
                min_genes=min_genes,
                min_cells=min_cells,
                max_pct_mito=max_pct_mito,
                n_pca_components=n_pca_components,
                pca_use_hvg=pca_use_hvg,
                n_latent=n_latent,
                n_layers_scvi=n_layers_scvi,
                n_hidden_scvi=n_hidden_scvi,
                max_epochs=max_epochs,
                early_stopping=early_stopping,
                batch_key=batch_key,
                labels_key=labels_key,
                protein_expression_obsm_key=protein_expression_obsm_key,
                batch_size=batch_size,
                obsm_prefix=obsm_prefix,
                copy=False,
                backend=backend,
            )
            if "embpy_cell_embeddings" in adata.uns:
                metadata.update(adata.uns["embpy_cell_embeddings"])

        # ==================================================================
        # Perturbation-level embeddings (gene / protein / molecule)
        # ==================================================================
        if perturbation_models is not None:
            if perturbation_column is None:
                raise ValueError(
                    "perturbation_column is required when perturbation_models "
                    "is specified. Set it to the .obs column containing "
                    "perturbation identifiers (e.g. 'perturbation', 'gene', "
                    "'smiles')."
                )
            if perturbation_column not in adata.obs.columns:
                raise ValueError(
                    f"Column '{perturbation_column}' not found in adata.obs. Available: {list(adata.obs.columns)}"
                )

            if isinstance(perturbation_models, str):
                perturbation_models = [perturbation_models]

            pert_ids = adata.obs[perturbation_column].astype(str).values
            unique_perts = list(dict.fromkeys(pert_ids))
            logging.info(
                "Perturbation embedding: %d cells, %d unique perturbations from column '%s'",
                len(pert_ids),
                len(unique_perts),
                perturbation_column,
            )

            # Determine id_type per perturbation if auto
            if perturbation_type == "auto":
                id_types = {p: detect_identifier_type(p) for p in unique_perts}
            else:
                id_types = dict.fromkeys(unique_perts, perturbation_type)

            for model_key in perturbation_models:
                obsm_key = f"{obsm_prefix}{model_key}"
                logging.info(
                    "Embedding perturbations with '%s' ...",
                    model_key,
                )

                try:
                    # Embed each unique perturbation
                    pert_embs: dict[str, np.ndarray | None] = {}
                    for pert in unique_perts:
                        id_t = id_types[pert]
                        try:
                            if id_t == "smiles":
                                emb = self.embed_molecule(
                                    identifier=pert,
                                    model=model_key,
                                    pooling_strategy=pooling_strategy,
                                )
                            else:
                                emb = self.embed_gene(
                                    identifier=pert,
                                    model=model_key,
                                    id_type=id_t,
                                    organism=perturbation_organism,
                                    pooling_strategy=pooling_strategy,
                                )
                            pert_embs[pert] = np.asarray(
                                emb,
                                dtype=np.float32,
                            ).ravel()
                        except Exception as e:  # noqa: BLE001
                            logging.warning(
                                "Failed to embed perturbation '%s': %s",
                                pert,
                                e,
                            )
                            pert_embs[pert] = None

                    # Determine embedding dim
                    emb_dim = 0
                    for v in pert_embs.values():
                        if v is not None:
                            emb_dim = v.shape[0]
                            break

                    if emb_dim == 0:
                        logging.error(
                            "No perturbation embeddings succeeded for '%s'",
                            model_key,
                        )
                        metadata[model_key] = {
                            "error": "no embeddings succeeded",
                        }
                        continue

                    n_ok = sum(1 for v in pert_embs.values() if v is not None)
                    successful_perts = [p for p in unique_perts if pert_embs.get(p) is not None]
                    entity_matrix = np.stack(
                        [np.asarray(pert_embs[p], dtype=np.float32).ravel() for p in successful_perts],
                        axis=0,
                    )
                    aliases = {
                        str(p): {
                            "perturbation_label": str(p),
                            "input_id_type": str(id_types.get(p, perturbation_type)),
                        }
                        for p in successful_perts
                    }
                    result = EmbeddingResult(
                        matrix=entity_matrix,
                        entity_ids=tuple(str(p) for p in successful_perts),
                        entity_type="perturbation",
                        id_scheme="perturbation_label",
                        provenance=EmbeddingProvenance.create(
                            model=model_key,
                            pooling=pooling_strategy,
                            extra={
                                "entity_type": "perturbation",
                                "input_kind": "anndata_obs_column",
                                "input_source": "AnnData.obs",
                                "input_id_column": perturbation_column,
                                "perturbation_column": perturbation_column,
                                "perturbation_type": perturbation_type,
                                "organism": perturbation_organism,
                                "n_requested_inputs": len(unique_perts),
                                "n_successfully_embedded_entities": n_ok,
                                "n_embedding_failures": len(unique_perts) - n_ok,
                                "n_cells": int(adata.n_obs),
                            },
                        ),
                        aliases=aliases,
                    )
                    to_anndata(result, target=adata, attach_to="uns", key=obsm_key)

                    metadata[model_key] = {
                        "uns_key": obsm_key,
                        "storage": "uns",
                        "embedding_dim": emb_dim,
                        "n_cells": int(adata.n_obs),
                        "n_perturbations_total": len(unique_perts),
                        "n_perturbations_embedded": n_ok,
                        "type": "perturbation",
                        "perturbation_column": perturbation_column,
                    }
                    logging.info(
                        "  -> stored in .uns perturbations[%r], shape=(%d, %d) (%d/%d perturbations embedded)",
                        obsm_key,
                        entity_matrix.shape[0],
                        entity_matrix.shape[1],
                        n_ok,
                        len(unique_perts),
                    )

                except Exception as e:  # noqa: BLE001
                    logging.error(
                        "Failed perturbation embedding with '%s': %s",
                        model_key,
                        e,
                    )
                    metadata[model_key] = {"error": str(e)}

        adata.uns["embpy_embeddings"] = metadata
        logging.info("embed_adata complete.")
        return adata

    def list_available_models(
        self,
        category: Literal[
            "all",
            "dna",
            "protein",
            "molecule",
            "text",
            "morphology",
            "single_cell",
            "static",
        ] = "all",
    ) -> list[str]:
        """Return available model names, optionally filtered by category.

        Parameters
        ----------
        category
            ``"all"`` returns every model (sequence + single-cell).
            ``"dna"``, ``"protein"``, ``"molecule"``, ``"text"``,
            ``"morphology"`` filter the sequence/structure model registry.
            ``"single_cell"`` returns single-cell foundation model keys.
            ``"static"`` returns fixed lookup embedding tables shipped via
            the default Hugging Face data repository.
        """
        from .models.singlecell_models import list_singlecell_models

        if category == "static":
            return sorted(DEFAULT_STATIC_EMBEDDING_MODELS)

        if category == "single_cell":
            return sorted(list_singlecell_models())

        if category == "all":
            seq_models = sorted(self._available_models.keys())
            sc_models = sorted(list_singlecell_models())
            static_models = sorted(DEFAULT_STATIC_EMBEDDING_MODELS)
            return sorted(set(seq_models + sc_models + static_models))

        result = []
        for name, (wrapper_cls, _) in self._available_models.items():
            if hasattr(wrapper_cls, "model_type") and wrapper_cls.model_type == category:
                result.append(name)
        return sorted(result)

    def model_catalog(
        self,
        category: str = "all",
        *,
        summary: bool = False,
    ) -> Any:
        """Return the available models as a table.

        The tabular counterpart to :meth:`list_available_models`, which returns bare
        strings. In a notebook this renders as a proper table rather than printed
        text, and being a DataFrame it can be filtered, sorted and exported.

        Parameters
        ----------
        category
            Restrict to one family (``"dna"``, ``"protein"``, ``"molecule"``,
            ``"text"``, ``"morphology"``, ``"single_cell"``, ``"static"``), or
            ``"all"`` for everything.
        summary
            ``False`` (default) gives one row per model with its family. ``True``
            gives one row per family with a count and a few example keys -- the
            quick "what can this package do?" view.

        Returns
        -------
        pandas.DataFrame
            With ``summary=False``: columns ``model``, ``family``.
            With ``summary=True``: columns ``family``, ``n_models``, ``examples``.

        Examples
        --------
        >>> embedder.model_catalog(summary=True)          # doctest: +SKIP
        >>> embedder.model_catalog("protein").head()      # doctest: +SKIP
        """
        import pandas as pd

        families = ["static", "dna", "protein", "molecule", "text", "single_cell", "morphology"]
        wanted = families if category == "all" else [category]

        rows: list[dict[str, Any]] = []
        for fam in wanted:
            for name in self.list_available_models(fam):  # type: ignore[arg-type]
                rows.append({"model": name, "family": fam})

        frame = pd.DataFrame(rows, columns=["model", "family"])
        # A model can appear under more than one family (e.g. a text encoder used
        # for genes); keep the first family so counts stay interpretable.
        frame = frame.drop_duplicates(subset="model", keep="first").reset_index(drop=True)

        if not summary:
            return frame.sort_values(["family", "model"]).reset_index(drop=True)

        grouped = (
            frame.groupby("family")["model"]
            .agg(n_models="size", examples=lambda s: ", ".join(sorted(s)[:4]))
            .reset_index()
            .sort_values("n_models", ascending=False)
            .reset_index(drop=True)
        )
        return grouped

    # ------------------------------------------------------------------
    # Morphological embedding API
    # ------------------------------------------------------------------

    def embed_morphological(
        self,
        image: str | np.ndarray | torch.Tensor,
        model: str = "subcell_mae_rybg",
        pooling_strategy: str = "attention_pool",
        **kwargs,
    ) -> np.ndarray:
        """Embed a single microscopy image using a morphology model.

        Parameters
        ----------
        image
            Path to a multi-channel image, or a ``(C, H, W)`` array/tensor.
        model
            Registered morphology model key (default ``"subcell_mae_rybg"``).
        pooling_strategy
            Pooling strategy for the model (``"cls"``, ``"mean"``,
            ``"attention_pool"``, ``"none"``).

        Returns
        -------
        np.ndarray
            1-D embedding vector.
        """
        wrapper = self._get_model(model)
        if getattr(wrapper, "model_type", None) != "morphology":
            raise ValueError(
                f"Model '{model}' is not a morphology embedder (type={getattr(wrapper, 'model_type', '?')})."
            )
        try:
            return wrapper.embed(image, pooling_strategy=pooling_strategy, **kwargs)
        except Exception as e:
            logging.error("Morphological embedding failed: %s", e)
            raise RuntimeError("Morphological embedding failed") from e

    def embed_morphological_batch(
        self,
        images: Sequence[str | np.ndarray | torch.Tensor],
        model: str = "subcell_mae_rybg",
        pooling_strategy: str = "attention_pool",
        **kwargs,
    ) -> list[np.ndarray]:
        """Embed a batch of microscopy images.

        Parameters
        ----------
        images
            Iterable of paths or ``(C, H, W)`` arrays/tensors.
        model
            Registered morphology model key.
        pooling_strategy
            Pooling strategy passed to the underlying model.

        Returns
        -------
        list[np.ndarray]
            One embedding per image.
        """
        wrapper = self._get_model(model)
        if getattr(wrapper, "model_type", None) != "morphology":
            raise ValueError(
                f"Model '{model}' is not a morphology embedder (type={getattr(wrapper, 'model_type', '?')})."
            )
        return wrapper.embed_batch(images, pooling_strategy=pooling_strategy, **kwargs)

    def embed_perturbation_morphology(
        self,
        perturbation: str,
        perturbation_type: str = "genetic",
        dataset: str = "jump",
        source: str = "subcell",
        model: str = "subcell_mae_rybg",
        pooling_strategy: str = "attention_pool",
        local_dir: str | None = None,
        aggregate: str = "mean",
        max_images: int | None = 5,
        plate_type: str | None = None,
        jump_profiles_dir: str | None = None,
        verbose: bool = True,
        **kwargs,
    ) -> np.ndarray:
        """Embed a perturbation from JUMP or HPA by resolving images automatically.

        The method handles identifier resolution (gene symbols via
        ``GeneResolver``, compound names via ``DrugResolver``), image
        fetching from CDN or local directories, preprocessing, and
        optional aggregation.

        Parameters
        ----------
        perturbation
            Gene symbol, compound name, JCP2022 ID, or Ensembl ID.
        perturbation_type
            ``"genetic"`` (CRISPR/ORF) or ``"compound"``.
        dataset
            ``"jump"`` or ``"hpa"``.
        source
            ``"subcell"`` -- run SubCell inference on images.
            ``"precomputed"`` -- load pre-computed JUMP CellProfiler profiles
            (only valid with ``dataset="jump"``).
        model
            Morphology model key (used when ``source="subcell"``).
        pooling_strategy
            Pooling strategy for the SubCell model.
        local_dir
            Path to a local image directory.  When images are already
            present they are loaded from disk; when they are fetched
            from the CDN they are cached here for future re-use.
            ``None`` = fetch from CDN without caching.
        aggregate
            ``"mean"`` averages embeddings across images; ``"none"``
            returns the first image embedding only.
        max_images
            Maximum number of images to download and embed.  Defaults
            to ``5`` to avoid long downloads.  Pass ``None`` to fetch
            all available images for the perturbation.
        plate_type
            JUMP plate type (``"crispr"``, ``"orf"``, ``"compound"``).
            Inferred from *perturbation_type* when ``None``.
        jump_profiles_dir
            Directory containing JUMP profile parquets (for
            ``source="precomputed"``).  When ``None``, defaults to
            ``data/embeddings/morphology_embeddings/JUMP/`` within the
            project.  If the parquet is not present it is automatically
            downloaded from the Cell Painting Gallery S3 bucket.
        verbose
            If ``True`` (default), prints a single summary message at the
            end indicating success/failure and the data source used.
            Set to ``False`` to suppress all output.

        Returns
        -------
        np.ndarray
            Embedding vector (shape depends on model and aggregation).
        """
        if source == "precomputed":
            if dataset != "jump":
                raise ValueError("source='precomputed' is only valid with dataset='jump'")
            result = self._load_jump_precomputed(
                perturbation,
                perturbation_type,
                plate_type,
                jump_profiles_dir,
                aggregate,
            )
            if verbose:
                print(
                    f"[embpy] {perturbation!r}: loaded precomputed JUMP profile (plate_type={plate_type or 'crispr'})"
                )
            return result

        # Use a resolution context to track steps without intermediate warnings
        resolution_ctx = _ResolutionContext()

        if dataset == "jump":
            images = self._resolve_jump_images(
                perturbation,
                perturbation_type,
                local_dir,
                plate_type,
                max_images,
                resolution_ctx,
            )
        elif dataset == "hpa":
            if perturbation_type == "compound":
                raise ValueError("perturbation_type='compound' is not supported for dataset='hpa'")
            images = self._resolve_hpa_images(perturbation, local_dir, max_images, resolution_ctx)
        else:
            raise ValueError(f"Unknown dataset: {dataset!r}")

        if not images:
            if verbose:
                # Print failure summary
                steps = resolution_ctx.steps
                attempted = ", ".join(steps) if steps else "direct lookup"
                print(f"[embpy] WARNING: {perturbation!r}: no images found (dataset={dataset}, tried: {attempted})")
            raise RuntimeError(
                f"No images found for perturbation={perturbation!r}, "
                f"dataset={dataset!r}, perturbation_type={perturbation_type!r}"
            )

        embeddings = self.embed_morphological_batch(
            images,
            model=model,
            pooling_strategy=pooling_strategy,
            **kwargs,
        )

        if verbose:
            # Print success summary
            source_desc = resolution_ctx.final_source or dataset
            n_images = len(images)
            print(f"[embpy] {perturbation!r}: embedded {n_images} image(s) (source: {source_desc})")

        if aggregate == "mean":
            return np.mean(embeddings, axis=0).astype(np.float32)
        return embeddings[0]

    def embed_perturbation_morphology_batch(
        self,
        perturbations: Sequence[str],
        perturbation_type: str = "genetic",
        dataset: str = "jump",
        source: str = "subcell",
        model: str = "subcell_mae_rybg",
        pooling_strategy: str = "attention_pool",
        local_dir: str | None = None,
        aggregate: str = "mean",
        max_images: int | None = 5,
        plate_type: str | None = None,
        jump_profiles_dir: str | None = None,
        verbose: bool = True,
        skip_failures: bool = True,
        n_workers: int = 8,
        **kwargs,
    ) -> tuple[np.ndarray, list[str]]:
        """Embed multiple perturbations from JUMP or HPA in batch.

        This is the batched version of :meth:`embed_perturbation_morphology`.
        Image resolution / download for all perturbations runs **in parallel**
        across a thread pool, and then a single GPU inference batch is run on
        the combined image stack.  This is dramatically faster than calling
        ``embed_perturbation_morphology`` in a Python loop, especially for
        HPA / JUMP CDN downloads which are network-bound.

        Parameters
        ----------
        perturbations
            Sequence of gene symbols, compound names, JCP2022 IDs, or Ensembl IDs.
        perturbation_type
            ``"genetic"`` (CRISPR/ORF) or ``"compound"``.
        dataset
            ``"jump"`` or ``"hpa"``.
        source
            ``"subcell"`` -- run SubCell inference on images.
            ``"precomputed"`` -- load pre-computed JUMP CellProfiler profiles
            (only valid with ``dataset="jump"``).
        model
            Morphology model key (used when ``source="subcell"``).
        pooling_strategy
            Pooling strategy for the SubCell model.
        local_dir
            Path to a local image directory.  When images are already
            present they are loaded from disk; when they are fetched
            from the CDN they are cached here for future re-use.
        aggregate
            ``"mean"`` averages embeddings across images per perturbation;
            ``"none"`` returns the first image embedding only.
        max_images
            Maximum number of images to download and embed per perturbation.
        plate_type
            JUMP plate type (``"crispr"``, ``"orf"``, ``"compound"``).
        jump_profiles_dir
            Directory containing JUMP profile parquets (for ``source="precomputed"``).
        verbose
            If ``True`` (default), prints progress and a summary at the end.
        skip_failures
            If ``True`` (default), perturbations that fail to resolve are
            skipped with a warning. If ``False``, raises on first failure.
        n_workers
            Number of parallel worker threads for resolving / downloading
            images across perturbations.  Defaults to ``8``.
        **kwargs
            Additional arguments passed to the embedding model.

        Returns
        -------
        tuple[np.ndarray, list[str]]
            A tuple of ``(embeddings, labels)`` where:

            - ``embeddings`` is a 2D array of shape ``(n_successful, embedding_dim)``
            - ``labels`` is a list of perturbation identifiers that succeeded,
              in the same order as the embedding rows.

        Examples
        --------
        >>> genes = ["TP53", "BRCA1", "EGFR", "UNKNOWN_GENE"]
        >>> embeddings, labels = embedder.embed_perturbation_morphology_batch(genes, dataset="hpa", max_images=3)
        >>> print(f"Embedded {len(labels)} / {len(genes)} genes")
        Embedded 3 / 4 genes
        >>> print(embeddings.shape)
        (3, 384)
        """
        import time
        from concurrent.futures import ThreadPoolExecutor, as_completed

        # ---- precomputed JUMP profiles: I/O-bound parquet reads only ----
        if source == "precomputed":
            return self._embed_perturbation_precomputed_batch(
                perturbations,
                perturbation_type=perturbation_type,
                plate_type=plate_type,
                jump_profiles_dir=jump_profiles_dir,
                aggregate=aggregate,
                verbose=verbose,
                skip_failures=skip_failures,
            )

        n_total = len(perturbations)
        t0 = time.time()
        hpa_catalog = None
        if dataset == "hpa":
            hpa_catalog = self._build_hpa_batch_catalog(
                perturbations,
                local_dir=local_dir,
                verbose=verbose,
            )

        # ── Stage 1: resolve images for ALL perturbations in parallel ──
        if verbose:
            print(
                f"[embpy] Stage 1/2: resolving images for {n_total} perturbations using {n_workers} parallel workers..."
            )

        def _resolve_one(pert: str) -> tuple[str, list[np.ndarray], str | None, Exception | None]:
            ctx = _ResolutionContext()
            try:
                if dataset == "jump":
                    imgs = self._resolve_jump_images(
                        pert,
                        perturbation_type,
                        local_dir,
                        plate_type,
                        max_images,
                        ctx,
                    )
                elif dataset == "hpa":
                    if perturbation_type == "compound":
                        raise ValueError("perturbation_type='compound' is not supported for dataset='hpa'")
                    imgs = self._resolve_hpa_images(
                        pert,
                        local_dir,
                        max_images,
                        ctx,
                        hpa_catalog=hpa_catalog,
                    )
                else:
                    raise ValueError(f"Unknown dataset: {dataset!r}")
                return pert, imgs, ctx.final_source, None
            except Exception as exc:  # noqa: BLE001
                return pert, [], ctx.final_source, exc

        per_gene_images: dict[str, list[np.ndarray]] = {}
        per_gene_source: dict[str, str | None] = {}
        failed_labels: list[str] = []
        first_exc: Exception | None = None

        with ThreadPoolExecutor(max_workers=n_workers) as pool:
            futures = {pool.submit(_resolve_one, p): p for p in perturbations}
            for i, fut in enumerate(as_completed(futures), start=1):
                pert, imgs, src, exc = fut.result()
                if exc is not None or not imgs:
                    failed_labels.append(pert)
                    if first_exc is None and exc is not None:
                        first_exc = exc
                    if verbose:
                        reason = str(exc) if exc else "no images found"
                        print(f"  [{i}/{n_total}] {pert!r:20s} FAILED ({reason})")
                else:
                    per_gene_images[pert] = imgs
                    per_gene_source[pert] = src
                    if verbose:
                        print(f"  [{i}/{n_total}] {pert!r:20s} {len(imgs)} image(s) from {src}")

        if not skip_failures and first_exc is not None:
            raise first_exc

        if not per_gene_images:
            if verbose:
                print("[embpy] No images resolved for any perturbation")
            return np.array([]).reshape(0, 0), []

        # ── Stage 2: stack all images and run a single GPU inference ──
        successful_labels = [p for p in perturbations if p in per_gene_images]
        all_images: list[np.ndarray] = []
        boundaries: list[int] = [0]
        for pert in successful_labels:
            imgs = per_gene_images[pert]
            all_images.extend(imgs)
            boundaries.append(boundaries[-1] + len(imgs))

        if verbose:
            t_resolve = time.time() - t0
            print(
                f"[embpy] Stage 2/2: running inference on "
                f"{len(all_images)} images (image resolution took "
                f"{t_resolve:.1f}s)..."
            )

        t1 = time.time()
        all_embeddings = self.embed_morphological_batch(
            all_images,
            model=model,
            pooling_strategy=pooling_strategy,
            **kwargs,
        )
        all_embeddings = np.asarray(all_embeddings)

        # Aggregate per gene using boundary slices
        per_gene_emb: list[np.ndarray] = []
        for j in range(len(successful_labels)):
            start, end = boundaries[j], boundaries[j + 1]
            sub = all_embeddings[start:end]
            if aggregate == "mean":
                per_gene_emb.append(np.mean(sub, axis=0).astype(np.float32))
            else:
                per_gene_emb.append(sub[0])

        if verbose:
            t_infer = time.time() - t1
            t_total = time.time() - t0
            print(
                f"[embpy] Done in {t_total:.1f}s "
                f"(resolve={t_resolve:.1f}s, inference={t_infer:.1f}s) -- "
                f"{len(successful_labels)}/{n_total} perturbations embedded"
            )
            if failed_labels:
                print(f"[embpy] Failed: {failed_labels}")

        return np.stack(per_gene_emb, axis=0), successful_labels

    def _embed_perturbation_precomputed_batch(
        self,
        perturbations: Sequence[str],
        perturbation_type: str,
        plate_type: str | None,
        jump_profiles_dir: str | None,
        aggregate: str,
        verbose: bool,
        skip_failures: bool,
    ) -> tuple[np.ndarray, list[str]]:
        """Batch-load pre-computed JUMP profiles for many perturbations.

        Reuses the parquet only once across all perturbations to avoid
        reloading it for every gene.
        """
        import pandas as pd

        from .resources.jump_metadata import get_jump_gene_mapper

        if plate_type is None:
            plate_type = "compound" if perturbation_type == "compound" else "crispr"

        pq_path = _resolve_jump_parquet_path(plate_type, jump_profiles_dir)
        df = pd.read_parquet(pq_path)
        feature_cols = [c for c in df.columns if not c.startswith("Metadata_")]
        jcp_col = "Metadata_JCP2022"

        mapper = get_jump_gene_mapper(plate_type) if perturbation_type != "compound" else {}
        inv_mapper = {v: k for k, v in mapper.items()}

        embeddings_list: list[np.ndarray] = []
        successful_labels: list[str] = []
        failed_labels: list[str] = []

        for pert in perturbations:
            if pert.startswith("JCP2022"):
                mask = df[jcp_col] == pert
            else:
                jcp_id = inv_mapper.get(pert)
                mask = df[jcp_col] == jcp_id if jcp_id else pd.Series(False, index=df.index)

            sub = df[mask]
            if len(sub) == 0:
                failed_labels.append(pert)
                if not skip_failures:
                    raise RuntimeError(f"No precomputed profile found for {pert!r}")
                continue

            features = sub[feature_cols].to_numpy(dtype=np.float32)
            emb = features.mean(axis=0) if aggregate == "mean" else features[0]
            embeddings_list.append(emb)
            successful_labels.append(pert)

        if verbose:
            print(
                f"[embpy] Loaded {len(successful_labels)}/{len(perturbations)} precomputed profiles from {pq_path.name}"
            )
            if failed_labels:
                print(f"[embpy] Failed: {failed_labels}")

        if not embeddings_list:
            return np.array([]).reshape(0, 0), []
        return np.stack(embeddings_list, axis=0), successful_labels

    # ------------------------------------------------------------------
    # Internal resolvers
    # ------------------------------------------------------------------

    @staticmethod
    def _hpa_shared_cache_root(local_dir: str | None) -> pathlib.Path | None:
        """Return a persistent shared HPA cache root when one is configured."""
        if os.environ.get("EMBPY_HPA_CACHE_DIR"):
            return pathlib.Path(os.environ["EMBPY_HPA_CACHE_DIR"])
        if local_dir is None:
            return None
        local_path = pathlib.Path(local_dir)
        if local_path.name in {"_hpa", "hpa"}:
            return local_path
        return local_path.parent / "_hpa"

    @staticmethod
    def _build_hpa_batch_catalog(
        perturbations: Sequence[str],
        *,
        local_dir: str | None,
        verbose: bool,
    ):
        """Build or load one HPA XML catalog slice for a batch of genes."""
        genes = sorted({str(p).strip().upper() for p in perturbations if str(p).strip()})
        if not genes:
            return None

        cache_root = BioEmbedder._hpa_shared_cache_root(local_dir)
        xml_source = None
        cache_path = None
        if cache_root is not None:
            cache_root.mkdir(parents=True, exist_ok=True)
            digest = sha1("\n".join(genes).encode("utf-8")).hexdigest()[:16]
            xml_source = cache_root / "proteinatlas.xml.gz"
            cache_path = cache_root / f"subcellular_catalog_{digest}.csv"

        from .resources.hpa_images import build_hpa_subcellular_catalog

        if verbose:
            cache_msg = f" (cache: {cache_path})" if cache_path is not None else ""
            print(f"[embpy] Building/loading one HPA XML catalog for {len(genes)} genes{cache_msg}...")
        try:
            return build_hpa_subcellular_catalog(
                xml_source=xml_source,
                cache_path=cache_path,
                genes=genes,
            )
        except Exception as exc:  # noqa: BLE001
            if verbose:
                print(
                    f"[embpy] WARNING: failed to build shared HPA XML catalog; falling back to per-gene lookup ({exc})"
                )
            return None

    @staticmethod
    def _hpa_catalog_rows_to_antibodies(catalog, perturbation: str) -> list[dict[str, Any]]:
        """Convert pre-filtered HPA catalog rows into fetchable image records."""
        if catalog is None or len(catalog) == 0:
            return []
        gene_upper = str(perturbation).upper()
        rows = catalog[catalog["gene"].astype(str).str.upper() == gene_upper]
        antibodies: list[dict[str, Any]] = []
        for _, row in rows.iterrows():
            antibodies.append(
                {
                    "id": row["antibody"],
                    "_plate": row["plate"],
                    "_position": row["position"],
                    "_sample": row["sample"],
                    "_url_prefix": row.get("image_url_prefix", ""),
                }
            )
        return antibodies

    @staticmethod
    def _resolve_jump_images(
        perturbation: str,
        perturbation_type: str,
        local_dir: str | None,
        plate_type: str | None,
        max_images: int | None,
        resolution_ctx: _ResolutionContext | None = None,
    ) -> list[np.ndarray]:
        """Resolve a JUMP perturbation to preprocessed SubCell-ready arrays.

        For compounds, uses ``DrugResolver`` to resolve alternate names
        when the initial lookup fails.  For genetic perturbations, uses
        ``GeneResolver`` to map symbols to canonical forms.
        """
        from .pp.morphology_preprocessing import (
            cell_painting_to_subcell,
            prepare_subcell_canvas,
        )
        from .resources.jump_metadata import (
            fetch_jump_fov,
            get_jump_item_location_metadata,
        )

        if plate_type is None:
            plate_type = "compound" if perturbation_type == "compound" else "crispr"

        is_jcp = perturbation.startswith("JCP2022")
        input_column = "JCP2022" if is_jcp else "standard_key"

        if resolution_ctx:
            resolution_ctx.add_step("JUMP direct lookup")

        rows = get_jump_item_location_metadata(
            perturbation,
            input_column=input_column,
        )

        used_fallback = False
        # Fallback resolution when initial lookup fails
        if not rows and not is_jcp:
            if resolution_ctx:
                resolution_ctx.add_step("JUMP fallback resolver")
            rows, fallback_source = _resolve_jump_perturbation_fallback(
                perturbation,
                perturbation_type,
                return_source=True,
            )
            if rows:
                used_fallback = True
                if resolution_ctx and fallback_source:
                    resolution_ctx.set_source(f"JUMP CDN via {fallback_source}")

        if not used_fallback and rows and resolution_ctx:
            resolution_ctx.set_source("JUMP CDN (direct)")

        if max_images is not None:
            rows = rows[:max_images]

        images: list[np.ndarray] = []
        failed_count = 0
        for row in rows:
            try:
                if local_dir is not None:
                    fov = _load_jump_fov_local(row, local_dir)
                else:
                    fov = fetch_jump_fov(row)
                subcell_4ch = cell_painting_to_subcell(fov)
                canvas = prepare_subcell_canvas(subcell_4ch, nm_per_pixel=325.0)
                images.append(canvas)
            except Exception:  # noqa: BLE001
                failed_count += 1
        if failed_count > 0 and resolution_ctx:
            resolution_ctx.add_step(f"{failed_count} image(s) failed to load")
        return images

    @staticmethod
    def _resolve_hpa_images(
        perturbation: str,
        local_dir: str | None,
        max_images: int | None,
        resolution_ctx: _ResolutionContext | None = None,
        hpa_catalog=None,
    ) -> list[np.ndarray]:
        """Resolve an HPA gene to SubCell-ready image arrays.

        When *local_dir* is given and already contains channel images for
        this gene the images are loaded from disk.  Otherwise images are
        fetched from the HPA CDN and -- if *local_dir* is provided --
        cached there as per-channel JPEGs for future re-use.

        Gene symbol resolution uses ``GeneResolver`` internally (via
        ``hpa_images._resolve_ensembl_id``).
        """
        import os
        import re
        from pathlib import Path as _Path

        from .resources.hpa_images import (
            fetch_hpa_if_image,
            fetch_hpa_if_image_by_prefix,
            get_hpa_antibodies_quiet,
            load_hpa_if_image,
        )

        # ── Try loading from local cache first ──────────────────────
        if local_dir is not None:
            gene_dir = _Path(local_dir) / perturbation
            if gene_dir.is_dir():
                if resolution_ctx:
                    resolution_ctx.add_step("local cache check")
                candidates = sorted(f for f in os.listdir(gene_dir) if f.endswith((".png", ".jpg")))
                prefixes: set[str] = set()
                for f in candidates:
                    if re.search(r"_blue\.\w+$", f):
                        prefixes.add(re.sub(r"_blue\.\w+$", "", f))
                if prefixes:
                    images: list[np.ndarray] = []
                    failed_local = 0
                    for prefix in sorted(prefixes):
                        try:
                            img = load_hpa_if_image(prefix=str(gene_dir / prefix))
                            images.append(img)
                        except Exception:  # noqa: BLE001
                            failed_local += 1
                    if max_images is not None:
                        images = images[:max_images]
                    if images:
                        if resolution_ctx:
                            resolution_ctx.set_source(f"HPA local cache ({gene_dir})")
                            if failed_local > 0:
                                resolution_ctx.add_step(f"{failed_local} local image(s) failed to load")
                        return images

        # ── Fetch from CDN ──────────────────────────────────────────
        gene_source = None
        if hpa_catalog is not None:
            if resolution_ctx:
                resolution_ctx.add_step("HPA shared XML catalog")
            antibodies = BioEmbedder._hpa_catalog_rows_to_antibodies(hpa_catalog, perturbation)
            if antibodies:
                gene_source = "HPA shared XML catalog"
            else:
                return []
        else:
            if resolution_ctx:
                resolution_ctx.add_step("HPA API lookup")
            antibodies, gene_source = get_hpa_antibodies_quiet(perturbation)
        if not antibodies:
            if resolution_ctx:
                resolution_ctx.add_step("HPA per-gene XML catalog fallback")
            try:
                from .resources.hpa_images import build_hpa_subcellular_catalog

                catalog = build_hpa_subcellular_catalog(genes=[perturbation])
                if len(catalog) == 0:
                    return []
                antibodies = BioEmbedder._hpa_catalog_rows_to_antibodies(catalog, perturbation)
                gene_source = "HPA XML catalog"
            except Exception:  # noqa: BLE001
                return []

        # Cap the antibody list *before* downloading so we only fetch
        # what the user actually requested.
        if max_images is not None:
            antibodies = antibodies[:max_images]

        cache_dir = None
        if local_dir is not None:
            cache_dir = _Path(local_dir) / perturbation
            cache_dir.mkdir(parents=True, exist_ok=True)

        def _unpack_ab(ab):
            if isinstance(ab, dict):
                return (
                    ab.get("id", ""),
                    ab.get("_plate", 1),
                    ab.get("_position", "A1"),
                    ab.get("_sample", 1),
                    ab.get("_url_prefix", ""),
                )
            return str(ab), 1, "A1", 1, ""

        failed_fetch = [0]  # Use list to allow mutation in nested function

        def _fetch_one(ab):
            ab_id, plate, position, sample, url_prefix = _unpack_ab(ab)
            try:
                # When we have a full URL prefix from the XML catalog, use
                # it directly -- it preserves cell-line subdirectories that
                # plate/position/sample alone cannot reconstruct.
                if url_prefix:
                    img = fetch_hpa_if_image_by_prefix(url_prefix)
                else:
                    img = fetch_hpa_if_image(ab_id, plate, position, sample)
                if cache_dir is not None:
                    _save_hpa_channels(
                        img,
                        cache_dir,
                        ab_id,
                        plate,
                        position,
                        sample,
                    )
                return img
            except Exception:  # noqa: BLE001
                failed_fetch[0] += 1
                return None

        from concurrent.futures import ThreadPoolExecutor

        with ThreadPoolExecutor(max_workers=min(8, len(antibodies))) as pool:
            results = list(pool.map(_fetch_one, antibodies))

        images = [r for r in results if r is not None]

        if resolution_ctx:
            if images:
                src = f"HPA CDN via {gene_source}" if gene_source else "HPA CDN"
                resolution_ctx.set_source(src)
            if failed_fetch[0] > 0:
                resolution_ctx.add_step(f"{failed_fetch[0]} CDN fetch(es) failed")

        return images

    @staticmethod
    def _load_jump_precomputed(
        perturbation: str,
        perturbation_type: str | None,
        plate_type: str | None,
        profiles_dir: str | None,
        aggregate: str,
    ) -> np.ndarray:
        """Load pre-computed JUMP profiles from parquet.

        When the parquet file is not found locally it is automatically
        downloaded from the Cell Painting Gallery S3 bucket into
        *profiles_dir* (or the default cache directory).
        """
        import pandas as pd

        from .resources.jump_metadata import get_jump_gene_mapper

        if plate_type is None:
            plate_type = "compound" if perturbation_type == "compound" else "crispr"

        pq_path = _resolve_jump_parquet_path(plate_type, profiles_dir)

        df = pd.read_parquet(pq_path)

        meta_cols = [c for c in df.columns if c.startswith("Metadata_")]
        jcp_col = "Metadata_JCP2022"
        feature_cols = [c for c in df.columns if not c.startswith("Metadata_")]

        if perturbation.startswith("JCP2022"):
            mask = df[jcp_col] == perturbation
        else:
            mapper = get_jump_gene_mapper(plate_type)
            inv_mapper = {v: k for k, v in mapper.items()}
            jcp_id = inv_mapper.get(perturbation)
            if jcp_id:
                mask = df[jcp_col] == jcp_id
            else:
                sk_col = next(
                    (c for c in meta_cols if c != jcp_col and "standard" in c.lower()),
                    None,
                )
                if sk_col and sk_col in df.columns:
                    mask = df[sk_col].astype(str).str.upper() == perturbation.upper()
                else:
                    jcp_to_name = pd.Series(mapper)
                    df["_mapped_name"] = df[jcp_col].map(jcp_to_name).fillna("")
                    mask = df["_mapped_name"].str.upper() == perturbation.upper()
                    df.drop(columns=["_mapped_name"], inplace=True)

        subset = df.loc[mask, feature_cols]
        if subset.empty:
            raise ValueError(f"Perturbation {perturbation!r} not found in JUMP {plate_type} profiles.")

        if aggregate == "mean":
            return subset.mean(axis=0).values.astype(np.float32)
        return subset.iloc[0].values.astype(np.float32)

    # ------------------------------------------------------------------
    # FASTA / FASTQ file embedding
    # ------------------------------------------------------------------

    _FASTA_EXTENSIONS: dict[str, str] = {
        ".fasta": "fasta",
        ".fa": "fasta",
        ".fna": "fasta",
        ".faa": "fasta",
        ".fastq": "fastq",
        ".fq": "fastq",
    }

    _DNA_CHARS = frozenset("ACGTNUacgtnu")

    def embed_fasta(
        self,
        path: str | os.PathLike[str],
        model: str,
        seq_type: Literal["dna", "protein"] | None = None,
        pooling_strategy: str = "mean",
        obsm_key: str | None = None,
        **kwargs: Any,
    ):
        """Embed all sequences from a FASTA or FASTQ file.

        Supports plain and gzip-compressed files.  The method
        auto-detects the file format from the extension and optionally
        infers whether sequences are DNA or protein when *seq_type*
        is not provided.

        Parameters
        ----------
        path
            Path to a FASTA/FASTQ file.  Recognised extensions:
            ``.fasta``, ``.fa``, ``.fna``, ``.faa`` (FASTA),
            ``.fastq``, ``.fq`` (FASTQ).  A trailing ``.gz`` is
            handled transparently.
        model
            Model name registered with this embedder.
        seq_type
            ``"dna"`` or ``"protein"``.  When ``None`` the type is
            resolved using a three-tier fallback:

            1. **Infer from model** -- if the model is a DNA or
               protein model the type is set accordingly.
            2. **Auto-detect from sequences** -- the first 100
               sequences are inspected; if they consist exclusively
               of ``{A, C, G, T, N, U}`` the type is set to
               ``"dna"``, otherwise ``"protein"``.
            3. **Fail** -- a ``ValueError`` is raised asking the
               caller to specify *seq_type* explicitly.
        pooling_strategy
            Passed to the underlying model.
        obsm_key
            Key used to store embeddings in ``adata.obsm``.
            Defaults to ``"X_{model}"``.
        **kwargs
            Forwarded to the model ``embed`` call.

        Returns
        -------
        anndata.AnnData
            An :class:`~anndata.AnnData` object with:

            * ``.obs`` -- ``sequence_id``, ``sequence_length``,
              ``description``, ``seq_type``
            * ``.obsm[obsm_key]`` -- embedding matrix
              (n_sequences x embedding_dim)
            * ``.uns`` -- ``model``, ``pooling_strategy``,
              ``source_file``, ``n_sequences``, ``n_failed``
        """
        import gzip
        from pathlib import Path

        import anndata as ad
        from Bio import SeqIO

        filepath = Path(path)

        # -- detect file format -----------------------------------------
        is_gz = filepath.suffix.lower() == ".gz"
        stem = filepath.with_suffix("") if is_gz else filepath
        fmt = self._FASTA_EXTENSIONS.get(stem.suffix.lower())
        if fmt is None:
            raise ValueError(
                f"Unrecognised file extension {stem.suffix!r}.  Supported: {', '.join(sorted(self._FASTA_EXTENSIONS))}"
            )

        # -- parse sequences --------------------------------------------
        handle = gzip.open(filepath, "rt") if is_gz else open(filepath)
        try:
            records = list(SeqIO.parse(handle, fmt))
        finally:
            handle.close()

        if not records:
            raise ValueError(f"No sequences found in {filepath}")

        seq_ids: list[str] = []
        descriptions: list[str] = []
        sequences: list[str] = []
        for rec in records:
            seq_ids.append(rec.id)
            descriptions.append(rec.description)
            sequences.append(str(rec.seq))

        # -- resolve seq_type -------------------------------------------
        resolved_seq_type = self._resolve_seq_type(
            seq_type,
            model,
            sequences,
        )

        # -- embed sequences --------------------------------------------
        inst = self._get_model(model)
        embeddings: list[np.ndarray | None] = []
        n_failed = 0
        for i, seq in enumerate(sequences):
            try:
                emb = inst.embed(
                    input=seq,
                    pooling_strategy=pooling_strategy,
                    **kwargs,
                )
                embeddings.append(emb)
            except Exception as e:  # noqa: BLE001
                logging.warning(
                    "Failed to embed sequence %s (%s): %s",
                    seq_ids[i],
                    seq[:30] + "...",
                    e,
                )
                embeddings.append(None)
                n_failed += 1

        # -- build AnnData ----------------------------------------------
        emb_dim = next(
            (e.shape[-1] for e in embeddings if e is not None),
            0,
        )
        emb_matrix = np.full(
            (len(sequences), emb_dim),
            np.nan,
            dtype=np.float32,
        )
        for i, emb in enumerate(embeddings):
            if emb is not None:
                emb_matrix[i] = emb

        import pandas as pd

        obs = pd.DataFrame(
            {
                "sequence_id": seq_ids,
                "sequence_length": [len(s) for s in sequences],
                "description": descriptions,
                "seq_type": resolved_seq_type,
            }
        )
        obs.index = obs["sequence_id"].astype(str)
        # ensure unique index
        if obs.index.duplicated().any():
            obs.index = pd.Index([f"{sid}_{i}" for i, sid in enumerate(obs.index)])

        key = obsm_key or f"X_{model}"
        adata = ad.AnnData(obs=obs)
        adata.obsm[key] = emb_matrix
        adata.uns["embed_fasta"] = {
            "model": model,
            "pooling_strategy": pooling_strategy,
            "source_file": str(filepath),
            "n_sequences": len(sequences),
            "n_failed": n_failed,
        }

        n_ok = len(sequences) - n_failed
        logging.info(
            "embed_fasta: embedded %d/%d sequences from %s (model=%s, seq_type=%s)",
            n_ok,
            len(sequences),
            filepath.name,
            model,
            resolved_seq_type,
        )
        return adata

    def _resolve_seq_type(
        self,
        explicit: Literal["dna", "protein"] | None,
        model: str,
        sequences: list[str],
    ) -> str:
        """Resolve sequence type with three-tier fallback."""
        if explicit is not None:
            logging.info("embed_fasta: using explicit seq_type=%r", explicit)
            return explicit

        # Tier 1: infer from model type
        try:
            inst = self._get_model(model)
            mtype = getattr(inst, "model_type", None)
            if mtype == "dna":
                logging.info(
                    "embed_fasta: inferred seq_type='dna' from model %r",
                    model,
                )
                return "dna"
            if mtype == "protein":
                logging.info(
                    "embed_fasta: inferred seq_type='protein' from model %r",
                    model,
                )
                return "protein"
        except Exception:  # noqa: BLE001
            pass

        # Tier 2: auto-detect from character set
        sample = sequences[:100]
        all_dna = all(set(seq.upper()).issubset(self._DNA_CHARS) for seq in sample if seq)
        if all_dna:
            logging.info(
                "embed_fasta: auto-detected seq_type='dna' from sequence character set",
            )
            return "dna"
        else:
            has_non_dna = any(not set(seq.upper()).issubset(self._DNA_CHARS) for seq in sample if seq)
            if has_non_dna:
                logging.info(
                    "embed_fasta: auto-detected seq_type='protein' from sequence character set",
                )
                return "protein"

        raise ValueError(
            "Could not determine sequence type automatically.  Please specify seq_type='dna' or seq_type='protein'."
        )

__init__(device='auto', organism='human', resolver_backend='api', mart_file=None, chromosome_folder=None)

Initializes the BioEmbedder.

Args: device: 'auto', 'cuda', 'mps', or 'cpu', or torch.device. organism: Default organism for sequence resolution and annotation (e.g. 'human', 'mouse', 'zebrafish'). Any species supported by Ensembl can be used. resolver_backend: 'api' to use online APIs, 'local' to use local FASTAs. mart_file: path to Mart CSV (required if resolver_backend='local'). chromosome_folder: path to folder with chr*.fa files (required if 'local').

Source code in src/embpy/embedder.py

def __init__(
    self,
    device: str | torch.device | None = "auto",  # type: ignore[name-defined]
    organism: str = "human",
    resolver_backend: Literal["api", "local"] = "api",
    mart_file: str | None = None,
    chromosome_folder: str | None = None,
):
    """
    Initializes the BioEmbedder.

    Args:
        device: 'auto', 'cuda', 'mps', or 'cpu', or torch.device.
        organism: Default organism for sequence resolution and annotation
            (e.g. 'human', 'mouse', 'zebrafish'). Any species supported
            by Ensembl can be used.
        resolver_backend: 'api' to use online APIs, 'local' to use local FASTAs.
        mart_file: path to Mart CSV (required if resolver_backend='local').
        chromosome_folder: path to folder with chr*.fa files (required if 'local').
    """
    self.organism = organism

    # Device setup
    if isinstance(device, str):
        if device == "auto":
            self.device = get_device()
        else:
            self.device = torch.device(device)  # type: ignore[attr-defined]
    elif isinstance(device, torch.device):  # type: ignore[attr-defined]
        self.device = device
    else:
        raise ConfigError("Invalid device; use 'auto','cpu','cuda','mps', or torch.device.")

    # GeneResolver setup
    self.resolver_backend = resolver_backend
    if resolver_backend == "local":
        if not mart_file or not chromosome_folder:
            raise ConfigError("mart_file and chromosome_folder must be provided for local resolver.")
        self.gene_resolver = GeneResolver(
            mart_file=mart_file,
            chromosome_folder=chromosome_folder,
        )
    else:
        self.gene_resolver = GeneResolver(species=organism)

    # Protein resolver
    self.protein_resolver = ProteinResolver(organism=organism)

    # Text resolver for description-based embeddings
    self.text_resolver = TextResolver(organism=organism)

    # Model cache and discovery.
    #
    # `model_cache` holds DNA/protein/molecule/text/morphology wrappers
    # that subclass `BaseModelWrapper`. Single-cell foundation-model
    # wrappers (scGPT, Geneformer, UCE, Tahoe, ...) subclass
    # `SingleCellWrapper` instead, so they live in a separate cache.
    # Both caches are keyed so that repeated calls to embed_cells /
    # embed_perturbation / etc. reuse an already-loaded torch model
    # instead of re-instantiating it on every call -- critical for
    # chunked inference over large datasets.
    self.model_cache: dict[str, BaseModelWrapper] = {}
    self._singlecell_cache: dict[tuple[str, str], Any] = {}
    self._available_models = self._discover_models()

    # Layer 2: per-call resolution report. Populated by
    # ``embed_genes_batch`` and exposed for callers (e.g.
    # ``BioEmbedderProvider`` and ``embed_perturbations.py``)
    # that want a per-identifier audit trail without parsing logs.
    # Each call to ``embed_genes_batch`` REPLACES this attribute,
    # so callers should grab it immediately after the call returns.
    self.last_report: ResolutionReport | None = None

    logging.info(
        "BioEmbedder initialized: device=%s, organism=%s, backend=%s",
        self.device,
        self.organism,
        self.resolver_backend,
    )

clear_model_cache(*, which='all')

Drop cached model wrappers and free associated GPU memory.

Parameters:

Name	Type	Description	Default
which	Literal['all', 'singlecell', 'other']	Which cache(s) to clear. "all" drops both the single-cell foundation-model cache and the DNA/protein/molecule/text/ morphology cache. "singlecell" clears only the single-cell cache (useful when switching from foundation-model inference to structure/morphology work). "other" clears only the non-single-cell cache.	'all'
After			required
the			required

Source code in src/embpy/embedder.py

def clear_model_cache(self, *, which: Literal["all", "singlecell", "other"] = "all") -> None:
    """Drop cached model wrappers and free associated GPU memory.

    Parameters
    ----------
    which
        Which cache(s) to clear. ``"all"`` drops both the single-cell
        foundation-model cache and the DNA/protein/molecule/text/
        morphology cache. ``"singlecell"`` clears only the single-cell
        cache (useful when switching from foundation-model inference
        to structure/morphology work). ``"other"`` clears only the
        non-single-cell cache.

    After dropping references, runs ``gc.collect()`` +
    ``torch.cuda.empty_cache()`` so the allocator actually releases
    the GPU memory back to the system.
    """
    import gc

    if which in ("all", "singlecell"):
        n = len(self._singlecell_cache)
        self._singlecell_cache.clear()
        logging.info("Cleared %d cached single-cell wrapper(s).", n)
    if which in ("all", "other"):
        n = len(self.model_cache)
        self.model_cache.clear()
        logging.info("Cleared %d cached model wrapper(s).", n)

    gc.collect()
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    except Exception:  # noqa: BLE001
        pass

decode_cells(adata=None, *, latent=None, model='scvi', obsm_key=None, gene_names=None, write_layer=None, **decode_kwargs)

Decode cell embeddings back to gene-expression space.

Companion to :meth:embed_cells that routes a latent matrix through the same fitted encoder-decoder wrapper that produced it. This is the decoder hook used by flow-matching / cellflow style perturbation pipelines: encode basal + perturbed cells to a shared latent, learn a transport map there, then decode the predicted latent back to counts with this method.

Requires a prior :meth:embed_cells call with the same model key so the trained wrapper is cached for this device. For STATE / Stack, instantiate the wrapper manually (checkpoint + gene-list paths cannot be supplied through :meth:embed_cells).

Parameters:

Name	Type	Description	Default
adata	AnnData	AnnData that holds the latent to decode. Used both to pick the latent from obsm[obsm_key] and to recover var_names for gene-parametric decoders (STATE).	None
latent	ndarray	Alternative to adata + obsm_key: pass the latent matrix directly.	None
model	str	Encoder-decoder registry key: "pca", "scvi", "scanvi", "totalvi", or "state".	'scvi'
obsm_key	str	Override key used to pull the latent from adata.obsm. Defaults to f"X_{model}".	None
gene_names	sequence of str	Target genes to decode to (STATE only). Defaults to adata.var_names when adata is supplied.	None
write_layer	str	If set and adata is supplied, the decoded expression is written to adata.layers[write_layer].	None
**decode_kwargs	Any	Forwarded to the wrapper's decode_cells (e.g. library_size, batch_index for scVI, read_depth for STATE).	{}

Returns:

Type	Description
np.ndarray of shape ``(n_cells, n_genes)``	Decoded expression (NB mean / log-probs / linear recon, depending on the model; see the wrapper docstring).

Source code in src/embpy/embedder.py

def decode_cells(
    self,
    adata=None,
    *,
    latent=None,
    model: str = "scvi",
    obsm_key: str | None = None,
    gene_names: Sequence[str] | None = None,
    write_layer: str | None = None,
    **decode_kwargs: Any,
):
    """Decode cell embeddings back to gene-expression space.

    Companion to :meth:`embed_cells` that routes a latent matrix
    through the *same* fitted encoder-decoder wrapper that produced
    it. This is the decoder hook used by flow-matching / cellflow
    style perturbation pipelines: encode basal + perturbed cells to
    a shared latent, learn a transport map there, then decode the
    predicted latent back to counts with this method.

    Requires a prior :meth:`embed_cells` call with the same
    ``model`` key so the trained wrapper is cached for this device.
    For STATE / Stack, instantiate the wrapper manually (checkpoint +
    gene-list paths cannot be supplied through :meth:`embed_cells`).

    Parameters
    ----------
    adata : anndata.AnnData, optional
        AnnData that holds the latent to decode. Used both to pick
        the latent from ``obsm[obsm_key]`` and to recover
        ``var_names`` for gene-parametric decoders (STATE).
    latent : np.ndarray, optional
        Alternative to ``adata`` + ``obsm_key``: pass the latent
        matrix directly.
    model : str
        Encoder-decoder registry key: ``"pca"``, ``"scvi"``,
        ``"scanvi"``, ``"totalvi"``, or ``"state"``.
    obsm_key : str, optional
        Override key used to pull the latent from ``adata.obsm``.
        Defaults to ``f"X_{model}"``.
    gene_names : sequence of str, optional
        Target genes to decode to (STATE only). Defaults to
        ``adata.var_names`` when ``adata`` is supplied.
    write_layer : str, optional
        If set and ``adata`` is supplied, the decoded expression is
        written to ``adata.layers[write_layer]``.
    **decode_kwargs
        Forwarded to the wrapper's ``decode_cells`` (e.g.
        ``library_size``, ``batch_index`` for scVI, ``read_depth``
        for STATE).

    Returns
    -------
    np.ndarray of shape ``(n_cells, n_genes)``
        Decoded expression (NB mean / log-probs / linear recon,
        depending on the model; see the wrapper docstring).
    """
    from .models.singlecell_models import singlecell_info

    card = singlecell_info(model)
    if not card.supports_decode:
        raise ValueError(
            f"Model '{model}' does not expose a decoder "
            f"(wrapper_class={card.wrapper_class_name}). "
            "Supported: pca, scvi, scanvi, totalvi, state."
        )

    if latent is None:
        if adata is None:
            raise ValueError("decode_cells needs either `latent` or `adata` (+ optional `obsm_key`).")
        key = obsm_key or f"X_{model}"
        if key not in adata.obsm:
            raise KeyError(
                f"{key!r} not present in adata.obsm. Run embed_cells "
                f"with models=[{model!r}] first, or pass `obsm_key`."
            )
        latent_arr = np.asarray(adata.obsm[key])
    else:
        latent_arr = np.asarray(latent)

    if gene_names is None and adata is not None:
        gene_names = list(adata.var_names)

    device_str = str(self.device)
    cache_key = (model, device_str)
    wrapper = self._singlecell_cache.get(cache_key)
    if wrapper is None:
        raise RuntimeError(
            f"No fitted '{model}' wrapper cached for device "
            f"{device_str!r}. Call `embed_cells(adata, models=["
            f"'{model}'])` first (for STATE/Stack, instantiate the "
            "wrapper manually with the checkpoint path and pass it to "
            "`wrapper.embed_cells` / `wrapper.decode_cells` directly)."
        )

    decoded = wrapper.decode_cells(
        latent_arr,
        gene_names=gene_names,
        adata=adata,
        **decode_kwargs,
    )

    if write_layer is not None and adata is not None:
        if decoded.shape[1] != adata.n_vars:
            raise ValueError(
                f"Decoded matrix has {decoded.shape[1]} genes but adata has "
                f"{adata.n_vars} var_names; cannot write to adata.layers."
            )
        adata.layers[write_layer] = decoded

    return decoded

embed(identifiers=None, *, entity_type=None, model, id_type=None, organism=None, pooling_strategy='mean', layer=None, output=None, target=None, input_path=None, output_path=None, attach_to='auto', harmonize_dim=None, path=None, fmt='npz', missing='error', id_column=None, identifier_column=None, anndata_axis=None, obs_column=None, var_column=None, is_perturbation=False, embedding_source='auto', hf_repo_id=DEFAULT_STATIC_EMBEDDING_REPO, hf_cache_dir=None, hf_token=None, hf_id_key=None, hf_id_type=None, whole_genome=False, biotype='protein_coding', show_progress=False, key=None, metadata_mode='full', include_matrix=True, random_state=0, **embed_kwargs)

Embed entities and return the standardized embpy output.

This is the package-wide output path. It accepts direct in-memory inputs (sequences, NumPy arrays, pandas Series/DataFrames, AnnData) and CSV/TSV/Parquet paths, normalizes them before inference, canonicalizes identifiers by entity type, then exports either AnnData or a table. Generated embeddings are never stored in AnnData .X; standalone AnnData uses sparse placeholder .X and stores matrices in .obsm, .varm or .uns. output="payload" returns the entity-aligned embedding payload directly, including canonical ids, model/provenance metadata and the requested embedding matrix.

Gene embeddings are feature embeddings by default and therefore route to .varm for AnnData output. Set is_perturbation=True when gene identifiers are perturbation/action labels; those embeddings route to .obsm and the flag is recorded in provenance. When the labels come from obs_column (or var_column), rows are aligned on that column of the target rather than on its obs_names/var_names, so every cell receives the vector of its label.

embedding_source="auto" routes known static lookup models such as genept and gene2vec to a Hugging Face dataset repository and routes all other models through normal inference. The default static repository is theislab/Embpy_Data and files are resolved via embeddings/{model}.npz / .parquet / .npy using :class:embpy.pp.HFHandler. Static lookup models return the same standardized outputs as model inference, but only for genes present in the downloaded table.

Defaults are explicit: if output is omitted, an output path chooses compact NPZ/CSV/Zarr file output, otherwise AnnData is returned. Multiple models or multiple entity types return the same output family with deterministic keys that include entity type and model name.

layer selects which transformer layer the embedding is pooled from, instead of the default final layer. Intermediate layers often transfer better to downstream tasks than the last one, which is specialised for the pretraining objective. Negative values count from the end, so layer=-1 reproduces the default and layer=-4 takes the fourth-from-last::

embedder.embed(adata, entity_type="protein", model="esm2_650M", layer=-4)

Not every architecture supports this. Models that cannot select a layer (fingerprints, static lookup tables, and wrappers that have not implemented it) raise a :class:ValueError rather than silently returning default-layer embeddings. Use :meth:~embpy.models.base.BaseModelWrapper.extract_hidden_states for finer-grained access to every layer at once.

Source code in src/embpy/embedder.py

def embed(
    self,
    identifiers: Any | None = None,
    *,
    entity_type: str | Sequence[str] | None = None,
    model: str | Sequence[str],
    id_type: str | Mapping[str, str] | None = None,
    organism: str | None = None,
    pooling_strategy: str = "mean",
    layer: int | None = None,
    output: Literal["anndata", "table", "payload"] | None = None,
    target: Any = None,
    input_path: str | os.PathLike[str] | None = None,
    output_path: str | os.PathLike[str] | None = None,
    attach_to: Literal["auto", "obs", "var", "uns"] = "auto",
    harmonize_dim: int | None = None,
    path: str | os.PathLike[str] | None = None,
    fmt: Literal["csv", "npz", "zarr"] = "npz",
    missing: Literal["error", "nan"] = "error",
    id_column: str | None = None,
    identifier_column: str | None = None,
    anndata_axis: Literal["obs", "var"] | None = None,
    obs_column: str | None = None,
    var_column: str | None = None,
    is_perturbation: bool = False,
    embedding_source: Literal["auto", "model", "hf", "huggingface", "static", "precomputed"] = "auto",
    hf_repo_id: str = DEFAULT_STATIC_EMBEDDING_REPO,
    hf_cache_dir: str | os.PathLike[str] | None = None,
    hf_token: str | None = None,
    hf_id_key: str | None = None,
    hf_id_type: str | None = None,
    whole_genome: bool = False,
    biotype: str = "protein_coding",
    show_progress: bool = False,
    key: str | None = None,
    metadata_mode: Literal["minimal", "full"] = "full",
    include_matrix: bool = True,
    random_state: int = 0,
    **embed_kwargs: Any,
):
    """Embed entities and return the standardized embpy output.

    This is the package-wide output path. It accepts direct in-memory
    inputs (sequences, NumPy arrays, pandas Series/DataFrames, AnnData)
    and CSV/TSV/Parquet paths, normalizes them before inference,
    canonicalizes identifiers by entity type, then exports either
    AnnData or a table. Generated embeddings are never stored in
    AnnData ``.X``; standalone AnnData uses sparse placeholder ``.X``
    and stores matrices in ``.obsm``, ``.varm`` or ``.uns``.
    ``output="payload"`` returns the entity-aligned embedding payload
    directly, including canonical ids, model/provenance metadata and
    the requested embedding matrix.

    Gene embeddings are feature embeddings by default and therefore
    route to ``.varm`` for AnnData output. Set ``is_perturbation=True``
    when gene identifiers are perturbation/action labels; those
    embeddings route to ``.obsm`` and the flag is recorded in
    provenance. When the labels come from ``obs_column`` (or
    ``var_column``), rows are aligned on that column of the target
    rather than on its ``obs_names``/``var_names``, so every cell
    receives the vector of its label.

    ``embedding_source="auto"`` routes known static lookup models
    such as ``genept`` and ``gene2vec`` to a Hugging Face dataset
    repository and routes all other models through normal inference.
    The default static repository is ``theislab/Embpy_Data`` and
    files are resolved via ``embeddings/{model}.npz`` / ``.parquet``
    / ``.npy`` using :class:`embpy.pp.HFHandler`. Static lookup
    models return the same standardized outputs as model inference,
    but only for genes present in the downloaded table.

    Defaults are explicit: if ``output`` is omitted, an output path
    chooses compact NPZ/CSV/Zarr file output, otherwise AnnData is
    returned. Multiple models or multiple entity types return the
    same output family with deterministic keys that include entity
    type and model name.

    ``layer`` selects which transformer layer the embedding is pooled
    from, instead of the default final layer. Intermediate layers often
    transfer better to downstream tasks than the last one, which is
    specialised for the pretraining objective. Negative values count
    from the end, so ``layer=-1`` reproduces the default and ``layer=-4``
    takes the fourth-from-last::

        embedder.embed(adata, entity_type="protein", model="esm2_650M", layer=-4)

    Not every architecture supports this. Models that cannot select a
    layer (fingerprints, static lookup tables, and wrappers that have not
    implemented it) raise a :class:`ValueError` rather than silently
    returning default-layer embeddings. Use
    :meth:`~embpy.models.base.BaseModelWrapper.extract_hidden_states` for
    finer-grained access to every layer at once.
    """
    from .io.exporters import route_output
    from .io.normalize import normalize_embedding_input

    # ``layer`` is the documented, discoverable spelling; ``target_layer``
    # is what the wrappers accept. Normalise here so both work and the
    # downstream support check sees a single canonical key.
    if layer is not None:
        embed_kwargs.setdefault("target_layer", layer)

    # Validate here, at the single public entry point, so static lookup
    # tables and other sources that bypass _embed_to_result are covered too.
    requested_layer = next(
        (
            embed_kwargs[k]
            for k in ("target_layer", "layer", "embedding_layer", "layer_name")
            if embed_kwargs.get(k) is not None
        ),
        None,
    )
    if requested_layer is not None:
        for _m in [model] if isinstance(model, str) else list(model):
            self._assert_layer_selection_supported(str(_m), requested_layer)

    org = organism or self.organism
    out_path = output_path if output_path is not None else path
    data = input_path if input_path is not None else identifiers

    if data is None and target is not None and not whole_genome:
        data = target
    if data is None and not whole_genome:
        raise ValueError(
            "input normalization: no identifiers or input_path were provided. "
            "Pass identifiers, an input file path, an AnnData object, or "
            "whole_genome=True."
        )

    if target is None and data is not None and self._is_anndata_like(data):
        target = data

    entity_types = self._resolve_entity_types(entity_type, data, whole_genome)
    if is_perturbation and any(et != "gene" for et in entity_types):
        raise ValueError(
            "input normalization: is_perturbation=True is currently only "
            "supported for entity_type='gene'. Use entity_type='perturbation' "
            "for morphology/action-image embeddings, or attach_to='obs' for "
            "other row-aligned entity embeddings."
        )
    models = self._as_string_list(model, arg_name="model")
    if "cell" in entity_types:
        if entity_types != ["cell"]:
            raise ValueError(
                "input normalization: entity_type='cell' cannot be mixed "
                "with other entity types in one BioEmbedder.embed(...) call. "
                "Run cell embeddings and feature/entity embeddings as "
                "separate calls so their AnnData alignment stays explicit."
            )
        return self._embed_cells_standardized(
            data,
            models=models,
            output=output or "anndata",
            target=target,
            path=out_path,
            fmt=fmt,
            missing=missing,
            key=key,
            metadata_mode=metadata_mode,
            include_matrix=include_matrix,
            **embed_kwargs,
        )

    normalized_by_entity = {}
    for et in entity_types:
        if whole_genome:
            normalized_by_entity[et] = self._whole_genome_input(
                entity_type=et,
                organism=org,
                biotype=biotype,
            )
            continue
        entity_data = data[et] if isinstance(data, Mapping) and et in data else data
        try:
            normalized_by_entity[et] = normalize_embedding_input(
                entity_data,
                entity_type=et,
                id_column=id_column,
                identifier_column=identifier_column,
                anndata_axis=anndata_axis,
                obs_column=obs_column,
                var_column=var_column,
            )
        except Exception as exc:
            if str(exc).startswith(("input loading:", "input normalization:")):
                raise
            raise ValueError(
                f"input normalization: failed for entity_type={et!r}: {type(exc).__name__}: {exc}"
            ) from exc

    tasks = [(et, m) for et in entity_types for m in models]
    iterator = tasks
    if show_progress:
        from tqdm.auto import tqdm

        iterator = tqdm(tasks, desc="embpy embed", unit="result")

    results = []
    for et, model_name in iterator:
        result = self._embed_to_result(
            normalized_by_entity[et],
            entity_type=et,
            model=model_name,
            id_type=self._id_type_for_entity(id_type, et),
            organism=org,
            pooling_strategy=pooling_strategy,
            is_perturbation=bool(is_perturbation and et == "gene"),
            embedding_source=embedding_source,
            hf_repo_id=hf_repo_id,
            hf_cache_dir=hf_cache_dir,
            hf_token=hf_token,
            hf_id_key=hf_id_key,
            hf_id_type=hf_id_type,
            show_progress=show_progress,
            **embed_kwargs,
        )
        results.append(result)

    resolved_output = output or ("table" if out_path is not None else "anndata")
    return route_output(
        results,
        output=resolved_output,
        target=target,
        attach_to=attach_to,
        harmonize_dim=harmonize_dim,
        path=out_path,
        fmt=fmt,
        missing=missing,
        key=key,
        metadata_mode=metadata_mode,
        include_matrix=include_matrix,
        random_state=random_state,
        obs_column=obs_column,
        var_column=var_column,
    )

embed_adata(adata, *, cell_models=None, preprocessing='auto', target_sum=10000.0, n_top_genes=2000, select_hvg=True, hvg_flavor='auto', log_transform=True, scale=False, max_value=10.0, min_genes=200, min_cells=3, max_pct_mito=None, n_pca_components=50, pca_use_hvg=True, n_latent=30, n_layers_scvi=2, n_hidden_scvi=128, max_epochs=200, early_stopping=True, batch_key=None, labels_key=None, protein_expression_obsm_key=None, perturbation_models=None, perturbation_column=None, perturbation_type='auto', perturbation_organism='human', pooling_strategy='mean', batch_size=32, obsm_prefix='X_', copy=True, backend='cpu')

Unified embedding of an AnnData -- cells and/or perturbations.

Combines cell-level embeddings (from expression via single-cell foundation models, PCA, or scVI) with perturbation-level embeddings (from gene/protein/molecule identifiers in .obs) in a single call.

Parameters:

Name	Type	Description	Default
adata	AnnData	Input AnnData with raw counts in .X and (optionally) perturbation annotations in .obs.	required
cell_models	str or list[str]	Single-cell model key(s) for expression-based cell embeddings (e.g. "scgpt", "pca", "scvi"). None skips cell embedding.	None
preprocessing	Literal['auto', 'raw', 'standard', 'none']	Preprocessing pipeline for cell models.	'auto'
target_sum	float | None	Normalization and scaling options forwarded to :func:~embpy.pp.preprocess_counts.	10000.0
n_top_genes	float | None	Normalization and scaling options forwarded to :func:~embpy.pp.preprocess_counts.	10000.0
log_transform	float | None	Normalization and scaling options forwarded to :func:~embpy.pp.preprocess_counts.	10000.0
scale	float | None	Normalization and scaling options forwarded to :func:~embpy.pp.preprocess_counts.	10000.0
max_value	float | None	Normalization and scaling options forwarded to :func:~embpy.pp.preprocess_counts.	10000.0
select_hvg	bool	Highly-variable-gene options forwarded to :func:~embpy.pp.preprocess_counts.	True
hvg_flavor	bool	Highly-variable-gene options forwarded to :func:~embpy.pp.preprocess_counts.	True
min_genes	int	QC filtering options forwarded to :func:~embpy.pp.preprocess_counts.	200
min_cells	int	QC filtering options forwarded to :func:~embpy.pp.preprocess_counts.	200
max_pct_mito	int	QC filtering options forwarded to :func:~embpy.pp.preprocess_counts.	200
n_pca_components	int	PCA-specific parameters.	50
pca_use_hvg	int	PCA-specific parameters.	50
n_latent	int		30
n_layers_scvi	int		30
n_hidden_scvi	int		30
max_epochs	int		30
early_stopping	bool		True
batch_key	bool		True
labels_key	bool		True
protein_expression_obsm_key	str | None	scvi-tools parameters.	None
perturbation_models	str or list[str]	Sequence/molecule model key(s) for perturbation embeddings (e.g. "esm2_650M", "enformer_human_rough", "chemberta2MTR"). None skips perturbation embedding.	None
perturbation_column	str	Column in adata.obs containing perturbation identifiers. Required when perturbation_models is set.	None
perturbation_type	Literal['auto', 'symbol', 'ensembl_id', 'uniprot_id', 'smiles']	Type of identifiers in the perturbation column. "auto" auto-detects per identifier.	'auto'
perturbation_organism	str	Organism for gene/protein resolution.	'human'
pooling_strategy	str	Pooling for sequence/molecule models.	'mean'
batch_size	int	Batch size for model inference.	32
obsm_prefix	str	Prefix for .obsm keys (default "X_").	'X_'
copy	bool	If True, operate on a copy of adata.	True

Returns:

Type	Description
AnnData	The enriched AnnData with: .X = original raw counts .layers["counts"] = raw counts copy .layers["log_normalized"] = processed expression (standard pipeline) .obsm["{prefix}{cell_model}"] = cell embeddings .uns["perturbations"]["{prefix}{pert_model}"] = perturbation embeddings, one row per unique perturbation .uns["embpy_embeddings"] = metadata dict

Examples:

>>> result = embedder.embed_adata(
...     adata,
...     cell_models=["pca", "scvi", "scgpt"],
...     perturbation_models=["esm2_650M", "chemberta2MTR"],
...     perturbation_column="perturbation",
...     perturbation_type="auto",
... )
>>> result.obsm["X_scgpt"].shape  # cell embeddings
(5000, 512)
>>> result.uns["perturbations"]["X_esm2_650M"]["matrix"].shape
(n_perturbations, 1280)

Source code in src/embpy/embedder.py

def embed_adata(
    self,
    adata,  # anndata.AnnData
    *,
    # --- Cell-level embeddings (from expression) ---
    cell_models: list[str] | str | None = None,
    preprocessing: Literal["auto", "raw", "standard", "none"] = "auto",
    target_sum: float | None = 1e4,
    n_top_genes: int = 2000,
    select_hvg: bool = True,
    hvg_flavor: Literal["auto", "seurat", "seurat_v3", "cell_ranger"] = "auto",
    log_transform: bool = True,
    scale: bool = False,
    max_value: float | None = 10.0,
    min_genes: int = 200,
    min_cells: int = 3,
    max_pct_mito: float | None = None,
    n_pca_components: int = 50,
    pca_use_hvg: bool = True,
    n_latent: int = 30,
    n_layers_scvi: int = 2,
    n_hidden_scvi: int = 128,
    max_epochs: int = 200,
    early_stopping: bool = True,
    batch_key: str | None = None,
    labels_key: str | None = None,
    protein_expression_obsm_key: str | None = None,
    # --- Perturbation-level embeddings (gene / protein / molecule) ---
    perturbation_models: list[str] | str | None = None,
    perturbation_column: str | None = None,
    perturbation_type: Literal[
        "auto",
        "symbol",
        "ensembl_id",
        "uniprot_id",
        "smiles",
    ] = "auto",
    perturbation_organism: str = "human",
    pooling_strategy: str = "mean",
    # --- General ---
    batch_size: int = 32,
    obsm_prefix: str = "X_",
    copy: bool = True,
    backend: Literal["cpu", "gpu"] = "cpu",
):
    """Unified embedding of an AnnData -- cells and/or perturbations.

    Combines cell-level embeddings (from expression via single-cell
    foundation models, PCA, or scVI) with perturbation-level
    embeddings (from gene/protein/molecule identifiers in ``.obs``)
    in a single call.

    Parameters
    ----------
    adata : anndata.AnnData
        Input AnnData with raw counts in ``.X`` and (optionally)
        perturbation annotations in ``.obs``.

    cell_models : str or list[str], optional
        Single-cell model key(s) for expression-based cell
        embeddings (e.g. ``"scgpt"``, ``"pca"``, ``"scvi"``).
        ``None`` skips cell embedding.
    preprocessing
        Preprocessing pipeline for cell models.
    target_sum, n_top_genes, log_transform, scale, max_value
        Normalization and scaling options forwarded to
        :func:`~embpy.pp.preprocess_counts`.
    select_hvg, hvg_flavor
        Highly-variable-gene options forwarded to
        :func:`~embpy.pp.preprocess_counts`.
    min_genes, min_cells, max_pct_mito
        QC filtering options forwarded to
        :func:`~embpy.pp.preprocess_counts`.
    n_pca_components, pca_use_hvg
        PCA-specific parameters.
    n_latent, n_layers_scvi, n_hidden_scvi, max_epochs,
    early_stopping, batch_key, labels_key,
    protein_expression_obsm_key
        scvi-tools parameters.

    perturbation_models : str or list[str], optional
        Sequence/molecule model key(s) for perturbation embeddings
        (e.g. ``"esm2_650M"``, ``"enformer_human_rough"``,
        ``"chemberta2MTR"``).  ``None`` skips perturbation embedding.
    perturbation_column : str, optional
        Column in ``adata.obs`` containing perturbation identifiers.
        Required when *perturbation_models* is set.
    perturbation_type
        Type of identifiers in the perturbation column.
        ``"auto"`` auto-detects per identifier.
    perturbation_organism
        Organism for gene/protein resolution.
    pooling_strategy
        Pooling for sequence/molecule models.

    batch_size
        Batch size for model inference.
    obsm_prefix
        Prefix for ``.obsm`` keys (default ``"X_"``).
    copy
        If ``True``, operate on a copy of adata.

    Returns
    -------
    anndata.AnnData
        The enriched AnnData with:

        - ``.X`` = original raw counts
        - ``.layers["counts"]`` = raw counts copy
        - ``.layers["log_normalized"]`` = processed expression
          (standard pipeline)
        - ``.obsm["{prefix}{cell_model}"]`` = cell embeddings
        - ``.uns["perturbations"]["{prefix}{pert_model}"]``
          = perturbation embeddings, one row per unique perturbation
        - ``.uns["embpy_embeddings"]`` = metadata dict

    Examples
    --------
    >>> result = embedder.embed_adata(
    ...     adata,
    ...     cell_models=["pca", "scvi", "scgpt"],
    ...     perturbation_models=["esm2_650M", "chemberta2MTR"],
    ...     perturbation_column="perturbation",
    ...     perturbation_type="auto",
    ... )
    >>> result.obsm["X_scgpt"].shape  # cell embeddings
    (5000, 512)
    >>> result.uns["perturbations"]["X_esm2_650M"]["matrix"].shape
    (n_perturbations, 1280)
    """
    from .io.exporters import to_anndata
    from .io.result import EmbeddingProvenance, EmbeddingResult
    from .resources.gene_resolver import detect_identifier_type

    if copy:
        adata = adata.copy()

    metadata: dict[str, dict[str, Any]] = {}

    # ==================================================================
    # Cell-level embeddings (expression-based)
    # ==================================================================
    if cell_models is not None:
        adata = self.embed_cells(
            adata,
            models=cell_models,
            preprocessing=preprocessing,
            target_sum=target_sum,
            n_top_genes=n_top_genes,
            select_hvg=select_hvg,
            hvg_flavor=hvg_flavor,
            log_transform=log_transform,
            scale=scale,
            max_value=max_value,
            min_genes=min_genes,
            min_cells=min_cells,
            max_pct_mito=max_pct_mito,
            n_pca_components=n_pca_components,
            pca_use_hvg=pca_use_hvg,
            n_latent=n_latent,
            n_layers_scvi=n_layers_scvi,
            n_hidden_scvi=n_hidden_scvi,
            max_epochs=max_epochs,
            early_stopping=early_stopping,
            batch_key=batch_key,
            labels_key=labels_key,
            protein_expression_obsm_key=protein_expression_obsm_key,
            batch_size=batch_size,
            obsm_prefix=obsm_prefix,
            copy=False,
            backend=backend,
        )
        if "embpy_cell_embeddings" in adata.uns:
            metadata.update(adata.uns["embpy_cell_embeddings"])

    # ==================================================================
    # Perturbation-level embeddings (gene / protein / molecule)
    # ==================================================================
    if perturbation_models is not None:
        if perturbation_column is None:
            raise ValueError(
                "perturbation_column is required when perturbation_models "
                "is specified. Set it to the .obs column containing "
                "perturbation identifiers (e.g. 'perturbation', 'gene', "
                "'smiles')."
            )
        if perturbation_column not in adata.obs.columns:
            raise ValueError(
                f"Column '{perturbation_column}' not found in adata.obs. Available: {list(adata.obs.columns)}"
            )

        if isinstance(perturbation_models, str):
            perturbation_models = [perturbation_models]

        pert_ids = adata.obs[perturbation_column].astype(str).values
        unique_perts = list(dict.fromkeys(pert_ids))
        logging.info(
            "Perturbation embedding: %d cells, %d unique perturbations from column '%s'",
            len(pert_ids),
            len(unique_perts),
            perturbation_column,
        )

        # Determine id_type per perturbation if auto
        if perturbation_type == "auto":
            id_types = {p: detect_identifier_type(p) for p in unique_perts}
        else:
            id_types = dict.fromkeys(unique_perts, perturbation_type)

        for model_key in perturbation_models:
            obsm_key = f"{obsm_prefix}{model_key}"
            logging.info(
                "Embedding perturbations with '%s' ...",
                model_key,
            )

            try:
                # Embed each unique perturbation
                pert_embs: dict[str, np.ndarray | None] = {}
                for pert in unique_perts:
                    id_t = id_types[pert]
                    try:
                        if id_t == "smiles":
                            emb = self.embed_molecule(
                                identifier=pert,
                                model=model_key,
                                pooling_strategy=pooling_strategy,
                            )
                        else:
                            emb = self.embed_gene(
                                identifier=pert,
                                model=model_key,
                                id_type=id_t,
                                organism=perturbation_organism,
                                pooling_strategy=pooling_strategy,
                            )
                        pert_embs[pert] = np.asarray(
                            emb,
                            dtype=np.float32,
                        ).ravel()
                    except Exception as e:  # noqa: BLE001
                        logging.warning(
                            "Failed to embed perturbation '%s': %s",
                            pert,
                            e,
                        )
                        pert_embs[pert] = None

                # Determine embedding dim
                emb_dim = 0
                for v in pert_embs.values():
                    if v is not None:
                        emb_dim = v.shape[0]
                        break

                if emb_dim == 0:
                    logging.error(
                        "No perturbation embeddings succeeded for '%s'",
                        model_key,
                    )
                    metadata[model_key] = {
                        "error": "no embeddings succeeded",
                    }
                    continue

                n_ok = sum(1 for v in pert_embs.values() if v is not None)
                successful_perts = [p for p in unique_perts if pert_embs.get(p) is not None]
                entity_matrix = np.stack(
                    [np.asarray(pert_embs[p], dtype=np.float32).ravel() for p in successful_perts],
                    axis=0,
                )
                aliases = {
                    str(p): {
                        "perturbation_label": str(p),
                        "input_id_type": str(id_types.get(p, perturbation_type)),
                    }
                    for p in successful_perts
                }
                result = EmbeddingResult(
                    matrix=entity_matrix,
                    entity_ids=tuple(str(p) for p in successful_perts),
                    entity_type="perturbation",
                    id_scheme="perturbation_label",
                    provenance=EmbeddingProvenance.create(
                        model=model_key,
                        pooling=pooling_strategy,
                        extra={
                            "entity_type": "perturbation",
                            "input_kind": "anndata_obs_column",
                            "input_source": "AnnData.obs",
                            "input_id_column": perturbation_column,
                            "perturbation_column": perturbation_column,
                            "perturbation_type": perturbation_type,
                            "organism": perturbation_organism,
                            "n_requested_inputs": len(unique_perts),
                            "n_successfully_embedded_entities": n_ok,
                            "n_embedding_failures": len(unique_perts) - n_ok,
                            "n_cells": int(adata.n_obs),
                        },
                    ),
                    aliases=aliases,
                )
                to_anndata(result, target=adata, attach_to="uns", key=obsm_key)

                metadata[model_key] = {
                    "uns_key": obsm_key,
                    "storage": "uns",
                    "embedding_dim": emb_dim,
                    "n_cells": int(adata.n_obs),
                    "n_perturbations_total": len(unique_perts),
                    "n_perturbations_embedded": n_ok,
                    "type": "perturbation",
                    "perturbation_column": perturbation_column,
                }
                logging.info(
                    "  -> stored in .uns perturbations[%r], shape=(%d, %d) (%d/%d perturbations embedded)",
                    obsm_key,
                    entity_matrix.shape[0],
                    entity_matrix.shape[1],
                    n_ok,
                    len(unique_perts),
                )

            except Exception as e:  # noqa: BLE001
                logging.error(
                    "Failed perturbation embedding with '%s': %s",
                    model_key,
                    e,
                )
                metadata[model_key] = {"error": str(e)}

    adata.uns["embpy_embeddings"] = metadata
    logging.info("embed_adata complete.")
    return adata

embed_cells(adata, models='scgpt', preprocessing='auto', *, target_sum=10000.0, n_top_genes=2000, select_hvg=True, hvg_flavor='auto', log_transform=True, scale=False, max_value=10.0, min_genes=200, min_cells=3, max_pct_mito=None, n_pca_components=50, pca_use_hvg=True, n_latent=30, n_layers_scvi=2, n_hidden_scvi=128, max_epochs=200, early_stopping=True, batch_key=None, labels_key=None, protein_expression_obsm_key=None, batch_size=32, obsm_prefix='X_', copy=True, backend='cpu', vocab_conversion='auto', model_kwargs=None)

Embed single cells from an AnnData object.

Orchestrates preprocessing and multi-model embedding in a single call. Each requested model's embeddings are stored in adata.obsm under "{obsm_prefix}{model_name}".

Parameters:

Name	Type	Description	Default
adata	AnnData	AnnData with raw counts in .X.	required
models	str or list[str]	Model key(s) from the single-cell model registry. Accepts any key from :func:~embpy.models.singlecell_models.list_singlecell_models, e.g. "scgpt", "geneformer_v2_12L", "pca", "scvi", "scanvi", "totalvi".	'scgpt'
preprocessing	('auto', 'raw', 'standard', 'none')	"auto" chooses a model-aware default from the single-cell registry: raw-count models get QC/counts-layer preparation, while models that consume processed expression (for example PCA) get the standard pipeline. "standard" runs log-normalize + HVG. "raw" applies only QC filtering. "none" skips preprocessing entirely.	"auto"
target_sum	float | None	Target total counts for normalization (standard pipeline).	10000.0
n_top_genes	int	Number of highly variable genes (standard pipeline).	2000
select_hvg	bool	Whether to compute highly variable genes in the standard pipeline.	True
hvg_flavor	Literal['auto', 'seurat', 'seurat_v3', 'cell_ranger']	Highly-variable-gene method. "auto" chooses a scanpy flavor compatible with log_transform.	'auto'
log_transform	bool	Whether to log1p-transform (standard pipeline).	True
scale	bool	Whether to scale to unit variance (standard pipeline).	False
max_value	float | None	Max value after scaling.	10.0
min_genes	int	Minimum genes per cell for QC.	200
min_cells	int	Minimum cells per gene for QC.	3
max_pct_mito	float | None	Maximum mitochondrial fraction (None to skip).	None
n_pca_components	int	Number of PCA components (when "pca" is in models).	50
pca_use_hvg	bool	Whether PCA restricts to HVGs.	True
n_latent	int	scvi-tools latent dimensionality.	30
n_layers_scvi	int	Number of hidden layers in scvi-tools encoder/decoder.	2
n_hidden_scvi	int	Number of nodes per hidden layer in scvi-tools.	128
max_epochs	int	Maximum training epochs for scvi-tools.	200
early_stopping	bool	Whether to use early stopping for scvi-tools.	True
batch_key	str | None	Batch column in adata.obs for scvi-tools batch correction.	None
labels_key	str | None	Label column in adata.obs for scANVI.	None
protein_expression_obsm_key	str | None	Key in adata.obsm with protein counts for totalVI.	None
batch_size	int	Batch size for foundation model inference.	32
obsm_prefix	str	Prefix for .obsm keys (default "X_").	'X_'
copy	bool	If True, operate on a copy of adata.	True
model_kwargs	dict[str, dict[str, Any]] | None	Per-model constructor arguments, keyed by model key, e.g. {"state": {"checkpoint": "/path/to/se600m.ckpt"}}. Needed by models whose weights are not fetched for you: state and stack take a checkpoint (state also accepts model_folder), and raise from load() without one, so before this existed they could not be reached through this method at all. The kwargs are part of the wrapper cache key, so two checkpoints of one model do not collide.	None

Returns:

Type	Description
AnnData	The AnnData with: .X = original raw counts .layers["counts"] = raw counts copy .layers["log_normalized"] = processed (standard pipeline) .obsm["{prefix}{model}"] = embeddings per model .uns["embpy_cell_embeddings"] = metadata dict

Source code in src/embpy/embedder.py

def embed_cells(
    self,
    adata,  # anndata.AnnData
    models: list[str] | str = "scgpt",
    preprocessing: Literal["auto", "raw", "standard", "none"] = "auto",
    *,
    # Preprocessing params (forwarded to preprocess_counts)
    target_sum: float | None = 1e4,
    n_top_genes: int = 2000,
    select_hvg: bool = True,
    hvg_flavor: Literal["auto", "seurat", "seurat_v3", "cell_ranger"] = "auto",
    log_transform: bool = True,
    scale: bool = False,
    max_value: float | None = 10.0,
    min_genes: int = 200,
    min_cells: int = 3,
    max_pct_mito: float | None = None,
    # PCA-specific
    n_pca_components: int = 50,
    pca_use_hvg: bool = True,
    # scVI-specific
    n_latent: int = 30,
    n_layers_scvi: int = 2,
    n_hidden_scvi: int = 128,
    max_epochs: int = 200,
    early_stopping: bool = True,
    batch_key: str | None = None,
    labels_key: str | None = None,
    protein_expression_obsm_key: str | None = None,
    # General
    batch_size: int = 32,
    obsm_prefix: str = "X_",
    copy: bool = True,
    backend: Literal["cpu", "gpu"] = "cpu",
    vocab_conversion: Literal["auto", "off"] = "auto",
    model_kwargs: dict[str, dict[str, Any]] | None = None,
):
    """Embed single cells from an AnnData object.

    Orchestrates preprocessing and multi-model embedding in a single
    call.  Each requested model's embeddings are stored in
    ``adata.obsm`` under ``"{obsm_prefix}{model_name}"``.

    Parameters
    ----------
    adata : anndata.AnnData
        AnnData with raw counts in ``.X``.
    models : str or list[str]
        Model key(s) from the single-cell model registry.  Accepts
        any key from :func:`~embpy.models.singlecell_models.list_singlecell_models`,
        e.g. ``"scgpt"``, ``"geneformer_v2_12L"``, ``"pca"``,
        ``"scvi"``, ``"scanvi"``, ``"totalvi"``.
    preprocessing : {"auto", "raw", "standard", "none"}
        ``"auto"`` chooses a model-aware default from the single-cell
        registry: raw-count models get QC/counts-layer preparation, while
        models that consume processed expression (for example PCA) get the
        standard pipeline.
        ``"standard"`` runs log-normalize + HVG.
        ``"raw"`` applies only QC filtering.
        ``"none"`` skips preprocessing entirely.
    target_sum
        Target total counts for normalization (standard pipeline).
    n_top_genes
        Number of highly variable genes (standard pipeline).
    select_hvg
        Whether to compute highly variable genes in the standard pipeline.
    hvg_flavor
        Highly-variable-gene method. ``"auto"`` chooses a scanpy flavor
        compatible with ``log_transform``.
    log_transform
        Whether to log1p-transform (standard pipeline).
    scale
        Whether to scale to unit variance (standard pipeline).
    max_value
        Max value after scaling.
    min_genes
        Minimum genes per cell for QC.
    min_cells
        Minimum cells per gene for QC.
    max_pct_mito
        Maximum mitochondrial fraction (``None`` to skip).
    n_pca_components
        Number of PCA components (when ``"pca"`` is in models).
    pca_use_hvg
        Whether PCA restricts to HVGs.
    n_latent
        scvi-tools latent dimensionality.
    n_layers_scvi
        Number of hidden layers in scvi-tools encoder/decoder.
    n_hidden_scvi
        Number of nodes per hidden layer in scvi-tools.
    max_epochs
        Maximum training epochs for scvi-tools.
    early_stopping
        Whether to use early stopping for scvi-tools.
    batch_key
        Batch column in ``adata.obs`` for scvi-tools batch correction.
    labels_key
        Label column in ``adata.obs`` for scANVI.
    protein_expression_obsm_key
        Key in ``adata.obsm`` with protein counts for totalVI.
    batch_size
        Batch size for foundation model inference.
    obsm_prefix
        Prefix for ``.obsm`` keys (default ``"X_"``).
    copy
        If ``True``, operate on a copy of adata.
    model_kwargs
        Per-model constructor arguments, keyed by model key, e.g.
        ``{"state": {"checkpoint": "/path/to/se600m.ckpt"}}``. Needed by
        models whose weights are not fetched for you: ``state`` and
        ``stack`` take a ``checkpoint`` (``state`` also accepts
        ``model_folder``), and raise from ``load()`` without one, so
        before this existed they could not be reached through this
        method at all. The kwargs are part of the wrapper cache key,
        so two checkpoints of one model do not collide.

    Returns
    -------
    anndata.AnnData
        The AnnData with:

        - ``.X`` = original raw counts
        - ``.layers["counts"]`` = raw counts copy
        - ``.layers["log_normalized"]`` = processed (standard pipeline)
        - ``.obsm["{prefix}{model}"]`` = embeddings per model
        - ``.uns["embpy_cell_embeddings"]`` = metadata dict
    """
    from .models.singlecell_models import (
        PCAEmbedding,
        ScVIToolsWrapper,
        list_singlecell_models,
        resolve_singlecell_preprocessing,
        singlecell_info,
    )
    from .pp.sc_preprocessing import preprocess_counts

    if isinstance(models, str):
        models = [models]

    available = list_singlecell_models()
    for m in models:
        if m not in available:
            raise ValueError(f"Unknown single-cell model '{m}'. Available: {available}")

    resolved_preprocessing, preprocessing_plan = resolve_singlecell_preprocessing(
        models,
        preprocessing,
    )

    if copy:
        adata = adata.copy()

    # ---- Preprocessing -----------------------------------------------
    preprocessing_options = {
        "target_sum": target_sum,
        "n_top_genes": n_top_genes,
        "select_hvg": select_hvg,
        "hvg_flavor": hvg_flavor,
        "log_transform": log_transform,
        "scale": scale,
        "max_value": max_value,
        "min_genes": min_genes,
        "min_cells": min_cells,
        "max_pct_mito": max_pct_mito,
        "backend": backend,
    }
    if resolved_preprocessing != "none":
        adata = preprocess_counts(
            adata,
            pipeline=resolved_preprocessing,
            min_genes=min_genes,
            min_cells=min_cells,
            max_pct_mito=max_pct_mito,
            target_sum=target_sum,
            n_top_genes=n_top_genes,
            select_hvg=select_hvg,
            hvg_flavor=hvg_flavor,
            log_transform=log_transform,
            scale=scale,
            max_value=max_value,
            copy=False,
            backend=backend,
        )
    else:
        logging.warning(
            "Skipping single-cell preprocessing because preprocessing='none'. "
            "Models that expect .layers['log_normalized'], .layers['counts'], "
            "or adata.var['highly_variable'] may fall back to .X or run "
            "without HVG restriction."
        )

    preprocessing_warnings = self._singlecell_preprocessing_warnings(
        adata,
        models=models,
        resolved_preprocessing=resolved_preprocessing,
        pca_use_hvg=pca_use_hvg,
    )
    preprocessing_metadata: dict[str, Any] = {
        **preprocessing_plan,
        "options": preprocessing_options,
        "warnings": preprocessing_warnings,
    }

    # ---- Embed with each model ---------------------------------------
    device_str = str(self.device)
    metadata: dict[str, dict[str, Any]] = {
        "__preprocessing__": preprocessing_metadata,
    }

    for model_key in models:
        obsm_key = f"{obsm_prefix}{model_key}"
        logging.info("Embedding cells with '%s' ...", model_key)

        try:
            if model_key == "pca":
                wrapper = PCAEmbedding(
                    n_components=n_pca_components,
                    use_hvg=pca_use_hvg,
                    layer="log_normalized",
                    scale=True,
                    backend=backend,
                )
                wrapper.load(device_str)
                embs = wrapper.embed_cells(adata)
                # Park the fitted wrapper so decode_cells() can reuse
                # the sklearn/cuml PCA + scaler for inverse_transform.
                self._singlecell_cache[(model_key, device_str)] = wrapper

            elif model_key in ("scvi", "scanvi", "totalvi"):
                model_cls_name = model_key.upper()
                if model_cls_name == "SCVI":
                    model_cls_name = "SCVI"
                elif model_cls_name == "SCANVI":
                    model_cls_name = "SCANVI"
                elif model_cls_name == "TOTALVI":
                    model_cls_name = "TOTALVI"

                wrapper = ScVIToolsWrapper(
                    model_class=model_cls_name,
                    n_latent=n_latent,
                    n_layers=n_layers_scvi,
                    n_hidden=n_hidden_scvi,
                    max_epochs=max_epochs,
                    early_stopping=early_stopping,
                    batch_key=batch_key,
                    labels_key=labels_key,
                    protein_expression_obsm_key=protein_expression_obsm_key,
                    layer="counts",
                    batch_size=batch_size,
                )
                wrapper.load(device_str)
                embs = wrapper.embed_cells(adata)
                # Park the trained scvi-tools model so decode_cells()
                # can route arbitrary latents through its generative
                # module without retraining.
                self._singlecell_cache[(model_key, device_str)] = wrapper

            else:
                # Foundation models (scGPT, Geneformer, UCE, Tahoe, ...)
                # are cached by (model_key, device) so repeated calls
                # to embed_cells -- e.g. chunked inference over a large
                # AnnData -- reuse the already-loaded torch model
                # instead of re-instantiating it on every call.
                wrapper = self._get_or_load_singlecell_wrapper(
                    model_key, batch_size, device_str,
                    (model_kwargs or {}).get(model_key),
                )
                # Auto-adapt var_names to the model's vocabulary.
                # This is the difference between a silent zero-match
                # failure (e.g. passing Ensembl IDs to scGPT) and a
                # working embedding. Set vocab_conversion="off" to
                # opt out of the automatic GeneResolver roundtrip.
                if vocab_conversion == "auto":
                    adata_for_model, report = self._ensure_singlecell_vocabulary(
                        adata, model_key, organism=self.organism
                    )
                    metadata[f"{model_key}__vocab"] = report
                else:
                    adata_for_model = adata
                embs = wrapper.embed_cells(adata_for_model)

            adata.obsm[obsm_key] = embs
            card = singlecell_info(model_key)
            metadata[model_key] = {
                "obsm_key": obsm_key,
                "embedding_dim": embs.shape[1],
                "n_cells": embs.shape[0],
                "wrapper_class": type(wrapper).__name__,
                "preprocessing": preprocessing_metadata,
                "model_requirements": {
                    "default_preprocessing": card.default_preprocessing,
                    "input_layer": card.input_layer,
                    "uses_hvg": card.uses_hvg,
                    "vocab_type": card.vocab_type,
                },
            }
            logging.info(
                "  -> stored in .obsm['%s'], shape %s",
                obsm_key,
                embs.shape,
            )

        except Exception as e:  # noqa: BLE001
            logging.error("Failed to embed with '%s': %s", model_key, e)
            metadata[model_key] = {
                "error": str(e),
                "preprocessing": preprocessing_metadata,
            }

    adata.uns["embpy_cell_embeddings"] = metadata
    logging.info(
        "embed_cells complete: %d models, %d cells",
        len(models),
        adata.n_obs,
    )
    return adata

embed_description(identifier, model='minilm_l6_v2', entity_type='auto', sources='all', pooling_strategy='mean', **kwargs)

Fetch text description(s) for a biological entity and embed them.

Resolves the identifier to rich text descriptions from public knowledge sources (MyGene, NCBI, Ensembl, UniProt, Wikipedia, PubChem), combines them, and embeds the result with a text model.

Parameters:

Name	Type	Description	Default
identifier	str	Gene symbol, Ensembl ID, UniProt accession, drug name, or SMILES string.	required
model	str	Text embedding model (default "minilm_l6_v2").	'minilm_l6_v2'
entity_type	Literal['gene', 'protein', 'molecule', 'auto']	"gene", "protein", "molecule", or "auto" to auto-detect.	'auto'
sources	list[str] | str	Knowledge sources to query. "all" uses the defaults for the detected entity type.	'all'
pooling_strategy	str	Pooling strategy for the text model.	'mean'
**kwargs	Any	Forwarded to the text model's embed method.	{}

Returns:

Type	Description
ndarray	The text embedding vector.

Examples:

>>> emb = embedder.embed_description("TP53", model="minilm_l6_v2")
>>> emb.shape
(384,)
>>> emb = embedder.embed_description("aspirin", entity_type="molecule")

Source code in src/embpy/embedder.py

def embed_description(
    self,
    identifier: str,
    model: str = "minilm_l6_v2",
    entity_type: Literal["gene", "protein", "molecule", "auto"] = "auto",
    sources: list[str] | str = "all",
    pooling_strategy: str = "mean",
    **kwargs: Any,
) -> np.ndarray:
    """Fetch text description(s) for a biological entity and embed them.

    Resolves the identifier to rich text descriptions from public
    knowledge sources (MyGene, NCBI, Ensembl, UniProt, Wikipedia,
    PubChem), combines them, and embeds the result with a text model.

    Parameters
    ----------
    identifier
        Gene symbol, Ensembl ID, UniProt accession, drug name,
        or SMILES string.
    model
        Text embedding model (default ``"minilm_l6_v2"``).
    entity_type
        ``"gene"``, ``"protein"``, ``"molecule"``, or ``"auto"``
        to auto-detect.
    sources
        Knowledge sources to query. ``"all"`` uses the defaults
        for the detected entity type.
    pooling_strategy
        Pooling strategy for the text model.
    **kwargs
        Forwarded to the text model's ``embed`` method.

    Returns
    -------
    np.ndarray
        The text embedding vector.

    Examples
    --------
    >>> emb = embedder.embed_description("TP53", model="minilm_l6_v2")
    >>> emb.shape
    (384,)
    >>> emb = embedder.embed_description("aspirin", entity_type="molecule")
    """
    description = self.text_resolver.get_combined_description(
        identifier,
        entity_type=entity_type,
        sources=sources,
    )
    logging.info(
        "Description for '%s' (%d chars): %s...",
        identifier,
        len(description),
        description[:100],
    )
    return self.embed_text(
        text=description,
        model=model,
        pooling_strategy=pooling_strategy,
        **kwargs,
    )

embed_descriptions_batch(identifiers, model='minilm_l6_v2', entity_type='auto', sources='all', pooling_strategy='mean', batch_size=None, **kwargs)

Fetch and embed text descriptions for a batch of entities.

Parameters:

Name	Type	Description	Default
identifiers	Sequence[str]	List of biological entity identifiers.	required
model	str	Text embedding model.	'minilm_l6_v2'
entity_type	Literal['gene', 'protein', 'molecule', 'auto']	Entity type (applied to all identifiers).	'auto'
sources	list[str] | str	Knowledge sources to query.	'all'
pooling_strategy	str	Pooling strategy for the text model.	'mean'
batch_size	int | None	Maximum batch size for embedding inference.	None
**kwargs	Any	Forwarded to the text model.	{}

Returns:

Type	Description
List of embedding arrays, with ``None`` for failed lookups.

Source code in src/embpy/embedder.py

def embed_descriptions_batch(
    self,
    identifiers: Sequence[str],
    model: str = "minilm_l6_v2",
    entity_type: Literal["gene", "protein", "molecule", "auto"] = "auto",
    sources: list[str] | str = "all",
    pooling_strategy: str = "mean",
    batch_size: int | None = None,
    **kwargs: Any,
) -> list[np.ndarray | None]:
    """Fetch and embed text descriptions for a batch of entities.

    Parameters
    ----------
    identifiers
        List of biological entity identifiers.
    model
        Text embedding model.
    entity_type
        Entity type (applied to all identifiers).
    sources
        Knowledge sources to query.
    pooling_strategy
        Pooling strategy for the text model.
    batch_size
        Maximum batch size for embedding inference.
    **kwargs
        Forwarded to the text model.

    Returns
    -------
    List of embedding arrays, with ``None`` for failed lookups.
    """
    texts: list[str] = []
    valid_indices: list[int] = []

    for i, ident in enumerate(identifiers):
        try:
            desc = self.text_resolver.get_combined_description(
                ident,
                entity_type=entity_type,
                sources=sources,
            )
            if desc:
                texts.append(desc)
                valid_indices.append(i)
            else:
                logging.warning("No description found for '%s'", ident)
        except Exception as e:  # noqa: BLE001
            logging.warning("Failed to fetch description for '%s': %s", ident, e)

    if not texts:
        return [None] * len(identifiers)

    batch_embs = self.embed_texts_batch(
        texts=texts,
        model=model,
        pooling_strategy=pooling_strategy,
        batch_size=batch_size,
        **kwargs,
    )

    results: list[np.ndarray | None] = [None] * len(identifiers)
    for idx, emb in zip(valid_indices, batch_embs, strict=False):
        results[idx] = emb

    return results

embed_fasta(path, model, seq_type=None, pooling_strategy='mean', obsm_key=None, **kwargs)

Embed all sequences from a FASTA or FASTQ file.

Supports plain and gzip-compressed files. The method auto-detects the file format from the extension and optionally infers whether sequences are DNA or protein when seq_type is not provided.

Parameters:

Name	Type	Description	Default
path	str | PathLike[str]	Path to a FASTA/FASTQ file. Recognised extensions: .fasta, .fa, .fna, .faa (FASTA), .fastq, .fq (FASTQ). A trailing .gz is handled transparently.	required
model	str	Model name registered with this embedder.	required
seq_type	Literal['dna', 'protein'] | None	"dna" or "protein". When None the type is resolved using a three-tier fallback: Infer from model -- if the model is a DNA or protein model the type is set accordingly. Auto-detect from sequences -- the first 100 sequences are inspected; if they consist exclusively of {A, C, G, T, N, U} the type is set to "dna", otherwise "protein". Fail -- a ValueError is raised asking the caller to specify seq_type explicitly.	None
pooling_strategy	str	Passed to the underlying model.	'mean'
obsm_key	str | None	Key used to store embeddings in adata.obsm. Defaults to "X_{model}".	None
**kwargs	Any	Forwarded to the model embed call.	{}

Returns:

Type	Description
AnnData	An :class:~anndata.AnnData object with: .obs -- sequence_id, sequence_length, description, seq_type .obsm[obsm_key] -- embedding matrix (n_sequences x embedding_dim) .uns -- model, pooling_strategy, source_file, n_sequences, n_failed

Source code in src/embpy/embedder.py

def embed_fasta(
    self,
    path: str | os.PathLike[str],
    model: str,
    seq_type: Literal["dna", "protein"] | None = None,
    pooling_strategy: str = "mean",
    obsm_key: str | None = None,
    **kwargs: Any,
):
    """Embed all sequences from a FASTA or FASTQ file.

    Supports plain and gzip-compressed files.  The method
    auto-detects the file format from the extension and optionally
    infers whether sequences are DNA or protein when *seq_type*
    is not provided.

    Parameters
    ----------
    path
        Path to a FASTA/FASTQ file.  Recognised extensions:
        ``.fasta``, ``.fa``, ``.fna``, ``.faa`` (FASTA),
        ``.fastq``, ``.fq`` (FASTQ).  A trailing ``.gz`` is
        handled transparently.
    model
        Model name registered with this embedder.
    seq_type
        ``"dna"`` or ``"protein"``.  When ``None`` the type is
        resolved using a three-tier fallback:

        1. **Infer from model** -- if the model is a DNA or
           protein model the type is set accordingly.
        2. **Auto-detect from sequences** -- the first 100
           sequences are inspected; if they consist exclusively
           of ``{A, C, G, T, N, U}`` the type is set to
           ``"dna"``, otherwise ``"protein"``.
        3. **Fail** -- a ``ValueError`` is raised asking the
           caller to specify *seq_type* explicitly.
    pooling_strategy
        Passed to the underlying model.
    obsm_key
        Key used to store embeddings in ``adata.obsm``.
        Defaults to ``"X_{model}"``.
    **kwargs
        Forwarded to the model ``embed`` call.

    Returns
    -------
    anndata.AnnData
        An :class:`~anndata.AnnData` object with:

        * ``.obs`` -- ``sequence_id``, ``sequence_length``,
          ``description``, ``seq_type``
        * ``.obsm[obsm_key]`` -- embedding matrix
          (n_sequences x embedding_dim)
        * ``.uns`` -- ``model``, ``pooling_strategy``,
          ``source_file``, ``n_sequences``, ``n_failed``
    """
    import gzip
    from pathlib import Path

    import anndata as ad
    from Bio import SeqIO

    filepath = Path(path)

    # -- detect file format -----------------------------------------
    is_gz = filepath.suffix.lower() == ".gz"
    stem = filepath.with_suffix("") if is_gz else filepath
    fmt = self._FASTA_EXTENSIONS.get(stem.suffix.lower())
    if fmt is None:
        raise ValueError(
            f"Unrecognised file extension {stem.suffix!r}.  Supported: {', '.join(sorted(self._FASTA_EXTENSIONS))}"
        )

    # -- parse sequences --------------------------------------------
    handle = gzip.open(filepath, "rt") if is_gz else open(filepath)
    try:
        records = list(SeqIO.parse(handle, fmt))
    finally:
        handle.close()

    if not records:
        raise ValueError(f"No sequences found in {filepath}")

    seq_ids: list[str] = []
    descriptions: list[str] = []
    sequences: list[str] = []
    for rec in records:
        seq_ids.append(rec.id)
        descriptions.append(rec.description)
        sequences.append(str(rec.seq))

    # -- resolve seq_type -------------------------------------------
    resolved_seq_type = self._resolve_seq_type(
        seq_type,
        model,
        sequences,
    )

    # -- embed sequences --------------------------------------------
    inst = self._get_model(model)
    embeddings: list[np.ndarray | None] = []
    n_failed = 0
    for i, seq in enumerate(sequences):
        try:
            emb = inst.embed(
                input=seq,
                pooling_strategy=pooling_strategy,
                **kwargs,
            )
            embeddings.append(emb)
        except Exception as e:  # noqa: BLE001
            logging.warning(
                "Failed to embed sequence %s (%s): %s",
                seq_ids[i],
                seq[:30] + "...",
                e,
            )
            embeddings.append(None)
            n_failed += 1

    # -- build AnnData ----------------------------------------------
    emb_dim = next(
        (e.shape[-1] for e in embeddings if e is not None),
        0,
    )
    emb_matrix = np.full(
        (len(sequences), emb_dim),
        np.nan,
        dtype=np.float32,
    )
    for i, emb in enumerate(embeddings):
        if emb is not None:
            emb_matrix[i] = emb

    import pandas as pd

    obs = pd.DataFrame(
        {
            "sequence_id": seq_ids,
            "sequence_length": [len(s) for s in sequences],
            "description": descriptions,
            "seq_type": resolved_seq_type,
        }
    )
    obs.index = obs["sequence_id"].astype(str)
    # ensure unique index
    if obs.index.duplicated().any():
        obs.index = pd.Index([f"{sid}_{i}" for i, sid in enumerate(obs.index)])

    key = obsm_key or f"X_{model}"
    adata = ad.AnnData(obs=obs)
    adata.obsm[key] = emb_matrix
    adata.uns["embed_fasta"] = {
        "model": model,
        "pooling_strategy": pooling_strategy,
        "source_file": str(filepath),
        "n_sequences": len(sequences),
        "n_failed": n_failed,
    }

    n_ok = len(sequences) - n_failed
    logging.info(
        "embed_fasta: embedded %d/%d sequences from %s (model=%s, seq_type=%s)",
        n_ok,
        len(sequences),
        filepath.name,
        model,
        resolved_seq_type,
    )
    return adata

embed_gene(identifier, model, id_type='symbol', organism='human', pooling_strategy='mean', region='full', protein_isoform='canonical', gene_description_format='Gene: {identifier}. Type: {id_type}. Organism: {organism}.', **kwargs)

Generates an embedding for a single gene using the specified model.

Determines whether to use sequence (DNA/protein) or text description based on the selected model's type.

Args: identifier (str): The gene identifier (e.g., "TP53", "ENSG00000141510") or a raw sequence when id_type="sequence". model (str): The user-facing name of the model to use (e.g., "enformer_human_rough", "esm2_650M", "minilm_l6_v2"). Must be a key in MODEL_REGISTRY. id_type (Literal): Type of the identifier provided. Use "sequence" to pass a raw DNA/protein/text string directly. Defaults to "symbol". organism (str): Organism name (e.g., "human", "mouse"). Used for sequence/description lookup. Defaults to "human". pooling_strategy (str): Pooling strategy if the model outputs per-token/residue embeddings. Defaults to "mean". Check model's available_pooling_strategies. region (Literal): Gene region to embed. "full" uses the complete genomic sequence (default), "exons" concatenates only exonic regions, and "introns" concatenates only intronic regions. Only applies to DNA models when id_type is "symbol" or "ensembl_id". protein_isoform (Literal): For protein models: "canonical" (default) embeds only the canonical UniProt sequence; "all" is not valid here -- use :meth:embed_protein_isoforms for all isoforms. gene_description_format (str): Format string used to generate input for text models. Defaults to "Gene: {identifier}. Type: {id_type}. Organism: {organism}.". **kwargs: Additional arguments passed to the specific model's embed method (e.g., target_layer for transformer models).

Returns:

Type	Description
np.ndarray: The computed gene embedding.

Raises:

Type	Description
ModelNotFoundError: If the requested model name is not registered or available.	IdentifierError: If the gene identifier cannot be resolved or required sequence/description fetched. ValueError: If the model type is ambiguous or incompatible inputs are generated. RuntimeError: If model loading or inference fails.

Source code in src/embpy/embedder.py

def embed_gene(
    self,
    identifier: str,
    model: str,
    id_type: Literal["symbol", "ensembl_id", "uniprot_id", "sequence"] = "symbol",
    organism: str = "human",
    pooling_strategy: str = "mean",
    region: Literal["full", "exons", "introns"] = "full",
    protein_isoform: Literal["canonical", "all"] = "canonical",
    gene_description_format: str = "Gene: {identifier}. Type: {id_type}. Organism: {organism}.",
    **kwargs: Any,
) -> np.ndarray:
    """
    Generates an embedding for a single gene using the specified model.

    Determines whether to use sequence (DNA/protein) or text description
    based on the selected model's type.

    Args:
        identifier (str): The gene identifier (e.g., "TP53", "ENSG00000141510")
            or a raw sequence when ``id_type="sequence"``.
        model (str): The user-facing name of the model to use (e.g., "enformer_human_rough",
                    "esm2_650M", "minilm_l6_v2"). Must be a key in MODEL_REGISTRY.
        id_type (Literal): Type of the identifier provided. Use ``"sequence"`` to
            pass a raw DNA/protein/text string directly. Defaults to "symbol".
        organism (str): Organism name (e.g., "human", "mouse"). Used for sequence/description lookup.
                        Defaults to "human".
        pooling_strategy (str): Pooling strategy if the model outputs per-token/residue embeddings.
                                Defaults to "mean". Check model's `available_pooling_strategies`.
        region (Literal): Gene region to embed. ``"full"`` uses the complete genomic
            sequence (default), ``"exons"`` concatenates only exonic regions, and
            ``"introns"`` concatenates only intronic regions. Only applies to DNA
            models when *id_type* is ``"symbol"`` or ``"ensembl_id"``.
        protein_isoform (Literal): For protein models: ``"canonical"`` (default)
            embeds only the canonical UniProt sequence; ``"all"`` is not valid
            here -- use :meth:`embed_protein_isoforms` for all isoforms.
        gene_description_format (str): Format string used to generate input for text models.
                                    Defaults to "Gene: {identifier}. Type: {id_type}. Organism: {organism}.".
        **kwargs: Additional arguments passed to the specific model's embed method
                (e.g., `target_layer` for transformer models).

    Returns
    -------
        np.ndarray: The computed gene embedding.

    Raises
    ------
        ModelNotFoundError: If the requested model name is not registered or available.
        IdentifierError: If the gene identifier cannot be resolved or required sequence/description fetched.
        ValueError: If the model type is ambiguous or incompatible inputs are generated.
        RuntimeError: If model loading or inference fails.
    """
    inst = self._get_model(model)
    mtype = inst.model_type

    if id_type != "sequence" and mtype == "protein":
        from embpy.resources.gene_resolver import detect_identifier_type

        detected = detect_identifier_type(identifier)
        if detected == "protein_sequence":
            logging.info(
                "Auto-detected raw protein sequence (id_type was %r); "
                "embedding directly. Pass id_type='sequence' explicitly "
                "to suppress this message.",
                id_type,
            )
            id_type = "sequence"

    if id_type == "sequence":
        input_data = identifier
    elif mtype == "dna":
        if id_type not in ("symbol", "ensembl_id"):
            raise IdentifierError(
                f"DNA models require id_type 'symbol', 'ensembl_id', or 'sequence', got '{id_type}'."
            )
        dna_id_type: Literal["symbol", "ensembl_id"] = id_type  # type: ignore[assignment]
        if region in ("exons", "introns"):
            seq = self.gene_resolver.get_gene_region_sequence(
                identifier,
                id_type=dna_id_type,
                organism=organism,
                region=region,
            )
        elif self.resolver_backend == "local":
            seq = self.gene_resolver.get_local_dna_sequence(identifier, dna_id_type)
        else:
            seq = self.gene_resolver.get_dna_sequence(identifier, dna_id_type, organism)
        if not seq:
            raise IdentifierError(f"DNA not found for {id_type}='{identifier}' (region={region})")
        input_data = seq
    elif mtype == "protein":
        prot = self.protein_resolver.get_canonical_sequence(identifier, id_type, organism)
        if not prot:
            raise IdentifierError(f"Protein not found for {id_type}='{identifier}'")
        input_data = prot
    elif mtype == "text":
        desc = self.gene_resolver.get_gene_description(
            identifier, id_type, organism, format_string=gene_description_format
        )
        if not desc:
            raise IdentifierError(f"Description not found for {id_type}='{identifier}'")
        input_data = desc
    elif mtype == "ppi":
        input_data = identifier
    else:
        raise ValueError(f"Unsupported model type '{mtype}' for embedding.")
    # Embed
    try:
        emb = inst.embed(input=input_data, pooling_strategy=pooling_strategy, **kwargs)
        return emb
    except Exception as e:
        logging.error(f"Embedding error for {identifier} with model {model}: {e}")
        raise RuntimeError(f"Embedding failed for {identifier}") from e

embed_genes_batch(model, identifiers=None, id_type='symbol', organism='human', pooling_strategy='mean', region='full', gene_description_format=None, fetch_all_dna=False, biotype='protein_coding', prefetched_sequences=None, **kwargs)

Generates embeddings for a batch of genes.

If identifiers is None and fetch_all_dna is True, it will automatically fetch ALL genes of the specified biotype and embed them.

Set region to "exons" or "introns" to embed only those gene regions (DNA models only).

Source code in src/embpy/embedder.py

def embed_genes_batch(
    self,
    model: str,
    identifiers: Sequence[str] | None = None,
    id_type: Literal["symbol", "ensembl_id", "uniprot_id", "sequence"] = "symbol",
    organism: str = "human",
    pooling_strategy: str = "mean",
    region: Literal["full", "exons", "introns"] = "full",
    gene_description_format: str | None = None,
    fetch_all_dna: bool = False,
    biotype: str = "protein_coding",
    prefetched_sequences: dict[str, str] | None = None,
    **kwargs: Any,
) -> list[np.ndarray | None]:
    """
    Generates embeddings for a batch of genes.

    If `identifiers` is None and `fetch_all_dna` is True, it will automatically
    fetch ALL genes of the specified `biotype` and embed them.

    Set ``region`` to ``"exons"`` or ``"introns"`` to embed only those
    gene regions (DNA models only).
    """
    inst = self._get_model(model)
    mtype = inst.model_type

    # Dictionary to hold pre-fetched sequences (if any)
    prefetched_data: dict[str, str] = dict(prefetched_sequences or {})
    if prefetched_data and identifiers is None:
        identifiers = list(prefetched_data.keys())
        id_type = "ensembl_id"
        logging.info("Using %d prefetched whole-genome sequences.", len(prefetched_data))

    # --- LOGIC CHANGE: Discovery Mode ---
    # If identifiers is None OR we explicitly want to pre-fetch
    if (fetch_all_dna or identifiers is None) and not prefetched_data:
        if self.resolver_backend == "api":
            logging.info(f"Fetching list of all '{biotype}' genes via API...")
            # This returns {ensembl_id: dna_sequence} or None
            fetched = self.gene_resolver.get_gene_sequences(biotype=biotype)
            if fetched is not None:
                prefetched_data = fetched

            if identifiers is None:
                # USE DISCOVERED GENES
                if prefetched_data:
                    identifiers = list(prefetched_data.keys())
                    logging.info(f"Auto-discovered {len(identifiers)} genes.")
                    # Force ID type to Ensembl ID as that's what get_gene_sequences returns
                    id_type = "ensembl_id"
                else:
                    logging.error(f"No genes found for biotype '{biotype}'.")
                    return []
        else:
            if identifiers is None:
                logging.error("Cannot auto-discover genes with 'local' backend. Provide identifiers.")
                return []
            logging.warning("fetch_all_dna is ignored in local mode.")

    if not identifiers:
        self.last_report = ResolutionReport(model_name=model, organism=organism)
        return []

    # Layer 2: fresh ResolutionReport for this batch. Every input
    # gets exactly one ``ResolutionRecord`` regardless of whether
    # it resolves, fails the resolver, fails the embedder, or hits
    # the OOM bisector. Replaces ``self.last_report`` so callers
    # can grab it after the call returns.
    report = ResolutionReport(model_name=model, organism=organism)
    self.last_report = report

    input_data_list: list[str | None] = []
    # Track per-input resolver state so we can attach reason codes
    # AFTER the embed_batch result is known (a resolver hit that
    # later fails to embed should be reported separately from a
    # resolver miss).
    per_input_state: list[dict[str, Any]] = []
    logging.info(f"Batch: {len(identifiers)} items for '{model}' ({mtype})...")

    # Narrow id_type for DNA resolver methods (only accept symbol/ensembl_id)
    dna_id: Literal["symbol", "ensembl_id"] | None = (
        id_type if id_type in ("symbol", "ensembl_id") else None  # type: ignore[assignment]
    )

    for ident in identifiers:
        data = None
        # Per-input audit state for Layer 2 reporting.
        state: dict[str, Any] = {
            "source": None,
            "reason": None,
            "attempted": [],
            "latency_ms": 0.0,
        }
        try:
            with time_block(
                "resolver_input",
                model=model,
                identifier=ident,
                id_type=id_type,
                mtype=mtype,
            ) as ev:
                if id_type == "sequence":
                    data = ident
                    state["source"] = "passthrough"
                elif mtype == "dna":
                    state["attempted"].append("gene_resolver_dna")
                    if region in ("exons", "introns") and dna_id is not None:
                        data = self.gene_resolver.get_gene_region_sequence(
                            ident,
                            id_type=dna_id,
                            organism=organism,
                            region=region,
                        )
                        state["source"] = f"gene_resolver_region:{region}"
                    else:
                        if prefetched_data:
                            if id_type == "ensembl_id":
                                data = prefetched_data.get(ident)
                                if data is not None:
                                    state["source"] = "prefetched"

                        if data is None:
                            if dna_id is None:
                                logging.warning(
                                    "DNA models require id_type 'symbol', 'ensembl_id', or 'sequence', "
                                    "got '%s'; skipping %s.",
                                    id_type,
                                    ident,
                                )
                                state["reason"] = f"resolver:unsupported_id_type:{id_type}"
                            elif self.resolver_backend == "local":
                                data = self.gene_resolver.get_local_dna_sequence(ident, dna_id)
                                state["source"] = "gene_resolver_local"
                            else:
                                data = self.gene_resolver.get_dna_sequence(ident, dna_id, organism)
                                state["source"] = "gene_resolver_api"

                elif mtype == "protein":
                    state["attempted"].append("protein_resolver")
                    data = self.protein_resolver.get_canonical_sequence(ident, id_type, organism)
                    state["source"] = "protein_resolver"

                elif mtype == "text":
                    # Default template combines the input identifier
                    # with MyGene.info fields. ``_SafeFormatDict`` in
                    # the resolver tolerates missing fields, so genes
                    # with no ``summary`` on MyGene still produce a
                    # thin but non-None description (e.g. just "Gene
                    # TP53 (human). TP53: tumor protein p53.").
                    state["attempted"].append("mygene")
                    fmt = gene_description_format or ("Gene {identifier} ({organism}). {symbol}: {name}. {summary}")
                    data = self.gene_resolver.get_gene_description(
                        ident,
                        id_type,
                        organism,
                        format_string=fmt,
                    )
                    state["source"] = "mygene_description"

                if data is None:
                    logging.warning(f"No data for {ident}; skipping.")
                    ev["status"] = "no_hit"
                    if state["reason"] is None:
                        state["reason"] = f"{mtype}_resolver:no_hit"
        except Exception as e:  # noqa: BLE001
            # We deliberately swallow resolver-side exceptions per
            # input (the resolver wraps its own errors in None
            # returns; anything still escaping here is unexpected
            # but should not nuke a 2000-gene batch). Layer 2
            # records the failure so callers can audit.
            logging.warning(f"Error fetching {ident}: {e}")
            state["reason"] = f"resolver_exception:{type(e).__name__}"
            log_event(
                "resolver_input_exception",
                level="warn",
                model=model,
                identifier=ident,
                error=type(e).__name__,
                error_msg=str(e)[:200],
            )
        input_data_list.append(data)
        per_input_state.append(state)

    valid_inputs = [d for d in input_data_list if d is not None]
    valid_indices = [i for i, d in enumerate(input_data_list) if d is not None]

    if not valid_inputs:
        # Every input failed the resolver -- record each one as
        # UNRESOLVED with the per-input reason captured above and
        # return early. This is the failure mode that used to be
        # silently swallowed under "BioEmbedder returned no
        # embeddings"; now the sidecar carries the actionable
        # detail.
        for ident, state in zip(identifiers, per_input_state, strict=False):
            report.record(
                ident,
                "unresolved",
                source=state["source"],
                reason=state["reason"] or "resolver:no_data",
                latency_ms=float(state["latency_ms"]),
                attempted_sources=list(state["attempted"]),
            )
        return [None] * len(identifiers)

    # Layer 3: wrap the model forward pass in an OOM-bisection
    # helper. If the wrapper already does its own batching (most
    # of them do, after the protein / DNA fixes), the outer call
    # succeeds and the wrapper is responsible for chunking. If
    # the wrapper raises an OOM-shaped exception, the bisector
    # halves the input list and retries, isolating the specific
    # input that genuinely cannot fit. The single-item case that
    # still OOMs propagates so it can be classified as
    # ContextOverflowError / ModelOOMError below.
    def _do_embed(batch: Sequence[str]) -> list[Any]:
        return list(
            inst.embed_batch(
                inputs=list(batch),
                pooling_strategy=pooling_strategy,
                **kwargs,
            )
        )

    try:
        with time_block(
            "model_forward",
            model=model,
            mtype=mtype,
            n_inputs=len(valid_inputs),
            device=str(self.device),
        ):
            batch_results = embed_batch_with_oom_recovery(
                _do_embed,
                valid_inputs,
                on_split=lambda left, right: log_event(
                    "oom_bisect",
                    level="warn",
                    model=model,
                    left=left,
                    right=right,
                ),
            )
    except EmbpyError:
        # Already a typed embpy exception (e.g. ModelOOMError raised by
        # a wrapper that classifies OOMs itself). Let it propagate
        # unchanged -- wrapping would only obscure the original fields.
        traceback.print_exc()
        raise
    except Exception as e:
        # Classify the failure so the upstream entry point can give a
        # category-specific exit code without parsing stderr. The
        # historical behaviour ("swallow + return [None]*N") masked
        # CUDA OOMs, context overflows, missing CUDA kernels, and
        # missing pip packages behind a single misleading "no
        # embeddings" error several layers up. We refuse to do that.
        logging.error(
            "Batch embed failed for model='%s' (n_inputs=%d): %s",
            model,
            len(valid_inputs),
            e,
        )
        traceback.print_exc()
        typed = _classify_embedder_exception(
            e,
            model_name=model,
            n_inputs=len(valid_inputs),
            device=str(self.device),
        )
        raise typed from e

    results: list[np.ndarray | None] = [None] * len(identifiers)
    for idx, emb in zip(valid_indices, batch_results, strict=False):
        results[idx] = emb

    # Layer 2: now that we know which inputs ended up with an
    # embedding (and which were resolver-misses), record one
    # ResolutionRecord per input. Embedder-side failures (OOM,
    # context overflow) raise typed errors above and never reach
    # this point, so anything in ``results[i] is None`` here is a
    # resolver miss with state already captured.
    valid_index_set = set(valid_indices)
    for i, (ident, state) in enumerate(zip(identifiers, per_input_state, strict=False)):
        if i in valid_index_set and results[i] is not None:
            report.record(
                ident,
                "resolved",
                source=state["source"],
                latency_ms=float(state["latency_ms"]),
                attempted_sources=list(state["attempted"]),
            )
        else:
            report.record(
                ident,
                "unresolved",
                source=state["source"],
                reason=state["reason"] or f"{mtype}_resolver:no_data",
                latency_ms=float(state["latency_ms"]),
                attempted_sources=list(state["attempted"]),
            )

    return results

embed_molecule(identifier, model, pooling_strategy='mean', **kwargs)

Generate an embedding for a single small molecule.

Validates that the input is a proper SMILES string before embedding.

Args: identifier (str): A SMILES string representing the molecule (e.g., "CCO" for ethanol). model (str): Key of the molecule embedding model to use (e.g., "chemberta_zinc_v1"). pooling_strategy (str, optional): Pooling strategy to aggregate token‐level embeddings. Defaults to "mean". **kwargs: Additional keyword arguments forwarded to the model’s embed method.

Returns:

Type	Description
np.ndarray:	A 1D NumPy array representing the molecule embedding.

Raises:

Type	Description
ModelNotFoundError:	If the specified model key is not registered. ValueError: If the chosen model is not of type "molecule", or if the SMILES string is invalid. RuntimeError: If the embedding process itself fails.

Source code in src/embpy/embedder.py

def embed_molecule(
    self,
    identifier: str,  # Expecting SMILES string
    model: str,  # e.g. "chemberta_zinc_v1"
    pooling_strategy: str = "mean",
    **kwargs: Any,
) -> np.ndarray:
    """
    Generate an embedding for a single small molecule.

    Validates that the input is a proper SMILES string before embedding.

    Args:
        identifier (str):
            A SMILES string representing the molecule (e.g., "CCO" for ethanol).
        model (str):
            Key of the molecule embedding model to use (e.g., "chemberta_zinc_v1").
        pooling_strategy (str, optional):
            Pooling strategy to aggregate token‐level embeddings.
            Defaults to "mean".
        **kwargs:
            Additional keyword arguments forwarded to the model’s `embed` method.

    Returns
    -------
        np.ndarray:
            A 1D NumPy array representing the molecule embedding.

    Raises
    ------
        ModelNotFoundError:
            If the specified model key is not registered.
        ValueError:
            If the chosen model is not of type "molecule", or if the SMILES string is invalid.
        RuntimeError:
            If the embedding process itself fails.
    """
    inst = self._get_model(model)
    if inst.model_type != "molecule":
        raise ValueError(f"Model '{model}' is not a molecule embedder.")

    smiles = identifier
    from rdkit import Chem
    # Validate SMILES
    if Chem.MolFromSmiles(smiles) is None:
        raise ValueError(f"Invalid SMILES string: '{smiles}'")

    logging.debug(f"Embedding molecule SMILES: {smiles} with model '{model}'")
    try:
        emb = inst.embed(input=smiles, pooling_strategy=pooling_strategy, **kwargs)
        logging.debug(f"Molecule embedding shape: {emb.shape}")
        return emb
    except Exception as e:
        logging.error(f"Embedding failed for SMILES '{smiles}': {e}")
        raise RuntimeError(f"Embedding failed for SMILES '{smiles}'") from e

embed_molecules_batch(identifiers, model, pooling_strategy='mean', **kwargs)

Generate embeddings for a batch of small molecules.

Invalid SMILES strings are skipped and return None in their positions.

Args: identifiers (Sequence[str]): A list of SMILES strings to embed. model (str): Key of the molecule embedding model to use. pooling_strategy (str, optional): Pooling strategy to apply. Defaults to "mean". **kwargs: Additional keyword arguments forwarded to embed_batch.

Returns:

Type	Description
List[Optional[np.ndarray]]:	A list of embeddings (NumPy arrays) for valid SMILES, with None for any invalid or failed entries. Order aligns with input list.

Raises:

Type	Description
ModelNotFoundError:	If the specified model key is not registered. ValueError: If the chosen model is not of type "molecule".

Source code in src/embpy/embedder.py

def embed_molecules_batch(
    self,
    identifiers: Sequence[str],
    model: str,
    pooling_strategy: str = "mean",
    **kwargs: Any,
) -> list[np.ndarray | None]:
    """
    Generate embeddings for a batch of small molecules.

    Invalid SMILES strings are skipped and return `None` in their positions.

    Args:
        identifiers (Sequence[str]):
            A list of SMILES strings to embed.
        model (str):
            Key of the molecule embedding model to use.
        pooling_strategy (str, optional):
            Pooling strategy to apply. Defaults to "mean".
        **kwargs:
            Additional keyword arguments forwarded to `embed_batch`.

    Returns
    -------
        List[Optional[np.ndarray]]:
            A list of embeddings (NumPy arrays) for valid SMILES, with `None`
            for any invalid or failed entries. Order aligns with input list.

    Raises
    ------
        ModelNotFoundError:
            If the specified model key is not registered.
        ValueError:
            If the chosen model is not of type "molecule".
    """
    inst = self._get_model(model)
    if inst.model_type != "molecule":
        raise ValueError(f"Model '{model}' is not a molecule embedder.")

    from rdkit import Chem
    valid_inputs: list[str] = []
    valid_indices: list[int] = []
    for idx, smi in enumerate(identifiers):
        if Chem.MolFromSmiles(smi) is not None:
            valid_inputs.append(smi)
            valid_indices.append(idx)
        else:
            logging.warning(f"Skipping invalid SMILES: '{smi}'")

    # Prepare output list
    results: list[np.ndarray | None] = [None] * len(identifiers)
    if not valid_inputs:
        logging.warning("No valid SMILES provided; returning all None.")
        return results

    logging.info(f"Embedding {len(valid_inputs)} valid SMILES with model '{model}'")
    try:
        # Wrappers disagree on the batch parameter name: most declare
        # ``inputs`` (base.py), while ChemBERTa declares ``input``. Try the
        # common spelling first and fall back, mirroring the generic embed
        # path -- otherwise every molecule model except ChemBERTa raises a
        # TypeError that is silently swallowed as "all inputs failed".
        try:
            batch_embs = inst.embed_batch(inputs=valid_inputs, pooling_strategy=pooling_strategy, **kwargs)
        except TypeError:
            batch_embs = inst.embed_batch(input=valid_inputs, pooling_strategy=pooling_strategy, **kwargs)
        for out_idx, emb in zip(valid_indices, batch_embs, strict=False):
            results[out_idx] = emb
    except Exception as e:  # noqa: BLE001
        logging.error(f"Batch embedding failed: {e}")
        # Leave all as None

    return results

embed_morphological(image, model='subcell_mae_rybg', pooling_strategy='attention_pool', **kwargs)

Embed a single microscopy image using a morphology model.

Parameters:

Name	Type	Description	Default
image	str | ndarray | Tensor	Path to a multi-channel image, or a (C, H, W) array/tensor.	required
model	str	Registered morphology model key (default "subcell_mae_rybg").	'subcell_mae_rybg'
pooling_strategy	str	Pooling strategy for the model ("cls", "mean", "attention_pool", "none").	'attention_pool'

Returns:

Type	Description
ndarray	1-D embedding vector.

Source code in src/embpy/embedder.py

def embed_morphological(
    self,
    image: str | np.ndarray | torch.Tensor,
    model: str = "subcell_mae_rybg",
    pooling_strategy: str = "attention_pool",
    **kwargs,
) -> np.ndarray:
    """Embed a single microscopy image using a morphology model.

    Parameters
    ----------
    image
        Path to a multi-channel image, or a ``(C, H, W)`` array/tensor.
    model
        Registered morphology model key (default ``"subcell_mae_rybg"``).
    pooling_strategy
        Pooling strategy for the model (``"cls"``, ``"mean"``,
        ``"attention_pool"``, ``"none"``).

    Returns
    -------
    np.ndarray
        1-D embedding vector.
    """
    wrapper = self._get_model(model)
    if getattr(wrapper, "model_type", None) != "morphology":
        raise ValueError(
            f"Model '{model}' is not a morphology embedder (type={getattr(wrapper, 'model_type', '?')})."
        )
    try:
        return wrapper.embed(image, pooling_strategy=pooling_strategy, **kwargs)
    except Exception as e:
        logging.error("Morphological embedding failed: %s", e)
        raise RuntimeError("Morphological embedding failed") from e

embed_morphological_batch(images, model='subcell_mae_rybg', pooling_strategy='attention_pool', **kwargs)

Embed a batch of microscopy images.

Parameters:

Name	Type	Description	Default
images	Sequence[str | ndarray | Tensor]	Iterable of paths or (C, H, W) arrays/tensors.	required
model	str	Registered morphology model key.	'subcell_mae_rybg'
pooling_strategy	str	Pooling strategy passed to the underlying model.	'attention_pool'

Returns:

Type	Description
list[ndarray]	One embedding per image.

Source code in src/embpy/embedder.py

def embed_morphological_batch(
    self,
    images: Sequence[str | np.ndarray | torch.Tensor],
    model: str = "subcell_mae_rybg",
    pooling_strategy: str = "attention_pool",
    **kwargs,
) -> list[np.ndarray]:
    """Embed a batch of microscopy images.

    Parameters
    ----------
    images
        Iterable of paths or ``(C, H, W)`` arrays/tensors.
    model
        Registered morphology model key.
    pooling_strategy
        Pooling strategy passed to the underlying model.

    Returns
    -------
    list[np.ndarray]
        One embedding per image.
    """
    wrapper = self._get_model(model)
    if getattr(wrapper, "model_type", None) != "morphology":
        raise ValueError(
            f"Model '{model}' is not a morphology embedder (type={getattr(wrapper, 'model_type', '?')})."
        )
    return wrapper.embed_batch(images, pooling_strategy=pooling_strategy, **kwargs)

embed_perturbation_morphology(perturbation, perturbation_type='genetic', dataset='jump', source='subcell', model='subcell_mae_rybg', pooling_strategy='attention_pool', local_dir=None, aggregate='mean', max_images=5, plate_type=None, jump_profiles_dir=None, verbose=True, **kwargs)

Embed a perturbation from JUMP or HPA by resolving images automatically.

The method handles identifier resolution (gene symbols via GeneResolver, compound names via DrugResolver), image fetching from CDN or local directories, preprocessing, and optional aggregation.

Parameters:

Name	Type	Description	Default
perturbation	str	Gene symbol, compound name, JCP2022 ID, or Ensembl ID.	required
perturbation_type	str	"genetic" (CRISPR/ORF) or "compound".	'genetic'
dataset	str	"jump" or "hpa".	'jump'
source	str	"subcell" -- run SubCell inference on images. "precomputed" -- load pre-computed JUMP CellProfiler profiles (only valid with dataset="jump").	'subcell'
model	str	Morphology model key (used when source="subcell").	'subcell_mae_rybg'
pooling_strategy	str	Pooling strategy for the SubCell model.	'attention_pool'
local_dir	str | None	Path to a local image directory. When images are already present they are loaded from disk; when they are fetched from the CDN they are cached here for future re-use. None = fetch from CDN without caching.	None
aggregate	str	"mean" averages embeddings across images; "none" returns the first image embedding only.	'mean'
max_images	int | None	Maximum number of images to download and embed. Defaults to 5 to avoid long downloads. Pass None to fetch all available images for the perturbation.	5
plate_type	str | None	JUMP plate type ("crispr", "orf", "compound"). Inferred from perturbation_type when None.	None
jump_profiles_dir	str | None	Directory containing JUMP profile parquets (for source="precomputed"). When None, defaults to data/embeddings/morphology_embeddings/JUMP/ within the project. If the parquet is not present it is automatically downloaded from the Cell Painting Gallery S3 bucket.	None
verbose	bool	If True (default), prints a single summary message at the end indicating success/failure and the data source used. Set to False to suppress all output.	True

Returns:

Type	Description
ndarray	Embedding vector (shape depends on model and aggregation).

Source code in src/embpy/embedder.py

def embed_perturbation_morphology(
    self,
    perturbation: str,
    perturbation_type: str = "genetic",
    dataset: str = "jump",
    source: str = "subcell",
    model: str = "subcell_mae_rybg",
    pooling_strategy: str = "attention_pool",
    local_dir: str | None = None,
    aggregate: str = "mean",
    max_images: int | None = 5,
    plate_type: str | None = None,
    jump_profiles_dir: str | None = None,
    verbose: bool = True,
    **kwargs,
) -> np.ndarray:
    """Embed a perturbation from JUMP or HPA by resolving images automatically.

    The method handles identifier resolution (gene symbols via
    ``GeneResolver``, compound names via ``DrugResolver``), image
    fetching from CDN or local directories, preprocessing, and
    optional aggregation.

    Parameters
    ----------
    perturbation
        Gene symbol, compound name, JCP2022 ID, or Ensembl ID.
    perturbation_type
        ``"genetic"`` (CRISPR/ORF) or ``"compound"``.
    dataset
        ``"jump"`` or ``"hpa"``.
    source
        ``"subcell"`` -- run SubCell inference on images.
        ``"precomputed"`` -- load pre-computed JUMP CellProfiler profiles
        (only valid with ``dataset="jump"``).
    model
        Morphology model key (used when ``source="subcell"``).
    pooling_strategy
        Pooling strategy for the SubCell model.
    local_dir
        Path to a local image directory.  When images are already
        present they are loaded from disk; when they are fetched
        from the CDN they are cached here for future re-use.
        ``None`` = fetch from CDN without caching.
    aggregate
        ``"mean"`` averages embeddings across images; ``"none"``
        returns the first image embedding only.
    max_images
        Maximum number of images to download and embed.  Defaults
        to ``5`` to avoid long downloads.  Pass ``None`` to fetch
        all available images for the perturbation.
    plate_type
        JUMP plate type (``"crispr"``, ``"orf"``, ``"compound"``).
        Inferred from *perturbation_type* when ``None``.
    jump_profiles_dir
        Directory containing JUMP profile parquets (for
        ``source="precomputed"``).  When ``None``, defaults to
        ``data/embeddings/morphology_embeddings/JUMP/`` within the
        project.  If the parquet is not present it is automatically
        downloaded from the Cell Painting Gallery S3 bucket.
    verbose
        If ``True`` (default), prints a single summary message at the
        end indicating success/failure and the data source used.
        Set to ``False`` to suppress all output.

    Returns
    -------
    np.ndarray
        Embedding vector (shape depends on model and aggregation).
    """
    if source == "precomputed":
        if dataset != "jump":
            raise ValueError("source='precomputed' is only valid with dataset='jump'")
        result = self._load_jump_precomputed(
            perturbation,
            perturbation_type,
            plate_type,
            jump_profiles_dir,
            aggregate,
        )
        if verbose:
            print(
                f"[embpy] {perturbation!r}: loaded precomputed JUMP profile (plate_type={plate_type or 'crispr'})"
            )
        return result

    # Use a resolution context to track steps without intermediate warnings
    resolution_ctx = _ResolutionContext()

    if dataset == "jump":
        images = self._resolve_jump_images(
            perturbation,
            perturbation_type,
            local_dir,
            plate_type,
            max_images,
            resolution_ctx,
        )
    elif dataset == "hpa":
        if perturbation_type == "compound":
            raise ValueError("perturbation_type='compound' is not supported for dataset='hpa'")
        images = self._resolve_hpa_images(perturbation, local_dir, max_images, resolution_ctx)
    else:
        raise ValueError(f"Unknown dataset: {dataset!r}")

    if not images:
        if verbose:
            # Print failure summary
            steps = resolution_ctx.steps
            attempted = ", ".join(steps) if steps else "direct lookup"
            print(f"[embpy] WARNING: {perturbation!r}: no images found (dataset={dataset}, tried: {attempted})")
        raise RuntimeError(
            f"No images found for perturbation={perturbation!r}, "
            f"dataset={dataset!r}, perturbation_type={perturbation_type!r}"
        )

    embeddings = self.embed_morphological_batch(
        images,
        model=model,
        pooling_strategy=pooling_strategy,
        **kwargs,
    )

    if verbose:
        # Print success summary
        source_desc = resolution_ctx.final_source or dataset
        n_images = len(images)
        print(f"[embpy] {perturbation!r}: embedded {n_images} image(s) (source: {source_desc})")

    if aggregate == "mean":
        return np.mean(embeddings, axis=0).astype(np.float32)
    return embeddings[0]

embed_perturbation_morphology_batch(perturbations, perturbation_type='genetic', dataset='jump', source='subcell', model='subcell_mae_rybg', pooling_strategy='attention_pool', local_dir=None, aggregate='mean', max_images=5, plate_type=None, jump_profiles_dir=None, verbose=True, skip_failures=True, n_workers=8, **kwargs)

Embed multiple perturbations from JUMP or HPA in batch.

This is the batched version of :meth:embed_perturbation_morphology. Image resolution / download for all perturbations runs in parallel across a thread pool, and then a single GPU inference batch is run on the combined image stack. This is dramatically faster than calling embed_perturbation_morphology in a Python loop, especially for HPA / JUMP CDN downloads which are network-bound.

Parameters:

Name	Type	Description	Default
perturbations	Sequence[str]	Sequence of gene symbols, compound names, JCP2022 IDs, or Ensembl IDs.	required
perturbation_type	str	"genetic" (CRISPR/ORF) or "compound".	'genetic'
dataset	str	"jump" or "hpa".	'jump'
source	str	"subcell" -- run SubCell inference on images. "precomputed" -- load pre-computed JUMP CellProfiler profiles (only valid with dataset="jump").	'subcell'
model	str	Morphology model key (used when source="subcell").	'subcell_mae_rybg'
pooling_strategy	str	Pooling strategy for the SubCell model.	'attention_pool'
local_dir	str | None	Path to a local image directory. When images are already present they are loaded from disk; when they are fetched from the CDN they are cached here for future re-use.	None
aggregate	str	"mean" averages embeddings across images per perturbation; "none" returns the first image embedding only.	'mean'
max_images	int | None	Maximum number of images to download and embed per perturbation.	5
plate_type	str | None	JUMP plate type ("crispr", "orf", "compound").	None
jump_profiles_dir	str | None	Directory containing JUMP profile parquets (for source="precomputed").	None
verbose	bool	If True (default), prints progress and a summary at the end.	True
skip_failures	bool	If True (default), perturbations that fail to resolve are skipped with a warning. If False, raises on first failure.	True
n_workers	int	Number of parallel worker threads for resolving / downloading images across perturbations. Defaults to 8.	8
**kwargs		Additional arguments passed to the embedding model.	{}

Returns:

Type	Description
tuple[ndarray, list[str]]	A tuple of (embeddings, labels) where: embeddings is a 2D array of shape (n_successful, embedding_dim) labels is a list of perturbation identifiers that succeeded, in the same order as the embedding rows.

Examples:

>>> genes = ["TP53", "BRCA1", "EGFR", "UNKNOWN_GENE"]
>>> embeddings, labels = embedder.embed_perturbation_morphology_batch(genes, dataset="hpa", max_images=3)
>>> print(f"Embedded {len(labels)} / {len(genes)} genes")
Embedded 3 / 4 genes
>>> print(embeddings.shape)
(3, 384)

Source code in src/embpy/embedder.py

def embed_perturbation_morphology_batch(
    self,
    perturbations: Sequence[str],
    perturbation_type: str = "genetic",
    dataset: str = "jump",
    source: str = "subcell",
    model: str = "subcell_mae_rybg",
    pooling_strategy: str = "attention_pool",
    local_dir: str | None = None,
    aggregate: str = "mean",
    max_images: int | None = 5,
    plate_type: str | None = None,
    jump_profiles_dir: str | None = None,
    verbose: bool = True,
    skip_failures: bool = True,
    n_workers: int = 8,
    **kwargs,
) -> tuple[np.ndarray, list[str]]:
    """Embed multiple perturbations from JUMP or HPA in batch.

    This is the batched version of :meth:`embed_perturbation_morphology`.
    Image resolution / download for all perturbations runs **in parallel**
    across a thread pool, and then a single GPU inference batch is run on
    the combined image stack.  This is dramatically faster than calling
    ``embed_perturbation_morphology`` in a Python loop, especially for
    HPA / JUMP CDN downloads which are network-bound.

    Parameters
    ----------
    perturbations
        Sequence of gene symbols, compound names, JCP2022 IDs, or Ensembl IDs.
    perturbation_type
        ``"genetic"`` (CRISPR/ORF) or ``"compound"``.
    dataset
        ``"jump"`` or ``"hpa"``.
    source
        ``"subcell"`` -- run SubCell inference on images.
        ``"precomputed"`` -- load pre-computed JUMP CellProfiler profiles
        (only valid with ``dataset="jump"``).
    model
        Morphology model key (used when ``source="subcell"``).
    pooling_strategy
        Pooling strategy for the SubCell model.
    local_dir
        Path to a local image directory.  When images are already
        present they are loaded from disk; when they are fetched
        from the CDN they are cached here for future re-use.
    aggregate
        ``"mean"`` averages embeddings across images per perturbation;
        ``"none"`` returns the first image embedding only.
    max_images
        Maximum number of images to download and embed per perturbation.
    plate_type
        JUMP plate type (``"crispr"``, ``"orf"``, ``"compound"``).
    jump_profiles_dir
        Directory containing JUMP profile parquets (for ``source="precomputed"``).
    verbose
        If ``True`` (default), prints progress and a summary at the end.
    skip_failures
        If ``True`` (default), perturbations that fail to resolve are
        skipped with a warning. If ``False``, raises on first failure.
    n_workers
        Number of parallel worker threads for resolving / downloading
        images across perturbations.  Defaults to ``8``.
    **kwargs
        Additional arguments passed to the embedding model.

    Returns
    -------
    tuple[np.ndarray, list[str]]
        A tuple of ``(embeddings, labels)`` where:

        - ``embeddings`` is a 2D array of shape ``(n_successful, embedding_dim)``
        - ``labels`` is a list of perturbation identifiers that succeeded,
          in the same order as the embedding rows.

    Examples
    --------
    >>> genes = ["TP53", "BRCA1", "EGFR", "UNKNOWN_GENE"]
    >>> embeddings, labels = embedder.embed_perturbation_morphology_batch(genes, dataset="hpa", max_images=3)
    >>> print(f"Embedded {len(labels)} / {len(genes)} genes")
    Embedded 3 / 4 genes
    >>> print(embeddings.shape)
    (3, 384)
    """
    import time
    from concurrent.futures import ThreadPoolExecutor, as_completed

    # ---- precomputed JUMP profiles: I/O-bound parquet reads only ----
    if source == "precomputed":
        return self._embed_perturbation_precomputed_batch(
            perturbations,
            perturbation_type=perturbation_type,
            plate_type=plate_type,
            jump_profiles_dir=jump_profiles_dir,
            aggregate=aggregate,
            verbose=verbose,
            skip_failures=skip_failures,
        )

    n_total = len(perturbations)
    t0 = time.time()
    hpa_catalog = None
    if dataset == "hpa":
        hpa_catalog = self._build_hpa_batch_catalog(
            perturbations,
            local_dir=local_dir,
            verbose=verbose,
        )

    # ── Stage 1: resolve images for ALL perturbations in parallel ──
    if verbose:
        print(
            f"[embpy] Stage 1/2: resolving images for {n_total} perturbations using {n_workers} parallel workers..."
        )

    def _resolve_one(pert: str) -> tuple[str, list[np.ndarray], str | None, Exception | None]:
        ctx = _ResolutionContext()
        try:
            if dataset == "jump":
                imgs = self._resolve_jump_images(
                    pert,
                    perturbation_type,
                    local_dir,
                    plate_type,
                    max_images,
                    ctx,
                )
            elif dataset == "hpa":
                if perturbation_type == "compound":
                    raise ValueError("perturbation_type='compound' is not supported for dataset='hpa'")
                imgs = self._resolve_hpa_images(
                    pert,
                    local_dir,
                    max_images,
                    ctx,
                    hpa_catalog=hpa_catalog,
                )
            else:
                raise ValueError(f"Unknown dataset: {dataset!r}")
            return pert, imgs, ctx.final_source, None
        except Exception as exc:  # noqa: BLE001
            return pert, [], ctx.final_source, exc

    per_gene_images: dict[str, list[np.ndarray]] = {}
    per_gene_source: dict[str, str | None] = {}
    failed_labels: list[str] = []
    first_exc: Exception | None = None

    with ThreadPoolExecutor(max_workers=n_workers) as pool:
        futures = {pool.submit(_resolve_one, p): p for p in perturbations}
        for i, fut in enumerate(as_completed(futures), start=1):
            pert, imgs, src, exc = fut.result()
            if exc is not None or not imgs:
                failed_labels.append(pert)
                if first_exc is None and exc is not None:
                    first_exc = exc
                if verbose:
                    reason = str(exc) if exc else "no images found"
                    print(f"  [{i}/{n_total}] {pert!r:20s} FAILED ({reason})")
            else:
                per_gene_images[pert] = imgs
                per_gene_source[pert] = src
                if verbose:
                    print(f"  [{i}/{n_total}] {pert!r:20s} {len(imgs)} image(s) from {src}")

    if not skip_failures and first_exc is not None:
        raise first_exc

    if not per_gene_images:
        if verbose:
            print("[embpy] No images resolved for any perturbation")
        return np.array([]).reshape(0, 0), []

    # ── Stage 2: stack all images and run a single GPU inference ──
    successful_labels = [p for p in perturbations if p in per_gene_images]
    all_images: list[np.ndarray] = []
    boundaries: list[int] = [0]
    for pert in successful_labels:
        imgs = per_gene_images[pert]
        all_images.extend(imgs)
        boundaries.append(boundaries[-1] + len(imgs))

    if verbose:
        t_resolve = time.time() - t0
        print(
            f"[embpy] Stage 2/2: running inference on "
            f"{len(all_images)} images (image resolution took "
            f"{t_resolve:.1f}s)..."
        )

    t1 = time.time()
    all_embeddings = self.embed_morphological_batch(
        all_images,
        model=model,
        pooling_strategy=pooling_strategy,
        **kwargs,
    )
    all_embeddings = np.asarray(all_embeddings)

    # Aggregate per gene using boundary slices
    per_gene_emb: list[np.ndarray] = []
    for j in range(len(successful_labels)):
        start, end = boundaries[j], boundaries[j + 1]
        sub = all_embeddings[start:end]
        if aggregate == "mean":
            per_gene_emb.append(np.mean(sub, axis=0).astype(np.float32))
        else:
            per_gene_emb.append(sub[0])

    if verbose:
        t_infer = time.time() - t1
        t_total = time.time() - t0
        print(
            f"[embpy] Done in {t_total:.1f}s "
            f"(resolve={t_resolve:.1f}s, inference={t_infer:.1f}s) -- "
            f"{len(successful_labels)}/{n_total} perturbations embedded"
        )
        if failed_labels:
            print(f"[embpy] Failed: {failed_labels}")

    return np.stack(per_gene_emb, axis=0), successful_labels

embed_protein(identifier, model, id_type='symbol', organism='human', pooling_strategy='mean', isoform='canonical', **kwargs)

Embed a protein sequence using a protein language model.

When isoform="canonical", returns a single embedding of the canonical UniProt sequence. When isoform="all", returns a dict mapping isoform accession IDs to their embeddings.

Parameters:

Name	Type	Description	Default
identifier	str	Gene symbol, Ensembl ID, UniProt accession, or raw amino acid sequence (when id_type="sequence").	required
model	str	Protein model name (e.g. "esm2_650M", "prot_t5_xl").	required
id_type	Literal['symbol', 'ensembl_id', 'uniprot_id', 'sequence']	Type of identifier.	'symbol'
organism	str	Organism name (default "human").	'human'
pooling_strategy	str	Pooling strategy ("mean", "max", "cls").	'mean'
isoform	Literal['canonical', 'all']	"canonical" for the canonical sequence only, "all" to embed every UniProt isoform.	'canonical'
**kwargs	Any	Forwarded to the model's embed method.	{}

Returns:

Type	Description
ndarray	When isoform="canonical".
dict[str, ndarray]	When isoform="all" -- maps isoform accession to embedding.

Source code in src/embpy/embedder.py

def embed_protein(
    self,
    identifier: str,
    model: str,
    id_type: Literal["symbol", "ensembl_id", "uniprot_id", "sequence"] = "symbol",
    organism: str = "human",
    pooling_strategy: str = "mean",
    isoform: Literal["canonical", "all"] = "canonical",
    **kwargs: Any,
) -> np.ndarray | dict[str, np.ndarray]:
    """Embed a protein sequence using a protein language model.

    When ``isoform="canonical"``, returns a single embedding of the
    canonical UniProt sequence.  When ``isoform="all"``, returns a
    dict mapping isoform accession IDs to their embeddings.

    Parameters
    ----------
    identifier
        Gene symbol, Ensembl ID, UniProt accession, or raw amino
        acid sequence (when ``id_type="sequence"``).
    model
        Protein model name (e.g. ``"esm2_650M"``, ``"prot_t5_xl"``).
    id_type
        Type of identifier.
    organism
        Organism name (default ``"human"``).
    pooling_strategy
        Pooling strategy (``"mean"``, ``"max"``, ``"cls"``).
    isoform
        ``"canonical"`` for the canonical sequence only, ``"all"``
        to embed every UniProt isoform.
    **kwargs
        Forwarded to the model's ``embed`` method.

    Returns
    -------
    np.ndarray
        When ``isoform="canonical"``.
    dict[str, np.ndarray]
        When ``isoform="all"`` -- maps isoform accession to embedding.
    """
    inst = self._get_model(model)
    if inst.model_type != "protein":
        raise ValueError(f"Model '{model}' is not a protein model (type={inst.model_type}).")

    if id_type == "sequence":
        emb = inst.embed(input=identifier, pooling_strategy=pooling_strategy, **kwargs)
        return emb

    if isoform == "canonical":
        seq = self.protein_resolver.get_canonical_sequence(identifier, id_type, organism)
        if not seq:
            raise IdentifierError(f"Canonical protein not found for {id_type}='{identifier}'")
        return inst.embed(input=seq, pooling_strategy=pooling_strategy, **kwargs)

    # isoform == "all"
    isoforms = self.protein_resolver.get_isoforms(
        identifier,
        id_type,
        organism,
        include_canonical=True,
    )
    if not isoforms:
        raise IdentifierError(f"No isoforms found for {id_type}='{identifier}'")

    results: dict[str, np.ndarray] = {}
    for iso_id, seq in isoforms.items():
        try:
            results[iso_id] = inst.embed(
                input=seq,
                pooling_strategy=pooling_strategy,
                **kwargs,
            )
        except Exception as e:  # noqa: BLE001
            logging.warning(f"Failed to embed isoform {iso_id}: {e}")
    return results

embed_proteins_batch(identifiers, model, id_type='symbol', organism='human', pooling_strategy='mean', isoform='canonical', **kwargs)

Embed proteins for a batch of gene identifiers.

Parameters:

Name	Type	Description	Default
identifiers	list[str]	List of gene identifiers or raw amino-acid sequences (when id_type="sequence").	required
model	str	Protein model name.	required
id_type	Literal['symbol', 'ensembl_id', 'uniprot_id', 'sequence']	Type of identifiers. Use "sequence" to embed raw amino-acid strings directly without resolver lookup.	'symbol'
organism	str	Organism name.	'human'
pooling_strategy	str	Pooling strategy.	'mean'
isoform	Literal['canonical', 'all']	"canonical" or "all". Ignored when id_type="sequence".	'canonical'
**kwargs	Any	Forwarded to the model.	{}

Returns:

Type	Description
dict[str, ndarray]	When isoform="canonical" or id_type="sequence" -- maps identifier (or sequence) to embedding.
dict[str, dict[str, ndarray]]	When isoform="all" -- maps identifier to {isoform_accession: embedding}.

Source code in src/embpy/embedder.py

def embed_proteins_batch(
    self,
    identifiers: list[str],
    model: str,
    id_type: Literal["symbol", "ensembl_id", "uniprot_id", "sequence"] = "symbol",
    organism: str = "human",
    pooling_strategy: str = "mean",
    isoform: Literal["canonical", "all"] = "canonical",
    **kwargs: Any,
) -> dict[str, np.ndarray] | dict[str, dict[str, np.ndarray]]:
    """Embed proteins for a batch of gene identifiers.

    Parameters
    ----------
    identifiers
        List of gene identifiers or raw amino-acid sequences
        (when ``id_type="sequence"``).
    model
        Protein model name.
    id_type
        Type of identifiers.  Use ``"sequence"`` to embed raw
        amino-acid strings directly without resolver lookup.
    organism
        Organism name.
    pooling_strategy
        Pooling strategy.
    isoform
        ``"canonical"`` or ``"all"``.  Ignored when
        ``id_type="sequence"``.
    **kwargs
        Forwarded to the model.

    Returns
    -------
    dict[str, np.ndarray]
        When ``isoform="canonical"`` or ``id_type="sequence"``
        -- maps identifier (or sequence) to embedding.
    dict[str, dict[str, np.ndarray]]
        When ``isoform="all"`` -- maps identifier to
        ``{isoform_accession: embedding}``.
    """
    inst = self._get_model(model)
    if inst.model_type != "protein":
        raise ValueError(f"Model '{model}' is not a protein model.")

    if id_type == "sequence":
        results: dict[str, np.ndarray] = {}
        for seq in identifiers:
            try:
                results[seq] = inst.embed(
                    input=seq,
                    pooling_strategy=pooling_strategy,
                    **kwargs,
                )
            except Exception as e:  # noqa: BLE001
                logging.warning(f"Failed to embed protein sequence: {e}")
        return results

    if isoform == "canonical":
        seqs = self.protein_resolver.get_canonical_sequences_batch(
            identifiers,
            id_type,
            organism,
        )
        results = {}
        for ident, seq in seqs.items():
            try:
                results[ident] = inst.embed(
                    input=seq,
                    pooling_strategy=pooling_strategy,
                    **kwargs,
                )
            except Exception as e:  # noqa: BLE001
                logging.warning(f"Failed to embed protein for {ident}: {e}")
        return results

    # isoform == "all"
    all_isoforms = self.protein_resolver.get_isoforms_batch(
        identifiers,
        id_type,
        organism,
        include_canonical=True,
    )
    results_iso: dict[str, dict[str, np.ndarray]] = {}
    for ident, isoforms_map in all_isoforms.items():
        results_iso[ident] = {}
        for iso_id, seq in isoforms_map.items():
            try:
                results_iso[ident][iso_id] = inst.embed(
                    input=seq,
                    pooling_strategy=pooling_strategy,
                    **kwargs,
                )
            except Exception as e:  # noqa: BLE001
                logging.warning(f"Failed to embed isoform {iso_id} for {ident}: {e}")
    return results_iso

embed_text(text, model, pooling_strategy='mean', **kwargs)

Generates an embedding for an arbitrary text string using a text model.

Args: text (str): The input text string. model (str): The name of the text embedding model. Can be either: - A predefined model name (e.g., "minilm_l6_v2", "bert_base_uncased") - Any Hugging Face model identifier (e.g., "sentence-transformers/all-MiniLM-L6-v2") pooling_strategy (str): Pooling strategy. Defaults to "mean". **kwargs: Additional arguments for the text model's embed method.

Returns:

Type	Description
np.ndarray: The computed text embedding.

Raises:

Type	Description
ModelNotFoundError: If the model cannot be loaded from HF or found in registry.	ValueError: If the loaded model is not a text model. RuntimeError: If model loading or inference fails.

Source code in src/embpy/embedder.py

def embed_text(
    self,
    text: str,
    model: str,  # Can be registry name OR any HF model identifier
    pooling_strategy: str = "mean",
    **kwargs: Any,
) -> np.ndarray:
    """
    Generates an embedding for an arbitrary text string using a text model.

    Args:
        text (str): The input text string.
        model (str): The name of the text embedding model. Can be either:
                    - A predefined model name (e.g., "minilm_l6_v2", "bert_base_uncased")
                    - Any Hugging Face model identifier (e.g., "sentence-transformers/all-MiniLM-L6-v2")
        pooling_strategy (str): Pooling strategy. Defaults to "mean".
        **kwargs: Additional arguments for the text model's embed method.

    Returns
    -------
        np.ndarray: The computed text embedding.

    Raises
    ------
        ModelNotFoundError: If the model cannot be loaded from HF or found in registry.
        ValueError: If the loaded model is not a text model.
        RuntimeError: If model loading or inference fails.
    """
    model_instance = self._get_model(model)
    if model_instance.model_type != "text":
        raise ValueError(f"Model '{model}' is not a text embedder. Use embed_gene or embed_molecule.")

    logging.debug(f"Embedding text: '{text[:100]}...' using model '{model}'")
    try:
        embedding = model_instance.embed(input=text, pooling_strategy=pooling_strategy, **kwargs)
        logging.debug(f"Text embedding generated with shape: {embedding.shape}")
        return embedding
    except Exception as e:
        logging.error(f"Error during text embedding generation for model '{model}': {e}")
        raise RuntimeError(f"Text embedding failed for input '{text[:50]}...'.") from e

embed_text_api(text, model='text-embedding-3-small', provider='openai', api_key=None, base_url=None, **kwargs)

Embed text using an API-based model with explicit provider config.

Convenience method for one-off API calls where you want to specify the provider, API key, and/or base URL directly without pre-registering the model.

Parameters:

Name	Type	Description	Default
text	str	Text string to embed.	required
model	str	Model name as expected by the API.	'text-embedding-3-small'
provider	str	"openai", "cohere", "voyage", "google", or "generic" (OpenAI-compatible endpoint).	'openai'
api_key	str | None	API key (or set the provider's env var).	None
base_url	str | None	Custom API endpoint URL.	None
**kwargs	Any	Forwarded to the API wrapper.	{}

Returns:

Type	Description
ndarray	Embedding vector.

Source code in src/embpy/embedder.py

def embed_text_api(
    self,
    text: str,
    model: str = "text-embedding-3-small",
    provider: str = "openai",
    api_key: str | None = None,
    base_url: str | None = None,
    **kwargs: Any,
) -> np.ndarray:
    """Embed text using an API-based model with explicit provider config.

    Convenience method for one-off API calls where you want to
    specify the provider, API key, and/or base URL directly
    without pre-registering the model.

    Parameters
    ----------
    text
        Text string to embed.
    model
        Model name as expected by the API.
    provider
        ``"openai"``, ``"cohere"``, ``"voyage"``, ``"google"``,
        or ``"generic"`` (OpenAI-compatible endpoint).
    api_key
        API key (or set the provider's env var).
    base_url
        Custom API endpoint URL.
    **kwargs
        Forwarded to the API wrapper.

    Returns
    -------
    np.ndarray
        Embedding vector.
    """
    from .models.api_models import APIEmbeddingWrapper

    cache_key = f"_api_{provider}_{model}"
    if cache_key not in self.model_cache:
        wrapper = APIEmbeddingWrapper(
            model_path_or_name=model,
            provider=provider,
            api_key=api_key,
            base_url=base_url,
        )
        wrapper.load(self.device)
        self.model_cache[cache_key] = wrapper

    return self.model_cache[cache_key].embed(input=text, **kwargs)

embed_texts_batch(texts, model, pooling_strategy='mean', batch_size=None, **kwargs)

Generates embeddings for a batch of arbitrary text strings.

Args: texts (Sequence[str]): A list or tuple of text strings. model (str): The name of the text embedding model. pooling_strategy (str): Pooling strategy. Defaults to "mean". batch_size (Optional[int]): Maximum batch size for processing. If None, processes all texts at once. Use smaller values to avoid OOM errors with large datasets. **kwargs: Additional arguments for the model's embed_batch method.

Returns:

Type	Description
list[Optional[np.ndarray]]: List of embeddings, with None for texts that failed.

Source code in src/embpy/embedder.py

def embed_texts_batch(
    self,
    texts: Sequence[str],
    model: str,
    pooling_strategy: str = "mean",
    batch_size: int | None = None,
    **kwargs: Any,
) -> list[np.ndarray | None]:  # Return None for errors
    """
    Generates embeddings for a batch of arbitrary text strings.

    Args:
        texts (Sequence[str]): A list or tuple of text strings.
        model (str): The name of the text embedding model.
        pooling_strategy (str): Pooling strategy. Defaults to "mean".
        batch_size (Optional[int]): Maximum batch size for processing. If None, processes all texts at once.
                                   Use smaller values to avoid OOM errors with large datasets.
        **kwargs: Additional arguments for the model's embed_batch method.

    Returns
    -------
        list[Optional[np.ndarray]]: List of embeddings, with None for texts that failed.
    """
    model_instance = self._get_model(model)
    if model_instance.model_type != "text":
        raise ValueError(f"Model '{model}' is not a text embedder.")

    valid_inputs = list(texts)
    logging.info(f"Embedding batch of {len(valid_inputs)} texts using model '{model}'...")

    try:
        batch_results = model_instance.embed_batch(
            inputs=valid_inputs,
            pooling_strategy=pooling_strategy,
            batch_size=batch_size,
            **kwargs,
        )
        if len(batch_results) != len(valid_inputs):
            logging.error(
                f"Batch embedding returned {len(batch_results)} results for {len(valid_inputs)} inputs. Mismatch!"
            )
            return [None] * len(texts)
        logging.info("Batch text embedding successful.")
        results: list[np.ndarray | None] = list(batch_results)

    except (ValueError, KeyError) as e:
        logging.error(f"Error during batch text embedding generation for model '{model}': {e}")
        results = [None] * len(texts)

    return results

generate_cells(base_adata, test_adata, *, model='stack', wrapper=None, **generate_kwargs)

Run in-context generation to synthesise new cell profiles.

Currently only supported by the Stack wrapper. Unlike :meth:decode_cells, this does NOT consume a latent matrix: it conditions on a donor-specific base_adata context and synthesises gene-expression for every cell in test_adata.

Parameters:

Name	Type	Description	Default
base_adata	AnnData	Donor-conditioned base context.	required
test_adata	AnnData	Cells to predict for.	required
model	str	Registry key (currently only "stack").	'stack'
wrapper	SingleCellWrapper	Already-constructed wrapper. Required for Stack because its checkpoint and gene-list paths cannot be expressed through embed_cells.	None
**generate_kwargs	Any	Forwarded to the wrapper's generate_cells (split_column, split_values, prompt_ratio, ...).	{}

Source code in src/embpy/embedder.py

def generate_cells(
    self,
    base_adata,
    test_adata,
    *,
    model: str = "stack",
    wrapper=None,
    **generate_kwargs: Any,
):
    """Run in-context generation to synthesise new cell profiles.

    Currently only supported by the Stack wrapper. Unlike
    :meth:`decode_cells`, this does NOT consume a latent matrix: it
    conditions on a donor-specific ``base_adata`` context and
    synthesises gene-expression for every cell in ``test_adata``.

    Parameters
    ----------
    base_adata : anndata.AnnData
        Donor-conditioned base context.
    test_adata : anndata.AnnData
        Cells to predict for.
    model : str
        Registry key (currently only ``"stack"``).
    wrapper : SingleCellWrapper, optional
        Already-constructed wrapper. Required for Stack because its
        checkpoint and gene-list paths cannot be expressed through
        ``embed_cells``.
    **generate_kwargs
        Forwarded to the wrapper's ``generate_cells`` (split_column,
        split_values, prompt_ratio, ...).
    """
    from .models.singlecell_models import singlecell_info

    card = singlecell_info(model)
    if not card.supports_generation:
        raise ValueError(f"Model '{model}' does not expose an in-context generation head. Supported: stack.")

    if wrapper is None:
        device_str = str(self.device)
        wrapper = self._singlecell_cache.get((model, device_str))
    if wrapper is None:
        raise RuntimeError(
            f"No '{model}' wrapper available. For Stack, instantiate "
            "it manually with `StackWrapper(checkpoint=..., "
            "genelist=...)`, call `.load(device)`, and pass it via "
            "the `wrapper=` kwarg."
        )

    return wrapper.generate_cells(base_adata, test_adata, **generate_kwargs)

get_model(model, *, load=True)

Return the model wrapper behind a model key, for direct introspection.

embed() covers the common path, but the wrappers expose more than a pooled vector -- per-layer hidden states, per-layer attention, and pooled embeddings for every layer at once. Reaching those used to require the private _get_model; this is the supported way.

Parameters:

Name	Type	Description	Default
model	str	A key from :meth:list_available_models.	required
load	bool	Load the weights onto this embedder's device if they are not loaded yet. Pass False to inspect class-level attributes (such as has_attention) without paying for a download.	True

Returns:

Type	Description
BaseModelWrapper	The wrapper, loaded unless load=False. Wrappers are cached, so repeated calls return the same instance.

Examples:

>>> embedder = BioEmbedder(device="cpu")
>>> wrapper = embedder.get_model("esm2_8M")
>>> ids = wrapper.tokenizer("MTEYKLVVVG", return_tensors="pt")["input_ids"]
>>> attn = wrapper.extract_attention(ids, layers=[-1])

See Also

embpy.models.base.BaseModelWrapper.extract_attention embpy.models.base.BaseModelWrapper.extract_hidden_states embpy.models.base.BaseModelWrapper.embed_all_layers

Source code in src/embpy/embedder.py

def get_model(self, model: str, *, load: bool = True) -> Any:
    """Return the model wrapper behind a model key, for direct introspection.

    ``embed()`` covers the common path, but the wrappers expose more than a
    pooled vector -- per-layer hidden states, per-layer attention, and pooled
    embeddings for every layer at once. Reaching those used to require the
    private ``_get_model``; this is the supported way.

    Parameters
    ----------
    model
        A key from :meth:`list_available_models`.
    load
        Load the weights onto this embedder's device if they are not loaded
        yet. Pass ``False`` to inspect class-level attributes (such as
        ``has_attention``) without paying for a download.

    Returns
    -------
    BaseModelWrapper
        The wrapper, loaded unless ``load=False``. Wrappers are cached, so
        repeated calls return the same instance.

    Examples
    --------
    >>> embedder = BioEmbedder(device="cpu")               # doctest: +SKIP
    >>> wrapper = embedder.get_model("esm2_8M")            # doctest: +SKIP
    >>> ids = wrapper.tokenizer("MTEYKLVVVG", return_tensors="pt")["input_ids"]
    >>> attn = wrapper.extract_attention(ids, layers=[-1])  # doctest: +SKIP

    See Also
    --------
    embpy.models.base.BaseModelWrapper.extract_attention
    embpy.models.base.BaseModelWrapper.extract_hidden_states
    embpy.models.base.BaseModelWrapper.embed_all_layers
    """
    inst = self._get_model(model, load=load)
    if load and getattr(inst, "model", None) is None:
        inst.load(self.device)
    return inst

list_available_models(category='all')

Return available model names, optionally filtered by category.

Parameters:

Name	Type	Description	Default
category	Literal['all', 'dna', 'protein', 'molecule', 'text', 'morphology', 'single_cell', 'static']	"all" returns every model (sequence + single-cell). "dna", "protein", "molecule", "text", "morphology" filter the sequence/structure model registry. "single_cell" returns single-cell foundation model keys. "static" returns fixed lookup embedding tables shipped via the default Hugging Face data repository.	'all'

Source code in src/embpy/embedder.py

def list_available_models(
    self,
    category: Literal[
        "all",
        "dna",
        "protein",
        "molecule",
        "text",
        "morphology",
        "single_cell",
        "static",
    ] = "all",
) -> list[str]:
    """Return available model names, optionally filtered by category.

    Parameters
    ----------
    category
        ``"all"`` returns every model (sequence + single-cell).
        ``"dna"``, ``"protein"``, ``"molecule"``, ``"text"``,
        ``"morphology"`` filter the sequence/structure model registry.
        ``"single_cell"`` returns single-cell foundation model keys.
        ``"static"`` returns fixed lookup embedding tables shipped via
        the default Hugging Face data repository.
    """
    from .models.singlecell_models import list_singlecell_models

    if category == "static":
        return sorted(DEFAULT_STATIC_EMBEDDING_MODELS)

    if category == "single_cell":
        return sorted(list_singlecell_models())

    if category == "all":
        seq_models = sorted(self._available_models.keys())
        sc_models = sorted(list_singlecell_models())
        static_models = sorted(DEFAULT_STATIC_EMBEDDING_MODELS)
        return sorted(set(seq_models + sc_models + static_models))

    result = []
    for name, (wrapper_cls, _) in self._available_models.items():
        if hasattr(wrapper_cls, "model_type") and wrapper_cls.model_type == category:
            result.append(name)
    return sorted(result)

model_catalog(category='all', *, summary=False)

Return the available models as a table.

The tabular counterpart to :meth:list_available_models, which returns bare strings. In a notebook this renders as a proper table rather than printed text, and being a DataFrame it can be filtered, sorted and exported.

Parameters:

Name	Type	Description	Default
category	str	Restrict to one family ("dna", "protein", "molecule", "text", "morphology", "single_cell", "static"), or "all" for everything.	'all'
summary	bool	False (default) gives one row per model with its family. True gives one row per family with a count and a few example keys -- the quick "what can this package do?" view.	False

Returns:

Type	Description
DataFrame	With summary=False: columns model, family. With summary=True: columns family, n_models, examples.

Examples:

>>> embedder.model_catalog(summary=True)
>>> embedder.model_catalog("protein").head()

Source code in src/embpy/embedder.py

def model_catalog(
    self,
    category: str = "all",
    *,
    summary: bool = False,
) -> Any:
    """Return the available models as a table.

    The tabular counterpart to :meth:`list_available_models`, which returns bare
    strings. In a notebook this renders as a proper table rather than printed
    text, and being a DataFrame it can be filtered, sorted and exported.

    Parameters
    ----------
    category
        Restrict to one family (``"dna"``, ``"protein"``, ``"molecule"``,
        ``"text"``, ``"morphology"``, ``"single_cell"``, ``"static"``), or
        ``"all"`` for everything.
    summary
        ``False`` (default) gives one row per model with its family. ``True``
        gives one row per family with a count and a few example keys -- the
        quick "what can this package do?" view.

    Returns
    -------
    pandas.DataFrame
        With ``summary=False``: columns ``model``, ``family``.
        With ``summary=True``: columns ``family``, ``n_models``, ``examples``.

    Examples
    --------
    >>> embedder.model_catalog(summary=True)          # doctest: +SKIP
    >>> embedder.model_catalog("protein").head()      # doctest: +SKIP
    """
    import pandas as pd

    families = ["static", "dna", "protein", "molecule", "text", "single_cell", "morphology"]
    wanted = families if category == "all" else [category]

    rows: list[dict[str, Any]] = []
    for fam in wanted:
        for name in self.list_available_models(fam):  # type: ignore[arg-type]
            rows.append({"model": name, "family": fam})

    frame = pd.DataFrame(rows, columns=["model", "family"])
    # A model can appear under more than one family (e.g. a text encoder used
    # for genes); keep the first family so counts stay interpretable.
    frame = frame.drop_duplicates(subset="model", keep="first").reset_index(drop=True)

    if not summary:
        return frame.sort_values(["family", "model"]).reset_index(drop=True)

    grouped = (
        frame.groupby("family")["model"]
        .agg(n_models="size", examples=lambda s: ", ".join(sorted(s)[:4]))
        .reset_index()
        .sort_values("n_models", ascending=False)
        .reset_index(drop=True)
    )
    return grouped

---

Model Wrappers

Protein Models

embpy.models.protein_models

ESM2Wrapper

Bases: BaseModelWrapper

Wrapper for ESM2 protein language models using Hugging Face Transformers.

Implements the ProteinEmbedder interface for ESM2 models. Supports tokenization, model inference, and pooling strategies to produce protein embeddings.

Source code in src/embpy/models/protein_models.py

class ESM2Wrapper(BaseModelWrapper):
    """
    Wrapper for ESM2 protein language models using Hugging Face Transformers.

    Implements the ProteinEmbedder interface for ESM2 models. Supports tokenization,
    model inference, and pooling strategies to produce protein embeddings.
    """

    model_type = "protein"
    available_pooling_strategies = ["mean", "max", "cls", "none"]

    def __init__(self, model_path_or_name: str = "facebook/esm2_t6_8M_UR50D", **kwargs):
        """
        Initialize the ESM2Wrapper.

        Args:
            model_path_or_name (str): Model identifier or path for the ESM2 model.
            **kwargs: Additional configuration parameters.
        """
        super().__init__(model_path_or_name, **kwargs)
        self.tokenizer = None

    def load(self, device: torch.device) -> None:
        """
        Load the ESM2 model and tokenizer onto the specified device.

        Args:
            device (torch.device): Device to load the model (e.g., cpu or cuda).

        Raises
        ------
            RuntimeError: If model or tokenizer loading fails.
        """
        if self.model is not None:
            logging.debug("ESM2 model already loaded.")
            return

        if not self.model_name:
            raise ValueError("model_path_or_name must be provided for ESM2Wrapper.")

        logging.info(f"Loading ESM2 model '{self.model_name}'...")
        from transformers import AutoTokenizer, AutoModel
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            self.model = AutoModel.from_pretrained(self.model_name)
            self.model.to(device)
            self.model.eval()
            self.device = device
            logging.info("ESM2 model loaded successfully.")
        except Exception as e:
            logging.error(f"Failed to load ESM2 model: {e}")
            raise RuntimeError(f"Could not load ESM2 model '{self.model_name}'") from e

    def _preprocess_sequence(self, sequence: str) -> Any:
        """
        Tokenize the protein sequence using the associated tokenizer.

        Args:
            sequence (str): The input protein sequence.

        Returns
        -------
            dict: Tokenized sequence inputs suitable for model inference.

        Raises
        ------
            RuntimeError: If the tokenizer is not loaded.
        """
        if not self.tokenizer:
            raise RuntimeError("Tokenizer is not loaded.")
        return self.tokenizer(sequence, return_tensors="pt", truncation=True, padding=True)

    def embed(
        self, input: str, pooling_strategy: str = "mean", target_layer: int | None = None, **kwargs
    ) -> np.ndarray:
        """
        Generate an embedding for a protein sequence using the ESM2 model.

        Args:
            sequence (str): The protein sequence.
            pooling_strategy (str): Pooling strategy to aggregate token-level embeddings.
                                    Options: 'mean', 'max', 'cls'.
            target_layer (Optional[int]): Specific layer from which to extract embeddings.
            **kwargs: Additional parameters.

        Returns
        -------
            np.ndarray: A pooled vector representing the protein embedding.

        Raises
        ------
            RuntimeError: If the model is not loaded.
            ValueError: If an invalid pooling strategy is provided.
        """
        if self.model is None or self.device is None:
            raise RuntimeError("ESM2 model not loaded. Please call load() first.")

        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(
                f"Invalid pooling strategy '{pooling_strategy}'. Available options: {self.available_pooling_strategies}"
            )

        # Preprocess the sequence and move inputs to the correct device:
        tokenized_input = self._preprocess_sequence(input)
        input_ids = tokenized_input["input_ids"].to(self.device)
        attention_mask = tokenized_input.get("attention_mask", None)
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)

        # Model inference:
        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=(target_layer is not None)
            )

            if target_layer is not None:
                hidden_states = outputs.hidden_states
                if not hidden_states:
                    raise ValueError("Model did not return hidden states.")
                if not (-len(hidden_states) <= target_layer < len(hidden_states)):
                    raise ValueError(f"Invalid target_layer index {target_layer}.")
                embeddings_tensor = hidden_states[target_layer]
            else:
                embeddings_tensor = outputs.last_hidden_state

        # If batch dimension is 1, squeeze it:
        if embeddings_tensor.dim() == 3 and embeddings_tensor.shape[0] == 1:
            embeddings_tensor = embeddings_tensor.squeeze(0)

        # Apply pooling strategy:
        if pooling_strategy == "none":
            pooled_embedding = embeddings_tensor.cpu().numpy()
        elif pooling_strategy == "cls":
            pooled_embedding = embeddings_tensor[0].cpu().numpy()
        elif pooling_strategy == "max":
            pooled_embedding = torch.max(embeddings_tensor, dim=0)[0].cpu().numpy()
        else:  # default to mean pooling:
            pooled_embedding = torch.mean(embeddings_tensor, dim=0).cpu().numpy()

        return pooled_embedding

    def _default_batch_size(self) -> int:
        """Pick a memory-safe default batch size based on the loaded model.

        ESM-2 attention activations scale as ``B * H * L^2``, which on an
        80 GB GPU caps the maximum safe batch size at roughly:

        * 35M  (T6,  hidden 320,  20 layers, 20 heads)   -> 64
        * 150M (T30, hidden 640,  30 layers, 20 heads)   -> 32
        * 650M (T33, hidden 1280, 33 layers, 20 heads)   -> 16
        * 3B   (T36, hidden 2560, 36 layers, 40 heads)   -> 4
        * 15B  (T48, hidden 5120, 48 layers, 40 heads)   -> 1 (rarely fits)

        Defaulting to a generic 32 OOMs the 3B model on long sequences
        (the observed failure: 33.87 GiB allocation for attention scores).
        We read ``model.config.hidden_size`` to scale automatically.
        """
        hidden = 1280
        cfg = getattr(self.model, "config", None)
        if cfg is not None:
            hidden = int(getattr(cfg, "hidden_size", 1280))
        if hidden >= 5000:
            return 1
        if hidden >= 2400:
            return 4
        if hidden >= 1200:
            return 16
        if hidden >= 600:
            return 32
        return 64

    def embed_batch(
        self,
        inputs: list[str],
        pooling_strategy: str = "mean",
        target_layer: int | None = None,
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """
        Compute embeddings for multiple protein sequences with chunking
        and a model-aware default batch size.

        Sequences longer than the model's context window are split into
        non-overlapping chunks, embedded separately, and mean-pooled
        across chunks. This matches the DNA path in
        ``embpy.models.dna_models._hf_batched_embed`` and was added
        because a) some proteins exceed ESM-2's 1024-token context (e.g.
        titin) and b) plain truncation would silently discard the entire
        C-terminus.

        Parameters
        ----------
        inputs
            A list of protein sequence strings.
        pooling_strategy
            One of {"mean", "max", "cls", "none"}. ``none`` is only
            supported when no input exceeds the context window (chunked
            inputs do not have a canonical per-token stitching).
        target_layer
            If specified, extract embeddings from a particular hidden
            state layer instead of the default ``last_hidden_state``.
        batch_size
            Optional override; otherwise picked from ``_default_batch_size()``.

        Returns
        -------
        list[np.ndarray]
            One 1D vector per input.

        Raises
        ------
        RuntimeError
            If the model has not been loaded.
        ValueError
            If ``pooling_strategy`` is unsupported.
        NotImplementedError
            If ``pooling_strategy='none'`` is requested with sequences
            that overflow the context window.
        """
        if self.model is None or self.device is None or self.tokenizer is None:
            raise RuntimeError("ESM2 model not loaded. Please call load() first.")
        if not inputs:
            return []
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(
                f"Invalid pooling strategy '{pooling_strategy}'. Choose from {self.available_pooling_strategies}"
            )

        batch_size = int(kwargs.pop("batch_size", self._default_batch_size()))

        cfg = getattr(self.model, "config", None)
        ctx_model = getattr(cfg, "max_position_embeddings", None) if cfg is not None else None
        ctx_tok = getattr(self.tokenizer, "model_max_length", None)
        candidates = [
            int(x) for x in (ctx_model, ctx_tok)
            if isinstance(x, int) and 0 < x < 1_000_000
        ]
        ctx = min(candidates) if candidates else 1024

        cls_id = getattr(self.tokenizer, "cls_token_id", None)
        sep_id = getattr(self.tokenizer, "sep_token_id", None)
        bos_id = getattr(self.tokenizer, "bos_token_id", None)
        eos_id = getattr(self.tokenizer, "eos_token_id", None)
        pad_id = getattr(self.tokenizer, "pad_token_id", None)
        if pad_id is None:
            pad_id = eos_id if eos_id is not None else 0
        head_id = cls_id if cls_id is not None else bos_id
        tail_id = sep_id if sep_id is not None else eos_id
        head_cost = 1 if head_id is not None else 0
        tail_cost = 1 if tail_id is not None else 0
        inner_max = ctx - head_cost - tail_cost
        if inner_max <= 0:
            raise RuntimeError(
                f"Context window {ctx} too small after reserving room for special tokens."
            )

        chunk_ids: list[list[int]] = []
        chunk_attn: list[list[int]] = []
        chunk_to_input_idx: list[int] = []
        overflow_warned = False
        for i, text in enumerate(inputs):
            enc = self.tokenizer(
                text,
                add_special_tokens=False,
                truncation=False,
                padding=False,
                return_tensors=None,
            )
            ids = enc["input_ids"]
            if ids and isinstance(ids[0], list):
                ids = ids[0]
            n_chunks = max(1, (len(ids) + inner_max - 1) // inner_max)
            if n_chunks > 1 and not overflow_warned:
                logging.info(
                    "ESM-2 input longer than context window (%d > %d) -- "
                    "chunking into %d pieces and mean-pooling per protein.",
                    len(ids), ctx, n_chunks,
                )
                overflow_warned = True
            for j in range(0, max(len(ids), 1), inner_max):
                piece = ids[j : j + inner_max]
                full = (
                    ([head_id] if head_id is not None else [])
                    + piece
                    + ([tail_id] if tail_id is not None else [])
                )
                chunk_ids.append(full)
                chunk_attn.append([1] * len(full))
                chunk_to_input_idx.append(i)

        chunk_pooled: list[np.ndarray | None] = [None] * len(chunk_ids)
        needs_sequence_out = pooling_strategy == "none"
        if needs_sequence_out and any(
            sum(1 for src in chunk_to_input_idx if src == k) > 1
            for k in range(len(inputs))
        ):
            raise NotImplementedError(
                "pooling_strategy='none' not supported when proteins exceed "
                "the model context window (would produce ragged outputs)."
            )

        for start in range(0, len(chunk_ids), batch_size):
            end = min(start + batch_size, len(chunk_ids))
            batch_ids = chunk_ids[start:end]
            batch_attn = chunk_attn[start:end]
            max_len = max(len(x) for x in batch_ids)
            input_ids = torch.tensor(
                [x + [pad_id] * (max_len - len(x)) for x in batch_ids],
                dtype=torch.long, device=self.device,
            )
            attention_mask = torch.tensor(
                [x + [0] * (max_len - len(x)) for x in batch_attn],
                dtype=torch.long, device=self.device,
            )

            with torch.no_grad():
                outputs = self.model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=(target_layer is not None),
                )
                if target_layer is not None:
                    emb_tensor = outputs.hidden_states[target_layer]
                else:
                    emb_tensor = outputs.last_hidden_state

            for k in range(emb_tensor.shape[0]):
                seq_emb = emb_tensor[k]
                mask = attention_mask[k].unsqueeze(-1).bool()
                if pooling_strategy == "none":
                    chunk_pooled[start + k] = seq_emb.cpu().numpy()
                    continue
                elif pooling_strategy == "cls":
                    pooled = seq_emb[0]
                elif pooling_strategy == "max":
                    pooled = seq_emb.masked_fill(~mask, float("-inf")).max(dim=0).values
                else:
                    pooled = (seq_emb * mask).sum(dim=0) / mask.sum(dim=0).clamp(min=1)
                chunk_pooled[start + k] = pooled.cpu().numpy()

            del input_ids, attention_mask, outputs, emb_tensor
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        chunks_per_input: list[list[int]] = [[] for _ in range(len(inputs))]
        for c, src in enumerate(chunk_to_input_idx):
            chunks_per_input[src].append(c)

        all_embeddings: list[np.ndarray] = []
        for i in range(len(inputs)):
            idxs = chunks_per_input[i]
            if not idxs:
                raise RuntimeError(f"No chunks produced for input {i}.")
            if len(idxs) == 1 or pooling_strategy == "none":
                all_embeddings.append(chunk_pooled[idxs[0]])  # type: ignore[arg-type]
                continue
            stacked = np.stack([chunk_pooled[c] for c in idxs], axis=0)  # type: ignore[misc]
            all_embeddings.append(stacked.mean(axis=0).astype(stacked.dtype))

        return all_embeddings

__init__(model_path_or_name='facebook/esm2_t6_8M_UR50D', **kwargs)

Initialize the ESM2Wrapper.

Args: model_path_or_name (str): Model identifier or path for the ESM2 model. **kwargs: Additional configuration parameters.

Source code in src/embpy/models/protein_models.py

def __init__(self, model_path_or_name: str = "facebook/esm2_t6_8M_UR50D", **kwargs):
    """
    Initialize the ESM2Wrapper.

    Args:
        model_path_or_name (str): Model identifier or path for the ESM2 model.
        **kwargs: Additional configuration parameters.
    """
    super().__init__(model_path_or_name, **kwargs)
    self.tokenizer = None

embed(input, pooling_strategy='mean', target_layer=None, **kwargs)

Generate an embedding for a protein sequence using the ESM2 model.

Args: sequence (str): The protein sequence. pooling_strategy (str): Pooling strategy to aggregate token-level embeddings. Options: 'mean', 'max', 'cls'. target_layer (Optional[int]): Specific layer from which to extract embeddings. **kwargs: Additional parameters.

Returns:

Type	Description
np.ndarray: A pooled vector representing the protein embedding.

Raises:

Type	Description
RuntimeError: If the model is not loaded.	ValueError: If an invalid pooling strategy is provided.

Source code in src/embpy/models/protein_models.py

def embed(
    self, input: str, pooling_strategy: str = "mean", target_layer: int | None = None, **kwargs
) -> np.ndarray:
    """
    Generate an embedding for a protein sequence using the ESM2 model.

    Args:
        sequence (str): The protein sequence.
        pooling_strategy (str): Pooling strategy to aggregate token-level embeddings.
                                Options: 'mean', 'max', 'cls'.
        target_layer (Optional[int]): Specific layer from which to extract embeddings.
        **kwargs: Additional parameters.

    Returns
    -------
        np.ndarray: A pooled vector representing the protein embedding.

    Raises
    ------
        RuntimeError: If the model is not loaded.
        ValueError: If an invalid pooling strategy is provided.
    """
    if self.model is None or self.device is None:
        raise RuntimeError("ESM2 model not loaded. Please call load() first.")

    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(
            f"Invalid pooling strategy '{pooling_strategy}'. Available options: {self.available_pooling_strategies}"
        )

    # Preprocess the sequence and move inputs to the correct device:
    tokenized_input = self._preprocess_sequence(input)
    input_ids = tokenized_input["input_ids"].to(self.device)
    attention_mask = tokenized_input.get("attention_mask", None)
    if attention_mask is not None:
        attention_mask = attention_mask.to(self.device)

    # Model inference:
    with torch.no_grad():
        outputs = self.model(
            input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=(target_layer is not None)
        )

        if target_layer is not None:
            hidden_states = outputs.hidden_states
            if not hidden_states:
                raise ValueError("Model did not return hidden states.")
            if not (-len(hidden_states) <= target_layer < len(hidden_states)):
                raise ValueError(f"Invalid target_layer index {target_layer}.")
            embeddings_tensor = hidden_states[target_layer]
        else:
            embeddings_tensor = outputs.last_hidden_state

    # If batch dimension is 1, squeeze it:
    if embeddings_tensor.dim() == 3 and embeddings_tensor.shape[0] == 1:
        embeddings_tensor = embeddings_tensor.squeeze(0)

    # Apply pooling strategy:
    if pooling_strategy == "none":
        pooled_embedding = embeddings_tensor.cpu().numpy()
    elif pooling_strategy == "cls":
        pooled_embedding = embeddings_tensor[0].cpu().numpy()
    elif pooling_strategy == "max":
        pooled_embedding = torch.max(embeddings_tensor, dim=0)[0].cpu().numpy()
    else:  # default to mean pooling:
        pooled_embedding = torch.mean(embeddings_tensor, dim=0).cpu().numpy()

    return pooled_embedding

embed_batch(inputs, pooling_strategy='mean', target_layer=None, **kwargs)

Compute embeddings for multiple protein sequences with chunking and a model-aware default batch size.

Sequences longer than the model's context window are split into non-overlapping chunks, embedded separately, and mean-pooled across chunks. This matches the DNA path in embpy.models.dna_models._hf_batched_embed and was added because a) some proteins exceed ESM-2's 1024-token context (e.g. titin) and b) plain truncation would silently discard the entire C-terminus.

Parameters:

Name	Type	Description	Default
inputs	list[str]	A list of protein sequence strings.	required
pooling_strategy	str	One of {"mean", "max", "cls", "none"}. none is only supported when no input exceeds the context window (chunked inputs do not have a canonical per-token stitching).	'mean'
target_layer	int | None	If specified, extract embeddings from a particular hidden state layer instead of the default last_hidden_state.	None
batch_size		Optional override; otherwise picked from _default_batch_size().	required

Returns:

Type	Description
list[ndarray]	One 1D vector per input.

Raises:

Type	Description
RuntimeError	If the model has not been loaded.
ValueError	If pooling_strategy is unsupported.
NotImplementedError	If pooling_strategy='none' is requested with sequences that overflow the context window.

Source code in src/embpy/models/protein_models.py

def embed_batch(
    self,
    inputs: list[str],
    pooling_strategy: str = "mean",
    target_layer: int | None = None,
    **kwargs: Any,
) -> list[np.ndarray]:
    """
    Compute embeddings for multiple protein sequences with chunking
    and a model-aware default batch size.

    Sequences longer than the model's context window are split into
    non-overlapping chunks, embedded separately, and mean-pooled
    across chunks. This matches the DNA path in
    ``embpy.models.dna_models._hf_batched_embed`` and was added
    because a) some proteins exceed ESM-2's 1024-token context (e.g.
    titin) and b) plain truncation would silently discard the entire
    C-terminus.

    Parameters
    ----------
    inputs
        A list of protein sequence strings.
    pooling_strategy
        One of {"mean", "max", "cls", "none"}. ``none`` is only
        supported when no input exceeds the context window (chunked
        inputs do not have a canonical per-token stitching).
    target_layer
        If specified, extract embeddings from a particular hidden
        state layer instead of the default ``last_hidden_state``.
    batch_size
        Optional override; otherwise picked from ``_default_batch_size()``.

    Returns
    -------
    list[np.ndarray]
        One 1D vector per input.

    Raises
    ------
    RuntimeError
        If the model has not been loaded.
    ValueError
        If ``pooling_strategy`` is unsupported.
    NotImplementedError
        If ``pooling_strategy='none'`` is requested with sequences
        that overflow the context window.
    """
    if self.model is None or self.device is None or self.tokenizer is None:
        raise RuntimeError("ESM2 model not loaded. Please call load() first.")
    if not inputs:
        return []
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(
            f"Invalid pooling strategy '{pooling_strategy}'. Choose from {self.available_pooling_strategies}"
        )

    batch_size = int(kwargs.pop("batch_size", self._default_batch_size()))

    cfg = getattr(self.model, "config", None)
    ctx_model = getattr(cfg, "max_position_embeddings", None) if cfg is not None else None
    ctx_tok = getattr(self.tokenizer, "model_max_length", None)
    candidates = [
        int(x) for x in (ctx_model, ctx_tok)
        if isinstance(x, int) and 0 < x < 1_000_000
    ]
    ctx = min(candidates) if candidates else 1024

    cls_id = getattr(self.tokenizer, "cls_token_id", None)
    sep_id = getattr(self.tokenizer, "sep_token_id", None)
    bos_id = getattr(self.tokenizer, "bos_token_id", None)
    eos_id = getattr(self.tokenizer, "eos_token_id", None)
    pad_id = getattr(self.tokenizer, "pad_token_id", None)
    if pad_id is None:
        pad_id = eos_id if eos_id is not None else 0
    head_id = cls_id if cls_id is not None else bos_id
    tail_id = sep_id if sep_id is not None else eos_id
    head_cost = 1 if head_id is not None else 0
    tail_cost = 1 if tail_id is not None else 0
    inner_max = ctx - head_cost - tail_cost
    if inner_max <= 0:
        raise RuntimeError(
            f"Context window {ctx} too small after reserving room for special tokens."
        )

    chunk_ids: list[list[int]] = []
    chunk_attn: list[list[int]] = []
    chunk_to_input_idx: list[int] = []
    overflow_warned = False
    for i, text in enumerate(inputs):
        enc = self.tokenizer(
            text,
            add_special_tokens=False,
            truncation=False,
            padding=False,
            return_tensors=None,
        )
        ids = enc["input_ids"]
        if ids and isinstance(ids[0], list):
            ids = ids[0]
        n_chunks = max(1, (len(ids) + inner_max - 1) // inner_max)
        if n_chunks > 1 and not overflow_warned:
            logging.info(
                "ESM-2 input longer than context window (%d > %d) -- "
                "chunking into %d pieces and mean-pooling per protein.",
                len(ids), ctx, n_chunks,
            )
            overflow_warned = True
        for j in range(0, max(len(ids), 1), inner_max):
            piece = ids[j : j + inner_max]
            full = (
                ([head_id] if head_id is not None else [])
                + piece
                + ([tail_id] if tail_id is not None else [])
            )
            chunk_ids.append(full)
            chunk_attn.append([1] * len(full))
            chunk_to_input_idx.append(i)

    chunk_pooled: list[np.ndarray | None] = [None] * len(chunk_ids)
    needs_sequence_out = pooling_strategy == "none"
    if needs_sequence_out and any(
        sum(1 for src in chunk_to_input_idx if src == k) > 1
        for k in range(len(inputs))
    ):
        raise NotImplementedError(
            "pooling_strategy='none' not supported when proteins exceed "
            "the model context window (would produce ragged outputs)."
        )

    for start in range(0, len(chunk_ids), batch_size):
        end = min(start + batch_size, len(chunk_ids))
        batch_ids = chunk_ids[start:end]
        batch_attn = chunk_attn[start:end]
        max_len = max(len(x) for x in batch_ids)
        input_ids = torch.tensor(
            [x + [pad_id] * (max_len - len(x)) for x in batch_ids],
            dtype=torch.long, device=self.device,
        )
        attention_mask = torch.tensor(
            [x + [0] * (max_len - len(x)) for x in batch_attn],
            dtype=torch.long, device=self.device,
        )

        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_hidden_states=(target_layer is not None),
            )
            if target_layer is not None:
                emb_tensor = outputs.hidden_states[target_layer]
            else:
                emb_tensor = outputs.last_hidden_state

        for k in range(emb_tensor.shape[0]):
            seq_emb = emb_tensor[k]
            mask = attention_mask[k].unsqueeze(-1).bool()
            if pooling_strategy == "none":
                chunk_pooled[start + k] = seq_emb.cpu().numpy()
                continue
            elif pooling_strategy == "cls":
                pooled = seq_emb[0]
            elif pooling_strategy == "max":
                pooled = seq_emb.masked_fill(~mask, float("-inf")).max(dim=0).values
            else:
                pooled = (seq_emb * mask).sum(dim=0) / mask.sum(dim=0).clamp(min=1)
            chunk_pooled[start + k] = pooled.cpu().numpy()

        del input_ids, attention_mask, outputs, emb_tensor
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    chunks_per_input: list[list[int]] = [[] for _ in range(len(inputs))]
    for c, src in enumerate(chunk_to_input_idx):
        chunks_per_input[src].append(c)

    all_embeddings: list[np.ndarray] = []
    for i in range(len(inputs)):
        idxs = chunks_per_input[i]
        if not idxs:
            raise RuntimeError(f"No chunks produced for input {i}.")
        if len(idxs) == 1 or pooling_strategy == "none":
            all_embeddings.append(chunk_pooled[idxs[0]])  # type: ignore[arg-type]
            continue
        stacked = np.stack([chunk_pooled[c] for c in idxs], axis=0)  # type: ignore[misc]
        all_embeddings.append(stacked.mean(axis=0).astype(stacked.dtype))

    return all_embeddings

load(device)

Load the ESM2 model and tokenizer onto the specified device.

Args: device (torch.device): Device to load the model (e.g., cpu or cuda).

Raises:

Type	Description
RuntimeError: If model or tokenizer loading fails.

Source code in src/embpy/models/protein_models.py

def load(self, device: torch.device) -> None:
    """
    Load the ESM2 model and tokenizer onto the specified device.

    Args:
        device (torch.device): Device to load the model (e.g., cpu or cuda).

    Raises
    ------
        RuntimeError: If model or tokenizer loading fails.
    """
    if self.model is not None:
        logging.debug("ESM2 model already loaded.")
        return

    if not self.model_name:
        raise ValueError("model_path_or_name must be provided for ESM2Wrapper.")

    logging.info(f"Loading ESM2 model '{self.model_name}'...")
    from transformers import AutoTokenizer, AutoModel
    try:
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModel.from_pretrained(self.model_name)
        self.model.to(device)
        self.model.eval()
        self.device = device
        logging.info("ESM2 model loaded successfully.")
    except Exception as e:
        logging.error(f"Failed to load ESM2 model: {e}")
        raise RuntimeError(f"Could not load ESM2 model '{self.model_name}'") from e

ESM3Wrapper

Bases: BaseModelWrapper

Wrapper for ESM3 protein language models (EvolutionaryScale SDK).

ESM3 is a multimodal generative model for proteins. This wrapper uses the sequence track to produce per-residue embeddings and pools them into a single vector.

Supported checkpoints:

• esm3-small-2024-08 (1.4B, open weights)
• esm3-medium-2024-08 (7B, Forge API)
• esm3-large-2024-03 (98B, Forge API)

For open weights, the model runs locally. For larger models, pass a Forge API token via forge_token.

Requires the esm package: pip install esm

Source code in src/embpy/models/protein_models.py

class ESM3Wrapper(BaseModelWrapper):
    """Wrapper for ESM3 protein language models (EvolutionaryScale SDK).

    ESM3 is a multimodal generative model for proteins. This wrapper
    uses the sequence track to produce per-residue embeddings and pools
    them into a single vector.

    Supported checkpoints:

    * ``esm3-small-2024-08`` (1.4B, open weights)
    * ``esm3-medium-2024-08`` (7B, Forge API)
    * ``esm3-large-2024-03`` (98B, Forge API)

    For open weights, the model runs locally. For larger models, pass
    a Forge API token via ``forge_token``.

    Requires the ``esm`` package: ``pip install esm``
    """

    model_type = "protein"
    available_pooling_strategies = ["mean", "max", "cls", "none"]
    # Same fused-kernel limit as ESM-C: the shared esm SDK attention layer calls
    # F.scaled_dot_product_attention (esm/layers/attention.py:70,76). The Forge
    # API checkpoints are remote in any case, so no local tensor exists to hook.
    has_attention = False

    def __init__(
        self,
        model_path_or_name: str = "esm3-small-2024-08",
        forge_token: str | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model_path_or_name, **kwargs)
        self.forge_token = forge_token
        self._client: Any = None

    def load(self, device: torch.device) -> None:
        """Load the ESM3 model locally or connect to Forge API."""
        if self._client is not None:
            logging.debug("ESM3 already loaded.")
            return

        if not _HAVE_ESM3:
            raise ImportError(
                "The 'esm' package is required for ESM3. "
                "Install with: pip install esm"
            )

        name = self.model_name or "esm3-small-2024-08"

        if self.forge_token:
            import esm.sdk.client as esm_client
            self._client = esm_client(name, token=self.forge_token)
            logging.info("Connected to ESM3 Forge API: %s", name)
        else:
            self._client = ESM3.from_pretrained(name).to(device).eval()
            logging.info("Loaded ESM3 '%s' locally on %s", name, device)

        self.device = device
        self.model = self._client

    def embed(
        self,
        input: str,
        pooling_strategy: str = "mean",
        **kwargs: Any,
    ) -> np.ndarray:
        """Embed a single protein sequence using ESM3.

        Uses the ESM3 encoder to get per-residue embeddings from the
        sequence track, then pools them.
        """
        if self._client is None:
            raise RuntimeError("ESM3 not loaded. Call load() first.")
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(
                f"Invalid pooling '{pooling_strategy}'. "
                f"Choose from {self.available_pooling_strategies}"
            )

        protein = ESMProtein(sequence=input)
        protein_tensor = self._client.encode(protein)
        cfg = LogitsConfig(sequence=True, return_embeddings=True)
        out = self._client.logits(protein_tensor, cfg)

        embs = out.embeddings
        if embs.dim() == 3 and embs.shape[0] == 1:
            embs = embs.squeeze(0)

        if pooling_strategy == "none":
            return embs.float().cpu().numpy()
        elif pooling_strategy == "cls":
            pooled = embs[0]
        elif pooling_strategy == "max":
            pooled = torch.max(embs, dim=0)[0]
        else:
            pooled = torch.mean(embs, dim=0)

        return pooled.float().cpu().numpy()

    def embed_batch(
        self,
        inputs: list[str],
        pooling_strategy: str = "mean",
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """Embed multiple protein sequences using ESM3."""
        if self._client is None:
            raise RuntimeError("ESM3 not loaded. Call load() first.")
        return [
            self.embed(seq, pooling_strategy=pooling_strategy, **kwargs)
            for seq in inputs
        ]

embed(input, pooling_strategy='mean', **kwargs)

Embed a single protein sequence using ESM3.

Uses the ESM3 encoder to get per-residue embeddings from the sequence track, then pools them.

Source code in src/embpy/models/protein_models.py

def embed(
    self,
    input: str,
    pooling_strategy: str = "mean",
    **kwargs: Any,
) -> np.ndarray:
    """Embed a single protein sequence using ESM3.

    Uses the ESM3 encoder to get per-residue embeddings from the
    sequence track, then pools them.
    """
    if self._client is None:
        raise RuntimeError("ESM3 not loaded. Call load() first.")
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(
            f"Invalid pooling '{pooling_strategy}'. "
            f"Choose from {self.available_pooling_strategies}"
        )

    protein = ESMProtein(sequence=input)
    protein_tensor = self._client.encode(protein)
    cfg = LogitsConfig(sequence=True, return_embeddings=True)
    out = self._client.logits(protein_tensor, cfg)

    embs = out.embeddings
    if embs.dim() == 3 and embs.shape[0] == 1:
        embs = embs.squeeze(0)

    if pooling_strategy == "none":
        return embs.float().cpu().numpy()
    elif pooling_strategy == "cls":
        pooled = embs[0]
    elif pooling_strategy == "max":
        pooled = torch.max(embs, dim=0)[0]
    else:
        pooled = torch.mean(embs, dim=0)

    return pooled.float().cpu().numpy()

embed_batch(inputs, pooling_strategy='mean', **kwargs)

Embed multiple protein sequences using ESM3.

Source code in src/embpy/models/protein_models.py

def embed_batch(
    self,
    inputs: list[str],
    pooling_strategy: str = "mean",
    **kwargs: Any,
) -> list[np.ndarray]:
    """Embed multiple protein sequences using ESM3."""
    if self._client is None:
        raise RuntimeError("ESM3 not loaded. Call load() first.")
    return [
        self.embed(seq, pooling_strategy=pooling_strategy, **kwargs)
        for seq in inputs
    ]

load(device)

Load the ESM3 model locally or connect to Forge API.

Source code in src/embpy/models/protein_models.py

def load(self, device: torch.device) -> None:
    """Load the ESM3 model locally or connect to Forge API."""
    if self._client is not None:
        logging.debug("ESM3 already loaded.")
        return

    if not _HAVE_ESM3:
        raise ImportError(
            "The 'esm' package is required for ESM3. "
            "Install with: pip install esm"
        )

    name = self.model_name or "esm3-small-2024-08"

    if self.forge_token:
        import esm.sdk.client as esm_client
        self._client = esm_client(name, token=self.forge_token)
        logging.info("Connected to ESM3 Forge API: %s", name)
    else:
        self._client = ESM3.from_pretrained(name).to(device).eval()
        logging.info("Loaded ESM3 '%s' locally on %s", name, device)

    self.device = device
    self.model = self._client

ESMCWrapper

Bases: BaseModelWrapper

Wrapper for the ESMC-SDK protein language models (e.g., 'esmc_300m', 'esmc_600m').

This class hides the details of the ESMC client and provides a simple interface for: - loading the model to a device, - embedding a single sequence with mean/max/cls pooling, - optionally returning hidden‐states from all layers, - embedding batches of sequences.

Attributes:

Name	Type	Description
model_name	str	The checkpoint identifier passed to ESMC.from_pretrained().
client	ESMC | None	The underlying ESMC inference client (None until loaded).
device	device | None	Device to which the client is moved (None until loaded).
available_pooling_strategies	List[str]	Allowed values for pooling_strategy: ['mean', 'max', 'cls'].

Source code in src/embpy/models/protein_models.py

class ESMCWrapper(BaseModelWrapper):
    """
    Wrapper for the ESMC-SDK protein language models (e.g., 'esmc_300m', 'esmc_600m').

    This class hides the details of the ESMC client and provides a
    simple interface for:
      - loading the model to a device,
      - embedding a single sequence with mean/max/cls pooling,
      - optionally returning hidden‐states from all layers,
      - embedding batches of sequences.

    Attributes
    ----------
    model_name : str
        The checkpoint identifier passed to ESMC.from_pretrained().
    client : ESMC | None
        The underlying ESMC inference client (None until loaded).
    device : torch.device | None
        Device to which the client is moved (None until loaded).
    available_pooling_strategies : List[str]
        Allowed values for pooling_strategy: ['mean', 'max', 'cls'].
    """

    model_type = "protein"
    available_pooling_strategies = ["mean", "max", "cls", "none"]
    # The ESM SDK computes attention with F.scaled_dot_product_attention
    # unconditionally (esm/layers/attention.py:70,76 -- both the masked and
    # unmasked branch), so the per-head matrix never exists as a tensor and no
    # hook can recover it. Declared False so extract_attention fails fast with an
    # explanation instead of running a forward pass that cannot produce weights.
    has_attention = False

    def __init__(self, model_path_or_name: str = "esmc_300m", **kwargs: Any):
        """
        Initialize the ESMCWrapper.

        Parameters
        ----------
        model_path_or_name : str, optional
            The ESMC checkpoint name, e.g. 'esmc_300m' or 'esmc_600m'.
            Defaults to 'esmc_300m'.
        """
        super().__init__(model_path_or_name=model_path_or_name, **kwargs)
        self.client: ESMC | None = None
        self.device: torch.device | None = None

    def load(self, device: torch.device) -> None:
        """
        Load the ESMC client onto the specified device and switch to eval mode.

        If already loaded, this is a no-op.

        Parameters
        ----------
        device : torch.device
            Target device for model inference (CPU, CUDA, or MPS).

        Raises
        ------
        ImportError
            If the ESMC SDK (the ``esm`` package) is not installed.
        RuntimeError
            If loading the client fails.
        """
        if self.client is not None:
            return
        if not _HAVE_ESMC:
            raise ImportError(
                "The 'esm' package is required for ESM-C. "
                "Recommended install: pip install esm --no-deps "
                "(the embpy[esm3] extra also works, but caps transformers below 4.48.2)."
            )
        self.device = device
        logging.info(f"Loading ESMC client '{self.model_name}' onto {device} …")
        try:
            self.client = ESMC.from_pretrained(self.model_name).to(device).eval()
            logging.info("ESMC client loaded successfully.")
        except Exception as e:
            logging.error(f"Failed to load ESMC client '{self.model_name}': {e}")
            raise RuntimeError(f"Could not load ESMC client '{self.model_name}'") from e

    def embed(
        self,
        input: str,
        pooling_strategy: str = "mean",
        max_length: int = 10000,
        **kwargs: Any,
    ) -> np.ndarray:
        """
        Embed a single protein sequence.

        Parameters
        ----------
        input : str
            Amino acid sequence string.
        pooling_strategy : str
            One of 'mean', 'max', 'cls'.
        max_length : int
            Safety truncation limit to prevent OOM on very large proteins.
        **kwargs : Any
            Additional arguments (unused, for interface consistency).

        Returns
        -------
        np.ndarray
            1D embedding vector.
        """
        if self.client is None or self.device is None:
            raise RuntimeError("ESMC client not loaded; call load() first.")
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(f"Invalid pooling '{pooling_strategy}'")

        sequence = input
        if len(sequence) > max_length:
            logging.warning(f"Sequence too long ({len(sequence)}). Truncating to {max_length}.")
            sequence = sequence[:max_length]

        prot = ESMProtein(sequence=sequence)
        tensor = self.client.encode(prot)

        cfg = LogitsConfig(sequence=True, return_embeddings=True, return_hidden_states=False)
        out: LogitsOutput = self.client.logits(tensor, cfg)

        embs = out.embeddings
        if embs.dim() == 3 and embs.shape[0] == 1:
            embs = embs.squeeze(0)

        if pooling_strategy == "none":
            return embs.cpu().numpy().astype(np.float32)
        elif pooling_strategy == "cls":
            pooled = embs[0]
        elif pooling_strategy == "max":
            pooled = torch.max(embs, dim=0)[0]
        else:
            pooled = torch.mean(embs, dim=0)

        return pooled.cpu().numpy()

    def embed_with_hidden_states(
        self,
        input: str,
        pooling_strategy: str = "mean",
        hidden_layers: Sequence[int] | None = None,
        max_length: int = 10000,
    ) -> dict[str, np.ndarray]:
        """
        Embed a single protein sequence and return hidden states.

        Parameters
        ----------
        input : str
            Amino acid sequence string.
        pooling_strategy : str
            One of 'mean', 'max', 'cls'.
        hidden_layers : Sequence[int], optional
            Specific layer indices to return. If None, returns all.
        max_length : int
            Safety truncation limit.

        Returns
        -------
        dict[str, np.ndarray]
            Dictionary with 'embedding' (1D) and 'hidden_states' (3D) arrays.
        """
        if self.client is None or self.device is None:
            raise RuntimeError("ESMC client not loaded; call load() first.")
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(f"Invalid pooling '{pooling_strategy}'")

        sequence = input
        if len(sequence) > max_length:
            logging.warning(f"Sequence too long ({len(sequence)}). Truncating to {max_length}.")
            sequence = sequence[:max_length]

        prot = ESMProtein(sequence=sequence)
        tensor = self.client.encode(prot)

        cfg = LogitsConfig(sequence=True, return_embeddings=True, return_hidden_states=True)
        out: LogitsOutput = self.client.logits(tensor, cfg)

        embs = out.embeddings
        if embs.dim() == 3 and embs.shape[0] == 1:
            embs = embs.squeeze(0)

        if pooling_strategy == "none":
            return embs.cpu().numpy().astype(np.float32)
        elif pooling_strategy == "cls":
            pooled = embs[0]
        elif pooling_strategy == "max":
            pooled = torch.max(embs, dim=0)[0]
        else:
            pooled = torch.mean(embs, dim=0)

        result: dict[str, np.ndarray] = {"embedding": pooled.cpu().numpy()}

        hidden = out.hidden_states
        if isinstance(hidden, tuple):
            hs = torch.stack(hidden, dim=0)
        else:
            hs = hidden
        while hs.dim() > 3 and hs.shape[0] == 1:
            hs = hs.squeeze(0)
        if hs.dim() == 4 and hs.shape[1] == 1:
            hs = hs.squeeze(1)

        if hidden_layers is not None:
            hs = hs[list(hidden_layers), :, :]
        result["hidden_states"] = hs.to(torch.float32).cpu().numpy()

        return result

    def embed_batch(
        self,
        inputs: list[str],
        pooling_strategy: str = "mean",
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """
        Embed a batch of protein sequences.

        Parameters
        ----------
        inputs : list[str]
            A list of amino-acid sequences.
        pooling_strategy : {'mean', 'max', 'cls'}, default='mean'
            Pooling to apply per sequence.
        **kwargs : Any
            Additional arguments forwarded to embed().

        Returns
        -------
        list[np.ndarray]
            A list of 1D embedding arrays, one per sequence.
        """
        return [
            self.embed(
                input=seq,
                pooling_strategy=pooling_strategy,
                **kwargs,
            )
            for seq in inputs
        ]

__init__(model_path_or_name='esmc_300m', **kwargs)

Initialize the ESMCWrapper.

Parameters:

Name	Type	Description	Default
model_path_or_name	str	The ESMC checkpoint name, e.g. 'esmc_300m' or 'esmc_600m'. Defaults to 'esmc_300m'.	'esmc_300m'

Source code in src/embpy/models/protein_models.py

def __init__(self, model_path_or_name: str = "esmc_300m", **kwargs: Any):
    """
    Initialize the ESMCWrapper.

    Parameters
    ----------
    model_path_or_name : str, optional
        The ESMC checkpoint name, e.g. 'esmc_300m' or 'esmc_600m'.
        Defaults to 'esmc_300m'.
    """
    super().__init__(model_path_or_name=model_path_or_name, **kwargs)
    self.client: ESMC | None = None
    self.device: torch.device | None = None

embed(input, pooling_strategy='mean', max_length=10000, **kwargs)

Embed a single protein sequence.

Parameters:

Name	Type	Description	Default
input	str	Amino acid sequence string.	required
pooling_strategy	str	One of 'mean', 'max', 'cls'.	'mean'
max_length	int	Safety truncation limit to prevent OOM on very large proteins.	10000
**kwargs	Any	Additional arguments (unused, for interface consistency).	{}

Returns:

Type	Description
ndarray	1D embedding vector.

Source code in src/embpy/models/protein_models.py

def embed(
    self,
    input: str,
    pooling_strategy: str = "mean",
    max_length: int = 10000,
    **kwargs: Any,
) -> np.ndarray:
    """
    Embed a single protein sequence.

    Parameters
    ----------
    input : str
        Amino acid sequence string.
    pooling_strategy : str
        One of 'mean', 'max', 'cls'.
    max_length : int
        Safety truncation limit to prevent OOM on very large proteins.
    **kwargs : Any
        Additional arguments (unused, for interface consistency).

    Returns
    -------
    np.ndarray
        1D embedding vector.
    """
    if self.client is None or self.device is None:
        raise RuntimeError("ESMC client not loaded; call load() first.")
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(f"Invalid pooling '{pooling_strategy}'")

    sequence = input
    if len(sequence) > max_length:
        logging.warning(f"Sequence too long ({len(sequence)}). Truncating to {max_length}.")
        sequence = sequence[:max_length]

    prot = ESMProtein(sequence=sequence)
    tensor = self.client.encode(prot)

    cfg = LogitsConfig(sequence=True, return_embeddings=True, return_hidden_states=False)
    out: LogitsOutput = self.client.logits(tensor, cfg)

    embs = out.embeddings
    if embs.dim() == 3 and embs.shape[0] == 1:
        embs = embs.squeeze(0)

    if pooling_strategy == "none":
        return embs.cpu().numpy().astype(np.float32)
    elif pooling_strategy == "cls":
        pooled = embs[0]
    elif pooling_strategy == "max":
        pooled = torch.max(embs, dim=0)[0]
    else:
        pooled = torch.mean(embs, dim=0)

    return pooled.cpu().numpy()

embed_batch(inputs, pooling_strategy='mean', **kwargs)

Embed a batch of protein sequences.

Parameters:

Name	Type	Description	Default
inputs	list[str]	A list of amino-acid sequences.	required
pooling_strategy	('mean', 'max', 'cls')	Pooling to apply per sequence.	'mean'
**kwargs	Any	Additional arguments forwarded to embed().	{}

Returns:

Type	Description
list[ndarray]	A list of 1D embedding arrays, one per sequence.

Source code in src/embpy/models/protein_models.py

def embed_batch(
    self,
    inputs: list[str],
    pooling_strategy: str = "mean",
    **kwargs: Any,
) -> list[np.ndarray]:
    """
    Embed a batch of protein sequences.

    Parameters
    ----------
    inputs : list[str]
        A list of amino-acid sequences.
    pooling_strategy : {'mean', 'max', 'cls'}, default='mean'
        Pooling to apply per sequence.
    **kwargs : Any
        Additional arguments forwarded to embed().

    Returns
    -------
    list[np.ndarray]
        A list of 1D embedding arrays, one per sequence.
    """
    return [
        self.embed(
            input=seq,
            pooling_strategy=pooling_strategy,
            **kwargs,
        )
        for seq in inputs
    ]

embed_with_hidden_states(input, pooling_strategy='mean', hidden_layers=None, max_length=10000)

Embed a single protein sequence and return hidden states.

Parameters:

Name	Type	Description	Default
input	str	Amino acid sequence string.	required
pooling_strategy	str	One of 'mean', 'max', 'cls'.	'mean'
hidden_layers	Sequence[int]	Specific layer indices to return. If None, returns all.	None
max_length	int	Safety truncation limit.	10000

Returns:

Type	Description
dict[str, ndarray]	Dictionary with 'embedding' (1D) and 'hidden_states' (3D) arrays.

Source code in src/embpy/models/protein_models.py

def embed_with_hidden_states(
    self,
    input: str,
    pooling_strategy: str = "mean",
    hidden_layers: Sequence[int] | None = None,
    max_length: int = 10000,
) -> dict[str, np.ndarray]:
    """
    Embed a single protein sequence and return hidden states.

    Parameters
    ----------
    input : str
        Amino acid sequence string.
    pooling_strategy : str
        One of 'mean', 'max', 'cls'.
    hidden_layers : Sequence[int], optional
        Specific layer indices to return. If None, returns all.
    max_length : int
        Safety truncation limit.

    Returns
    -------
    dict[str, np.ndarray]
        Dictionary with 'embedding' (1D) and 'hidden_states' (3D) arrays.
    """
    if self.client is None or self.device is None:
        raise RuntimeError("ESMC client not loaded; call load() first.")
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(f"Invalid pooling '{pooling_strategy}'")

    sequence = input
    if len(sequence) > max_length:
        logging.warning(f"Sequence too long ({len(sequence)}). Truncating to {max_length}.")
        sequence = sequence[:max_length]

    prot = ESMProtein(sequence=sequence)
    tensor = self.client.encode(prot)

    cfg = LogitsConfig(sequence=True, return_embeddings=True, return_hidden_states=True)
    out: LogitsOutput = self.client.logits(tensor, cfg)

    embs = out.embeddings
    if embs.dim() == 3 and embs.shape[0] == 1:
        embs = embs.squeeze(0)

    if pooling_strategy == "none":
        return embs.cpu().numpy().astype(np.float32)
    elif pooling_strategy == "cls":
        pooled = embs[0]
    elif pooling_strategy == "max":
        pooled = torch.max(embs, dim=0)[0]
    else:
        pooled = torch.mean(embs, dim=0)

    result: dict[str, np.ndarray] = {"embedding": pooled.cpu().numpy()}

    hidden = out.hidden_states
    if isinstance(hidden, tuple):
        hs = torch.stack(hidden, dim=0)
    else:
        hs = hidden
    while hs.dim() > 3 and hs.shape[0] == 1:
        hs = hs.squeeze(0)
    if hs.dim() == 4 and hs.shape[1] == 1:
        hs = hs.squeeze(1)

    if hidden_layers is not None:
        hs = hs[list(hidden_layers), :, :]
    result["hidden_states"] = hs.to(torch.float32).cpu().numpy()

    return result

load(device)

Load the ESMC client onto the specified device and switch to eval mode.

If already loaded, this is a no-op.

Parameters:

Name	Type	Description	Default
device	device	Target device for model inference (CPU, CUDA, or MPS).	required

Raises:

Type	Description
ImportError	If the ESMC SDK (the esm package) is not installed.
RuntimeError	If loading the client fails.

Source code in src/embpy/models/protein_models.py

def load(self, device: torch.device) -> None:
    """
    Load the ESMC client onto the specified device and switch to eval mode.

    If already loaded, this is a no-op.

    Parameters
    ----------
    device : torch.device
        Target device for model inference (CPU, CUDA, or MPS).

    Raises
    ------
    ImportError
        If the ESMC SDK (the ``esm`` package) is not installed.
    RuntimeError
        If loading the client fails.
    """
    if self.client is not None:
        return
    if not _HAVE_ESMC:
        raise ImportError(
            "The 'esm' package is required for ESM-C. "
            "Recommended install: pip install esm --no-deps "
            "(the embpy[esm3] extra also works, but caps transformers below 4.48.2)."
        )
    self.device = device
    logging.info(f"Loading ESMC client '{self.model_name}' onto {device} …")
    try:
        self.client = ESMC.from_pretrained(self.model_name).to(device).eval()
        logging.info("ESMC client loaded successfully.")
    except Exception as e:
        logging.error(f"Failed to load ESMC client '{self.model_name}': {e}")
        raise RuntimeError(f"Could not load ESMC client '{self.model_name}'") from e

ProtT5Wrapper

Bases: BaseModelWrapper

Wrapper for ProtTrans ProtT5-XL-UniRef50 protein language model.

Produces 1024-dimensional per-protein embeddings by extracting residue-level representations from the T5 encoder and mean-pooling over the sequence length.

This is the same embedding used as initial node features in the STRING-SPACE (sequence) and STRING-GNN frameworks (Hu et al., 2024).

Supported model identifiers (Hugging Face):

• Rostlab/prot_t5_xl_uniref50 – full-precision (float32)
• Rostlab/prot_t5_xl_half_uniref50-enc – half-precision encoder-only

Parameters:

Name	Type	Description	Default
model_path_or_name	str	Hugging Face model name or local path.	'Rostlab/prot_t5_xl_half_uniref50-enc'

Source code in src/embpy/models/protein_models.py

class ProtT5Wrapper(BaseModelWrapper):
    """Wrapper for ProtTrans ProtT5-XL-UniRef50 protein language model.

    Produces **1024-dimensional** per-protein embeddings by extracting
    residue-level representations from the T5 *encoder* and mean-pooling
    over the sequence length.

    This is the same embedding used as initial node features in the
    STRING-SPACE (sequence) and STRING-GNN frameworks (Hu et al., 2024).

    Supported model identifiers (Hugging Face):

    * ``Rostlab/prot_t5_xl_uniref50``  – full-precision (float32)
    * ``Rostlab/prot_t5_xl_half_uniref50-enc``  – half-precision encoder-only

    Parameters
    ----------
    model_path_or_name
        Hugging Face model name or local path.
    """

    model_type = "protein"
    available_pooling_strategies = ["mean", "max", "cls", "none"]

    # ProtT5-XL hidden dimension
    EMBEDDING_DIM = 1024

    def __init__(
        self,
        model_path_or_name: str = "Rostlab/prot_t5_xl_half_uniref50-enc",
        **kwargs: Any,
    ) -> None:
        super().__init__(model_path_or_name, **kwargs)
        self.tokenizer: T5Tokenizer | None = None

    # ------------------------------------------------------------------
    # Loading
    # ------------------------------------------------------------------

    def load(self, device: torch.device) -> None:
        """Load the ProtT5 encoder and tokenizer onto *device*."""
        if self.model is not None:
            logging.debug("ProtT5 model already loaded.")
            return

        if not self.model_name:
            raise ValueError("model_path_or_name must be provided for ProtT5Wrapper.")

        logging.info("Loading ProtT5 model '%s' …", self.model_name)
        from transformers import T5Tokenizer, T5EncoderModel
        try:
            self.tokenizer = T5Tokenizer.from_pretrained(self.model_name, do_lower_case=False)
            # Use T5EncoderModel to avoid decoder overhead
            self.model = T5EncoderModel.from_pretrained(self.model_name)
            self.model.to(device).eval()  # type: ignore[union-attr]
            self.device = device
            logging.info("ProtT5 model loaded successfully (dim=%d).", self.EMBEDDING_DIM)
        except Exception as e:
            logging.error("Failed to load ProtT5 model: %s", e)
            raise RuntimeError(f"Could not load ProtT5 model '{self.model_name}'") from e

    # ------------------------------------------------------------------
    # Preprocessing
    # ------------------------------------------------------------------

    @staticmethod
    def _prepare_sequence(sequence: str) -> str:
        """Prepare a protein sequence for ProtT5 tokenization.

        ProtT5 expects **spaces between each amino acid** and all rare /
        ambiguous residues (B, J, O, U, X, Z) mapped to ``X``.
        """
        # Remove whitespace
        sequence = sequence.strip().upper()
        # Replace rare amino acids with X
        sequence = re.sub(r"[UZOB]", "X", sequence)
        # Insert spaces between every amino acid
        return " ".join(list(sequence))

    # ------------------------------------------------------------------
    # Embedding
    # ------------------------------------------------------------------

    def embed(
        self,
        input: str,
        pooling_strategy: str = "mean",
        target_layer: int | None = None,
        **kwargs: Any,
    ) -> np.ndarray:
        """Embed a single protein sequence.

        Parameters
        ----------
        input
            Amino-acid sequence (e.g. ``"MTEYKLVVVG..."``).
        pooling_strategy
            One of ``"mean"`` (default), ``"max"``, or ``"cls"``.
        target_layer
            If given, extract hidden states from a specific encoder layer
            instead of the last one.

        Returns
        -------
        np.ndarray
            1-D array of shape ``(1024,)`` (or model hidden dim).
        """
        if self.model is None or self.device is None or self.tokenizer is None:
            raise RuntimeError("ProtT5 model not loaded. Call load() first.")
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(f"Invalid pooling '{pooling_strategy}'. Choose from {self.available_pooling_strategies}")

        prepared = self._prepare_sequence(input)
        ids = self.tokenizer(
            prepared,
            add_special_tokens=True,
            padding="longest",
            return_tensors="pt",
        )
        input_ids = ids["input_ids"].to(self.device)
        attention_mask = ids.get("attention_mask")
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)

        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_hidden_states=(target_layer is not None),
            )

            if target_layer is not None:
                hidden_states = outputs.hidden_states
                if not hidden_states:
                    raise ValueError("Model did not return hidden states.")
                if not (-len(hidden_states) <= target_layer < len(hidden_states)):
                    raise ValueError(f"Invalid target_layer {target_layer}.")
                emb = hidden_states[target_layer]
            else:
                emb = outputs.last_hidden_state  # (1, seq_len, 1024)

        # Squeeze batch dim
        if emb.dim() == 3 and emb.shape[0] == 1:
            emb = emb.squeeze(0)  # (seq_len, 1024)

        # Mask out padding tokens for mean/max pooling
        if attention_mask is not None and pooling_strategy not in ("cls", "none"):
            mask = attention_mask.squeeze(0).unsqueeze(-1).bool()  # (seq_len, 1)
            emb = emb.masked_fill(~mask, 0.0)

        if pooling_strategy == "none":
            return emb.float().cpu().numpy()
        elif pooling_strategy == "cls":
            pooled = emb[0]
        elif pooling_strategy == "max":
            if attention_mask is not None:
                emb = emb.masked_fill(~mask, float("-inf"))
            pooled = torch.max(emb, dim=0)[0]
        else:  # mean
            if attention_mask is not None:
                pooled = emb.sum(dim=0) / mask.sum(dim=0).clamp(min=1)
            else:
                pooled = emb.mean(dim=0)

        return pooled.float().cpu().numpy()

    def embed_batch(
        self,
        inputs: list[str],
        pooling_strategy: str = "mean",
        target_layer: int | None = None,
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """Embed multiple protein sequences.

        Parameters
        ----------
        inputs
            List of amino-acid sequence strings.
        pooling_strategy
            One of ``"mean"`` (default), ``"max"``, ``"cls"``.
        target_layer
            Optional encoder layer index.

        Returns
        -------
        list[np.ndarray]
            One 1-D array of shape ``(1024,)`` per input sequence.
        """
        if self.model is None or self.device is None or self.tokenizer is None:
            raise RuntimeError("ProtT5 model not loaded. Call load() first.")
        if not inputs:
            return []
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(
                f"Invalid pooling '{pooling_strategy}'. Choose from {self.available_pooling_strategies}"
            )

        batch_size = kwargs.pop("batch_size", 16)
        all_embeddings: list[np.ndarray] = []

        for start in range(0, len(inputs), batch_size):
            chunk = inputs[start : start + batch_size]
            prepared = [self._prepare_sequence(seq) for seq in chunk]
            tokenized = self.tokenizer(
                prepared,
                add_special_tokens=True,
                padding="longest",
                return_tensors="pt",
            )
            input_ids = tokenized["input_ids"].to(self.device)
            attention_mask = tokenized.get("attention_mask")
            if attention_mask is not None:
                attention_mask = attention_mask.to(self.device)

            with torch.no_grad():
                outputs = self.model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=(target_layer is not None),
                )
                if target_layer is not None:
                    emb_tensor = outputs.hidden_states[target_layer]
                else:
                    emb_tensor = outputs.last_hidden_state

            for i in range(emb_tensor.shape[0]):
                seq_emb = emb_tensor[i]
                if attention_mask is not None:
                    mask = attention_mask[i].unsqueeze(-1).bool()
                else:
                    mask = None

                if pooling_strategy == "none":
                    all_embeddings.append(seq_emb.float().cpu().numpy())
                    continue
                elif pooling_strategy == "cls":
                    pooled = seq_emb[0]
                elif pooling_strategy == "max":
                    if mask is not None:
                        seq_emb = seq_emb.masked_fill(~mask, float("-inf"))
                    pooled = torch.max(seq_emb, dim=0)[0]
                else:
                    if mask is not None:
                        pooled = (seq_emb * mask).sum(dim=0) / mask.sum(dim=0).clamp(min=1)
                    else:
                        pooled = seq_emb.mean(dim=0)

                all_embeddings.append(pooled.float().cpu().numpy())

        return all_embeddings

embed(input, pooling_strategy='mean', target_layer=None, **kwargs)

Embed a single protein sequence.

Parameters:

Name	Type	Description	Default
input	str	Amino-acid sequence (e.g. "MTEYKLVVVG...").	required
pooling_strategy	str	One of "mean" (default), "max", or "cls".	'mean'
target_layer	int | None	If given, extract hidden states from a specific encoder layer instead of the last one.	None

Returns:

Type	Description
ndarray	1-D array of shape (1024,) (or model hidden dim).

Source code in src/embpy/models/protein_models.py

def embed(
    self,
    input: str,
    pooling_strategy: str = "mean",
    target_layer: int | None = None,
    **kwargs: Any,
) -> np.ndarray:
    """Embed a single protein sequence.

    Parameters
    ----------
    input
        Amino-acid sequence (e.g. ``"MTEYKLVVVG..."``).
    pooling_strategy
        One of ``"mean"`` (default), ``"max"``, or ``"cls"``.
    target_layer
        If given, extract hidden states from a specific encoder layer
        instead of the last one.

    Returns
    -------
    np.ndarray
        1-D array of shape ``(1024,)`` (or model hidden dim).
    """
    if self.model is None or self.device is None or self.tokenizer is None:
        raise RuntimeError("ProtT5 model not loaded. Call load() first.")
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(f"Invalid pooling '{pooling_strategy}'. Choose from {self.available_pooling_strategies}")

    prepared = self._prepare_sequence(input)
    ids = self.tokenizer(
        prepared,
        add_special_tokens=True,
        padding="longest",
        return_tensors="pt",
    )
    input_ids = ids["input_ids"].to(self.device)
    attention_mask = ids.get("attention_mask")
    if attention_mask is not None:
        attention_mask = attention_mask.to(self.device)

    with torch.no_grad():
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=(target_layer is not None),
        )

        if target_layer is not None:
            hidden_states = outputs.hidden_states
            if not hidden_states:
                raise ValueError("Model did not return hidden states.")
            if not (-len(hidden_states) <= target_layer < len(hidden_states)):
                raise ValueError(f"Invalid target_layer {target_layer}.")
            emb = hidden_states[target_layer]
        else:
            emb = outputs.last_hidden_state  # (1, seq_len, 1024)

    # Squeeze batch dim
    if emb.dim() == 3 and emb.shape[0] == 1:
        emb = emb.squeeze(0)  # (seq_len, 1024)

    # Mask out padding tokens for mean/max pooling
    if attention_mask is not None and pooling_strategy not in ("cls", "none"):
        mask = attention_mask.squeeze(0).unsqueeze(-1).bool()  # (seq_len, 1)
        emb = emb.masked_fill(~mask, 0.0)

    if pooling_strategy == "none":
        return emb.float().cpu().numpy()
    elif pooling_strategy == "cls":
        pooled = emb[0]
    elif pooling_strategy == "max":
        if attention_mask is not None:
            emb = emb.masked_fill(~mask, float("-inf"))
        pooled = torch.max(emb, dim=0)[0]
    else:  # mean
        if attention_mask is not None:
            pooled = emb.sum(dim=0) / mask.sum(dim=0).clamp(min=1)
        else:
            pooled = emb.mean(dim=0)

    return pooled.float().cpu().numpy()

embed_batch(inputs, pooling_strategy='mean', target_layer=None, **kwargs)

Embed multiple protein sequences.

Parameters:

Name	Type	Description	Default
inputs	list[str]	List of amino-acid sequence strings.	required
pooling_strategy	str	One of "mean" (default), "max", "cls".	'mean'
target_layer	int | None	Optional encoder layer index.	None

Returns:

Type	Description
list[ndarray]	One 1-D array of shape (1024,) per input sequence.

Source code in src/embpy/models/protein_models.py

def embed_batch(
    self,
    inputs: list[str],
    pooling_strategy: str = "mean",
    target_layer: int | None = None,
    **kwargs: Any,
) -> list[np.ndarray]:
    """Embed multiple protein sequences.

    Parameters
    ----------
    inputs
        List of amino-acid sequence strings.
    pooling_strategy
        One of ``"mean"`` (default), ``"max"``, ``"cls"``.
    target_layer
        Optional encoder layer index.

    Returns
    -------
    list[np.ndarray]
        One 1-D array of shape ``(1024,)`` per input sequence.
    """
    if self.model is None or self.device is None or self.tokenizer is None:
        raise RuntimeError("ProtT5 model not loaded. Call load() first.")
    if not inputs:
        return []
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(
            f"Invalid pooling '{pooling_strategy}'. Choose from {self.available_pooling_strategies}"
        )

    batch_size = kwargs.pop("batch_size", 16)
    all_embeddings: list[np.ndarray] = []

    for start in range(0, len(inputs), batch_size):
        chunk = inputs[start : start + batch_size]
        prepared = [self._prepare_sequence(seq) for seq in chunk]
        tokenized = self.tokenizer(
            prepared,
            add_special_tokens=True,
            padding="longest",
            return_tensors="pt",
        )
        input_ids = tokenized["input_ids"].to(self.device)
        attention_mask = tokenized.get("attention_mask")
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)

        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_hidden_states=(target_layer is not None),
            )
            if target_layer is not None:
                emb_tensor = outputs.hidden_states[target_layer]
            else:
                emb_tensor = outputs.last_hidden_state

        for i in range(emb_tensor.shape[0]):
            seq_emb = emb_tensor[i]
            if attention_mask is not None:
                mask = attention_mask[i].unsqueeze(-1).bool()
            else:
                mask = None

            if pooling_strategy == "none":
                all_embeddings.append(seq_emb.float().cpu().numpy())
                continue
            elif pooling_strategy == "cls":
                pooled = seq_emb[0]
            elif pooling_strategy == "max":
                if mask is not None:
                    seq_emb = seq_emb.masked_fill(~mask, float("-inf"))
                pooled = torch.max(seq_emb, dim=0)[0]
            else:
                if mask is not None:
                    pooled = (seq_emb * mask).sum(dim=0) / mask.sum(dim=0).clamp(min=1)
                else:
                    pooled = seq_emb.mean(dim=0)

            all_embeddings.append(pooled.float().cpu().numpy())

    return all_embeddings

load(device)

Load the ProtT5 encoder and tokenizer onto device.

Source code in src/embpy/models/protein_models.py

def load(self, device: torch.device) -> None:
    """Load the ProtT5 encoder and tokenizer onto *device*."""
    if self.model is not None:
        logging.debug("ProtT5 model already loaded.")
        return

    if not self.model_name:
        raise ValueError("model_path_or_name must be provided for ProtT5Wrapper.")

    logging.info("Loading ProtT5 model '%s' …", self.model_name)
    from transformers import T5Tokenizer, T5EncoderModel
    try:
        self.tokenizer = T5Tokenizer.from_pretrained(self.model_name, do_lower_case=False)
        # Use T5EncoderModel to avoid decoder overhead
        self.model = T5EncoderModel.from_pretrained(self.model_name)
        self.model.to(device).eval()  # type: ignore[union-attr]
        self.device = device
        logging.info("ProtT5 model loaded successfully (dim=%d).", self.EMBEDDING_DIM)
    except Exception as e:
        logging.error("Failed to load ProtT5 model: %s", e)
        raise RuntimeError(f"Could not load ProtT5 model '{self.model_name}'") from e

Text & Description Models

embpy.models.text_models

LlamaEmbeddingWrapper

Bases: TextLLMWrapper

Extract embeddings from LLaMA decoder-only models.

LLaMA models are autoregressive (decoder-only), so the natural embedding is the last non-padding token representation. This wrapper loads the model via AutoModelForCausalLM in bfloat16 for memory efficiency and defaults to last_token pooling.

Requires a HuggingFace token for gated models -- set HF_TOKEN environment variable or run huggingface-cli login.

Parameters:

Name	Type	Description	Default
model_path_or_name	str	HuggingFace model ID (e.g. "meta-llama/Llama-3.2-3B").	'meta-llama/Llama-3.2-3B'
max_length	int | None	Maximum sequence length. Defaults to 4096.	4096

Source code in src/embpy/models/text_models.py

class LlamaEmbeddingWrapper(TextLLMWrapper):
    """Extract embeddings from LLaMA decoder-only models.

    LLaMA models are autoregressive (decoder-only), so the natural
    embedding is the **last non-padding token** representation. This
    wrapper loads the model via ``AutoModelForCausalLM`` in bfloat16
    for memory efficiency and defaults to ``last_token`` pooling.

    Requires a HuggingFace token for gated models -- set ``HF_TOKEN``
    environment variable or run ``huggingface-cli login``.

    Parameters
    ----------
    model_path_or_name
        HuggingFace model ID (e.g. ``"meta-llama/Llama-3.2-3B"``).
    max_length
        Maximum sequence length. Defaults to 4096.
    """

    available_pooling_strategies = ["last_token", "mean", "max", "none"]

    def __init__(
        self,
        model_path_or_name: str = "meta-llama/Llama-3.2-3B",
        max_length: int | None = 4096,
        **kwargs,
    ):
        super().__init__(model_path_or_name, max_length=max_length, **kwargs)

    def load(self, device: torch.device):
        """Load a LLaMA model in bfloat16 for embedding extraction."""
        if self.model is not None:
            return

        if self.model_name is None:
            raise ValueError("model_path_or_name must be set.")

        import os

        logging.info("Loading LLaMA '%s' in bfloat16...", self.model_name)
        from transformers import AutoModelForCausalLM, AutoTokenizer

        try:

            token = os.environ.get("HF_TOKEN")

            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_name, token=token,
            )
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token
                self.tokenizer.pad_token_id = self.tokenizer.eos_token_id

            # No ``device_map``: it needs ``accelerate`` (not in any embpy extra).
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                torch_dtype=torch.bfloat16,
                token=token,
            )
            self.model = self.model.to(device)
            self.model.eval()
            self.device = device

            if self.max_length is None:
                self.max_length = 4096
            self.tokenizer.model_max_length = self.max_length

            logging.info(
                "LLaMA loaded: %s, dtype=bfloat16, max_length=%d",
                self.model_name, self.max_length,
            )
        except Exception as e:
            logging.error("Failed to load LLaMA '%s': %s", self.model_name, e)
            raise RuntimeError(f"Could not load LLaMA '{self.model_name}'") from e

    def embed(
        self,
        input: str,
        pooling_strategy: str = "last_token",
        target_layer: int | None = None,
        **kwargs: Any,
    ) -> np.ndarray:
        """Compute embedding from a LLaMA model.

        Default pooling is ``last_token`` (last non-padding token),
        which is the natural choice for decoder-only models.
        """
        text = input
        if self.model is None or self.device is None:
            raise RuntimeError("Model not loaded. Call load() first.")
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(
                f"Invalid pooling '{pooling_strategy}'. "
                f"Available: {self.available_pooling_strategies}"
            )

        return self._process_batch_chunk([text], pooling_strategy, target_layer)[0]

    def embed_batch(
        self,
        inputs: Sequence[str],
        pooling_strategy: str = "last_token",
        target_layer: int | None = None,
        batch_size: int | None = None,
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """Compute embeddings for several texts from a LLaMA model.

        Same pooling, ``target_layer`` and ``max_length`` handling as
        :meth:`embed` (default pooling is ``last_token``). The texts are
        tokenized together with padding and an attention mask, so padding
        never contributes to the pooled vector. ``batch_size`` splits
        ``inputs`` into chunks to bound memory.
        """
        return super().embed_batch(
            inputs,
            pooling_strategy=pooling_strategy,
            target_layer=target_layer,
            batch_size=batch_size,
            **kwargs,
        )

    def _process_batch_chunk(
        self,
        inputs: Sequence[str],
        pooling_strategy: str,
        target_layer: int | None = None,
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """Embed one chunk of texts; shared by :meth:`embed` and :meth:`embed_batch`.

        Decoder-only models expose ``hidden_states`` rather than
        ``last_hidden_state``, so the layer is always taken from
        ``hidden_states`` (the last one unless ``target_layer`` is given).
        """
        tokenized = self.tokenizer(
            list(inputs),
            return_tensors="pt",
            truncation=True,
            padding=True,
            max_length=self.max_length,
        )
        input_ids = tokenized["input_ids"].to(self.device)
        attention_mask = tokenized["attention_mask"].to(self.device)

        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_hidden_states=True,
            )
            if target_layer is not None:
                hidden = outputs.hidden_states[target_layer]  # (B, L, D)
            else:
                hidden = outputs.hidden_states[-1]  # (B, L, D)

        if pooling_strategy == "none":
            # Drop padding so each result matches what embed() gives for that text alone.
            pooled = [h[m.bool()].float().cpu().numpy() for h, m in zip(hidden, attention_mask, strict=True)]
        else:
            if pooling_strategy == "last_token":
                batch_pooled = self._last_token_pool(hidden, attention_mask)
            elif pooling_strategy == "mean":
                mask = attention_mask.unsqueeze(-1).float()  # (B, L, 1)
                batch_pooled = torch.sum(hidden * mask, dim=1) / torch.sum(mask, dim=1)
            else:  # max
                batch_pooled = hidden.masked_fill(attention_mask.unsqueeze(-1) == 0, float("-inf")).max(dim=1).values
            pooled = list(batch_pooled.cpu().float().numpy())

        del input_ids, attention_mask, hidden, outputs
        if self.device.type == "cuda":
            torch.cuda.empty_cache()

        return pooled

embed(input, pooling_strategy='last_token', target_layer=None, **kwargs)

Compute embedding from a LLaMA model.

Default pooling is last_token (last non-padding token), which is the natural choice for decoder-only models.

Source code in src/embpy/models/text_models.py

def embed(
    self,
    input: str,
    pooling_strategy: str = "last_token",
    target_layer: int | None = None,
    **kwargs: Any,
) -> np.ndarray:
    """Compute embedding from a LLaMA model.

    Default pooling is ``last_token`` (last non-padding token),
    which is the natural choice for decoder-only models.
    """
    text = input
    if self.model is None or self.device is None:
        raise RuntimeError("Model not loaded. Call load() first.")
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(
            f"Invalid pooling '{pooling_strategy}'. "
            f"Available: {self.available_pooling_strategies}"
        )

    return self._process_batch_chunk([text], pooling_strategy, target_layer)[0]

embed_batch(inputs, pooling_strategy='last_token', target_layer=None, batch_size=None, **kwargs)

Compute embeddings for several texts from a LLaMA model.

Same pooling, target_layer and max_length handling as :meth:embed (default pooling is last_token). The texts are tokenized together with padding and an attention mask, so padding never contributes to the pooled vector. batch_size splits inputs into chunks to bound memory.

Source code in src/embpy/models/text_models.py

def embed_batch(
    self,
    inputs: Sequence[str],
    pooling_strategy: str = "last_token",
    target_layer: int | None = None,
    batch_size: int | None = None,
    **kwargs: Any,
) -> list[np.ndarray]:
    """Compute embeddings for several texts from a LLaMA model.

    Same pooling, ``target_layer`` and ``max_length`` handling as
    :meth:`embed` (default pooling is ``last_token``). The texts are
    tokenized together with padding and an attention mask, so padding
    never contributes to the pooled vector. ``batch_size`` splits
    ``inputs`` into chunks to bound memory.
    """
    return super().embed_batch(
        inputs,
        pooling_strategy=pooling_strategy,
        target_layer=target_layer,
        batch_size=batch_size,
        **kwargs,
    )

load(device)

Load a LLaMA model in bfloat16 for embedding extraction.

Source code in src/embpy/models/text_models.py

def load(self, device: torch.device):
    """Load a LLaMA model in bfloat16 for embedding extraction."""
    if self.model is not None:
        return

    if self.model_name is None:
        raise ValueError("model_path_or_name must be set.")

    import os

    logging.info("Loading LLaMA '%s' in bfloat16...", self.model_name)
    from transformers import AutoModelForCausalLM, AutoTokenizer

    try:

        token = os.environ.get("HF_TOKEN")

        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_name, token=token,
        )
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id

        # No ``device_map``: it needs ``accelerate`` (not in any embpy extra).
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            torch_dtype=torch.bfloat16,
            token=token,
        )
        self.model = self.model.to(device)
        self.model.eval()
        self.device = device

        if self.max_length is None:
            self.max_length = 4096
        self.tokenizer.model_max_length = self.max_length

        logging.info(
            "LLaMA loaded: %s, dtype=bfloat16, max_length=%d",
            self.model_name, self.max_length,
        )
    except Exception as e:
        logging.error("Failed to load LLaMA '%s': %s", self.model_name, e)
        raise RuntimeError(f"Could not load LLaMA '{self.model_name}'") from e

TextLLMWrapper

Bases: BaseModelWrapper

Wrapper for general-purpose text embedding models

Source code in src/embpy/models/text_models.py

class TextLLMWrapper(BaseModelWrapper):
    """Wrapper for general-purpose text embedding models"""

    model_type = "text"
    available_pooling_strategies = ["mean", "max", "cls", "last_token", "none"]

    def __init__(
        self,
        model_path_or_name: str = "sentence-transformers/all-MiniLM-L6-v2",
        max_length: int | None = None,
        cls_token_position: int = 0,
        **kwargs,
    ):
        """
        Initializes the Text LLM wrapper.

        Args:
            model_path_or_name (str): Hugging Face model name or Sentence-Transformer name.
            max_length (Optional[int]): Maximum sequence length. If None, will be determined
                                      from model config or default to 512.
            cls_token_position (int): Position of CLS token (default: 0 for BERT-style models).
            **kwargs: Additional config.
        """
        super().__init__(model_path_or_name, **kwargs)
        self.tokenizer = None
        self.max_length = max_length
        self.cls_token_position = cls_token_position

    def load(self, device: torch.device):
        """Loads the text model and potentially tokenizer."""
        if self.model is not None:
            logging.debug("Text model already loaded.")
            return

        if self.model_name is None:
            raise ValueError("model_path_or_name must be set for TextLLMWrapper.")

        logging.info(f"Loading text model '{self.model_name}'...")
        from transformers import AutoTokenizer, AutoModel

        try:
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            self.model = AutoModel.from_pretrained(self.model_name)
            self.model = self.model.to(device)
            self.model.eval()
            self.device = device

            # Determine max_length if not provided
            if self.max_length is None:
                model_max = getattr(self.model.config, "max_position_embeddings", None)
                if model_max is None:
                    model_max = getattr(self.tokenizer, "model_max_length", None)

                if model_max is None or model_max > 1024:
                    model_max = 512

                self.max_length = int(model_max)

            # Ensure tokenizer respects our max_length
            self.tokenizer.model_max_length = self.max_length
            logging.info(f"Using max_length={self.max_length}, cls_token_position={self.cls_token_position}")
        except Exception as e:
            logging.error(f"Failed to load text model '{self.model_name}': {e}")
            raise RuntimeError(f"Could not load text model '{self.model_name}'") from e

    def _get_cls_embedding(self, embeddings_tensor: torch.Tensor) -> torch.Tensor:
        """
        Extracts CLS token embedding from the sequence.

        Args:
            embeddings_tensor: Tensor of shape (seq_len, hidden_dim) or (batch, seq_len, hidden_dim)

        Returns
        -------
            torch.Tensor: CLS token embedding
        """
        if embeddings_tensor.dim() == 3:  # Batch dimension present
            return embeddings_tensor[:, self.cls_token_position, :]
        else:  # Single sequence
            return embeddings_tensor[self.cls_token_position, :]

    def _preprocess_text_hf(self, text: str) -> Any:
        """Tokenizes text using Hugging Face tokenizer."""
        if not self.tokenizer:
            raise RuntimeError("Hugging Face Tokenizer not loaded.")
        return self.tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=self.max_length)

    def embed(
        self,
        input: str,
        pooling_strategy: str = "mean",
        target_layer: int | None = None,  # For HF models if needed
        **kwargs: Any,
    ) -> np.ndarray:
        """
        Computes the embedding for a single text string.

        Args:
            input (str): The input text string.
            pooling_strategy (str): Pooling strategy ('mean', 'max', 'cls'). Sentence-Transformers typically uses 'mean'.
            target_layer (Optional[int]): For HF models, extract embeddings from this hidden layer index.
            **kwargs: Additional arguments (e.g., `normalize_embeddings` for SentenceTransformer).

        Returns
        -------
            np.ndarray: The resulting embedding vector.
        """
        text = input
        if self.model is None or self.device is None:
            raise RuntimeError("Text model not loaded. Call load() first.")
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(
                f"Invalid pooling strategy '{pooling_strategy}'. Available: {self.available_pooling_strategies}"
            )

        # Preprocess
        tokenized_input = self._preprocess_text_hf(text)
        input_ids = tokenized_input["input_ids"].to(self.device)
        attention_mask = tokenized_input.get("attention_mask", None)
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)

        # Inference
        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=(target_layer is not None)
            )
            if target_layer is not None:
                embeddings_tensor = outputs.hidden_states[target_layer]  # (B, L, D)
            else:
                embeddings_tensor = outputs.last_hidden_state  # (B, L, D)

        # Squeeze batch dim
        if embeddings_tensor.dim() == 3 and embeddings_tensor.shape[0] == 1:
            embeddings_tensor = embeddings_tensor.squeeze(0)  # (L, D)

        # Apply pooling (Needs attention mask awareness for HF mean pooling)
        if pooling_strategy == "none":
            pooled_embedding = embeddings_tensor.cpu().numpy()
        elif pooling_strategy == "cls":
            pooled_embedding = self._get_cls_embedding(embeddings_tensor).cpu().numpy()
        elif pooling_strategy == "last_token":
            pooled_embedding = self._last_token_pool(embeddings_tensor, attention_mask).cpu().numpy()
        else:
            # Apply proper pooling with attention mask
            if pooling_strategy == "mean" and attention_mask is not None:
                # Mask-aware mean pooling
                mask = attention_mask.squeeze(0).unsqueeze(-1)  # (L, 1)
                sum_embeddings = torch.sum(embeddings_tensor * mask, dim=0)
                sum_mask = torch.sum(mask, dim=0)
                pooled_embedding = (sum_embeddings / sum_mask).cpu().numpy()
            else:
                # Simple mean/max over sequence length
                pooled_embedding = self._apply_pooling(embeddings_tensor, pooling_strategy)

        # Cleanup
        del input_ids, attention_mask, embeddings_tensor, outputs
        if self.device.type == "cuda":
            torch.cuda.empty_cache()

        return pooled_embedding

    def embed_batch(
        self,
        inputs: Sequence[str],
        pooling_strategy: str = "mean",
        target_layer: int | None = None,
        batch_size: int | None = None,
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """
        Computes embeddings for a batch of text strings efficiently.

        Args:
            inputs (Sequence[str]): A list/tuple of text strings.
            pooling_strategy (str): Pooling strategy ('mean', 'max', 'cls').
            target_layer (Optional[int]): For HF models, extract embeddings from this hidden layer index.
            batch_size (Optional[int]): Maximum batch size for processing. If None, processes all inputs at once.
                                   Use smaller values (e.g., 16-64) to avoid OOM errors.
            **kwargs: Additional arguments for the model.

        Returns
        -------
            list[np.ndarray]: A list of resulting embedding vectors.
        """
        if self.model is None or self.device is None:
            raise RuntimeError("Text model not loaded. Call load() first.")
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(
                f"Invalid pooling strategy '{pooling_strategy}'. Available: {self.available_pooling_strategies}"
            )
        if not inputs:
            return []

        # If batch_size is specified and smaller than input size, process in chunks
        if batch_size is not None and len(inputs) > batch_size:
            logging.info(f"Processing {len(inputs)} inputs in chunks of {batch_size}")
            results = []
            for i in range(0, len(inputs), batch_size):
                chunk = inputs[i : i + batch_size]
                chunk_results = self._process_batch_chunk(chunk, pooling_strategy, target_layer, **kwargs)
                results.extend(chunk_results)
            return results
        else:
            # Process all inputs at once (original behavior)
            return self._process_batch_chunk(inputs, pooling_strategy, target_layer, **kwargs)

    def _process_batch_chunk(
        self,
        inputs: Sequence[str],
        pooling_strategy: str,
        target_layer: int | None = None,
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """Process a single batch chunk."""
        logging.debug(f"Processing batch chunk of {len(inputs)} text inputs")

        # Tokenize all inputs in the chunk
        tokenized_batch = self.tokenizer(
            list(inputs),
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.max_length,
        )

        input_ids = tokenized_batch["input_ids"].to(self.device)
        attention_mask = tokenized_batch.get("attention_mask", None)
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)

        # Batch inference
        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=(target_layer is not None)
            )
            if target_layer is not None:
                embeddings_tensor = outputs.hidden_states[target_layer]  # (B, L, D)
            else:
                embeddings_tensor = outputs.last_hidden_state  # (B, L, D)

        # Apply pooling for each item in batch
        results = []
        batch_size = embeddings_tensor.shape[0]

        for i in range(batch_size):
            seq_embeddings = embeddings_tensor[i]
            seq_mask = None if attention_mask is None else attention_mask[i]

            if pooling_strategy == "none":
                pooled_embedding = seq_embeddings.cpu().numpy()
            elif pooling_strategy == "cls":
                pooled_embedding = seq_embeddings[self.cls_token_position, :].cpu().numpy()
            elif pooling_strategy == "last_token":
                # For batch processing, we need to handle single sequence
                single_hidden = seq_embeddings.unsqueeze(0)  # Add batch dim
                single_mask = seq_mask.unsqueeze(0) if seq_mask is not None else None
                pooled_embedding = self._last_token_pool(single_hidden, single_mask).cpu().numpy()
            else:
                if pooling_strategy == "mean" and seq_mask is not None:
                    mask = seq_mask.unsqueeze(-1)
                    sum_embeddings = torch.sum(seq_embeddings * mask, dim=0)
                    sum_mask = torch.sum(mask, dim=0)
                    pooled_embedding = (sum_embeddings / sum_mask).cpu().numpy()
                else:
                    pooled_embedding = self._apply_pooling(seq_embeddings, pooling_strategy)

            results.append(pooled_embedding)

        # Cleanup
        del input_ids, attention_mask, embeddings_tensor, outputs
        if self.device.type == "cuda":
            torch.cuda.empty_cache()

        return results

    def _last_token_pool(self, last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        """
        Pool by taking the last token's embedding, handling both left and right padding.

        Args:
            last_hidden_states: Tensor of shape (batch_size, seq_len, hidden_dim) or (seq_len, hidden_dim)
            attention_mask: Tensor of shape (batch_size, seq_len) or (seq_len,)

        Returns
        -------
            torch.Tensor: Last token embeddings of shape (batch_size, hidden_dim) or (hidden_dim,)
        """
        # Handle single sequence case
        if last_hidden_states.dim() == 2:
            last_hidden_states = last_hidden_states.unsqueeze(0)  # Add batch dim
            attention_mask = attention_mask.unsqueeze(0)
            squeeze_batch = True
        else:
            squeeze_batch = False

        left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
        if left_padding:
            result = last_hidden_states[:, -1]
        else:
            sequence_lengths = attention_mask.sum(dim=1) - 1
            batch_size = last_hidden_states.shape[0]
            result = last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]

        if squeeze_batch:
            result = result.squeeze(0)

        return result

__init__(model_path_or_name='sentence-transformers/all-MiniLM-L6-v2', max_length=None, cls_token_position=0, **kwargs)

Initializes the Text LLM wrapper.

Args: model_path_or_name (str): Hugging Face model name or Sentence-Transformer name. max_length (Optional[int]): Maximum sequence length. If None, will be determined from model config or default to 512. cls_token_position (int): Position of CLS token (default: 0 for BERT-style models). **kwargs: Additional config.

Source code in src/embpy/models/text_models.py

def __init__(
    self,
    model_path_or_name: str = "sentence-transformers/all-MiniLM-L6-v2",
    max_length: int | None = None,
    cls_token_position: int = 0,
    **kwargs,
):
    """
    Initializes the Text LLM wrapper.

    Args:
        model_path_or_name (str): Hugging Face model name or Sentence-Transformer name.
        max_length (Optional[int]): Maximum sequence length. If None, will be determined
                                  from model config or default to 512.
        cls_token_position (int): Position of CLS token (default: 0 for BERT-style models).
        **kwargs: Additional config.
    """
    super().__init__(model_path_or_name, **kwargs)
    self.tokenizer = None
    self.max_length = max_length
    self.cls_token_position = cls_token_position

embed(input, pooling_strategy='mean', target_layer=None, **kwargs)

Computes the embedding for a single text string.

Args: input (str): The input text string. pooling_strategy (str): Pooling strategy ('mean', 'max', 'cls'). Sentence-Transformers typically uses 'mean'. target_layer (Optional[int]): For HF models, extract embeddings from this hidden layer index. **kwargs: Additional arguments (e.g., normalize_embeddings for SentenceTransformer).

Returns:

Type	Description
np.ndarray: The resulting embedding vector.

Source code in src/embpy/models/text_models.py

def embed(
    self,
    input: str,
    pooling_strategy: str = "mean",
    target_layer: int | None = None,  # For HF models if needed
    **kwargs: Any,
) -> np.ndarray:
    """
    Computes the embedding for a single text string.

    Args:
        input (str): The input text string.
        pooling_strategy (str): Pooling strategy ('mean', 'max', 'cls'). Sentence-Transformers typically uses 'mean'.
        target_layer (Optional[int]): For HF models, extract embeddings from this hidden layer index.
        **kwargs: Additional arguments (e.g., `normalize_embeddings` for SentenceTransformer).

    Returns
    -------
        np.ndarray: The resulting embedding vector.
    """
    text = input
    if self.model is None or self.device is None:
        raise RuntimeError("Text model not loaded. Call load() first.")
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(
            f"Invalid pooling strategy '{pooling_strategy}'. Available: {self.available_pooling_strategies}"
        )

    # Preprocess
    tokenized_input = self._preprocess_text_hf(text)
    input_ids = tokenized_input["input_ids"].to(self.device)
    attention_mask = tokenized_input.get("attention_mask", None)
    if attention_mask is not None:
        attention_mask = attention_mask.to(self.device)

    # Inference
    with torch.no_grad():
        outputs = self.model(
            input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=(target_layer is not None)
        )
        if target_layer is not None:
            embeddings_tensor = outputs.hidden_states[target_layer]  # (B, L, D)
        else:
            embeddings_tensor = outputs.last_hidden_state  # (B, L, D)

    # Squeeze batch dim
    if embeddings_tensor.dim() == 3 and embeddings_tensor.shape[0] == 1:
        embeddings_tensor = embeddings_tensor.squeeze(0)  # (L, D)

    # Apply pooling (Needs attention mask awareness for HF mean pooling)
    if pooling_strategy == "none":
        pooled_embedding = embeddings_tensor.cpu().numpy()
    elif pooling_strategy == "cls":
        pooled_embedding = self._get_cls_embedding(embeddings_tensor).cpu().numpy()
    elif pooling_strategy == "last_token":
        pooled_embedding = self._last_token_pool(embeddings_tensor, attention_mask).cpu().numpy()
    else:
        # Apply proper pooling with attention mask
        if pooling_strategy == "mean" and attention_mask is not None:
            # Mask-aware mean pooling
            mask = attention_mask.squeeze(0).unsqueeze(-1)  # (L, 1)
            sum_embeddings = torch.sum(embeddings_tensor * mask, dim=0)
            sum_mask = torch.sum(mask, dim=0)
            pooled_embedding = (sum_embeddings / sum_mask).cpu().numpy()
        else:
            # Simple mean/max over sequence length
            pooled_embedding = self._apply_pooling(embeddings_tensor, pooling_strategy)

    # Cleanup
    del input_ids, attention_mask, embeddings_tensor, outputs
    if self.device.type == "cuda":
        torch.cuda.empty_cache()

    return pooled_embedding

embed_batch(inputs, pooling_strategy='mean', target_layer=None, batch_size=None, **kwargs)

Computes embeddings for a batch of text strings efficiently.

Args: inputs (Sequence[str]): A list/tuple of text strings. pooling_strategy (str): Pooling strategy ('mean', 'max', 'cls'). target_layer (Optional[int]): For HF models, extract embeddings from this hidden layer index. batch_size (Optional[int]): Maximum batch size for processing. If None, processes all inputs at once. Use smaller values (e.g., 16-64) to avoid OOM errors. **kwargs: Additional arguments for the model.

Returns:

Type	Description
list[np.ndarray]: A list of resulting embedding vectors.

Source code in src/embpy/models/text_models.py

def embed_batch(
    self,
    inputs: Sequence[str],
    pooling_strategy: str = "mean",
    target_layer: int | None = None,
    batch_size: int | None = None,
    **kwargs: Any,
) -> list[np.ndarray]:
    """
    Computes embeddings for a batch of text strings efficiently.

    Args:
        inputs (Sequence[str]): A list/tuple of text strings.
        pooling_strategy (str): Pooling strategy ('mean', 'max', 'cls').
        target_layer (Optional[int]): For HF models, extract embeddings from this hidden layer index.
        batch_size (Optional[int]): Maximum batch size for processing. If None, processes all inputs at once.
                               Use smaller values (e.g., 16-64) to avoid OOM errors.
        **kwargs: Additional arguments for the model.

    Returns
    -------
        list[np.ndarray]: A list of resulting embedding vectors.
    """
    if self.model is None or self.device is None:
        raise RuntimeError("Text model not loaded. Call load() first.")
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(
            f"Invalid pooling strategy '{pooling_strategy}'. Available: {self.available_pooling_strategies}"
        )
    if not inputs:
        return []

    # If batch_size is specified and smaller than input size, process in chunks
    if batch_size is not None and len(inputs) > batch_size:
        logging.info(f"Processing {len(inputs)} inputs in chunks of {batch_size}")
        results = []
        for i in range(0, len(inputs), batch_size):
            chunk = inputs[i : i + batch_size]
            chunk_results = self._process_batch_chunk(chunk, pooling_strategy, target_layer, **kwargs)
            results.extend(chunk_results)
        return results
    else:
        # Process all inputs at once (original behavior)
        return self._process_batch_chunk(inputs, pooling_strategy, target_layer, **kwargs)

load(device)

Loads the text model and potentially tokenizer.

Source code in src/embpy/models/text_models.py

def load(self, device: torch.device):
    """Loads the text model and potentially tokenizer."""
    if self.model is not None:
        logging.debug("Text model already loaded.")
        return

    if self.model_name is None:
        raise ValueError("model_path_or_name must be set for TextLLMWrapper.")

    logging.info(f"Loading text model '{self.model_name}'...")
    from transformers import AutoTokenizer, AutoModel

    try:
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModel.from_pretrained(self.model_name)
        self.model = self.model.to(device)
        self.model.eval()
        self.device = device

        # Determine max_length if not provided
        if self.max_length is None:
            model_max = getattr(self.model.config, "max_position_embeddings", None)
            if model_max is None:
                model_max = getattr(self.tokenizer, "model_max_length", None)

            if model_max is None or model_max > 1024:
                model_max = 512

            self.max_length = int(model_max)

        # Ensure tokenizer respects our max_length
        self.tokenizer.model_max_length = self.max_length
        logging.info(f"Using max_length={self.max_length}, cls_token_position={self.cls_token_position}")
    except Exception as e:
        logging.error(f"Failed to load text model '{self.model_name}': {e}")
        raise RuntimeError(f"Could not load text model '{self.model_name}'") from e

Molecule Models

embpy.models.molecule_models

ChembertaWrapper

Bases: BaseModelWrapper

Wrapper for ChemBERTa models for SMILES strings.

Implements: - embed(smiles, pooling_strategy, use_pooler) → np.ndarray - embed_batch([...], pooling_strategy, use_pooler) → list[np.ndarray]

Pooling strategies: • cls → take token 0 from last_hidden_state • mean → mean over all tokens (mask‐aware) • max → max over all tokens (mask‐aware)

You can also set use_pooler=True to get outputs.pooler_output (if your model has one).

Source code in src/embpy/models/molecule_models.py

class ChembertaWrapper(BaseModelWrapper):
    """
    Wrapper for ChemBERTa models for SMILES strings.

    Implements:
      - embed(smiles, pooling_strategy, use_pooler) → np.ndarray
      - embed_batch([...], pooling_strategy, use_pooler) → list[np.ndarray]

    Pooling strategies:
      • cls  → take token 0 from last_hidden_state
      • mean → mean over all tokens (mask‐aware)
      • max  → max  over all tokens (mask‐aware)

    You can also set `use_pooler=True` to get `outputs.pooler_output` (if your model has one).
    """

    model_type = "molecule"
    available_pooling_strategies = ["cls", "mean", "max", "none"]

    def __init__(self, model_path_or_name: str = "DeepChem/ChemBERTa-77M-MTR", **kwargs):
        super().__init__(model_path_or_name, **kwargs)
        self.tokenizer: Any = None
        self.model: Any = None
        self.device: torch.device | None = None
        self.max_len: int | None = None

    def load(self, device: torch.device):
        """Load the dataset"""
        if self.model is not None:
            return
        logging.info(f"Loading ChemBERTa '{self.model_name}'…")
        from transformers import AutoTokenizer, AutoModel
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        import transformers
        prev_level = transformers.logging.get_verbosity()
        transformers.logging.set_verbosity_error()
        self.model = AutoModel.from_pretrained(self.model_name).to(device).eval()
        transformers.logging.set_verbosity(prev_level)
        self.device = device

        # derive a safe max sequence length (RoBERTa-style ≈512)
        mmax = getattr(self.model.config, "max_position_embeddings", None)
        if mmax is None or mmax > 512:
            mmax = 512
        self.max_len = int(mmax)
        # make tokenizer aware (some HF tokenizers have "no max" by default)
        self.tokenizer.model_max_length = self.max_len
        logging.info(f"Using model_max_length={self.max_len}")

    def _token_length(self, smiles: str) -> int:
        """Tokenized length incl. special tokens, without truncation/padding."""
        if self.tokenizer is None:
            raise RuntimeError("Call load() first.")
        toks: BatchEncoding = self.tokenizer(  # type: ignore[operator]
            smiles,
            add_special_tokens=True,
            padding=False,
            truncation=False,
            return_attention_mask=False,
            return_token_type_ids=False,
        )
        input_ids: list[int] = toks["input_ids"]  # type: ignore[assignment]
        return len(input_ids)

    def _preprocess_smiles(self, smiles: str) -> dict[str, torch.Tensor]:
        if self.tokenizer is None:
            raise RuntimeError("Call load() first.")
        # we’ll still pass max_length for safety even though we pre-check length
        return self.tokenizer(  # type: ignore[operator]
            smiles,
            return_tensors="pt",
            padding=False,
            truncation=True,
            max_length=self.max_len,
        )

    def embed(
        self,
        input: str,
        pooling_strategy: str = "mean",
        use_pooler: bool = False,
        **kwargs: Any,
    ) -> np.ndarray | None:  # << allow None for skipped items
        """Embeddings function, which generates the embeddings for a specific provided model"""
        if self.model is None or self.device is None:
            raise RuntimeError("Call load() first.")
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(f"pooling_strategy must be one of {self.available_pooling_strategies}")

        # --- Skip too-long SMILES ---
        L = self._token_length(input)
        if self.max_len is not None and L > self.max_len:
            logging.warning(f"Skipping SMILES ({L} tokens > max {self.max_len}): {input[:80]}...")
            return None

        # 1) tokenize & move to device
        enc = self._preprocess_smiles(input)
        input_ids = enc["input_ids"].to(self.device)
        attention_mask = enc["attention_mask"].to(self.device)

        # 2) forward
        with torch.no_grad():
            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
            last_hidden = outputs.last_hidden_state  # (1, seq_len, hidden_dim)
            pooler_out = getattr(outputs, "pooler_output", None)

        # 3) use pooler if requested
        if use_pooler and pooler_out is not None:
            return pooler_out[0].cpu().numpy()

        # 4) squeeze batch dim
        hidden = last_hidden.squeeze(0)  # (seq_len, hidden_dim)

        # 5) masking-aware pooling
        if pooling_strategy == "none":
            return hidden.cpu().numpy()
        elif pooling_strategy == "cls":
            vec = hidden[0]
        else:
            mask = attention_mask.squeeze(0).unsqueeze(-1)  # (seq_len, 1)
            if pooling_strategy == "mean":
                summed = (hidden * mask).sum(dim=0)
                vec = summed / mask.sum(dim=0).clamp(min=1)  # safe denom
            else:  # 'max'
                neg_inf = torch.finfo(hidden.dtype).min
                masked = hidden.masked_fill(mask == 0, neg_inf)
                vec = masked.max(dim=0).values

        return vec.cpu().numpy()

    def embed_batch(
        self,
        input: list[str],
        pooling_strategy: str = "mean",
        use_pooler: bool = False,
        batch_size: int = 64,
        **kwargs: Any,
    ) -> list[np.ndarray | None]:
        """
        Embed a batch of SMILES strings with true batched inference.

        Filters out too-long SMILES, batches the valid ones for a
        single forward pass per chunk, then reassembles the results
        preserving original order.

        Returns a list of 1D numpy arrays (or None for skipped SMILES).
        """
        if self.model is None or self.device is None or self.tokenizer is None:
            raise RuntimeError("Call load() first.")
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(f"pooling_strategy must be one of {self.available_pooling_strategies}")

        n = len(input)
        results: list[np.ndarray | None] = [None] * n

        valid_indices: list[int] = []
        valid_smiles: list[str] = []
        for idx, smi in enumerate(input):
            tok_len = self._token_length(smi)
            if self.max_len is not None and tok_len > self.max_len:
                logging.warning(f"Skipping SMILES ({tok_len} tokens > max {self.max_len}): {smi[:80]}...")
                continue
            valid_indices.append(idx)
            valid_smiles.append(smi)

        for start in range(0, len(valid_smiles), batch_size):
            chunk_smiles = valid_smiles[start : start + batch_size]
            chunk_indices = valid_indices[start : start + batch_size]

            enc = self.tokenizer(
                chunk_smiles,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=self.max_len,
            )
            input_ids = enc["input_ids"].to(self.device)
            attention_mask = enc["attention_mask"].to(self.device)

            with torch.no_grad():
                outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
                last_hidden = outputs.last_hidden_state
                pooler_out = getattr(outputs, "pooler_output", None)

            for i, orig_idx in enumerate(chunk_indices):
                if use_pooler and pooler_out is not None:
                    results[orig_idx] = pooler_out[i].cpu().numpy()
                    continue

                hidden = last_hidden[i]
                mask = attention_mask[i].unsqueeze(-1)

                if pooling_strategy == "none":
                    results[orig_idx] = hidden.cpu().numpy()
                    continue
                elif pooling_strategy == "cls":
                    vec = hidden[0]
                elif pooling_strategy == "mean":
                    vec = (hidden * mask).sum(dim=0) / mask.sum(dim=0).clamp(min=1)
                else:
                    neg_inf = torch.finfo(hidden.dtype).min
                    masked = hidden.masked_fill(mask == 0, neg_inf)
                    vec = masked.max(dim=0).values

                results[orig_idx] = vec.cpu().numpy()

        return results

embed(input, pooling_strategy='mean', use_pooler=False, **kwargs)

Embeddings function, which generates the embeddings for a specific provided model

Source code in src/embpy/models/molecule_models.py

def embed(
    self,
    input: str,
    pooling_strategy: str = "mean",
    use_pooler: bool = False,
    **kwargs: Any,
) -> np.ndarray | None:  # << allow None for skipped items
    """Embeddings function, which generates the embeddings for a specific provided model"""
    if self.model is None or self.device is None:
        raise RuntimeError("Call load() first.")
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(f"pooling_strategy must be one of {self.available_pooling_strategies}")

    # --- Skip too-long SMILES ---
    L = self._token_length(input)
    if self.max_len is not None and L > self.max_len:
        logging.warning(f"Skipping SMILES ({L} tokens > max {self.max_len}): {input[:80]}...")
        return None

    # 1) tokenize & move to device
    enc = self._preprocess_smiles(input)
    input_ids = enc["input_ids"].to(self.device)
    attention_mask = enc["attention_mask"].to(self.device)

    # 2) forward
    with torch.no_grad():
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        last_hidden = outputs.last_hidden_state  # (1, seq_len, hidden_dim)
        pooler_out = getattr(outputs, "pooler_output", None)

    # 3) use pooler if requested
    if use_pooler and pooler_out is not None:
        return pooler_out[0].cpu().numpy()

    # 4) squeeze batch dim
    hidden = last_hidden.squeeze(0)  # (seq_len, hidden_dim)

    # 5) masking-aware pooling
    if pooling_strategy == "none":
        return hidden.cpu().numpy()
    elif pooling_strategy == "cls":
        vec = hidden[0]
    else:
        mask = attention_mask.squeeze(0).unsqueeze(-1)  # (seq_len, 1)
        if pooling_strategy == "mean":
            summed = (hidden * mask).sum(dim=0)
            vec = summed / mask.sum(dim=0).clamp(min=1)  # safe denom
        else:  # 'max'
            neg_inf = torch.finfo(hidden.dtype).min
            masked = hidden.masked_fill(mask == 0, neg_inf)
            vec = masked.max(dim=0).values

    return vec.cpu().numpy()

embed_batch(input, pooling_strategy='mean', use_pooler=False, batch_size=64, **kwargs)

Embed a batch of SMILES strings with true batched inference.

Filters out too-long SMILES, batches the valid ones for a single forward pass per chunk, then reassembles the results preserving original order.

Returns a list of 1D numpy arrays (or None for skipped SMILES).

Source code in src/embpy/models/molecule_models.py

def embed_batch(
    self,
    input: list[str],
    pooling_strategy: str = "mean",
    use_pooler: bool = False,
    batch_size: int = 64,
    **kwargs: Any,
) -> list[np.ndarray | None]:
    """
    Embed a batch of SMILES strings with true batched inference.

    Filters out too-long SMILES, batches the valid ones for a
    single forward pass per chunk, then reassembles the results
    preserving original order.

    Returns a list of 1D numpy arrays (or None for skipped SMILES).
    """
    if self.model is None or self.device is None or self.tokenizer is None:
        raise RuntimeError("Call load() first.")
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(f"pooling_strategy must be one of {self.available_pooling_strategies}")

    n = len(input)
    results: list[np.ndarray | None] = [None] * n

    valid_indices: list[int] = []
    valid_smiles: list[str] = []
    for idx, smi in enumerate(input):
        tok_len = self._token_length(smi)
        if self.max_len is not None and tok_len > self.max_len:
            logging.warning(f"Skipping SMILES ({tok_len} tokens > max {self.max_len}): {smi[:80]}...")
            continue
        valid_indices.append(idx)
        valid_smiles.append(smi)

    for start in range(0, len(valid_smiles), batch_size):
        chunk_smiles = valid_smiles[start : start + batch_size]
        chunk_indices = valid_indices[start : start + batch_size]

        enc = self.tokenizer(
            chunk_smiles,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.max_len,
        )
        input_ids = enc["input_ids"].to(self.device)
        attention_mask = enc["attention_mask"].to(self.device)

        with torch.no_grad():
            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
            last_hidden = outputs.last_hidden_state
            pooler_out = getattr(outputs, "pooler_output", None)

        for i, orig_idx in enumerate(chunk_indices):
            if use_pooler and pooler_out is not None:
                results[orig_idx] = pooler_out[i].cpu().numpy()
                continue

            hidden = last_hidden[i]
            mask = attention_mask[i].unsqueeze(-1)

            if pooling_strategy == "none":
                results[orig_idx] = hidden.cpu().numpy()
                continue
            elif pooling_strategy == "cls":
                vec = hidden[0]
            elif pooling_strategy == "mean":
                vec = (hidden * mask).sum(dim=0) / mask.sum(dim=0).clamp(min=1)
            else:
                neg_inf = torch.finfo(hidden.dtype).min
                masked = hidden.masked_fill(mask == 0, neg_inf)
                vec = masked.max(dim=0).values

            results[orig_idx] = vec.cpu().numpy()

    return results

load(device)

Load the dataset

Source code in src/embpy/models/molecule_models.py

def load(self, device: torch.device):
    """Load the dataset"""
    if self.model is not None:
        return
    logging.info(f"Loading ChemBERTa '{self.model_name}'…")
    from transformers import AutoTokenizer, AutoModel
    self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
    import transformers
    prev_level = transformers.logging.get_verbosity()
    transformers.logging.set_verbosity_error()
    self.model = AutoModel.from_pretrained(self.model_name).to(device).eval()
    transformers.logging.set_verbosity(prev_level)
    self.device = device

    # derive a safe max sequence length (RoBERTa-style ≈512)
    mmax = getattr(self.model.config, "max_position_embeddings", None)
    if mmax is None or mmax > 512:
        mmax = 512
    self.max_len = int(mmax)
    # make tokenizer aware (some HF tokenizers have "no max" by default)
    self.tokenizer.model_max_length = self.max_len
    logging.info(f"Using model_max_length={self.max_len}")

MHGGNNWrapper

Bases: BaseModelWrapper

MHG-GNN molecular hypergraph autoencoder wrapper.

MHG-GNN combines a GIN-based graph encoder with a sequential decoder based on Molecular Hypergraph Grammar (MHG). It was pre-trained on ~1.34M molecules from PubChem.

The encoder produces latent-space embeddings from SMILES strings. The decoder can reconstruct structurally valid SMILES from those latent vectors, which is unique among molecular GNN models.

Requires

• mhg-gnn model package (from IBM Research / GT4SD)
• torch-geometric
• huggingface-hub

Install::

pip install torch-geometric huggingface-hub
pip install git+https://github.com/GT4SD/mhg-gnn.git

HuggingFace model: ibm-research/materials.mhg-ged

References

• Paper: https://arxiv.org/abs/2309.16374
• HF: https://huggingface.co/ibm-research/materials.mhg-ged

Source code in src/embpy/models/molecule_models.py

class MHGGNNWrapper(BaseModelWrapper):
    """MHG-GNN molecular hypergraph autoencoder wrapper.

    MHG-GNN combines a GIN-based graph encoder with a sequential decoder
    based on Molecular Hypergraph Grammar (MHG).  It was pre-trained on
    ~1.34M molecules from PubChem.

    The **encoder** produces latent-space embeddings from SMILES strings.
    The **decoder** can reconstruct structurally valid SMILES from those
    latent vectors, which is unique among molecular GNN models.

    Requires
    --------
    * ``mhg-gnn`` model package (from IBM Research / GT4SD)
    * ``torch-geometric``
    * ``huggingface-hub``

    Install::

        pip install torch-geometric huggingface-hub
        pip install git+https://github.com/GT4SD/mhg-gnn.git

    HuggingFace model: ``ibm-research/materials.mhg-ged``

    References
    ----------
    * Paper: https://arxiv.org/abs/2309.16374
    * HF: https://huggingface.co/ibm-research/materials.mhg-ged
    """

    model_type: str = "molecule"  # type: ignore[assignment]
    available_pooling_strategies: list[str] = ["flat"]
    # MHG-GNN is a GIN-based graph autoencoder over molecular hypergraphs, not a
    # transformer, so it has no attention weights to extract.
    has_attention = False

    def __init__(
        self,
        model_path_or_name: str = "ibm-research/materials.mhg-ged",
        **kwargs: Any,
    ):
        super().__init__(model_path_or_name, **kwargs)
        self._wrapper: Any = None
        self._loaded = False

    def load(self, device: torch.device) -> None:
        """Load MHG-GNN pretrained weights from HuggingFace Hub.

        The model package's built-in ``load()`` function is used, which
        downloads the pretrained pickle from HuggingFace Hub and
        constructs the ``GrammarGINVAE`` model automatically.
        """
        if self._loaded:
            return

        import importlib

        _load_fn: Any = None
        for mod_path in ("mhg_model.load", "mhg_gnn.load"):
            try:
                mod = importlib.import_module(mod_path)
                _load_fn = getattr(mod, "load", None)
                if _load_fn is not None:
                    break
            except ImportError:
                continue

        if _load_fn is None:
            raise ImportError(
                "MHG-GNN model code not found.  Install one of:\n"
                "  pip install git+https://github.com/GT4SD/mhg-gnn.git\n"
                "  (or from the IBM internal mhg-gnn repository)\n"
                "Also requires: pip install torch-geometric huggingface-hub\n"
                "See: https://huggingface.co/ibm-research/materials.mhg-ged"
            )

        logging.info(f"Loading MHG-GNN from HuggingFace: {self.model_name}…")
        self._wrapper = _load_fn()
        self._wrapper.to(device)
        if hasattr(self._wrapper, "eval"):
            self._wrapper.eval()
        self.device = device
        self._loaded = True
        logging.info("MHG-GNN loaded successfully.")

    def embed(
        self,
        input: str,
        pooling_strategy: str = "flat",
        **kwargs: Any,
    ) -> np.ndarray:
        """Encode a single SMILES string into its latent representation."""
        if not self._loaded or self._wrapper is None:
            raise RuntimeError("MHG-GNN not loaded. Call load() first.")

        with torch.no_grad():
            embeddings = self._wrapper.encode([input])
        if not embeddings:
            raise ValueError(f"MHG-GNN returned no embeddings for SMILES: {input}")
        return embeddings[0].detach().cpu().numpy().astype(np.float32)

    def embed_batch(
        self,
        inputs: Sequence[str],
        pooling_strategy: str = "flat",
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """Encode a batch of SMILES strings into latent representations."""
        if not self._loaded or self._wrapper is None:
            raise RuntimeError("MHG-GNN not loaded. Call load() first.")

        with torch.no_grad():
            embeddings = self._wrapper.encode(list(inputs))
        return [e.detach().cpu().numpy().astype(np.float32) for e in embeddings]

    def decode(self, embeddings: list[np.ndarray]) -> list[str | None]:
        """Decode latent vectors back to SMILES strings.

        This leverages MHG-GNN's decoder, which is guaranteed to produce
        structurally valid molecules.

        Parameters
        ----------
        embeddings : list[np.ndarray]
            Latent vectors (e.g. from :meth:`embed` or :meth:`embed_batch`).

        Returns
        -------
        list[str | None]
            Reconstructed SMILES strings (``None`` for failed decodings).
        """
        if not self._loaded or self._wrapper is None:
            raise RuntimeError("MHG-GNN not loaded. Call load() first.")

        tensors = [torch.tensor(e, dtype=torch.float32) for e in embeddings]
        dev = self.device if self.device is not None else torch.device("cpu")
        tensors = [t.to(dev) for t in tensors]

        with torch.no_grad():
            decoded: list[str | None] = self._wrapper.decode(tensors)
        return decoded

decode(embeddings)

Decode latent vectors back to SMILES strings.

This leverages MHG-GNN's decoder, which is guaranteed to produce structurally valid molecules.

Parameters:

Name	Type	Description	Default
embeddings	list[ndarray]	Latent vectors (e.g. from :meth:embed or :meth:embed_batch).	required

Returns:

Type	Description
list[str | None]	Reconstructed SMILES strings (None for failed decodings).

Source code in src/embpy/models/molecule_models.py

def decode(self, embeddings: list[np.ndarray]) -> list[str | None]:
    """Decode latent vectors back to SMILES strings.

    This leverages MHG-GNN's decoder, which is guaranteed to produce
    structurally valid molecules.

    Parameters
    ----------
    embeddings : list[np.ndarray]
        Latent vectors (e.g. from :meth:`embed` or :meth:`embed_batch`).

    Returns
    -------
    list[str | None]
        Reconstructed SMILES strings (``None`` for failed decodings).
    """
    if not self._loaded or self._wrapper is None:
        raise RuntimeError("MHG-GNN not loaded. Call load() first.")

    tensors = [torch.tensor(e, dtype=torch.float32) for e in embeddings]
    dev = self.device if self.device is not None else torch.device("cpu")
    tensors = [t.to(dev) for t in tensors]

    with torch.no_grad():
        decoded: list[str | None] = self._wrapper.decode(tensors)
    return decoded

embed(input, pooling_strategy='flat', **kwargs)

Encode a single SMILES string into its latent representation.

Source code in src/embpy/models/molecule_models.py

def embed(
    self,
    input: str,
    pooling_strategy: str = "flat",
    **kwargs: Any,
) -> np.ndarray:
    """Encode a single SMILES string into its latent representation."""
    if not self._loaded or self._wrapper is None:
        raise RuntimeError("MHG-GNN not loaded. Call load() first.")

    with torch.no_grad():
        embeddings = self._wrapper.encode([input])
    if not embeddings:
        raise ValueError(f"MHG-GNN returned no embeddings for SMILES: {input}")
    return embeddings[0].detach().cpu().numpy().astype(np.float32)

embed_batch(inputs, pooling_strategy='flat', **kwargs)

Encode a batch of SMILES strings into latent representations.

Source code in src/embpy/models/molecule_models.py

def embed_batch(
    self,
    inputs: Sequence[str],
    pooling_strategy: str = "flat",
    **kwargs: Any,
) -> list[np.ndarray]:
    """Encode a batch of SMILES strings into latent representations."""
    if not self._loaded or self._wrapper is None:
        raise RuntimeError("MHG-GNN not loaded. Call load() first.")

    with torch.no_grad():
        embeddings = self._wrapper.encode(list(inputs))
    return [e.detach().cpu().numpy().astype(np.float32) for e in embeddings]

load(device)

Load MHG-GNN pretrained weights from HuggingFace Hub.

The model package's built-in load() function is used, which downloads the pretrained pickle from HuggingFace Hub and constructs the GrammarGINVAE model automatically.

Source code in src/embpy/models/molecule_models.py

def load(self, device: torch.device) -> None:
    """Load MHG-GNN pretrained weights from HuggingFace Hub.

    The model package's built-in ``load()`` function is used, which
    downloads the pretrained pickle from HuggingFace Hub and
    constructs the ``GrammarGINVAE`` model automatically.
    """
    if self._loaded:
        return

    import importlib

    _load_fn: Any = None
    for mod_path in ("mhg_model.load", "mhg_gnn.load"):
        try:
            mod = importlib.import_module(mod_path)
            _load_fn = getattr(mod, "load", None)
            if _load_fn is not None:
                break
        except ImportError:
            continue

    if _load_fn is None:
        raise ImportError(
            "MHG-GNN model code not found.  Install one of:\n"
            "  pip install git+https://github.com/GT4SD/mhg-gnn.git\n"
            "  (or from the IBM internal mhg-gnn repository)\n"
            "Also requires: pip install torch-geometric huggingface-hub\n"
            "See: https://huggingface.co/ibm-research/materials.mhg-ged"
        )

    logging.info(f"Loading MHG-GNN from HuggingFace: {self.model_name}…")
    self._wrapper = _load_fn()
    self._wrapper.to(device)
    if hasattr(self._wrapper, "eval"):
        self._wrapper.eval()
    self.device = device
    self._loaded = True
    logging.info("MHG-GNN loaded successfully.")

MiniMolWrapper

Bases: BaseModelWrapper

MiniMol pre-trained GINE molecular fingerprint wrapper.

MiniMol is a 10M-parameter message-passing GNN (GIN with edge features) pre-trained on 6M molecules across 3,300+ bio/quantum tasks from the Graphium LargeMix dataset. It produces fixed 512-dimensional molecular fingerprints from SMILES strings.

The model and weights ship together in the minimol pip package, so no separate download step is required.

Requires

pip install minimol (also installs graphium and torch-geometric automatically)

References

• Repository: https://github.com/graphcore-research/minimol
• MiniMol outperforms models 10x its size on 17/22 TDC ADMET tasks.

Notes

MiniMol manages its own device selection (CUDA if available, else CPU). The device parameter in :meth:load is stored for interface consistency but does not override MiniMol's internal device selection.

Source code in src/embpy/models/molecule_models.py

class MiniMolWrapper(BaseModelWrapper):
    """MiniMol pre-trained GINE molecular fingerprint wrapper.

    MiniMol is a 10M-parameter message-passing GNN (GIN with edge features)
    pre-trained on 6M molecules across 3,300+ bio/quantum tasks from the
    Graphium LargeMix dataset.  It produces fixed **512-dimensional**
    molecular fingerprints from SMILES strings.

    The model and weights ship together in the ``minimol`` pip package,
    so no separate download step is required.

    Requires
    --------
    ``pip install minimol``
    (also installs ``graphium`` and ``torch-geometric`` automatically)

    References
    ----------
    * Repository: https://github.com/graphcore-research/minimol
    * MiniMol outperforms models 10x its size on 17/22 TDC ADMET tasks.

    Notes
    -----
    MiniMol manages its own device selection (CUDA if available, else CPU).
    The ``device`` parameter in :meth:`load` is stored for interface
    consistency but does not override MiniMol's internal device selection.
    """

    model_type: str = "molecule"  # type: ignore[assignment]
    available_pooling_strategies: list[str] = ["flat"]
    EMBEDDING_DIM: int = 512
    # MiniMol is a message-passing GNN (GIN with edge features), not a transformer,
    # so it has no attention weights to extract.
    has_attention = False

    def __init__(
        self,
        model_path_or_name: str = "minimol",
        batch_size: int = 100,
        **kwargs: Any,
    ):
        super().__init__(model_path_or_name, **kwargs)
        self._batch_size = batch_size
        self._minimol: Any = None
        self._loaded = False

    def load(self, device: torch.device) -> None:
        """Load the MiniMol model.

        Weights are bundled with the ``minimol`` package.  The first call
        initialises the underlying Graphium predictor and loads the
        pre-trained checkpoint.
        """
        if self._loaded:
            return

        try:
            from minimol import Minimol  # type: ignore[import-not-found]
        except ImportError as e:
            raise ImportError(
                "MiniMol is not installed.  Install with:\n"
                "  pip install minimol\n"
                "See: https://github.com/graphcore-research/minimol"
            ) from e

        logging.info("Loading MiniMol (GINE, 512-dim fingerprints)…")
        self._minimol = Minimol(batch_size=self._batch_size)
        self.device = device
        self._loaded = True
        logging.info("MiniMol loaded successfully.")

    def embed(
        self,
        input: str,
        pooling_strategy: str = "flat",
        **kwargs: Any,
    ) -> np.ndarray:
        """Compute a 512-dim fingerprint for a single SMILES string."""
        if not self._loaded or self._minimol is None:
            raise RuntimeError("MiniMol not loaded. Call load() first.")

        results = self._minimol([input])
        if not results:
            raise ValueError(f"MiniMol returned no results for SMILES: {input}")
        return results[0].detach().cpu().numpy().astype(np.float32)

    def embed_batch(
        self,
        inputs: Sequence[str],
        pooling_strategy: str = "flat",
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """Compute 512-dim fingerprints for a batch of SMILES strings."""
        if not self._loaded or self._minimol is None:
            raise RuntimeError("MiniMol not loaded. Call load() first.")

        smiles_list = list(inputs)
        results = self._minimol(smiles_list)
        return [r.detach().cpu().numpy().astype(np.float32) for r in results]

embed(input, pooling_strategy='flat', **kwargs)

Compute a 512-dim fingerprint for a single SMILES string.

Source code in src/embpy/models/molecule_models.py

def embed(
    self,
    input: str,
    pooling_strategy: str = "flat",
    **kwargs: Any,
) -> np.ndarray:
    """Compute a 512-dim fingerprint for a single SMILES string."""
    if not self._loaded or self._minimol is None:
        raise RuntimeError("MiniMol not loaded. Call load() first.")

    results = self._minimol([input])
    if not results:
        raise ValueError(f"MiniMol returned no results for SMILES: {input}")
    return results[0].detach().cpu().numpy().astype(np.float32)

embed_batch(inputs, pooling_strategy='flat', **kwargs)

Compute 512-dim fingerprints for a batch of SMILES strings.

Source code in src/embpy/models/molecule_models.py

def embed_batch(
    self,
    inputs: Sequence[str],
    pooling_strategy: str = "flat",
    **kwargs: Any,
) -> list[np.ndarray]:
    """Compute 512-dim fingerprints for a batch of SMILES strings."""
    if not self._loaded or self._minimol is None:
        raise RuntimeError("MiniMol not loaded. Call load() first.")

    smiles_list = list(inputs)
    results = self._minimol(smiles_list)
    return [r.detach().cpu().numpy().astype(np.float32) for r in results]

load(device)

Load the MiniMol model.

Weights are bundled with the minimol package. The first call initialises the underlying Graphium predictor and loads the pre-trained checkpoint.

Source code in src/embpy/models/molecule_models.py

def load(self, device: torch.device) -> None:
    """Load the MiniMol model.

    Weights are bundled with the ``minimol`` package.  The first call
    initialises the underlying Graphium predictor and loads the
    pre-trained checkpoint.
    """
    if self._loaded:
        return

    try:
        from minimol import Minimol  # type: ignore[import-not-found]
    except ImportError as e:
        raise ImportError(
            "MiniMol is not installed.  Install with:\n"
            "  pip install minimol\n"
            "See: https://github.com/graphcore-research/minimol"
        ) from e

    logging.info("Loading MiniMol (GINE, 512-dim fingerprints)…")
    self._minimol = Minimol(batch_size=self._batch_size)
    self.device = device
    self._loaded = True
    logging.info("MiniMol loaded successfully.")

MolEWrapper

Bases: BaseModelWrapper

MolE (Molecular Embeddings) graph transformer wrapper.

MolE is a DeBERTa-based graph transformer with disentangled attention, pre-trained on ~842M molecules via a two-step strategy:

1. Self-supervised pre-training on molecular graph representations.
2. Multi-task supervised training to assimilate biological information.

It produces CLS-pooled embeddings from molecular graphs and achieves SOTA on 9/22 ADMET tasks in the TDC benchmark.

.. note::

Pretrained weights are not yet publicly released. You must provide a checkpoint_path obtained through your own training or directly from the authors.

Requires

pip install git+https://github.com/recursionpharma/mole_public.git

References

• Repository: https://github.com/recursionpharma/mole_public
• Published in Nature Communications (2024).

Source code in src/embpy/models/molecule_models.py

class MolEWrapper(BaseModelWrapper):
    """MolE (Molecular Embeddings) graph transformer wrapper.

    MolE is a DeBERTa-based graph transformer with disentangled attention,
    pre-trained on ~842M molecules via a two-step strategy:

    1. **Self-supervised** pre-training on molecular graph representations.
    2. **Multi-task supervised** training to assimilate biological information.

    It produces CLS-pooled embeddings from molecular graphs and achieves
    SOTA on 9/22 ADMET tasks in the TDC benchmark.

    .. note::

       Pretrained weights are **not** yet publicly released.  You must
       provide a ``checkpoint_path`` obtained through your own training
       or directly from the authors.

    Requires
    --------
    ``pip install git+https://github.com/recursionpharma/mole_public.git``

    References
    ----------
    * Repository: https://github.com/recursionpharma/mole_public
    * Published in *Nature Communications* (2024).
    """

    model_type: str = "molecule"  # type: ignore[assignment]
    available_pooling_strategies: list[str] = ["cls"]

    def __init__(
        self,
        model_path_or_name: str = "mole",
        checkpoint_path: str | None = None,
        **kwargs: Any,
    ):
        super().__init__(model_path_or_name, **kwargs)
        self._checkpoint_path = checkpoint_path
        self._mole_predict: Any = None
        self._loaded = False

    def load(self, device: torch.device) -> None:
        """Load the MolE model.

        Uses the ``mole.mole_predict`` high-level API internally.  Requires
        a valid ``checkpoint_path`` pointing to a pre-trained ``.ckpt`` file.
        """
        if self._loaded:
            return

        _mole_predict: Any = None
        for mod_path in ("mole.mole_predict", "mole.cli.mole_predict"):
            try:
                import importlib

                _mole_predict = importlib.import_module(mod_path)
                break
            except ImportError:
                continue

        if _mole_predict is None:
            raise ImportError(
                "MolE is not installed.  Install with:\n"
                "  pip install git+https://github.com/recursionpharma/mole_public.git\n"
                "See: https://github.com/recursionpharma/mole_public"
            )

        if self._checkpoint_path is None:
            raise ValueError(
                "MolE requires a pretrained checkpoint path.  Pass "
                "checkpoint_path='path/to/checkpoint.ckpt' to the constructor.\n"
                "Note: pretrained weights are not yet publicly released."
            )

        self._mole_predict = _mole_predict
        self.device = device
        self._loaded = True
        logging.info(f"MolE wrapper ready (checkpoint: {self._checkpoint_path}).")

    def embed(
        self,
        input: str,
        pooling_strategy: str = "cls",
        **kwargs: Any,
    ) -> np.ndarray:
        """Compute the CLS-pooled embedding for a single SMILES string."""
        if not self._loaded or self._mole_predict is None:
            raise RuntimeError("MolE not loaded. Call load() first.")

        batch_size: int = kwargs.get("batch_size", 32)  # type: ignore[assignment]
        num_workers: int = kwargs.get("num_workers", 0)  # type: ignore[assignment]

        embeddings = self._mole_predict.encode(
            smiles=[input],
            pretrained_model=self._checkpoint_path,
            batch_size=batch_size,
            num_workers=num_workers,
        )
        return np.asarray(embeddings[0], dtype=np.float32)

    def embed_batch(
        self,
        inputs: Sequence[str],
        pooling_strategy: str = "cls",
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """Compute CLS-pooled embeddings for a batch of SMILES strings."""
        if not self._loaded or self._mole_predict is None:
            raise RuntimeError("MolE not loaded. Call load() first.")

        batch_size: int = kwargs.get("batch_size", 32)  # type: ignore[assignment]
        num_workers: int = kwargs.get("num_workers", 0 if sys.platform == "win32" else 4)  # type: ignore[assignment]

        embeddings = self._mole_predict.encode(
            smiles=list(inputs),
            pretrained_model=self._checkpoint_path,
            batch_size=batch_size,
            num_workers=num_workers,
        )
        return [np.asarray(e, dtype=np.float32) for e in embeddings]

embed(input, pooling_strategy='cls', **kwargs)

Compute the CLS-pooled embedding for a single SMILES string.

Source code in src/embpy/models/molecule_models.py

def embed(
    self,
    input: str,
    pooling_strategy: str = "cls",
    **kwargs: Any,
) -> np.ndarray:
    """Compute the CLS-pooled embedding for a single SMILES string."""
    if not self._loaded or self._mole_predict is None:
        raise RuntimeError("MolE not loaded. Call load() first.")

    batch_size: int = kwargs.get("batch_size", 32)  # type: ignore[assignment]
    num_workers: int = kwargs.get("num_workers", 0)  # type: ignore[assignment]

    embeddings = self._mole_predict.encode(
        smiles=[input],
        pretrained_model=self._checkpoint_path,
        batch_size=batch_size,
        num_workers=num_workers,
    )
    return np.asarray(embeddings[0], dtype=np.float32)

embed_batch(inputs, pooling_strategy='cls', **kwargs)

Compute CLS-pooled embeddings for a batch of SMILES strings.

Source code in src/embpy/models/molecule_models.py

def embed_batch(
    self,
    inputs: Sequence[str],
    pooling_strategy: str = "cls",
    **kwargs: Any,
) -> list[np.ndarray]:
    """Compute CLS-pooled embeddings for a batch of SMILES strings."""
    if not self._loaded or self._mole_predict is None:
        raise RuntimeError("MolE not loaded. Call load() first.")

    batch_size: int = kwargs.get("batch_size", 32)  # type: ignore[assignment]
    num_workers: int = kwargs.get("num_workers", 0 if sys.platform == "win32" else 4)  # type: ignore[assignment]

    embeddings = self._mole_predict.encode(
        smiles=list(inputs),
        pretrained_model=self._checkpoint_path,
        batch_size=batch_size,
        num_workers=num_workers,
    )
    return [np.asarray(e, dtype=np.float32) for e in embeddings]

load(device)

Load the MolE model.

Uses the mole.mole_predict high-level API internally. Requires a valid checkpoint_path pointing to a pre-trained .ckpt file.

Source code in src/embpy/models/molecule_models.py

def load(self, device: torch.device) -> None:
    """Load the MolE model.

    Uses the ``mole.mole_predict`` high-level API internally.  Requires
    a valid ``checkpoint_path`` pointing to a pre-trained ``.ckpt`` file.
    """
    if self._loaded:
        return

    _mole_predict: Any = None
    for mod_path in ("mole.mole_predict", "mole.cli.mole_predict"):
        try:
            import importlib

            _mole_predict = importlib.import_module(mod_path)
            break
        except ImportError:
            continue

    if _mole_predict is None:
        raise ImportError(
            "MolE is not installed.  Install with:\n"
            "  pip install git+https://github.com/recursionpharma/mole_public.git\n"
            "See: https://github.com/recursionpharma/mole_public"
        )

    if self._checkpoint_path is None:
        raise ValueError(
            "MolE requires a pretrained checkpoint path.  Pass "
            "checkpoint_path='path/to/checkpoint.ckpt' to the constructor.\n"
            "Note: pretrained weights are not yet publicly released."
        )

    self._mole_predict = _mole_predict
    self.device = device
    self._loaded = True
    logging.info(f"MolE wrapper ready (checkpoint: {self._checkpoint_path}).")

MolformerWrapper

Bases: BaseModelWrapper

Wrapper for MoLFormer molecular language models (via HuggingFace Transformers).

Loads a MolFormer checkpoint and exposes: - embed(...) for a single SMILES - embed_batch(...) for a list of SMILES

Pooling strategies: • cls → model.pooler_output • mean → average over token embeddings • max → max over token embeddings

Source code in src/embpy/models/molecule_models.py

class MolformerWrapper(BaseModelWrapper):
    """
    Wrapper for MoLFormer molecular language models (via HuggingFace Transformers).

    Loads a MolFormer checkpoint and exposes:
      - `embed(...)` for a single SMILES
      - `embed_batch(...)` for a list of SMILES

    Pooling strategies:
      • cls  → model.pooler_output
      • mean → average over token embeddings
      • max  → max over token embeddings
    """

    model_type = "molecule"
    available_pooling_strategies = ["cls", "mean", "max", "none"]

    def __init__(self, model_path_or_name: str = "ibm/MoLFormer-XL-both-10pct", **kwargs):
        """
        Initialize the MolformerWrapper.

        Parameters
        ----------
        model_path_or_name : str, optional
            Hugging Face model identifier or local path for the MolFormer weights.
            Defaults to `"ibm/MoLFormer-XL-both-10pct"`.
        **kwargs
            Additional keyword arguments passed to `BaseModelWrapper`.
            `revision` overrides the commit pinned in `embpy.models.hf_revisions` (which keeps
            the repo's `trust_remote_code` files at a version that loads on transformers 4.x).
        """
        super().__init__(model_path_or_name, **kwargs)
        self.tokenizer = None

    def load(self, device: torch.device) -> None:
        """
        Load the MolFormer model and tokenizer onto `device`.

        Raises
        ------
        RuntimeError
            If the model or tokenizer fail to load.
        """
        if self.model is not None:
            logging.debug("MolFormer already loaded.")
            return

        logging.info(f"Loading MolFormer '{self.model_name}' (trust_remote_code)…")
        from transformers import AutoTokenizer, AutoModel
        revision = resolve_revision(self.model_name, self.config.get("revision"))
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_name, trust_remote_code=True, revision=revision, model_max_length=512,
            )
            model = AutoModel.from_pretrained(
                self.model_name,
                deterministic_eval=True,
                trust_remote_code=True,
                revision=revision,
            )
            self.model = model.to(device).eval()  # type: ignore[union-attr]
            self.device = device
            logging.info("MolFormer loaded successfully.")
        except Exception as e:
            logging.error(f"MolFormer load error: {e}")
            raise RuntimeError(f"Could not load MolFormer '{self.model_name}'") from e

    def _preprocess_smiles(self, smiles: str | list[str]) -> BatchEncoding:
        """
        Tokenize one or more SMILES strings into model-ready input tensors.

        This wraps the HuggingFace tokenizer to convert raw SMILES into
        `input_ids` and `attention_mask`, handling batching, padding, and
        truncation automatically.

        Parameters
        ----------
        smiles : str or list of str
            A single SMILES string (e.g. `"CCO"`) or a list of them (e.g.
            `["CCO", "Cn1c(=O)c2c(ncn2C)n(C)c1=O"]`).

        Returns
        -------
        transformers.BatchEncoding
            A BatchEncoding containing at least:

            - `input_ids` (`torch.LongTensor` of shape `(batch_size, seq_len)`)
            - `attention_mask` (`torch.LongTensor` of shape `(batch_size, seq_len)`)

        Raises
        ------
        RuntimeError
            If called before `load()`, i.e. when `self.tokenizer` is still `None`.
        """
        if self.tokenizer is None:
            raise RuntimeError("Tokenizer not loaded. Call load() first.")
        seqs = smiles if isinstance(smiles, (list | tuple)) else [smiles]
        return self.tokenizer(seqs, padding=True, truncation=True, return_tensors="pt")

    def embed(
        self,
        input: str,
        pooling_strategy: str = "cls",
        **kwargs: Any,
    ) -> np.ndarray:
        """
        Embed a single SMILES string.

        Parameters
        ----------
        input : str
            The SMILES string.
        pooling_strategy : str, default "cls"
            One of "cls", "mean", or "max".

        Returns
        -------
        np.ndarray
            1-D embedding vector.
        """
        if self.model is None or self.device is None:
            raise RuntimeError("MolFormer model not loaded. Call load() first.")
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(f"Invalid pooling '{pooling_strategy}'. Choose from {self.available_pooling_strategies}")

        tokens = self._preprocess_smiles(input)
        tokens = {k: v.to(self.device) for k, v in tokens.items()}

        with torch.no_grad():
            out = self.model(**tokens)

        # out.pooler_output: (B, hidden_dim)
        # out.last_hidden_state: (B, seq_len, hidden_dim)
        if pooling_strategy == "none":
            return out.last_hidden_state[0].cpu().numpy()
        elif pooling_strategy == "cls":
            vec = out.pooler_output[0]
        else:
            seq_emb = out.last_hidden_state[0]  # (seq_len, hidden_dim)
            if pooling_strategy == "mean":
                vec = seq_emb.mean(dim=0)
            else:  # max
                vec = seq_emb.max(dim=0).values

        return vec.cpu().numpy()

    def embed_batch(
        self,
        inputs: Sequence[str],
        pooling_strategy: str = "cls",
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """
        Embed a batch of SMILES strings.

        Parameters
        ----------
        inputs : Sequence[str]
            List of SMILES strings.
        pooling_strategy : str, default "cls"
            One of "cls", "mean", or "max".

        Returns
        -------
        List[np.ndarray]
            A list of 1-D embedding vectors.
        """
        if self.model is None or self.device is None:
            raise RuntimeError("MolFormer model not loaded. Call load() first.")
        if pooling_strategy not in self.available_pooling_strategies:
            raise ValueError(f"Invalid pooling '{pooling_strategy}'. Choose from {self.available_pooling_strategies}")

        tokens = self._preprocess_smiles(list(inputs))
        tokens = {k: v.to(self.device) for k, v in tokens.items()}

        with torch.no_grad():
            out = self.model(**tokens)

        embeddings: list[np.ndarray] = []
        if pooling_strategy == "none":
            for seq_emb in out.last_hidden_state:
                embeddings.append(seq_emb.cpu().numpy())
        elif pooling_strategy == "cls":
            for vec in out.pooler_output:
                embeddings.append(vec.cpu().numpy())
        else:
            seqs = out.last_hidden_state
            if pooling_strategy == "mean":
                pooled = seqs.mean(dim=1)
            else:
                pooled = seqs.max(dim=1).values
            for vec in pooled:
                embeddings.append(vec.cpu().numpy())

        return embeddings

__init__(model_path_or_name='ibm/MoLFormer-XL-both-10pct', **kwargs)

Initialize the MolformerWrapper.

Parameters:

Name	Type	Description	Default
model_path_or_name	str	Hugging Face model identifier or local path for the MolFormer weights. Defaults to "ibm/MoLFormer-XL-both-10pct".	'ibm/MoLFormer-XL-both-10pct'
**kwargs		Additional keyword arguments passed to BaseModelWrapper. revision overrides the commit pinned in embpy.models.hf_revisions (which keeps the repo's trust_remote_code files at a version that loads on transformers 4.x).	{}

Source code in src/embpy/models/molecule_models.py

def __init__(self, model_path_or_name: str = "ibm/MoLFormer-XL-both-10pct", **kwargs):
    """
    Initialize the MolformerWrapper.

    Parameters
    ----------
    model_path_or_name : str, optional
        Hugging Face model identifier or local path for the MolFormer weights.
        Defaults to `"ibm/MoLFormer-XL-both-10pct"`.
    **kwargs
        Additional keyword arguments passed to `BaseModelWrapper`.
        `revision` overrides the commit pinned in `embpy.models.hf_revisions` (which keeps
        the repo's `trust_remote_code` files at a version that loads on transformers 4.x).
    """
    super().__init__(model_path_or_name, **kwargs)
    self.tokenizer = None

embed(input, pooling_strategy='cls', **kwargs)

Embed a single SMILES string.

Parameters:

Name	Type	Description	Default
input	str	The SMILES string.	required
pooling_strategy	str	One of "cls", "mean", or "max".	"cls"

Returns:

Type	Description
ndarray	1-D embedding vector.

Source code in src/embpy/models/molecule_models.py

def embed(
    self,
    input: str,
    pooling_strategy: str = "cls",
    **kwargs: Any,
) -> np.ndarray:
    """
    Embed a single SMILES string.

    Parameters
    ----------
    input : str
        The SMILES string.
    pooling_strategy : str, default "cls"
        One of "cls", "mean", or "max".

    Returns
    -------
    np.ndarray
        1-D embedding vector.
    """
    if self.model is None or self.device is None:
        raise RuntimeError("MolFormer model not loaded. Call load() first.")
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(f"Invalid pooling '{pooling_strategy}'. Choose from {self.available_pooling_strategies}")

    tokens = self._preprocess_smiles(input)
    tokens = {k: v.to(self.device) for k, v in tokens.items()}

    with torch.no_grad():
        out = self.model(**tokens)

    # out.pooler_output: (B, hidden_dim)
    # out.last_hidden_state: (B, seq_len, hidden_dim)
    if pooling_strategy == "none":
        return out.last_hidden_state[0].cpu().numpy()
    elif pooling_strategy == "cls":
        vec = out.pooler_output[0]
    else:
        seq_emb = out.last_hidden_state[0]  # (seq_len, hidden_dim)
        if pooling_strategy == "mean":
            vec = seq_emb.mean(dim=0)
        else:  # max
            vec = seq_emb.max(dim=0).values

    return vec.cpu().numpy()

embed_batch(inputs, pooling_strategy='cls', **kwargs)

Embed a batch of SMILES strings.

Parameters:

Name	Type	Description	Default
inputs	Sequence[str]	List of SMILES strings.	required
pooling_strategy	str	One of "cls", "mean", or "max".	"cls"

Returns:

Type	Description
List[ndarray]	A list of 1-D embedding vectors.

Source code in src/embpy/models/molecule_models.py

def embed_batch(
    self,
    inputs: Sequence[str],
    pooling_strategy: str = "cls",
    **kwargs: Any,
) -> list[np.ndarray]:
    """
    Embed a batch of SMILES strings.

    Parameters
    ----------
    inputs : Sequence[str]
        List of SMILES strings.
    pooling_strategy : str, default "cls"
        One of "cls", "mean", or "max".

    Returns
    -------
    List[np.ndarray]
        A list of 1-D embedding vectors.
    """
    if self.model is None or self.device is None:
        raise RuntimeError("MolFormer model not loaded. Call load() first.")
    if pooling_strategy not in self.available_pooling_strategies:
        raise ValueError(f"Invalid pooling '{pooling_strategy}'. Choose from {self.available_pooling_strategies}")

    tokens = self._preprocess_smiles(list(inputs))
    tokens = {k: v.to(self.device) for k, v in tokens.items()}

    with torch.no_grad():
        out = self.model(**tokens)

    embeddings: list[np.ndarray] = []
    if pooling_strategy == "none":
        for seq_emb in out.last_hidden_state:
            embeddings.append(seq_emb.cpu().numpy())
    elif pooling_strategy == "cls":
        for vec in out.pooler_output:
            embeddings.append(vec.cpu().numpy())
    else:
        seqs = out.last_hidden_state
        if pooling_strategy == "mean":
            pooled = seqs.mean(dim=1)
        else:
            pooled = seqs.max(dim=1).values
        for vec in pooled:
            embeddings.append(vec.cpu().numpy())

    return embeddings

load(device)

Load the MolFormer model and tokenizer onto device.

Raises:

Type	Description
RuntimeError	If the model or tokenizer fail to load.

Source code in src/embpy/models/molecule_models.py

def load(self, device: torch.device) -> None:
    """
    Load the MolFormer model and tokenizer onto `device`.

    Raises
    ------
    RuntimeError
        If the model or tokenizer fail to load.
    """
    if self.model is not None:
        logging.debug("MolFormer already loaded.")
        return

    logging.info(f"Loading MolFormer '{self.model_name}' (trust_remote_code)…")
    from transformers import AutoTokenizer, AutoModel
    revision = resolve_revision(self.model_name, self.config.get("revision"))
    try:
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_name, trust_remote_code=True, revision=revision, model_max_length=512,
        )
        model = AutoModel.from_pretrained(
            self.model_name,
            deterministic_eval=True,
            trust_remote_code=True,
            revision=revision,
        )
        self.model = model.to(device).eval()  # type: ignore[union-attr]
        self.device = device
        logging.info("MolFormer loaded successfully.")
    except Exception as e:
        logging.error(f"MolFormer load error: {e}")
        raise RuntimeError(f"Could not load MolFormer '{self.model_name}'") from e

RDKitWrapper

Bases: BaseModelWrapper

RDKit molecular fingerprint wrapper.

Generates fixed-length fingerprint vectors from SMILES strings. Supports both binary (0/1 bit vectors) and count-based (continuous integer) representations.

Fingerprint types

• "rdkit" -- RDKit topological (Daylight-like), binary
• "morgan" -- Morgan / ECFP bit vector, binary
• "morgan_count" -- Morgan / ECFP count vector, continuous
• "maccs" -- MACCS keys (166 public keys), binary
• "atom_pair" -- Atom-pair bit fingerprint, binary
• "atom_pair_count" -- Atom-pair count fingerprint, continuous
• "topological_torsion" -- Topological torsion bits, binary
• "topological_torsion_count" -- Topological torsion counts, continuous

Source code in src/embpy/models/molecule_models.py

class RDKitWrapper(BaseModelWrapper):
    """RDKit molecular fingerprint wrapper.

    Generates fixed-length fingerprint vectors from SMILES strings.
    Supports both **binary** (0/1 bit vectors) and **count-based**
    (continuous integer) representations.

    Fingerprint types
    -----------------
    * ``"rdkit"``  -- RDKit topological (Daylight-like), binary
    * ``"morgan"`` -- Morgan / ECFP bit vector, binary
    * ``"morgan_count"`` -- Morgan / ECFP **count** vector, continuous
    * ``"maccs"``  -- MACCS keys (166 public keys), binary
    * ``"atom_pair"`` -- Atom-pair bit fingerprint, binary
    * ``"atom_pair_count"`` -- Atom-pair count fingerprint, continuous
    * ``"topological_torsion"`` -- Topological torsion bits, binary
    * ``"topological_torsion_count"`` -- Topological torsion counts, continuous
    """

    model_type = "molecule"
    available_pooling_strategies = ["flat"]
    # Structural fingerprints (Morgan, MACCS, atom-pair) are hand-computed bit
    # vectors, not neural networks, so there is nothing to attend with.
    has_attention = False

    _VALID_FP_TYPES = (
        "morgan",
        "morgan_count",
        "rdkit",
        "maccs",
        "atom_pair",
        "atom_pair_count",
        "topological_torsion",
        "topological_torsion_count",
    )

    def __init__(
        self,
        model_path_or_name: str = "rdkit_fingerprint",
        fingerprint_type: str | None = None,
        radius: int = 2,
        n_bits: int = 2048,
        **kwargs: Any,
    ):
        super().__init__(model_path_or_name, **kwargs)
        # Auto-detect: when model_path_or_name is itself a valid fingerprint
        # type (e.g. loaded from MODEL_REGISTRY), use it as fingerprint_type.
        if fingerprint_type is None:
            if model_path_or_name in self._VALID_FP_TYPES:
                fingerprint_type = model_path_or_name
            else:
                fingerprint_type = "rdkit"
        if fingerprint_type not in self._VALID_FP_TYPES:
            raise ValueError(f"Unknown fingerprint_type '{fingerprint_type}'. Choose from {self._VALID_FP_TYPES}")
        self.fingerprint_type = fingerprint_type
        self.radius = radius
        # MACCS keys always have a fixed size (167 bits)
        self.n_bits = 167 if fingerprint_type == "maccs" else n_bits
        self._loaded = False

    @property
    def is_count_fingerprint(self) -> bool:
        """``True`` when the fingerprint produces continuous counts."""
        return self.fingerprint_type.endswith("_count")

    @property
    def is_binary_fingerprint(self) -> bool:
        """``True`` when the fingerprint produces binary 0/1 vectors."""
        return not self.is_count_fingerprint

    def load(self, device: torch.device) -> None:
        """Mark the wrapper as loaded (RDKit is CPU-only)."""
        if self._loaded:
            return
        logging.info(
            "Loading RDKit wrapper (type=%s, radius=%d, n_bits=%d)",
            self.fingerprint_type,
            self.radius,
            self.n_bits,
        )
        self.device = device
        self._loaded = True

    def embed(
        self,
        input: str,
        pooling_strategy: str = "flat",
        **kwargs: Any,
    ) -> np.ndarray:
        """Compute the fingerprint for a single SMILES string.

        Returns
        -------
        np.ndarray
            1-D ``float32`` array of length :attr:`n_bits`.
            Binary fingerprints contain only 0.0 / 1.0.
            Count fingerprints contain non-negative integers (as float32).
        """
        if not self._loaded:
            raise RuntimeError("RDKit model not loaded. Call load() first.")

        from rdkit import Chem, DataStructs
        mol = Chem.MolFromSmiles(input)
        if mol is None:
            raise ValueError(f"RDKit failed to parse SMILES: {input}")

        return self._compute_fingerprint(mol)

    def embed_batch(
        self,
        inputs: Sequence[str],
        pooling_strategy: str = "flat",
        **kwargs: Any,
    ) -> list[np.ndarray]:
        """Compute fingerprints for a batch of SMILES strings."""
        return [self.embed(s, pooling_strategy=pooling_strategy, **kwargs) for s in inputs]

    # ------------------------------------------------------------------
    # Internal fingerprint computation
    # ------------------------------------------------------------------

    def _compute_fingerprint(self, mol: Any) -> np.ndarray:
        """Dispatch to the correct RDKit fingerprint generator."""
        from rdkit.Chem import rdFingerprintGenerator

        fp_type = self.fingerprint_type

        # --- Morgan (ECFP-like) ---
        if fp_type == "morgan":
            gen = rdFingerprintGenerator.GetMorganGenerator(
                radius=self.radius, fpSize=self.n_bits,
            )
            fp = gen.GetFingerprintAsNumPy(mol)
            return fp.astype(np.float32)

        if fp_type == "morgan_count":
            gen = rdFingerprintGenerator.GetMorganGenerator(
                radius=self.radius, fpSize=self.n_bits,
            )
            fp = gen.GetCountFingerprintAsNumPy(mol)
            return fp.astype(np.float32)

        # --- RDKit topological ---
        if fp_type == "rdkit":
            gen = rdFingerprintGenerator.GetRDKitFPGenerator(fpSize=self.n_bits)
            fp = gen.GetFingerprintAsNumPy(mol)
            return fp.astype(np.float32)

        # --- MACCS keys (167 bits) ---
        if fp_type == "maccs":
            from rdkit.Chem import MACCSkeys

            fp = MACCSkeys.GenMACCSKeys(mol)  # type: ignore[attr-defined]
            return self._bitvect_to_array(fp)

        # --- Atom pair ---
        if fp_type == "atom_pair":
            gen = rdFingerprintGenerator.GetAtomPairGenerator(fpSize=self.n_bits)
            fp = gen.GetFingerprintAsNumPy(mol)
            return fp.astype(np.float32)

        if fp_type == "atom_pair_count":
            gen = rdFingerprintGenerator.GetAtomPairGenerator(fpSize=self.n_bits)
            fp = gen.GetCountFingerprintAsNumPy(mol)
            return fp.astype(np.float32)

        # --- Topological torsion ---
        if fp_type == "topological_torsion":
            gen = rdFingerprintGenerator.GetTopologicalTorsionGenerator(fpSize=self.n_bits)
            fp = gen.GetFingerprintAsNumPy(mol)
            return fp.astype(np.float32)

        if fp_type == "topological_torsion_count":
            gen = rdFingerprintGenerator.GetTopologicalTorsionGenerator(fpSize=self.n_bits)
            fp = gen.GetCountFingerprintAsNumPy(mol)
            return fp.astype(np.float32)

        raise ValueError(f"Unknown fingerprint type: {fp_type}")

    # ------------------------------------------------------------------
    # Conversion helpers
    # ------------------------------------------------------------------

    def _bitvect_to_array(self, fp: Any) -> np.ndarray:
        """Convert an RDKit ``ExplicitBitVect`` to a float32 numpy array."""
        from rdkit import DataStructs
        arr = np.zeros(self.n_bits, dtype=np.float32)
        DataStructs.ConvertToNumpyArray(fp, arr)
        return arr

    def _sparse_intvect_to_array(self, fp: Any) -> np.ndarray:
        """Convert an RDKit ``(U)IntSparseIntVect`` to a dense float32 array."""
        arr = np.zeros(self.n_bits, dtype=np.float32)
        for idx, count in fp.GetNonzeroElements().items():
            arr[idx % self.n_bits] += count
        return arr

is_binary_fingerprint property

True when the fingerprint produces binary 0/1 vectors.

is_count_fingerprint property

True when the fingerprint produces continuous counts.

embed(input, pooling_strategy='flat', **kwargs)

Compute the fingerprint for a single SMILES string.

Returns:

Type	Description
ndarray	1-D float32 array of length :attr:n_bits. Binary fingerprints contain only 0.0 / 1.0. Count fingerprints contain non-negative integers (as float32).

Source code in src/embpy/models/molecule_models.py

def embed(
    self,
    input: str,
    pooling_strategy: str = "flat",
    **kwargs: Any,
) -> np.ndarray:
    """Compute the fingerprint for a single SMILES string.

    Returns
    -------
    np.ndarray
        1-D ``float32`` array of length :attr:`n_bits`.
        Binary fingerprints contain only 0.0 / 1.0.
        Count fingerprints contain non-negative integers (as float32).
    """
    if not self._loaded:
        raise RuntimeError("RDKit model not loaded. Call load() first.")

    from rdkit import Chem, DataStructs
    mol = Chem.MolFromSmiles(input)
    if mol is None:
        raise ValueError(f"RDKit failed to parse SMILES: {input}")

    return self._compute_fingerprint(mol)

embed_batch(inputs, pooling_strategy='flat', **kwargs)

Compute fingerprints for a batch of SMILES strings.

Source code in src/embpy/models/molecule_models.py

def embed_batch(
    self,
    inputs: Sequence[str],
    pooling_strategy: str = "flat",
    **kwargs: Any,
) -> list[np.ndarray]:
    """Compute fingerprints for a batch of SMILES strings."""
    return [self.embed(s, pooling_strategy=pooling_strategy, **kwargs) for s in inputs]

load(device)

Mark the wrapper as loaded (RDKit is CPU-only).

Source code in src/embpy/models/molecule_models.py

def load(self, device: torch.device) -> None:
    """Mark the wrapper as loaded (RDKit is CPU-only)."""
    if self._loaded:
        return
    logging.info(
        "Loading RDKit wrapper (type=%s, radius=%d, n_bits=%d)",
        self.fingerprint_type,
        self.radius,
        self.n_bits,
    )
    self.device = device
    self._loaded = True