MicroGlot / modeling_microglot.py
MicroGlot v1.0
6e50d10 15 days ago
82.1 kB
"""MicroGlot: a taxonomy-informed sparse (mixture-of-experts) decoder-only DNA language model.
One code file serves the three models of the repository ``athanzli/MicroGlot``:
* the repository root  - MicroGlot, the species-conditioned language model (``MicroGlotForCausalLM``)
* ``plain/``            - MicroGlot-plain, the same architecture without species conditioning
* ``species_encoder/``  - the Species-encoder, which predicts a species embedding from DNA (``MicroGlotSpeciesEncoder``)
Species inputs follow the ``input_ids`` / ``inputs_embeds`` convention:
* ``species_ids``    (LongTensor ``[batch]``): rows of the model's table of the 99,700 pretraining species
                     (produced by the tokenizer's ``species=`` argument); ``-1`` marks a sequence whose
                     species is unknown and must be inferred by an attached species encoder.
* ``species_embeds`` (FloatTensor ``[batch, 32]`` or ``[32]``): your own 32-d species vectors.
MicroGlot runs on NVIDIA GPUs with flash-attn installed, as it was trained and benchmarked. The numerical
code (attention, mixture of experts, decoder layers) is the code MicroGlot was trained with, unchanged.
"""

import math
import warnings
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import os as _os
_os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")

import torch
import torch.nn as nn
import torch.nn.functional as F
import transformers
from packaging import version
from torch.utils.checkpoint import checkpoint

# Supported versions (see the model card): PyTorch >= 2.7, transformers >= 4.51.3, flash-attn 2.7.4-2.8.3 tested.
# PyTorch 2.5 and 2.6 compute nn.RMSNorm(eps=None) of bf16 inputs with the bf16 epsilon, which changes the outputs.
_INSTALL = "https://huggingface.co/athanzli/MicroGlot#installation"
if version.parse(torch.__version__).release < (2, 7):
    raise ImportError(f"MicroGlot requires PyTorch 2.7 or later (found {torch.__version__}). See {_INSTALL}.")
if version.parse(transformers.__version__).release < (4, 51, 3):
    raise ImportError(f"MicroGlot requires transformers 4.51.3 or later (found {transformers.__version__}). "
                      f"See {_INSTALL}.")
# MicroGlot was trained and benchmarked with flash-attn's rotary position-embedding kernel, so it
# is required: every installation then computes the same embeddings.
try:
    from flash_attn.layers.rotary import RotaryEmbedding
except ImportError as _flash_attn_error:
    raise ImportError(
        "MicroGlot requires FlashAttention-2 (the flash-attn package, whose rotary position-embedding kernel was "
        f"used to train and benchmark the model) and an NVIDIA GPU. Install it as described at {_INSTALL}. "
        f"(Importing flash_attn failed with: {_flash_attn_error})"
    ) from _flash_attn_error

from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import ModelOutput
from transformers.utils import logging

logger = logging.get_logger(__name__)

# from_pretrained(torch_dtype=...) is called dtype= from transformers 4.56 on
_DTYPE_KWARG = "dtype" if version.parse(transformers.__version__).release >= (4, 56) else "torch_dtype"


@dataclass
class MicroGlotModelOutput(ModelOutput):
    """Backbone output. The species position is removed, so every tensor lines up with ``input_ids``.
    last_hidden_state: ``[batch, length, hidden]``, the final-norm output.
    hidden_states: with ``output_hidden_states=True``, ``num_hidden_layers + 1`` tensors (24): ``[0]`` the
        token embeddings, ``[k]`` the output of decoder layer k (1..22), ``[-1]`` the final-norm output
        (equal to ``last_hidden_state``).
    moe_loss: the router load-balancing loss summed over layers (already scaled by
        ``router_aux_loss_coef``), or None with ``output_router_logits=False``.
    species_embeds: ``[batch, 32]`` unit-norm species vectors actually used, in the model dtype
        (None for MicroGlot-plain).
    """

    last_hidden_state: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    moe_loss: Optional[torch.FloatTensor] = None
    species_embeds: Optional[torch.FloatTensor] = None


@dataclass
class MicroGlotCausalLMOutput(ModelOutput):
    """Next-token output of ``MicroGlotForCausalLM``.
    loss: with ``labels``, the mean next-token cross-entropy (plus ``moe_loss`` with ``output_router_logits=True``).
    logits: ``[batch, length, vocab]`` next-token logits.
    hidden_states: as in ``MicroGlotModelOutput`` (with ``skip_lm_head=True``, the final hidden state only).
    species_embeds: ``[batch, 32]`` unit-norm species vectors actually used (None for MicroGlot-plain).
    moe_loss: with ``output_router_logits=True``, the router load-balancing loss (as in ``MicroGlotModelOutput``).
    """

    loss: Optional[torch.FloatTensor] = None
    logits: Optional[torch.FloatTensor] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    species_embeds: Optional[torch.FloatTensor] = None
    moe_loss: Optional[torch.FloatTensor] = None


@dataclass
class MicroGlotSpeciesEncoderOutput(ModelOutput):
    """Species-encoder output, in the style of ``CLIPTextModelOutput``.
    species_embeds: ``[batch, 32]`` unit-norm predicted species vectors, normalised in the model dtype: the
        vectors MicroGlot is conditioned on (listed first, so pipelines return it).
    embeddings: ``[batch, 32]`` the projection output before normalisation.
    pooler_output: ``[batch, hidden]`` the pooled state (last real token).
    last_hidden_state, hidden_states, moe_loss: as in ``MicroGlotModelOutput``.
    """

    species_embeds: torch.FloatTensor = None
    embeddings: Optional[torch.FloatTensor] = None
    pooler_output: Optional[torch.FloatTensor] = None
    last_hidden_state: Optional[torch.FloatTensor] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    moe_loss: Optional[torch.FloatTensor] = None


@dataclass
class MoEParameterCounts:
    """Parameter counts for MoE models."""
    total_params: int
    activated_params: int
    non_embedding_total: int
    non_embedding_activated: int
    embedding_params: int
    moe_total_params: int
    moe_activated_params: int
    dense_params: int
    num_moe_layers: int
    num_dense_layers: int
    experts_per_layer: List[int]
    top_k: int
    
    def __str__(self) -> str:
        """Human-readable summary string."""
        return (
            f"MoE Parameter Counts:\n"
            f"  Total Parameters: {self.total_params:,} ({self.total_params / 1e9:.2f}B)\n"
            f"  Activated Parameters: {self.activated_params:,} ({self.activated_params / 1e6:.1f}M)\n"
            f"  Non-Embedding Total: {self.non_embedding_total:,} ({self.non_embedding_total / 1e9:.2f}B)\n"
            f"  Non-Embedding Activated: {self.non_embedding_activated:,} ({self.non_embedding_activated / 1e6:.1f}M)\n"
            f"  Embedding Parameters: {self.embedding_params:,} ({self.embedding_params / 1e6:.1f}M)\n"
            f"  MoE Layers: {self.num_moe_layers} (total: {self.moe_total_params:,}, activated: {self.moe_activated_params:,})\n"
            f"  Dense Layers: {self.num_dense_layers} ({self.dense_params:,})\n"
            f"  Top-K Routing: {self.top_k}\n"
            f"  Experts per Layer: {self.experts_per_layer}"
        )
    
    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        return {
            "total_params": self.total_params,
            "activated_params": self.activated_params,
            "non_embedding_total": self.non_embedding_total,
            "non_embedding_activated": self.non_embedding_activated,
            "embedding_params": self.embedding_params,
            "moe_total_params": self.moe_total_params,
            "moe_activated_params": self.moe_activated_params,
            "dense_params": self.dense_params,
            "num_moe_layers": self.num_moe_layers,
            "num_dense_layers": self.num_dense_layers,
            "experts_per_layer": self.experts_per_layer,
            "top_k": self.top_k,
        }


def compute_moe_parameter_counts(config: "GenomicModelConfig") -> MoEParameterCounts:
    """Compute total and activated parameter counts for MoE models from config."""
    hidden_size = config.hidden_size
    intermediate_size = config.intermediate_size
    vocab_size = config.vocab_size
    num_layers = config.num_hidden_layers
    num_query_heads = config.num_query_heads
    num_kv_heads = config.num_kv_heads
    head_dim = config.head_dim
    top_k = config.num_experts_per_tok if config.use_moe else 0
    use_residual = config.moe_use_residual if config.use_moe else False
    
    if config.use_moe and config.moe_layer_experts is not None:
        experts_per_layer = config.moe_layer_experts
    else:
        experts_per_layer = [0] * num_layers
    
    embedding_params = vocab_size * hidden_size
    
    attn_params = (
        hidden_size * (num_query_heads * head_dim) +
        hidden_size * (num_kv_heads * head_dim) +
        hidden_size * (num_kv_heads * head_dim) +
        (num_query_heads * head_dim) * hidden_size
    )
    
    layernorm_params = 2 * hidden_size
    
    ffn_params = (
        hidden_size * (2 * intermediate_size) +
        intermediate_size * hidden_size
    )
    
    moe_total_params = 0
    moe_activated_params = 0
    dense_params = 0
    num_moe_layers = 0
    num_dense_layers = 0
    
    for layer_idx, num_experts in enumerate(experts_per_layer):
        layer_base_params = attn_params + layernorm_params
        
        if num_experts > 0:
            num_moe_layers += 1
            
            expert_total = num_experts * ffn_params
            
            gate_params = hidden_size * num_experts
            
            expert_activated = (top_k / num_experts) * expert_total + gate_params
            
            species_film_params = 0
            if config.species_emb_dim is not None:
                species_film_params = 2 * (config.species_emb_dim * num_experts + num_experts)
            
            residual_params = (ffn_params + hidden_size * 2) if use_residual else 0
            
            layer_moe_total = expert_total + gate_params + species_film_params + residual_params
            layer_moe_activated = expert_activated + species_film_params + residual_params
            
            moe_total_params += layer_base_params + layer_moe_total
            moe_activated_params += layer_base_params + layer_moe_activated
        else:
            num_dense_layers += 1
            layer_total = layer_base_params + ffn_params
            dense_params += layer_total
    
    final_norm_params = hidden_size
    
    species_encoder_params = 0
    if config.use_species_encoder and config.species_emb_dim is not None:
        species_encoder_params = config.species_emb_dim
    
    non_embedding_total = moe_total_params + dense_params + final_norm_params + species_encoder_params
    non_embedding_activated = moe_activated_params + dense_params + final_norm_params + species_encoder_params
    
    total_params = embedding_params + non_embedding_total
    activated_params = embedding_params + non_embedding_activated
    
    return MoEParameterCounts(
        total_params=total_params,
        activated_params=activated_params,
        non_embedding_total=non_embedding_total,
        non_embedding_activated=non_embedding_activated,
        embedding_params=embedding_params,
        moe_total_params=moe_total_params,
        moe_activated_params=moe_activated_params,
        dense_params=dense_params,
        num_moe_layers=num_moe_layers,
        num_dense_layers=num_dense_layers,
        experts_per_layer=experts_per_layer,
        top_k=top_k,
    )


def count_parameters(model: nn.Module, requires_grad_only: bool = False) -> int:
    """Count the number of parameters in a model."""
    if requires_grad_only:
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
    return sum(p.numel() for p in model.parameters())


def get_model_parameter_counts(model: "GenomicModel") -> MoEParameterCounts:
    """Get parameter counts from an instantiated model."""
    if hasattr(model, 'model'):
        base_model = model.model
        config = model.config
    else:
        base_model = model
        config = model.config
    
    counts = compute_moe_parameter_counts(config)
    
    actual_total = count_parameters(model)
    
    expected_diff = 0
    if hasattr(model, 'lm_head'):
        expected_diff = config.vocab_size * config.hidden_size
    
    if abs(actual_total - counts.total_params - expected_diff) > 1000:
        import warnings
        warnings.warn(
            f"Parameter count mismatch: analytical={counts.total_params:,}, "
            f"actual={actual_total:,}, expected_diff={expected_diff:,}. "
            f"Difference: {actual_total - counts.total_params - expected_diff:,}"
        )
    
    return counts


def format_parameter_count(params: int) -> str:
    """Format parameter count in human-readable form."""
    if params >= 1e9:
        return f"{params / 1e9:.2f}B"
    elif params >= 1e6:
        return f"{params / 1e6:.1f}M"
    elif params >= 1e3:
        return f"{params / 1e3:.1f}K"
    else:
        return str(params)


def print_model_parameter_summary(config: "GenomicModelConfig") -> None:
    """Print a detailed parameter summary for a model configuration."""
    counts = compute_moe_parameter_counts(config)
    
    print("=" * 60)
    print("Model Parameter Summary")
    print("=" * 60)
    print(f"Total Parameters:      {format_parameter_count(counts.total_params):>10} ({counts.total_params:,})")
    print(f"Activated Parameters:  {format_parameter_count(counts.activated_params):>10} ({counts.activated_params:,})")
    print("-" * 60)
    print(f"Embedding Parameters:  {format_parameter_count(counts.embedding_params):>10}")
    print(f"Non-Embedding Total:   {format_parameter_count(counts.non_embedding_total):>10}")
    print(f"Non-Embedding Active:  {format_parameter_count(counts.non_embedding_activated):>10}")
    print("-" * 60)
    print(f"Dense Layers:          {counts.num_dense_layers:>10} ({format_parameter_count(counts.dense_params)})")
    print(f"MoE Layers:            {counts.num_moe_layers:>10} (total: {format_parameter_count(counts.moe_total_params)}, "
          f"active: {format_parameter_count(counts.moe_activated_params)})")
    print(f"Top-K Routing:         {counts.top_k:>10}")
    print(f"Experts per Layer:     {counts.experts_per_layer}")
    print("=" * 60)



class MicroGlotConfig(PretrainedConfig):
    """Configuration shared by MicroGlot, MicroGlot-plain and the Species-encoder.
    The architecture fields (vocab_size, hidden_size, num_hidden_layers, num_query_heads, num_kv_heads,
    intermediate_size, rope_theta, moe_layer_experts, num_experts_per_tok, moe_use_residual,
    gate_softmax_over_all_experts, router_aux_loss_coef, moe_dispatch, ...) are those of the pretrained models.
    `moe_dispatch`: "sorted" (default), "loop" or "grouped".
    Species-related fields:
        species_emb_dim (`int`, *optional*): size of the species vector (32). `None` -> no species input
            (MicroGlot-plain, and the backbone of the species encoder).
        prepend_species_token (`bool`): insert the projected species vector as a token after `[BOS]`.
        num_species (`int`): rows in the built-in species table addressed by `species_ids` (99,700; 0 = none).
        species_vocab_sha256 (`str`, *optional*): sha256 of the species names in table order, joined by
            newlines; the tokenizer's `species_vocab_sha256` must match it.
        species_encoder_repo / species_encoder_revision / species_encoder_subfolder (`str`, *optional*): where
            `load_species_encoder()` and `from_pretrained(..., species_encoder=True)` find the Species-encoder.
            `species_encoder_repo=None`: the repository or folder this model was loaded from, at the same
            revision; `species_encoder_subfolder` is the encoder's folder there (`"species_encoder"`).
        species_encoder_embedding_dim (`int`): output size of `MicroGlotSpeciesEncoder` (32).
        species_encoder_pooling_strategy (`str`): `"eos"` (last real token) or `"mean"`.
    """

    model_type = "microglot"
    keys_to_ignore_at_inference = ["moe_loss", "species_embeds"]

    def __init__(
        self,
        vocab_size: int = 8192,
        hidden_size: int = 1024,
        num_hidden_layers: int = 23,
        num_query_heads: int = 16,
        num_kv_heads: int = 8,
        intermediate_size: Optional[int] = None,
        attention_dropout_prob: float = 0.0,
        classifier_dropout_prob: float = 0.0,
        attn_output_dropout_prob: float = 0.0,
        moe_output_dropout_prob: float = 0.0,
        ffn_dropout_prob: float = 0.0,
        moe_expert_dropout_prob: float = 0.0,
        rope_theta: float = 500000.0,
        max_trained_length: int = 8192,
        init_scale: float = 0.1,
        use_swiglu: bool = True,
        gradient_checkpointing: bool = False,
        use_flash_attention: bool = True,
        pad_token_id: int = 0,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        num_experts_per_tok: int = 1,
        router_aux_loss_coef: float = 0.001,
        moe_layer_experts: Optional[List[int]] = None,
        moe_use_residual: bool = True,
        moe_capacity_factor: float = 1.25,
        moe_eval_capacity_factor: float = 2.0,
        moe_min_capacity: int = 4,
        moe_noisy_gate_policy: Optional[str] = None,
        moe_drop_tokens: bool = True,
        moe_use_rts: bool = True,
        gate_softmax_over_all_experts: bool = True,
        moe_dispatch: str = "sorted",
        species_emb_dim: Optional[int] = None,
        prepend_species_token: bool = False,
        num_species: int = 0,
        species_encoder_embedding_dim: int = 32,
        species_encoder_pooling_strategy: str = "eos",
        species_encoder_repo: Optional[str] = None,
        species_encoder_revision: Optional[str] = None,
        species_encoder_subfolder: Optional[str] = None,
        species_vocab_sha256: Optional[str] = None,
        use_species_encoder: bool = False,
        **kwargs,
    ):
        if use_species_encoder:  # a training-code config: the species encoder is bundled inside the checkpoint
            logger.warning_once(
                "This checkpoint bundles a copy of the species encoder, which this code ignores. Species names "
                "and vectors work unchanged. To infer species, attach the Species-encoder with "
                "model.load_species_encoder()."
            )
            if species_encoder_repo is None:
                species_encoder_repo, species_encoder_subfolder = "athanzli/MicroGlot", "species_encoder"
        super().__init__(pad_token_id=pad_token_id, eos_token_id=eos_token_id, bos_token_id=bos_token_id, **kwargs)
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_query_heads = num_query_heads
        self.num_kv_heads = num_kv_heads
        if intermediate_size is None:
            if use_swiglu:
                hidden_dim = (8 * hidden_size) // 3
                intermediate_size = 256 * ((hidden_dim + 255) // 256)
            else:
                intermediate_size = 4 * hidden_size
        self.intermediate_size = intermediate_size
        self.attention_dropout_prob = attention_dropout_prob
        self.classifier_dropout_prob = classifier_dropout_prob
        self.attn_output_dropout_prob = attn_output_dropout_prob
        self.moe_output_dropout_prob = moe_output_dropout_prob
        self.ffn_dropout_prob = ffn_dropout_prob
        self.moe_expert_dropout_prob = moe_expert_dropout_prob
        self.rope_theta = rope_theta
        self.max_trained_length = max_trained_length
        self.init_scale = init_scale
        self.use_swiglu = use_swiglu
        self.gradient_checkpointing = gradient_checkpointing
        self.use_flash_attention = use_flash_attention
        self.num_experts_per_tok = num_experts_per_tok
        self.router_aux_loss_coef = router_aux_loss_coef
        self.moe_layer_experts = moe_layer_experts
        self.use_moe = moe_layer_experts is not None and any(n > 0 for n in moe_layer_experts)
        self.moe_use_residual = moe_use_residual
        self.moe_capacity_factor = moe_capacity_factor
        self.moe_eval_capacity_factor = moe_eval_capacity_factor
        self.moe_min_capacity = moe_min_capacity
        self.moe_noisy_gate_policy = moe_noisy_gate_policy
        self.moe_drop_tokens = moe_drop_tokens
        self.moe_use_rts = moe_use_rts
        self.gate_softmax_over_all_experts = gate_softmax_over_all_experts
        self.moe_dispatch = moe_dispatch
        self.species_emb_dim = species_emb_dim
        self.use_species_encoder = False  # a training-code flag; the encoder is attached separately
        self.prepend_species_token = prepend_species_token
        self.num_species = num_species
        self.species_encoder_embedding_dim = species_encoder_embedding_dim
        self.species_encoder_pooling_strategy = species_encoder_pooling_strategy
        self.species_encoder_repo = species_encoder_repo
        self.species_encoder_revision = species_encoder_revision
        self.species_encoder_subfolder = species_encoder_subfolder
        self.species_vocab_sha256 = species_vocab_sha256
        if prepend_species_token and species_emb_dim is None:
            raise ValueError("species_emb_dim must be set when prepend_species_token=True.")
        if num_species and species_emb_dim is None:
            raise ValueError("num_species > 0 needs species_emb_dim (the width of the species table).")
        assert hidden_size % num_query_heads == 0
        self.head_dim = hidden_size // num_query_heads
        assert num_query_heads % num_kv_heads == 0, "num_query_heads must be divisible by num_kv_heads"
        self.num_kv_groups = num_query_heads // num_kv_heads
        self._validate_moe_config()

    def _validate_moe_config(self):
        if not self.use_moe:
            return
        if len(self.moe_layer_experts) != self.num_hidden_layers:
            raise ValueError("moe_layer_experts must have one entry per layer (0 = dense).")
        for i, n in enumerate(self.moe_layer_experts):
            if not isinstance(n, int) or n < 0:
                raise ValueError(f"moe_layer_experts[{i}] = {n} is invalid.")
            if 0 < n < self.num_experts_per_tok:
                raise ValueError(f"moe_layer_experts[{i}] = {n} < num_experts_per_tok.")

    def get_num_experts_for_layer(self, layer_idx: int) -> int:
        """Number of experts in layer `layer_idx` (0 = dense feed-forward)."""
        if not self.use_moe:
            return 0
        return self.moe_layer_experts[layer_idx]

    def get_parameter_counts(self) -> "MoEParameterCounts":
        """Total and activated parameter counts for this configuration."""
        return compute_moe_parameter_counts(self)

    def print_parameter_summary(self) -> None:
        """Print a detailed parameter summary for this configuration."""
        print_model_parameter_summary(self)

    def to_dict(self):
        # save_pretrained() copies this code file next to config.json (see register_for_auto_class below), so
        # keep every auto_map entry local: a saved model then runs the code it was saved with, not the Hub's
        # latest (transformers rewrites Hub-loaded entries to "athanzli/MicroGlot--modeling_microglot...").
        out = super().to_dict()
        if isinstance(out.get("auto_map"), dict):
            out["auto_map"] = {k: v.split("--")[-1] if isinstance(v, str) else v for k, v in out["auto_map"].items()}
        return out

    @classmethod
    def get_config_dict(cls, *args, **kwargs):
        config_dict, kwargs = super().get_config_dict(*args, **kwargs)
        if config_dict.get("model_type") == "genomic_lm":  # the training code's name of this model type
            config_dict["model_type"] = cls.model_type
        return config_dict, kwargs


GenomicModelConfig = MicroGlotConfig  # the training code's name, used in the layer code's type hints


class MicroGlotSpeciesEncoderConfig(MicroGlotConfig):
    """Configuration of the Species-encoder: `MicroGlotConfig` under its own model type, so that the three
    models of the repository, which share this code file, keep distinct Auto classes."""

    model_type = "microglot_species_encoder"

    def __init__(self, **kwargs):  # explicit: transformers 5 would otherwise turn the subclass into a dataclass
        super().__init__(**kwargs)

    @classmethod
    def get_config_dict(cls, *args, **kwargs):
        config_dict, kwargs = super().get_config_dict(*args, **kwargs)
        if config_dict.get("model_type") in ("microglot", "genomic_lm"):  # encoder folders of the training code
            config_dict["model_type"] = cls.model_type
        return config_dict, kwargs


class GenomicAttention(nn.Module):
    """Grouped Query Attention (GQA) with RoPE and Flash Attention 2.0."""
    
    def __init__(self, config: GenomicModelConfig):
        super().__init__()
        self.config = config
        
        self.hidden_size = config.hidden_size
        self.num_query_heads = config.num_query_heads
        self.num_kv_heads = config.num_kv_heads
        self.num_kv_groups = config.num_kv_groups
        self.head_dim = config.head_dim
        self.attention_dropout = config.attention_dropout_prob
        
        self.q_proj = nn.Linear(self.hidden_size, self.num_query_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_query_heads * self.head_dim, self.hidden_size, bias=False)
        
        self.attn_output_dropout_prob = config.attn_output_dropout_prob
        self.attn_output_dropout = nn.Dropout(config.attn_output_dropout_prob) if config.attn_output_dropout_prob > 0 else nn.Identity()
        
        self.rotary_emb = RotaryEmbedding(
            dim=self.head_dim,
            base=config.rope_theta,
            interleaved=False
        )
    
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Apply causal (autoregressive) Flash Attention with RoPE."""
        batch_size, seq_len, _ = hidden_states.shape
        
        if attention_mask is None:
            attention_mask = torch.ones((batch_size, seq_len), device=hidden_states.device, dtype=torch.bool)
        else:
            attention_mask = attention_mask.to(dtype=torch.bool)

        query_states = self.q_proj(hidden_states).view(batch_size, seq_len, self.num_query_heads, self.head_dim).contiguous()
        key_states = self.k_proj(hidden_states).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).contiguous()
        value_states = self.v_proj(hidden_states).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).contiguous()
        
        qkv_states = self.rotary_emb(
            qkv = torch.cat([query_states, key_states, value_states], dim=2),
            num_heads_q = self.num_query_heads,
        )
        query_states = qkv_states[:, :, :self.num_query_heads, :].contiguous()
        key_states = qkv_states[:, :, self.num_query_heads:self.num_query_heads + self.num_kv_heads, :].contiguous()
        value_states = qkv_states[:, :, self.num_query_heads + self.num_kv_heads:, :].contiguous()
        
        query_states = query_states.transpose(1, 2)
        key_states = key_states.transpose(1, 2)
        value_states = value_states.transpose(1, 2)
        
        # SDPA picks the fastest kernel the GPU and dtype support: flash attention for bf16/fp16 on
        # Ampere or newer, another kernel otherwise (forcing flash left fp32 and older GPUs without one).
        attn_output = F.scaled_dot_product_attention(
            query_states,
            key_states,
            value_states,
            dropout_p=self.attention_dropout if self.training else 0.0,
            is_causal=True,
            enable_gqa=True,
        )
        attn_output = attn_output.transpose(1, 2)
        attn_output = attn_output.reshape(batch_size, seq_len, self.num_query_heads * self.head_dim)
        attn_output = self.o_proj(attn_output)
        attn_output = self.attn_output_dropout(attn_output)
        return attn_output


class SwiGLU(nn.Module):
    """SwiGLU activation function."""
    
    def __init__(self, config: GenomicModelConfig, dropout_prob: float = None):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size

        self.gate_up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        
        self.dropout_prob = dropout_prob if dropout_prob is not None else getattr(config, 'ffn_dropout_prob', 0.0)
        self.dropout = nn.Dropout(self.dropout_prob) if self.dropout_prob > 0 else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate_up = self.gate_up_proj(x)
        gate, up = gate_up.chunk(2, dim=-1)
        out = self.down_proj(F.silu(gate) * up)
        return self.dropout(out)


class TopKGate(nn.Module):
    """Top-k gating network for Mixture of Experts."""
    
    def __init__(
        self,
        hidden_size: int,
        num_experts: int,
        top_k: int = 1,
        noisy_gate_policy: Optional[str] = None,
        species_emb_dim: Optional[int] = None,
        gate_softmax_over_all_experts: bool = False,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_experts = num_experts
        self.top_k = top_k
        self.noisy_gate_policy = noisy_gate_policy
        self.gate_softmax_over_all_experts = gate_softmax_over_all_experts
        
        self.wg = nn.Linear(hidden_size, num_experts, bias=False)
        
        self.use_species_film = species_emb_dim is not None
        if self.use_species_film:
            self.film_gamma = nn.Linear(species_emb_dim, num_experts, bias=True)
            self.film_beta = nn.Linear(species_emb_dim, num_experts, bias=True)
            nn.init.trunc_normal_(self.film_gamma.weight, std=0.02)
            nn.init.zeros_(self.film_gamma.bias)
            nn.init.trunc_normal_(self.film_beta.weight, std=0.02)
            nn.init.zeros_(self.film_beta.bias)
        else:
            self.film_gamma = None
            self.film_beta = None
    
    def _compute_aux_loss(
        self,
        logits: torch.Tensor,
        top_k_indices: torch.Tensor,
    ) -> torch.Tensor:
        """Compute GShard-style load balancing auxiliary loss."""
        num_tokens = logits.shape[0]
        
        flat_indices = top_k_indices.reshape(-1)
        expert_counts = torch.zeros(self.num_experts, device=logits.device, dtype=logits.dtype)
        expert_counts.scatter_add_(0, flat_indices, torch.ones_like(flat_indices, dtype=logits.dtype))
        f = expert_counts / num_tokens
        
        probs = F.softmax(logits, dim=-1)
        p = probs.mean(dim=0)
        
        if self.top_k == 1:
            aux_loss = self.num_experts * torch.sum(f * p)
        else:
            aux_loss = (self.num_experts ** 2) * torch.mean(f * p) / self.top_k
        
        return aux_loss
    
    def forward(
        self,
        hidden_states: torch.Tensor,
        species_emb: Optional[torch.Tensor] = None,
        seq_len: Optional[int] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Compute top-k gate scores."""
        if self.training and self.noisy_gate_policy == 'Jitter':
            noise = torch.empty_like(hidden_states).uniform_(1.0 - 0.01, 1.0 + 0.01)
            hidden_states = hidden_states * noise
        
        logits = F.linear(hidden_states, self.wg.weight, self.wg.bias).float()

        if self.use_species_film and species_emb is not None and self.film_gamma is not None:
            gamma = F.linear(species_emb, self.film_gamma.weight, self.film_gamma.bias).float()
            beta = F.linear(species_emb, self.film_beta.weight, self.film_beta.bias).float()
            
            gamma_expanded = gamma.repeat_interleave(seq_len, dim=0)
            beta_expanded = beta.repeat_interleave(seq_len, dim=0)
            
            logits = (1.0 + gamma_expanded) * logits + beta_expanded
        
        if self.gate_softmax_over_all_experts:
            all_probs = F.softmax(logits, dim=-1)
            gate_scores, top_k_indices = torch.topk(all_probs, self.top_k, dim=-1)
            if self.top_k > 1:
                gate_scores = gate_scores / gate_scores.sum(dim=-1, keepdim=True)
        else:
            top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)
            gate_scores = F.softmax(top_k_logits, dim=-1)

        aux_loss = self._compute_aux_loss(logits, top_k_indices)

        return gate_scores, top_k_indices, aux_loss


class MoELayer(nn.Module):
    """Native PyTorch Mixture of Experts layer."""
    
    def __init__(self, config: GenomicModelConfig, num_experts: int):
        """Initialize native MoE layer."""
        super().__init__()
        self.config = config
        self.num_experts = num_experts
        self.top_k = config.num_experts_per_tok
        self.aux_loss_coef = config.router_aux_loss_coef
        self.moe_dispatch = getattr(config, "moe_dispatch", "loop")
        
        use_species = config.use_species_encoder or (config.species_emb_dim is not None)
        species_dim = config.species_emb_dim if use_species else None
        
        self.gate = TopKGate(
            hidden_size=config.hidden_size,
            num_experts=num_experts,
            top_k=self.top_k,
            noisy_gate_policy=config.moe_noisy_gate_policy,
            species_emb_dim=species_dim,
            gate_softmax_over_all_experts=getattr(config, "gate_softmax_over_all_experts", False),
        )
        
        self.experts = nn.ModuleList([
            SwiGLU(config, dropout_prob=config.moe_expert_dropout_prob)
            for _ in range(num_experts)
        ])
        
        self.use_residual = config.moe_use_residual
        if self.use_residual:
            self.residual_mlp = SwiGLU(config, dropout_prob=config.moe_expert_dropout_prob)
            self.residual_mix = nn.Linear(config.hidden_size, 2, bias=True)
        
        self.moe_output_dropout_prob = config.moe_output_dropout_prob
        self.moe_output_dropout = (
            nn.Dropout(config.moe_output_dropout_prob)
            if config.moe_output_dropout_prob > 0
            else nn.Identity()
        )
    
    def _dispatch_experts(self, flat_hidden, gate_scores, expert_indices):
        mode = self.moe_dispatch
        if mode == "grouped":
            try:
                return self._dispatch_grouped(flat_hidden, gate_scores, expert_indices)
            except (RuntimeError, NotImplementedError) as e:
                if not getattr(self, "_grouped_warned", False):
                    print(f"[MoELayer] torch._grouped_mm unusable ({type(e).__name__}: "
                          f"{str(e)[:80]}); falling back to moe_dispatch='sorted'.")
                    self._grouped_warned = True
                self.moe_dispatch = "sorted"
                return self._dispatch_sorted(flat_hidden, gate_scores, expert_indices)
        if mode == "sorted":
            return self._dispatch_sorted(flat_hidden, gate_scores, expert_indices)
        return self._dispatch_loop(flat_hidden, gate_scores, expert_indices)

    def _dispatch_loop(self, flat_hidden, gate_scores, expert_indices):
        """Legacy per-expert Python loop."""
        output = torch.zeros_like(flat_hidden)
        for k in range(self.top_k):
            indices_k = expert_indices[:, k]
            scores_k = gate_scores[:, k]
            for expert_idx in range(self.num_experts):
                mask = (indices_k == expert_idx)
                if not mask.any():
                    continue
                token_indices = mask.nonzero(as_tuple=True)[0]
                expert_input = flat_hidden[token_indices]
                expert_output = self.experts[expert_idx](expert_input)
                weighted_output = expert_output * scores_k[token_indices].unsqueeze(-1)
                output[token_indices] += weighted_output
        output = output + sum(
            p.flatten()[0] * 0 for expert in self.experts for p in expert.parameters()
        )
        return output

    def _flatten_assignments(self, gate_scores, expert_indices):
        """(token, slot) pairs sorted by expert; shared by sorted/grouped dispatch."""
        num_tokens, top_k = expert_indices.shape
        flat_expert = expert_indices.reshape(-1)
        flat_score = gate_scores.reshape(-1)
        flat_token = torch.arange(num_tokens, device=expert_indices.device).repeat_interleave(top_k)
        order = torch.argsort(flat_expert, stable=True)
        sorted_token = flat_token[order]
        sorted_score = flat_score[order]
        counts = torch.bincount(flat_expert, minlength=self.num_experts)
        return sorted_token, sorted_score, counts

    def _dispatch_sorted(self, flat_hidden, gate_scores, expert_indices):
        """Sort tokens by expert, run each expert on its contiguous slice, scatter-add."""
        sorted_token, sorted_score, counts = self._flatten_assignments(gate_scores, expert_indices)
        sorted_input = flat_hidden.index_select(0, sorted_token)
        sizes = counts.tolist()
        chunks = torch.split(sorted_input, sizes, dim=0)
        expert_out = torch.cat(
            [self.experts[e](chunks[e]) for e in range(self.num_experts)], dim=0
        )
        expert_out = expert_out * sorted_score.to(expert_out.dtype).unsqueeze(-1)
        output = torch.zeros_like(flat_hidden)
        output.index_add_(0, sorted_token, expert_out.to(output.dtype))
        return output

    def _dispatch_grouped(self, flat_hidden, gate_scores, expert_indices):
        """Fully batched dispatch via torch._grouped_mm over stacked expert weights."""
        sorted_token, sorted_score, counts = self._flatten_assignments(gate_scores, expert_indices)
        sorted_input = flat_hidden.index_select(0, sorted_token).contiguous()
        offs = torch.cumsum(counts, 0).to(torch.int32)
        w_gate_up = torch.stack([e.gate_up_proj.weight for e in self.experts])
        w_down = torch.stack([e.down_proj.weight for e in self.experts])
        gate_up = torch._grouped_mm(
            sorted_input, w_gate_up.transpose(-2, -1), offs=offs
        )
        gate, up = gate_up.chunk(2, dim=-1)
        act = (F.silu(gate) * up).contiguous()
        down = torch._grouped_mm(
            act, w_down.transpose(-2, -1), offs=offs
        )
        down = down * sorted_score.to(down.dtype).unsqueeze(-1)
        output = torch.zeros_like(flat_hidden)
        output.index_add_(0, sorted_token, down.to(output.dtype))
        return output

    def forward(
        self,
        hidden_states: torch.Tensor,
        species_emb: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Apply MoE layer with sparse index-based routing."""
        batch_size, seq_len, hidden_size = hidden_states.shape
        input_dtype = hidden_states.dtype
        
        num_tokens = batch_size * seq_len
        flat_hidden = hidden_states.reshape(num_tokens, hidden_size)
        
        gate_scores, expert_indices, aux_loss = self.gate(
            flat_hidden, species_emb=species_emb, seq_len=seq_len
        )
        
        gate_scores = gate_scores.to(input_dtype)
        
        output = self._dispatch_experts(flat_hidden, gate_scores, expert_indices)

        if self.use_residual:
            residual_output = self.residual_mlp(flat_hidden)
            
            mix_logits = F.linear(
                flat_hidden,
                self.residual_mix.weight,
                self.residual_mix.bias,
            ).float()
            mix_coef = F.softmax(mix_logits, dim=-1).to(input_dtype)
            
            output = mix_coef[:, 0:1] * output + mix_coef[:, 1:2] * residual_output
        
        output = output.reshape(batch_size, seq_len, hidden_size)
        
        if self.moe_output_dropout_prob > 0 and self.training:
            output = self.moe_output_dropout(output)
        
        l_aux = aux_loss.float() * self.aux_loss_coef
        
        return output, l_aux


class GenomicDecoderLayer(nn.Module):
    """Single decoder layer for the Genomic Language Model."""
    def __init__(self, config: GenomicModelConfig, layer_idx: int = 0):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx

        self.self_attn = GenomicAttention(config)
        self.input_layernorm = nn.RMSNorm(normalized_shape=config.hidden_size)
        self.post_attention_layernorm = nn.RMSNorm(normalized_shape=config.hidden_size)
        
        if not config.use_swiglu:
            raise NotImplementedError("Only SwiGLU is supported.")
        
        self.num_experts = config.get_num_experts_for_layer(layer_idx)
        self.use_moe = self.num_experts > 0
        
        if self.use_moe:
            self.ffn = MoELayer(config, num_experts=self.num_experts)
        else:
            self.ffn = SwiGLU(config)
    
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        species_emb: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(hidden_states, attention_mask=attention_mask)
        hidden_states = residual + hidden_states
        
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        
        moe_loss = None
        if self.use_moe:
            hidden_states, moe_loss = self.ffn(hidden_states, species_emb=species_emb)
        else:
            hidden_states = self.ffn(hidden_states)
        
        hidden_states = residual + hidden_states
        
        return hidden_states, moe_loss


class MicroGlotSpeciesTable(nn.Module):
    """The pretraining species' 32-d Poincare vectors, addressed by `species_ids` (cf. VITS's speaker table).
    `weight` is a float32 buffer `[num_species, 32]` that follows device moves but never dtype casts
    (`torch_dtype=...`, `.to(torch.bfloat16)`, `.half()`): the prior is normalised in float32 and then cast.
    Row i belongs to line i of the tokenizer's species_vocab.txt.
    """

    def __init__(self, num_species: int, dim: int):
        super().__init__()
        self.register_buffer("weight", torch.zeros(num_species, dim, dtype=torch.float32))

    def _apply(self, fn, recurse=True):
        moved = fn(self.weight)
        self._buffers["weight"] = moved if moved.dtype == torch.float32 else self.weight.to(moved.device)
        return self

    def forward(self, species_ids: torch.LongTensor) -> torch.Tensor:
        return self.weight[species_ids]


class MicroGlotHead(nn.Module):
    """dense -> GELU -> out_proj: the Species-encoder's projection."""

    def __init__(self, in_dim: int, hidden: int, out_dim: int):
        super().__init__()
        self.dense = nn.Linear(in_dim, hidden)
        self.act = nn.GELU()
        self.out_proj = nn.Linear(hidden, out_dim)

    def forward(self, x):
        return self.out_proj(self.act(self.dense(x)))


def _rope_modules(module: nn.Module):
    """The flash-attn rotary-embedding modules inside `module`."""
    return [m for m in module.modules() if isinstance(m, RotaryEmbedding)]


class MicroGlotPreTrainedModel(PreTrainedModel):
    """Base class of the MicroGlot models: weight initialisation, loading, and the species-encoder calls
    (delegated to the `MicroGlotModel` backbone)."""

    config_class = MicroGlotConfig
    base_model_prefix = "model"
    main_input_name = "input_ids"
    supports_gradient_checkpointing = True
    _no_split_modules = ["GenomicDecoderLayer"]
    _keys_to_ignore_on_load_missing = [r"rotary_emb"]
    # training-code checkpoints bundle the species encoder (here a separate model); base loads of a causal-LM
    # checkpoint carry lm_head.
    _keys_to_ignore_on_load_unexpected = [r"(^|\.)species_encoder\.", r"^lm_head\."]

    def _init_weights(self, module: nn.Module):
        if isinstance(module, nn.Linear):
            fan_in = module.weight.shape[1]
            std = math.sqrt(self.config.init_scale / fan_in)
            nn.init.trunc_normal_(module.weight, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.trunc_normal_(module.weight, std=0.02)
        elif isinstance(module, RotaryEmbedding):
            # inv_freq is a non-persistent buffer (not in the checkpoint): transformers>=5 re-creates it
            # uninitialised after loading and relies on _init_weights to fill it.
            dim, base = module.dim, float(module.base)
            module.inv_freq.copy_(1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)))

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, species_encoder=None, **kwargs):
        """The standard `from_pretrained`, plus `species_encoder`:
        species_encoder (`bool` or `str`, *optional*): `True` also loads and attaches the Species-encoder, as
            `model.load_species_encoder()` does; a repo id or a local path loads the encoder found there. To
            attach an encoder object you already loaded, call `model.set_species_encoder(encoder)`.
        """
        if species_encoder is not None and not isinstance(species_encoder, (bool, str, _os.PathLike)):
            raise TypeError(
                "species_encoder= takes True (this model's Species-encoder), a repo id or a local path. "
                "To attach an encoder object you already loaded, call model.set_species_encoder(encoder)."
            )
        if _DTYPE_KWARG == "torch_dtype" and "dtype" in kwargs:  # accept the newer `dtype=` before transformers 4.56
            kwargs.setdefault("torch_dtype", kwargs.pop("dtype"))
        loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
        model = loaded[0] if isinstance(loaded, tuple) else loaded  # output_loading_info=True: (model, info)
        # where this model came from: load_species_encoder() looks for the encoder there, at the same revision
        if pretrained_model_name_or_path is not None:  # None: built from config= and state_dict=
            source = _os.fspath(pretrained_model_name_or_path)
            if _os.path.isdir(source):
                source = _os.path.abspath(source)
            revision = getattr(model.config, "_commit_hash", None) or kwargs.get("revision")
            for m in (model, getattr(model, "model", None)):
                if isinstance(m, MicroGlotModel):
                    m.__dict__["_microglot_source"] = (source, revision)
        if species_encoder:
            hub = {k: kwargs[k] for k in ("cache_dir", "force_download", "local_files_only", "token", "proxies")
                   if k in kwargs}
            model.load_species_encoder(None if species_encoder is True else species_encoder, **hub)
        return loaded

    def to_bfloat16(self) -> "MicroGlotPreTrainedModel":
        """Convert model weights to bfloat16 (the species table stays float32)."""
        return self.to(torch.bfloat16)

    # -- species-encoder composition, delegated to the backbone (MicroGlotModel implements it) ----------
    def _microglot_backbone(self):
        return getattr(self, "model", None)

    @property
    def species_encoder(self):
        """The attached species encoder, or None."""
        return self._microglot_backbone().species_encoder

    def set_species_encoder(self, encoder, match_backbone: bool = True):
        """Attach a species encoder you loaded yourself (None detaches); see `MicroGlotModel.set_species_encoder`."""
        self._microglot_backbone().set_species_encoder(encoder, match_backbone=match_backbone)

    def get_species_encoder(self):
        """The attached species encoder, or None."""
        return self._microglot_backbone().species_encoder

    def load_species_encoder(self, *args, **kwargs):
        """Load and attach the species encoder; see `MicroGlotModel.load_species_encoder`."""
        return self._microglot_backbone().load_species_encoder(*args, **kwargs)

    def get_species_embeds(self, *args, **kwargs):
        """The species vectors the model would use; see `MicroGlotModel.get_species_embeds`."""
        return self._microglot_backbone().get_species_embeds(*args, **kwargs)

    def encode_species(self, species, unknown_species: str = "raise") -> torch.LongTensor:
        """Species name(s) -> `species_ids` on the model's device: a 0-d tensor for one name, `[n]` for a list.
        A thin wrapper over this repository's tokenizer (`tokenizer.convert_species_to_ids`; `None` -> -1,
        loose matching, unknown names raise `KeyError` unless `unknown_species="infer"`), loaded once from
        where the model was loaded and checked against `config.species_vocab_sha256`.
        """
        if not self.config.num_species:
            raise ValueError("This checkpoint has no species table (MicroGlot-plain or the species encoder).")
        tok = self.__dict__.get("_species_tokenizer")
        if tok is None:
            from transformers import AutoTokenizer

            source = self.config._name_or_path
            commit = getattr(self.config, "_commit_hash", None)
            try:
                tok = AutoTokenizer.from_pretrained(source, trust_remote_code=True,
                                                    **({"revision": commit} if commit else {}))
            except (OSError, ValueError):
                tok = None
            if getattr(tok, "species_vocab_sha256", None) is None:
                raise ValueError(
                    f"Could not load the MicroGlot tokenizer from {source!r}. Load it with AutoTokenizer."
                    "from_pretrained('athanzli/MicroGlot', trust_remote_code=True) and call "
                    "tokenizer.convert_species_to_ids(names)."
                )
            if tok.species_vocab_sha256 != self.config.species_vocab_sha256:
                raise ValueError("The tokenizer's species list does not match this model's species table.")
            self.__dict__["_species_tokenizer"] = tok
        ids = tok.convert_species_to_ids(species, unknown_species=unknown_species)
        return torch.tensor(ids, dtype=torch.long, device=self.device)


class MicroGlotModel(MicroGlotPreTrainedModel):
    """Decoder-only backbone (`AutoModel`). Returns token states aligned with `input_ids`.
    Species inputs: `species_ids` (rows of the built-in table; -1 = infer), `species_embeds` (your own
    vectors), or nothing (every row inferred by the attached species encoder). The encoder is held by
    reference: not a submodule, so it is never in `parameters()`, `state_dict()` or a saved checkpoint and
    never trained, but `.to()/.half()/.cuda()/.float()` reach it.
    """

    def __init__(self, config: MicroGlotConfig, species_encoder=None):
        super().__init__(config)
        self.gradient_checkpointing = False
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            [GenomicDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)]
        )
        self.norm = nn.RMSNorm(normalized_shape=config.hidden_size)
        self.embed_species = (
            MicroGlotSpeciesTable(config.num_species, config.species_emb_dim) if config.num_species else None
        )
        self.species_token_proj = (
            nn.Linear(config.species_emb_dim, config.hidden_size, bias=False)
            if config.prepend_species_token
            else None
        )
        self.__dict__["_species_encoder"] = None
        self.__dict__["_species_encoder_follow"] = True
        self.post_init()
        if species_encoder is not None:
            self.set_species_encoder(species_encoder)

    def _microglot_backbone(self):
        return self

    @property
    def species_encoder(self):
        """The attached species encoder, or None."""
        return self.__dict__.get("_species_encoder")

    def set_species_encoder(self, encoder, match_backbone: bool = True):
        """Attach (or with None detach) the encoder that infers species for rows without one.
        The encoder is put in eval mode with `requires_grad_(False)`.
        match_backbone=True (default): move/cast it to the backbone's device and dtype and give it exact rotary
        frequencies in the backbone buffer's dtype, so the pair computes exactly what the training code's
        combined model computed; later .to()/.half()/.cuda() calls on this model reach it too. match_backbone=False: leave it
        where it is (e.g. another GPU); inputs are moved to it and its output is cast back.
        """
        if encoder is not None:
            if self.config.species_emb_dim is None:
                raise ValueError(
                    "This checkpoint takes no species information (MicroGlot-plain, or the species encoder's own "
                    "backbone), so it has no use for a species encoder."
                )
            if not hasattr(encoder, "projection") or getattr(
                encoder.config, "species_encoder_embedding_dim", None
            ) != self.config.species_emb_dim:
                raise ValueError(
                    f"Expected the Species-encoder (athanzli/MicroGlot, subfolder='species_encoder') or another "
                    f"MicroGlotSpeciesEncoder "
                    f"with {self.config.species_emb_dim}-d output."
                )
            encoder.eval().requires_grad_(False)
            if match_backbone:
                ref = self.embed_tokens.weight
                encoder.to(device=ref.device, dtype=ref.dtype)
                ref_rope = _rope_modules(self)
                if ref_rope:
                    # The cast above also rounds the encoder's rotary frequencies, which flash-attn reads as
                    # they are when they are float32. Rebuild them exactly and give them the backbone buffer's
                    # dtype, as the encoder inside the training code's combined model had.
                    ref_inv = ref_rope[0].inv_freq
                    for m in _rope_modules(encoder):
                        inv_freq = 1.0 / (m.base ** (torch.arange(0, m.dim, 2, dtype=torch.float32) / m.dim))
                        m.inv_freq = inv_freq.to(device=ref_inv.device, dtype=ref_inv.dtype)
                        m._cos_cached, m._seq_len_cached = None, 0
        self.__dict__["_species_encoder"] = encoder
        self.__dict__["_species_encoder_follow"] = bool(match_backbone)

    def get_species_encoder(self):
        """The attached species encoder, or None."""
        return self.species_encoder

    def load_species_encoder(self, pretrained_model_name_or_path=None, *, subfolder=None, revision=None,
                             match_backbone=True, **kwargs):
        """Load the Species-encoder and attach it (see `set_species_encoder`); returns the encoder.
        Default: the folder `config.species_encoder_subfolder` ("species_encoder") of the repository or local
        folder this model was loaded from, at the same revision (or `config.species_encoder_repo` at
        `config.species_encoder_revision` if set), in the model's dtype. Other keyword arguments go to
        `from_pretrained` (e.g. `cache_dir`, `token`). It uses the encoder class in this code file.
        """
        if torch.is_inference_mode_enabled():
            raise RuntimeError(
                "Load the species encoder outside torch.inference_mode(): weights created inside it are "
                "inference tensors and cannot be moved or cast afterwards."
            )
        if self.config.species_emb_dim is None:
            raise ValueError("MicroGlot-plain takes no species information, so it has no species encoder.")
        path = pretrained_model_name_or_path
        if path is None:
            if self.config.species_encoder_repo is not None:
                path, revision = self.config.species_encoder_repo, revision or self.config.species_encoder_revision
            else:
                source, source_revision = self.__dict__.get("_microglot_source") or (self.config._name_or_path, None)
                path, revision = source, revision or source_revision
            subfolder = subfolder if subfolder is not None else self.config.species_encoder_subfolder
            if not path:
                raise ValueError("Pass the encoder's location, e.g. "
                                 "model.load_species_encoder('athanzli/MicroGlot', subfolder='species_encoder').")
        if subfolder:
            kwargs["subfolder"] = subfolder
        if revision is not None:
            kwargs["revision"] = revision
        if "torch_dtype" not in kwargs and "dtype" not in kwargs:
            kwargs[_DTYPE_KWARG] = self.embed_tokens.weight.dtype
        try:
            encoder = MicroGlotSpeciesEncoder.from_pretrained(path, **kwargs)
        except OSError as e:
            raise OSError(
                f"No Species-encoder found at {path!r}" + (f" (subfolder {subfolder!r})" if subfolder else "")
                + ". Load it from the Hub: model.load_species_encoder('athanzli/MicroGlot', subfolder='species_encoder')."
            ) from e
        self.set_species_encoder(encoder, match_backbone=match_backbone)
        return encoder

    def _apply(self, fn, recurse=True):
        super()._apply(fn, recurse)
        encoder = self.__dict__.get("_species_encoder")
        if encoder is not None and self.__dict__.get("_species_encoder_follow", True):
            encoder._apply(fn, recurse)
        return self

    def get_input_embeddings(self) -> nn.Embedding:
        return self.embed_tokens

    def set_input_embeddings(self, value: nn.Embedding):
        self.embed_tokens = value

    # -- species resolution -------------------------------------------------------------------------
    @torch.no_grad()
    def get_species_embeds(self, input_ids=None, attention_mask=None, species_ids=None, species_embeds=None,
                           normalize_species_embeds: bool = True):
        """The [batch, 32] unit vectors this model would condition on, in its dtype, without running the
        backbone (prior rows from the table, custom vectors, inferred rows from the attached encoder).
        Takes the same species arguments as `forward`; `input_ids` (and `attention_mask`) are only needed for
        rows the encoder infers. Use it to inspect inferred species or to precompute a `species_embeds`
        column (pass it back with `normalize_species_embeds=False` to reuse the exact vectors).
        """
        if input_ids is None:
            ref = species_ids if species_ids is not None else species_embeds
            if ref is None:
                raise ValueError("Pass input_ids, species_ids or species_embeds.")
            ref = torch.as_tensor(ref)
            n = (ref.shape[0] if ref.dim() > 0 else 1) if species_ids is not None else (
                ref.shape[0] if ref.dim() > 1 else 1)
            if species_ids is not None and bool((ref < 0).any()):
                raise ValueError("Rows to infer (species_ids == -1) need input_ids.")
            input_ids = torch.full((n, 1), self.config.bos_token_id, device=self.embed_tokens.weight.device)
        return self._species_vectors(input_ids, attention_mask, species_ids, species_embeds,
                                     normalize_species_embeds)

    def _species_vectors(self, input_ids, attention_mask, species_ids, species_embeds, normalize=True):
        """Unit-norm species vectors [batch, species_emb_dim] in the model dtype, or None (plain)."""
        dim = self.config.species_emb_dim
        if dim is None:
            if species_ids is not None or species_embeds is not None:
                raise ValueError(
                    "This checkpoint (MicroGlot-plain) takes no species information: drop the species "
                    "inputs, or load athanzli/MicroGlot to condition on a species."
                )
            return None
        if species_ids is not None and species_embeds is not None:
            raise ValueError("You cannot specify both species_ids and species_embeds at the same time.")
        batch, device, dtype = input_ids.shape[0], input_ids.device, self.embed_tokens.weight.dtype

        if species_embeds is not None:  # the training code's `species_emb=` path (plus a no-op device move)
            species_embeds = torch.as_tensor(species_embeds)
            if species_embeds.dim() == 1:
                species_embeds = species_embeds.unsqueeze(0)
            if species_embeds.dim() != 2 or species_embeds.shape[-1] != dim:
                raise ValueError(
                    f"species_embeds must be [{dim}] or [batch, {dim}], got {tuple(species_embeds.shape)}."
                )
            if species_embeds.shape[0] == 1 and batch > 1:
                species_embeds = species_embeds.expand(batch, -1)
            elif species_embeds.shape[0] != batch:
                raise ValueError(f"species_embeds has {species_embeds.shape[0]} rows for {batch} sequences.")
            species_embeds = species_embeds.to(device=device)
            if normalize:
                return F.normalize(species_embeds.to(torch.float32), dim=-1).to(dtype)
            norms = species_embeds.float().norm(dim=-1)
            if bool(((norms - 1).abs() > 1e-2).any()):
                raise ValueError(
                    "normalize_species_embeds=False expects unit vectors that are used as they are (e.g. "
                    "out.species_embeds or encoder(...).species_embeds)."
                )
            return species_embeds.to(dtype)

        vectors, unknown = None, None
        if species_ids is not None:
            species_ids = torch.as_tensor(species_ids, device=device)
            if species_ids.dim() == 0:
                species_ids = species_ids.expand(batch)
            if species_ids.shape != (batch,) or species_ids.dtype not in (torch.int64, torch.int32):
                raise ValueError(f"species_ids must be an integer tensor of shape [{batch}].")
            species_ids = species_ids.long()
            if self.embed_species is None:
                raise ValueError("This checkpoint has no species table; pass species_embeds instead.")
            n = self.embed_species.weight.shape[0]
            if bool(((species_ids < -1) | (species_ids >= n)).any()):
                raise ValueError(f"species_ids must lie in [-1, {n}); -1 marks an unknown species.")
            prior = self.embed_species(species_ids.clamp(min=0))
            vectors = F.normalize(prior.to(torch.float32), dim=-1).to(dtype)
            unknown = species_ids < 0
            if not bool(unknown.any()):
                return vectors

        encoder = self.species_encoder
        if encoder is None:
            raise ValueError(
                "No species given for some sequences and no species encoder is attached. Pass species= to "
                "the tokenizer (or species_ids / species_embeds), or attach the encoder with "
                "model.load_species_encoder(). For no species information, use MicroGlot-plain "
                "(subfolder='plain')."
            )
        rows = slice(None) if unknown is None else unknown
        enc_ids = input_ids[rows].to(encoder.device)
        enc_mask = None if attention_mask is None else attention_mask[rows].to(encoder.device)
        with torch.no_grad():
            inferred = encoder(input_ids=enc_ids, attention_mask=enc_mask).species_embeds
        inferred = inferred.to(device=device, dtype=dtype)
        if unknown is None:
            return inferred
        vectors = vectors.clone()
        vectors[unknown] = inferred
        return vectors

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        species_embeds: Optional[torch.FloatTensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_router_logits: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        species_ids: Optional[torch.LongTensor] = None,
        normalize_species_embeds: bool = True,
        species_emb: Optional[torch.FloatTensor] = None,
    ) -> Union[Tuple, MicroGlotModelOutput]:
        """Token states for right-padded `input_ids` `[batch, length]` (each row starts with [BOS]).
        Species (MicroGlot only; at most one of the two):
            species_ids (`LongTensor` `[batch]`, or 0-d for the whole batch): rows of the species table
                (`tokenizer(dna, species=...)` returns them); -1 = infer this row with the attached encoder.
            species_embeds (`FloatTensor` `[batch, 32]` or `[32]`): your own vectors, any float dtype and norm;
                normalised in float32 and cast to the model dtype. With `normalize_species_embeds=False` they
                must already be unit vectors and are used as they are (exact replay of `out.species_embeds`).
                `species_emb` (the training code's name) is also accepted.
            Neither: every row is inferred by the attached species encoder (ValueError if none is attached).
        output_hidden_states: also return the 24 hidden states. output_router_logits: compute `moe_loss`
            (default: on). return_dict=False returns the tuple `(last_hidden_state, hidden_states, moe_loss)`.
        The positional order is the training code's (the third argument is the species vector).
        """
        if not (input_ids.is_cuda and self.embed_tokens.weight.is_cuda):
            raise RuntimeError(
                "MicroGlot runs on an NVIDIA GPU (flash-attn's rotary kernel is CUDA-only): move the "
                "model and the inputs to a CUDA device, e.g. model.to('cuda')."
            )
        if species_emb is not None:
            if species_embeds is not None:
                raise ValueError("Pass species_embeds or species_emb (the same argument), not both.")
            species_embeds = species_emb
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None
            else self.config.output_hidden_states if hasattr(self.config, 'output_hidden_states') else False
        )
        output_router_logits = (
            output_router_logits if output_router_logits is not None
            else self.config.use_moe
        )
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        if attention_mask is not None and not bool(attention_mask[:, 0].all()):
            raise ValueError(
                "MicroGlot needs right padding (every sequence starts with [BOS] at position 0); "
                "tokenize with padding_side='right'."
            )
        if attention_mask is None and self.config.pad_token_id is not None and bool(
            (input_ids == self.config.pad_token_id).any()
        ):
            raise ValueError("input_ids contain [PAD] but no attention_mask was given; pass the tokenizer's mask.")
        if (
            self.species_token_proj is not None
            and (species_ids is not None or species_embeds is not None or self.species_encoder is not None)
            and bool((input_ids[:, 0] != self.config.bos_token_id).any())
        ):
            raise ValueError(
                "With a species, every sequence must start with [BOS] (the species token goes right after it); "
                "tokenize with add_special_tokens=True."
            )
        seq_len = input_ids.shape[1]
        max_trained = getattr(self.config, "max_trained_length", 8192)
        if seq_len > max_trained:
            warnings.warn(
                f"Input sequence length ({seq_len} tokens) exceeds the maximum length seen during "
                f"pretraining ({max_trained} tokens); split long sequences into windows as shown under "
                "\"Long sequences\" at https://huggingface.co/athanzli/MicroGlot.",
                UserWarning,
                stacklevel=3,
            )
        if torch.is_grad_enabled():
            # flash-attn rebuilds a cos/sin cache made under torch.inference_mode() only in training mode;
            # in eval mode it cannot be saved for backward. Drop it (it is rebuilt with the same values).
            for layer in self.layers:
                rope = layer.self_attn.rotary_emb
                if rope._cos_cached is not None and rope._cos_cached.is_inference():
                    rope._cos_cached = None

        species_emb = self._species_vectors(input_ids, attention_mask, species_ids, species_embeds,
                                            normalize_species_embeds)

        # ---- from here on: the forward used in training and benchmarking, unchanged ---------------
        hidden_states = self.embed_tokens(input_ids).to(self.embed_tokens.weight.dtype)

        if self.species_token_proj is not None and species_emb is not None:
            batch_size = hidden_states.size(0)
            device = hidden_states.device

            species_token = self.species_token_proj(species_emb).unsqueeze(1).to(hidden_states.dtype)

            bos_embedding = hidden_states[:, :1, :]
            rest_embeddings = hidden_states[:, 1:, :]

            hidden_states = torch.cat([bos_embedding, species_token, rest_embeddings], dim=1)

            if attention_mask is not None:
                species_mask = torch.ones(batch_size, 1, dtype=attention_mask.dtype, device=device)
                bos_mask = attention_mask[:, :1]
                rest_mask = attention_mask[:, 1:]
                attention_mask = torch.cat([bos_mask, species_mask, rest_mask], dim=1)

        all_hidden_states = () if output_hidden_states else None
        total_moe_loss = torch.tensor(0.0, device=hidden_states.device, dtype=torch.float32) if output_router_logits else None

        for layer in self.layers:
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if self.gradient_checkpointing and self.training:
                layer_outputs = checkpoint(
                    layer,
                    hidden_states,
                    attention_mask,
                    species_emb,
                    use_reentrant=False,
                )
            else:
                layer_outputs = layer(hidden_states, attention_mask=attention_mask, species_emb=species_emb)

            hidden_states, moe_loss = layer_outputs

            if output_router_logits and moe_loss is not None:
                total_moe_loss = total_moe_loss + moe_loss

        hidden_states = self.norm(hidden_states)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        species_was_prepended = (
            self.species_token_proj is not None and species_emb is not None
        )
        if species_was_prepended:
            hidden_states = torch.cat(
                [hidden_states[:, :1, :], hidden_states[:, 2:, :]], dim=1
            )
            if all_hidden_states is not None:
                all_hidden_states = tuple(
                    torch.cat([hs[:, :1, :], hs[:, 2:, :]], dim=1)
                    for hs in all_hidden_states
                )

        if not return_dict:
            return (hidden_states, all_hidden_states, total_moe_loss)

        return MicroGlotModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            moe_loss=total_moe_loss,
            species_embeds=species_emb,
        )


class MicroGlotForCausalLM(MicroGlotPreTrainedModel):
    """MicroGlot with its next-token head (`AutoModelForCausalLM`, the released checkpoints' class).
    `lm_head` is tied to the input embeddings. Sequence generation is not supported.
    """

    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}  # dict: required by transformers 5

    def __init__(self, config: MicroGlotConfig, species_encoder=None):
        super().__init__(config)
        self.model = MicroGlotModel(config, species_encoder=species_encoder)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.model.embed_tokens

    def set_input_embeddings(self, value: nn.Embedding):
        self.model.embed_tokens = value

    def get_output_embeddings(self) -> nn.Linear:
        return self.lm_head

    def set_output_embeddings(self, new_embeddings: nn.Linear):
        self.lm_head = new_embeddings

    def tie_weights(self, missing_keys=None, recompute_mapping=True, **kwargs):
        """Always tie lm_head to the input embeddings, as in training (also for configs without
        tie_word_embeddings); under transformers 5 also clear the tied key from the load report."""
        self.lm_head.weight = self.model.embed_tokens.weight
        if missing_keys is not None:
            missing_keys.discard("lm_head.weight")

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        species_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_router_logits: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        skip_lm_head: bool = False,
        species_ids: Optional[torch.LongTensor] = None,
        normalize_species_embeds: bool = True,
        species_emb: Optional[torch.FloatTensor] = None,
    ) -> Union[Tuple, MicroGlotCausalLMOutput]:
        """Next-token logits `[batch, length, vocab]`; the species arguments are those of `MicroGlotModel.forward`.
        labels (`LongTensor` `[batch, length]`, -100 = ignored): `loss` = mean next-token cross-entropy. With a
            species token the first prediction (made at [BOS], before the model has seen the species) is excluded,
            as in pretraining.
        output_router_logits: also compute the router load-balancing loss (`moe_loss`) and add it to `loss`, as in
            pretraining (default: off). Pass True when training this model, so that the experts stay balanced.
        skip_lm_head: return only the final hidden state (in `hidden_states`).
        return_dict=False returns the tuple `([loss,] logits, hidden_states, moe_loss)`.
        """
        return_dict = return_dict if return_dict is not None else self.config.return_dict
        # Off by default, as in transformers' mixture-of-experts language models: `loss` is the next-token
        # cross-entropy only. The hidden states and logits do not depend on this flag.
        output_router_logits = bool(output_router_logits)
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            species_ids=species_ids,
            species_embeds=species_embeds,
            species_emb=species_emb,
            normalize_species_embeds=normalize_species_embeds,
            output_hidden_states=output_hidden_states,
            output_router_logits=output_router_logits,
            return_dict=True,
        )
        hidden_states = outputs.last_hidden_state

        if skip_lm_head:
            return MicroGlotCausalLMOutput(loss=None, logits=None, hidden_states=hidden_states)

        logits = self.lm_head(hidden_states)
        moe_aux_loss = outputs.moe_loss

        loss = None
        if labels is not None:
            species_was_prepended = self.config.prepend_species_token and outputs.species_embeds is not None
            if species_was_prepended:
                shift_logits = logits[..., 1:-1, :].contiguous()
                shift_labels = labels[..., 2:].contiguous()
            else:
                shift_logits = logits[..., :-1, :].contiguous()
                shift_labels = labels[..., 1:].contiguous()

            loss_fn = nn.CrossEntropyLoss(ignore_index=-100)
            loss = loss_fn(
                shift_logits.view(-1, self.config.vocab_size),
                shift_labels.view(-1)
            ).float()

            if self.config.use_moe and moe_aux_loss is not None:
                loss = loss + moe_aux_loss

        if not return_dict:
            output = (logits, outputs.hidden_states, outputs.moe_loss)
            return (loss,) + output if loss is not None else output

        return MicroGlotCausalLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            species_embeds=outputs.species_embeds,
            moe_loss=outputs.moe_loss,
        )



class MicroGlotSpeciesEncoder(MicroGlotPreTrainedModel):
    """Predicts a species embedding from DNA (folder species_encoder/ of athanzli/MicroGlot, `AutoModel`).
    A MicroGlot backbone without species input, pooled at the last real token (needs right padding), then a
    Linear -> GELU -> Linear projection to 32 dimensions (`embeddings`), L2-normalised in the model dtype
    (`species_embeds`, the vector MicroGlot is conditioned on).
    """

    config_class = MicroGlotSpeciesEncoderConfig

    def __init__(self, config: MicroGlotSpeciesEncoderConfig):
        super().__init__(config)
        if config.species_emb_dim is not None or config.prepend_species_token:
            raise ValueError("The species encoder's backbone takes no species input (species_emb_dim=None).")
        self.embedding_dim = config.species_encoder_embedding_dim
        self.pooling_strategy = config.species_encoder_pooling_strategy
        self.model = MicroGlotModel(config)
        self.projection_dropout = nn.Dropout(config.classifier_dropout_prob)
        self.projection = MicroGlotHead(config.hidden_size, config.hidden_size, config.species_encoder_embedding_dim)
        self.post_init()

    def load_species_encoder(self, *args, **kwargs):
        """Not available: this is the species encoder itself."""
        raise TypeError("This is the species encoder; it takes no species encoder of its own.")

    def get_input_embeddings(self) -> nn.Embedding:
        return self.model.embed_tokens

    def set_input_embeddings(self, value: nn.Embedding):
        self.model.embed_tokens = value

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, MicroGlotSpeciesEncoderOutput]:
        """Species vectors for right-padded `input_ids` `[batch, length]` (one per sequence).
        return_dict=False returns `(species_embeds, embeddings, pooler_output, last_hidden_state,
        hidden_states)`.
        """
        return_dict = return_dict if return_dict is not None else self.config.return_dict
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=output_hidden_states,
            output_router_logits=self.config.use_moe,
            return_dict=True,
        )
        hidden_states = outputs.last_hidden_state

        if self.pooling_strategy == "eos":
            if attention_mask is not None:
                seq_lengths = attention_mask.sum(dim=1) - 1
                seq_lengths = seq_lengths.clamp(min=0)
                batch_indices = torch.arange(hidden_states.size(0), device=hidden_states.device)
                pooled = hidden_states[batch_indices, seq_lengths]
            else:
                pooled = hidden_states[:, -1, :]
        else:
            if attention_mask is not None:
                mask = attention_mask.unsqueeze(-1).float()
                pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
            else:
                pooled = hidden_states.mean(dim=1)

        pooled = self.projection_dropout(pooled)
        raw = self.projection(pooled)
        species_embeds = F.normalize(raw, dim=-1)

        if not return_dict:
            return (species_embeds, raw, pooled, hidden_states, outputs.hidden_states)
        return MicroGlotSpeciesEncoderOutput(
            species_embeds=species_embeds,
            embeddings=raw,
            pooler_output=pooled,
            last_hidden_state=hidden_states,
            hidden_states=outputs.hidden_states,
            moe_loss=outputs.moe_loss,
        )


# Training-code names (configs and auto_maps written by the training code keep resolving).
GenomicPreTrainedModel = MicroGlotPreTrainedModel
GenomicModel = MicroGlotModel
GenomicLMForCausalLM = MicroGlotForCausalLM
GenomicSpeciesEmbeddingEncoder = MicroGlotSpeciesEncoder
GenomicModelOutput = MicroGlotModelOutput

# save_pretrained() then copies this file and writes local auto_map entries, also for Hub-loaded models.
MicroGlotConfig.register_for_auto_class("AutoConfig")
MicroGlotSpeciesEncoderConfig.register_for_auto_class("AutoConfig")
MicroGlotModel.register_for_auto_class("AutoModel")
MicroGlotForCausalLM.register_for_auto_class("AutoModelForCausalLM")
MicroGlotSpeciesEncoder.register_for_auto_class("AutoModel")