import torch.nn as nn
import torch
import torch.nn.functional as F
from torch.distributions import Normal, Distribution
import numpy as np
from typing import (
    Dict,
    Optional,
    Union,
    Tuple,
    Any,
    Callable,
    List,
)
from ..cvae.module import CVAE
from ..cvae.base import LossOutput
from ..cvae.utils import init_library_size
from . import prepare_cell_text_llm
from transformers import (
    PreTrainedModel,
    PretrainedConfig,
    AutoConfig,
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    AutoModelForCausalLM,
    CONFIG_MAPPING,
    PreTrainedTokenizer,
    PreTrainedTokenizerFast,
    Blip2QFormerConfig,
    Blip2QFormerModel,
)
from transformers.generation.streamers import BaseStreamer
from transformers.modeling_outputs import (
    Seq2SeqLMOutput,
    CausalLMOutput,
)

ACTIVATION_FAMILY = {
    "gelu": F.gelu,
    "relu": F.relu,
    "tanh": F.tanh,
    "sigmoid": F.sigmoid,
    "selu": F.selu,
    "silu": F.silu,
}


def get_torch_activation(name: str) -> Callable[[torch.Tensor], torch.Tensor]:
    """Get the torch activation function by name (e.g, ``get_torch_activation("gelu")``)."""
    func = ACTIVATION_FAMILY.get(name, None)
    if callable(func):
        return func

    raise ValueError(f"Unknown activation function: {name}.")


class ResidualBlock(nn.Module):
    """
    A PyTorch module implementing a residual block. It is composed of a linear layer, an activation
    function, and a dropout layer.

    Parameters
    ----------
    n_inputs: int
        The number of input features (input dimension) for the linear layer.
    n_outputs: int
        The number of output features (output dimension) for the linear layer. If a residual
        connection is enabled, ``n_inputs`` must be equal to ``n_outputs``.
    activation_func: str, default "gelu"
        The name of the activation function to apply after the linear transformation. It should
        be one of "gelu", "relu", "tanh", "sigmoid", "selu", or "silu".
    dropout_rate: float, default 0.1
        The probability of an element to be zeroed in the dropout layer applied after the
        activation function.
    residual: bool, default True
        If True, the input will be added to the output (residual connection), assuming the
        input and output dimensions are the same. If False, no residual connection is used.
    """

    def __init__(
        self,
        n_inputs: int,
        n_outputs: int,
        activation_func: str = "gelu",
        dropout_rate: float = 0.1,
        residual: bool = True,
    ) -> "ResidualBlock":
        super(ResidualBlock, self).__init__()
        if residual:
            assert n_inputs == n_outputs, "Residual connection requires the same input and output dimensions"

        self.residual = residual
        self.linear = nn.Linear(n_inputs, n_outputs)
        self.activation = get_torch_activation(activation_func)
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Given the input ``x``, get the corresponding output of the residual block."""
        if self.residual:
            return x + self.dropout(self.activation(self.linear(x)))
        else:
            return self.dropout(self.activation(self.linear(x)))


class SCQFormer(Blip2QFormerModel):
    """
    A variant of the BLIP2 Q-Former model designed for encoding cell counts.

    The SCQFormer takes in a set of input cell counts, passes them through a series of residual blocks
    (cell encoders) to generate key-value pairs. Q-Former is used to pool the key-value pairs via query
    tokens to generate the final output.

    Parameters
    ----------
    count_dim: int
        The dimensionality of the input cell counts. This corresponds to the number of input features.
    num_query_tokens: int
        The number of query tokens generated by the model. These tokens interact with the hidden states
        from residual blocks through cross-attention layers and interact with each other through
        self-attention layers.
    num_key_value_tokens: int
        The number of key-value tokens used in the cross-attention mechanism.
    config: transformers.Blip2QFormerConfig
        The configuration object that contains hyperparameters for the ``transformers.Blip2QFormer`` model.
    num_hidden_layers: int, default 2
        The number of ``ResidualBlock`` in the cell encoder.
    """

    def __init__(
        self,
        count_dim: int,
        num_query_tokens: int,
        num_key_value_tokens: int,
        config: Blip2QFormerConfig,
        num_hidden_layers: int = 2,
    ) -> "SCQFormer":
        # initialize the Blip2QFormerModel and model weights
        super().__init__(config)

        hidden_dropout_prob = config.hidden_dropout_prob
        hidden_act = config.hidden_act
        input_dim, output_dim = count_dim, num_key_value_tokens * config.hidden_size

        self.cell_encoder = []
        for i in range(num_hidden_layers):
            if i == 0:
                self.cell_encoder.append(
                    ResidualBlock(
                        input_dim,
                        output_dim,
                        activation_func=hidden_act,
                        dropout_rate=hidden_dropout_prob,
                        residual=False,
                    ))
            elif i == num_hidden_layers - 1:
                self.cell_encoder.append(
                    ResidualBlock(
                        output_dim,
                        output_dim,
                        activation_func=hidden_act,
                        dropout_rate=0.0,
                        residual=True,
                    ))
            else:
                self.cell_encoder.append(
                    ResidualBlock(
                        output_dim,
                        output_dim,
                        activation_func=hidden_act,
                        dropout_rate=hidden_dropout_prob,
                        residual=True,
                    ))
        self.cell_encoder = nn.Sequential(*self.cell_encoder)
        self.query_tokens = nn.Parameter(
            torch.zeros(1, num_query_tokens, config.hidden_size))

        # initialize the weights of the model again
        self.init_weights()
        nn.init.normal_(self.query_tokens, std=config.initializer_range)

    def forward(self, input_counts: torch.Tensor) -> torch.Tensor:
        """Encode the input cell counts and generate the final output."""
        encoder_hidden_states = self.cell_encoder(input_counts)
        batch_size = encoder_hidden_states.size(0)
        encoder_hidden_states = torch.reshape(
            encoder_hidden_states, (batch_size, -1, self.config.hidden_size))

        query_tokens = self.query_tokens.expand(batch_size, -1, -1)
        outputs = super().forward(
            query_tokens, encoder_hidden_states=encoder_hidden_states)[0]

        return outputs


class Generator(nn.Module):
    """
    A Generator model for conditional generation of gene expression profiles based on
    conditional embeddings and latent variables. This model is built upon a Conditional
    Variational Autoencoder (CVAE), which integrates continuous and categorical covariates
    to generate synthetic data.

    Parameters
    ----------
    n_input: int
        The dimensionality of the input data (i.e., the number of genes).
    condition_dim: int, default 4096
        The dimensionality of the condition embedding input, which represents the condition
        that controls the generative process (e.g., cell type embeddings).
    condition_input_dim: int, default 128
        The dimensionality of the mapped condition after processing by the mapping layer.
    **kwargs: dict, optional
        Additional arguments passed to the base model, typically the ``scvi.module.CVAE`` model.
    """

    def __init__(self,
                 n_input: int,
                 condition_dim: int = 4096,
                 condition_input_dim: int = 128,
                 **kwargs) -> "Generator":

        super(Generator, self).__init__()
        self.mapping_layer = nn.Linear(condition_dim, condition_input_dim)
        kwargs["n_continuous_cov"] = condition_input_dim
        self.base_model = CVAE(n_input, **kwargs)

    def forward(
        self,
        condition_embeds: torch.Tensor,
        gene_counts: torch.Tensor,
        batch_ids: Optional[torch.Tensor] = None,
        cat_covs: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        loss_kwargs: Optional[Dict[str, Any]] = None,
        compute_loss: bool = True,
    ) -> Union[Tuple[Dict[str, torch.Tensor | Distribution | None], Dict[
            str, Distribution | None]], LossOutput]:
        """
        Perform a forward pass through the generator.

        Parameters
        ----------
        condition_embeds: torch.Tensor
            The condition embeddings control the generation process.
        gene_counts: torch.Tensor
            A tensor representing observed gene counts, which serves as input for training.
        batch_ids: torch.Tensor, optional, default None
            The batch identifiers for each sample in the input data. If None, all samples are assigned
            to batch 0. Another way to consider batch information is to use  ``condition_embeds`` to encode
            batch information.
        cat_covs: torch.Tensor, optional, default None
            Categorical covariates related to the sample.
        labels: torch.Tensor, optional, default None
            Class labels associated with the samples.
        loss_kwargs: dict, optional, default None
            Additional keyword arguments passed to the loss function ``scvi.module.CVAE.loss`` (e.g., the weight of
            the KL divergence term between latent variables and the prior distribution).
        compute_loss: bool, default True
            Whether to compute and return the loss during the forward pass. If False, the outputs of
            inference process and generative process are returned. If True, the loss is returned.

        Returns
        -------
        outputs: Union[Tuple[Dict[str, torch.Tensor | Distribution | None], Dict[str, Distribution | None]], LossOutput]
            The outputs of the model. If ``compute_loss`` is True, the loss is returned. Otherwise, the outputs
            of the inference process and generative process are returned.
        """
        if condition_embeds.ndim > 2:
            condition_embeds = condition_embeds.squeeze(1)
        if batch_ids is None:
            batch_ids = torch.zeros(
                (gene_counts.shape[0], 1)).to(dtype=torch.long,
                                              device=gene_counts.device)
        inputs = {
            "gene_counts": gene_counts,
            "batch_ids": batch_ids,
            "cat_covs": cat_covs,
            "y": labels,
        }
        cont_covs = self.mapping_layer(condition_embeds)
        inputs["cont_covs"] = cont_covs

        outputs = self.base_model(inputs,
                                  loss_kwargs=loss_kwargs,
                                  compute_loss=compute_loss)
        if compute_loss:
            return outputs[-1]

        return outputs

    @torch.inference_mode()
    def generate(
        self,
        condition_embeds: torch.Tensor,
        batch_ids: Optional[torch.Tensor] = None,
        library_size: Optional[torch.Tensor] = None,
        batch_size: int = 128,
        to_numpy: bool = True,
    ) -> torch.Tensor | np.ndarray:
        """
        Generate synthetic gene expression profiles from the latent space.

        Parameters
        ----------
        condition_embeds: torch.Tensor
            The condition embeddings used to control the generation process.
        batch_ids: torch.Tensor, optional, default None
            The batch identifiers for each sample in the input data. If None, all samples are assigned
            to batch 0. Another way to consider batch information is to use  ``condition_embeds`` to encode
            batch information.
        library_size: torch.Tensor, optional, default None
            The library size vector predefined for each sample. It is used to control the library size of each
            generated cell. If None, a prior distribution over library size will be used to sample the library
            size for each cell. Note that, in general, the prior distribution is set by
            ``scvi.utils.init_library_size``.
        batch_size: int, default 128
            The batch size for generating data. Control the number of samples processed at once.
        to_numpy: bool, default True
            Whether to return the output as a NumPy array. If False, a PyTorch tensor is returned.

        Returns
        -------
        outputs: torch.Tensor | np.ndarray
            The generated gene expression data. The shape of the output depends on the input shape
            and batch size. If ``to_numpy`` is True, the output is a NumPy array, otherwise it is a
            PyTorch tensor.
        """
        if condition_embeds.ndim > 2:
            condition_embeds = condition_embeds.squeeze(1)
        cont_covs = self.mapping_layer(condition_embeds)
        n_samples = condition_embeds.shape[0]
        if batch_ids is None:
            batch_ids = torch.zeros(
                (n_samples, 1)).to(dtype=torch.long,
                                   device=condition_embeds.device)
        device = next(self.base_model.parameters()).device

        outputs = []
        for i in range(0, n_samples, batch_size):
            inputs = {}
            cont_covs_batch = cont_covs[
                i:i + batch_size] if cont_covs is not None else None
            batch_ids_batch = batch_ids[i:i + batch_size]
            # the prior distribution of latent vectors
            p_z = Normal(
                torch.zeros(
                    (batch_ids_batch.size(0), self.base_model.n_latent)),
                torch.ones((batch_ids_batch.size(0), self.base_model.n_latent)))
            inputs = {
                "z": p_z.sample(),
                "batch_ids": batch_ids_batch,
                "cont_covs": cont_covs_batch,
            }
            if library_size is not None:
                inputs["library"] = library_size[i:i + batch_size]
            else:
                p_library = self.base_model.get_prior_library_distribution(
                    batch_ids_batch)
                assert p_library is not None, "The model does not have a prior distribution for library size, please provide it"
                inputs["library"] = p_library.sample()
            inputs = {k: v.to(device) for k, v in inputs.items()}
            p_gene = self.base_model.generative(**inputs)["p_gene"]
            fake_samples = p_gene.sample()
            if fake_samples.device.type == "cuda":
                fake_samples = fake_samples.cpu()
            outputs.append(fake_samples)
        outputs = torch.concat(outputs,
                               dim=0) if len(outputs) > 1 else outputs[0]

        return outputs if not to_numpy else outputs.numpy()


class CellTextLLM(nn.Module):
    """
    A model class that integrates cell encoders, language models, and feature decoders.

    The model can merge encoded cell features with tokenized input text, generating both text and cell data.

    Parameters
    ----------
    base_model: transformers.PreTrainedModel or torch.nn.Module
        The underlying language model that will be used for text generation.
    tokenizer: transformers.PreTrainedTokenizer or transformers.PreTrainedTokenizerFast
        The tokenizer to tokenize the input text sequences.
    feature_encoder: transformers.PreTrainedModel or torch.nn.Module or None, optional, default None
        A model used to encode input cell features, such as cell counts.
    feature_decoder: transformers.PreTrainedModel or torhc.nn.Module or None, optional, default None
        A model used to decode cell features from the hidden states of the language model.
    normalize_total: bool, default True
        Whether to normalize the total input cell counts by dividing by the total sum of the counts.
    log_variational: bool, default True
        Whether to apply a log transformation (log1p) to the input cell counts before encoding. If
        both normalization and log transformation are enabled, the log transformation is applied after
        normalization.
    """

    def __init__(
        self,
        base_model: PreTrainedModel | nn.Module,
        tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
        feature_encoder: PreTrainedModel | nn.Module | None = None,
        feature_decoder: PreTrainedModel | nn.Module | None = None,
        normalize_total: bool = True,
        log_variational: bool = True,
    ) -> "CellTextLLM":
        super(CellTextLLM, self).__init__()
        self.base_model = base_model
        self.tokenizer = tokenizer
        if feature_encoder is not None:
            self.feature_encoder = feature_encoder
        else:
            self.register_buffer("feature_encoder", None)
        if feature_decoder is not None:
            self.feature_decoder = feature_decoder
        else:
            self.register_buffer("feature_decoder", None)

        assert hasattr(self.base_model, "config"),  \
            "Base model must have a config attribute. " + \
            "Please use prepare_cell_text_llm to prepare the model."
        assert hasattr(self.base_model.config, "eos_token_id"), \
            "Base model config must have eos_token_id attribute. " + \
            "Please use prepare_cell_text_llm to prepare the model."

        if self.feature_encoder is not None or self.feature_decoder is not None:
            assert hasattr(self.base_model.config, "special_tokens_index_dict"), \
                "Base model config must have special_tokens_index_dict attribute. " + \
                "Please use prepare_cell_text_llm to prepare the model."
            if self.feature_encoder is not None:
                assert hasattr(self.base_model.config, "pad_token_id"), \
                    "Base model config must have pad_token_id attribute. " + \
                    "Please use prepare_cell_text_llm to prepare the model."
                assert hasattr(self.base_model.config, "ignore_index"), \
                    "Base model config must have ignore_index attribute. " + \
                    "Please use prepare_cell_text_llm to prepare the model."
            if self.feature_decoder is not None:
                assert hasattr(self.base_model.config, "num_signal_tokens"), \
                    "Base model config must have num_signal_tokens attribute. " + \
                    "Please use prepare_cell_text_llm to prepare the model."

        self.normalize_total = normalize_total
        self.log_variational = log_variational

    def _merge_input_ids_with_features(
        self,
        input_ids: torch.Tensor,
        inputs_embeds: torch.Tensor,
        features: torch.Tensor,
        attention_mask: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
    ) -> Dict[str, torch.Tensor | Dict[str, torch.Tensor] | None]:
        """
        Merge input token embeddings with encoded cell features and updates attention masks and labels accordingly.

        This method is designed to handle a mix of text and cell feature tokens. It identifies placeholder tokens in the
        input text sequence (denoted by special tokens) and replaces them with corresponding cell features. It then constructs
        the final sequence of embeddings that contains both text and feature embeddings, ensuring that the attention masks and
        labels are updated accordingly for a consistent model input.

        Adapted from https://github.com/huggingface/transformers/blob/v4.40.1/src/transformers/models/llava/modeling_llava.py#L355.

        Parameters
        ----------
        input_ids: torch.Tensor
            The input token IDs of shape (batch_size, sequence_length), where some tokens may be special placeholders
            indicating the positions to insert cell features.
        inputs_embeds: torch.Tensor
            Precomputed embeddings corresponding to the ``input_ids`` of shape (batch_size, sequence_length, embed_dim).
        features: torch.Tensor
            Encoded cell features of shape (num_cells, num_features, embed_dim). These features will be inserted into
            ``input_embeds`` based on the placeholders in ``input_ids``. If the features are 2D, they will be unsqueezed
            to 3D with an additional dimension.
        attention_mask: torch.Tensor
            Attention mask indicating which tokens should be attended to, of shape (batch_size, sequence_length).
        labels: torch.Tensor, optional, default None
            Ground truth labels for token prediction tasks of shape (batch_size, sequence_length). These labels will also
            be updated to match the new token positions after merging.

        Returns
        -------
        outputs: dict
            A dict containing:
            - "inputs_embeds": torch.Tensor
              The final merged embeddings with both text and cell feature embeddings of shape (batch_size, max_sequence_length, embed_dim).
            - "attention_mask": torch.Tensor
              The updated attention mask reflecting the new merged token sequence of shape (batch_size, max_sequence_length).
            - "labels": torch.Tensor, optional
              The updated labels of shape (batch_size, max_sequence_length), None if labels are not provided.

        Notes
        -----
        - This function assumes that the placeholders for cell features in the ``input_ids`` are identified by a special token.
          The exact special token index is retrieved from the ``base_model.config.special_tokens_index_dict``. The final merged
          sequence is padded to the maximum sequence length derived from the maximum number of input embeddings in the batch.
        - The number of placeholders in the input text sequence should match the number of cell features provided.
        - For encoder-decoder models, the labels are not updated because the labels do not contain input tokens.
        """
        # both MLP-based and transformer-based encoders are considered
        if len(features.shape) == 2:
            features = features.unsqueeze(dim=1)

        num_features, num_feature_legnth, embed_dim = features.shape
        placeholder_index = self.base_model.config.special_tokens_index_dict[
            "placeholder"]
        batch_size, sequence_length = input_ids.shape
        left_padding = not torch.sum(
            input_ids[:,
                      -1] == torch.tensor(self.base_model.config.pad_token_id))

        # 1. create a mask to know where special placeholders are
        special_token_mask = input_ids == placeholder_index
        # for each sample, we calculate the number of placeholders for this modality
        num_special_tokens = torch.sum(special_token_mask, dim=-1)
        # compute the maximum embed dimension
        # if the size of feature is equal to 1, the sequence length after merging is the same
        # we just replace the placeholder with the corresponding features
        max_embed_dim = (num_special_tokens.max() *
                         (num_feature_legnth - 1)) + sequence_length
        batch_indices, non_modality_indices = torch.where(
            input_ids != placeholder_index)

        # 2. compute the positions where text should be written
        # calculate new positions for text tokens in merged modality-text sequence
        # `special_token_mask` identifies placeholders
        # each placeholder will be replaced by `nb_text_tokens_per_feature` text tokens
        # `torch.cumsum` computes how each modality token shifts subsequent text token positions
        # - 1 to adjust for zero-based indexing, as `cumsum` inherently increases indices by one
        new_token_positions = torch.cumsum(
            (special_token_mask * (num_feature_legnth - 1) + 1), -1) - 1
        nb_pad = max_embed_dim - 1 - new_token_positions[:, -1]
        if left_padding:
            new_token_positions += nb_pad[:, None]  # offset for left padding
        text_to_overwrite = new_token_positions[batch_indices,
                                                non_modality_indices]

        # 3. create the full embedding, already padded to the maximum position
        final_embedding = torch.zeros(batch_size,
                                      max_embed_dim,
                                      embed_dim,
                                      dtype=inputs_embeds.dtype,
                                      device=inputs_embeds.device)
        final_attention_mask = torch.zeros(batch_size,
                                           max_embed_dim,
                                           dtype=attention_mask.dtype,
                                           device=inputs_embeds.device)
        if labels is not None and self.base_model.config.is_decoder:
            final_labels = torch.full(
                (batch_size, max_embed_dim),
                self.base_model.config.ignore_index,
                dtype=input_ids.dtype,
                device=input_ids.device,
            )
        # in case the modality encoder or the Language model has been offloaded to CPU, we need to manually
        # set the corresponding tensors into their correct target device
        target_device = inputs_embeds.device
        batch_indices, non_modality_indices, text_to_overwrite = (
            batch_indices.to(target_device),
            non_modality_indices.to(target_device),
            text_to_overwrite.to(target_device),
        )
        attention_mask = attention_mask.to(target_device)

        # 4. fill the embeddings based on the mask
        # if we have ["hey" "<CELL>", "how", "are"]
        # we need to index copy on [0, 577, 578, 579] for the text and [1:576] for the cell features
        final_embedding[batch_indices,
                        text_to_overwrite] = inputs_embeds[batch_indices,
                                                           non_modality_indices]
        final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[
            batch_indices, non_modality_indices]
        if labels is not None and self.base_model.config.is_decoder:
            final_labels[batch_indices,
                         text_to_overwrite] = labels[batch_indices,
                                                     non_modality_indices]

        # 5. fill the embeddings corresponding to the modality. Anything that is still zeros needs filling
        # B * L
        modality_to_overwrite = torch.all(final_embedding == 0, dim=-1)
        modality_to_overwrite &= modality_to_overwrite.cumsum(
            -1) - 1 >= nb_pad[:, None].to(target_device)
        if modality_to_overwrite.sum() != features.shape[:-1].numel():
            raise ValueError(
                f"The input provided to the model are wrong. The number of cell tokens is {torch.sum(special_token_mask)} while"
                f" the number of cells given to the model is {num_features}. This prevents correct indexing and breaks batch generation."
            )
        final_embedding[modality_to_overwrite] = features.contiguous().reshape(
            -1, embed_dim).to(target_device)
        final_attention_mask |= modality_to_overwrite

        # 6. mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens
        batch_indices, pad_indices = torch.where(
            input_ids == self.base_model.config.pad_token_id)
        indices_to_mask = new_token_positions[batch_indices, pad_indices]
        final_embedding[batch_indices, indices_to_mask] = 0
        if labels is None or self.base_model.config.is_encoder_decoder:
            # for encoder-decoder models, we don't need to update the labels
            # because the labels don't contain input tokens
            final_labels = None
        labels = final_labels
        attention_mask = final_attention_mask
        inputs_embeds = final_embedding

        return {
            "inputs_embeds": inputs_embeds,
            "attention_mask": attention_mask,
            "labels": labels,
        }

    def _get_condition_embeds(
        self,
        last_hidden_states: torch.Tensor,
        input_ids: torch.Tensor,
        output_counts: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Extract conditional embeddings from the hidden states of the language model.

        This function is used to retrieve the embeddings corresponding to a set of signal tokens
        in the input sequence. These signal tokens are used to indicate the presence of specific conditional information
        (e.g., cell-related information). The embeddings of these tokens are extracted and used to generate gene profiles.

        Parameters
        ----------
        last_hidden_states: torch.Tensor
            The last hidden states of the language model's forward pass, of shape
            (batch_size, sequence_length, hidden_size).
        input_ids: torch.Tensor
            The input token IDs of shape (batch_size, sequence_length), which include the special
            tokens used to identify the locations of the conditional information.
        output_counts: torch.Tensor, optional, default None
            If provided, this tensor contains the expected output cell counts. It is used to verify
            that the number of samples containing signal tokens matches the number of expected outputs.

        Returns
        -------
        condition_embeds: torch.Tensor
            A tensor containing the extracted conditional embeddings from the hidden states,
            of shape (num_cells, num_signal_tokens, hidden_size). These embeddings can be used
            for generating gene profiles.

        Notes
        -----
        - This function assumes that the signal tokens in an input sequence are consecutive.
        - The number of samples containing signal tokens must match the number of expected conditional outputs
        (e.g., gene counts) when provided.
        """
        # get conditional embeddings
        num_signal_tokens = self.base_model.config.num_signal_tokens
        first_signal_token_index = self.base_model.config.special_tokens_index_dict[
            "first_signal_token"]
        output_ids = input_ids
        batch_indices, signal_token_indices = torch.where(
            output_ids == first_signal_token_index)
        if output_counts is not None and len(batch_indices) != len(
                output_counts):
            raise ValueError(
                f"The number of first signal tokens in inputs is {len(batch_indices)} "
                f"but the number of output counts you provide is {len(output_counts)}. "
                f"The row index of each signal token listed as follows:\n{batch_indices}"
            )
        indexer = signal_token_indices[:, None] + torch.arange(
            num_signal_tokens)[None].to(device=output_ids.device,
                                        dtype=output_ids.dtype)
        target_ids = torch.arange(
            first_signal_token_index, first_signal_token_index +
            num_signal_tokens).to(device=output_ids.device,
                                  dtype=output_ids.dtype)
        if not torch.all(
                torch.gather(output_ids[batch_indices], 1, indexer) ==
                target_ids[None]):
            raise ValueError(
                "The arrangement of signal tokens is not consecutive.")
        batch_indices = batch_indices[:, None].expand_as(indexer)
        condition_embeds = last_hidden_states[batch_indices, indexer]

        return condition_embeds

    @torch.inference_mode()
    def generate(
        self,
        input_ids: torch.Tensor,
        input_counts: Optional[torch.Tensor] = None,
        return_conditions: bool = False,
        attention_mask: Optional[torch.Tensor] = None,
        prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor],
                                                    List[int]]] = None,
        synced_gpus: Optional[bool] = None,
        assistant_model: Optional["PreTrainedModel"] = None,
        streamer: Optional["BaseStreamer"] = None,
        **kwargs,
    ) -> Dict[str, List[str | None | np.ndarray]]:
        """
        Generate texts or cells from text sequences ``input_ids`` and input counts ``input_counts`` (optionally).

        This method combines tokenized text input with encoded cell features to generate both text
        sequences and optionally cell data. It handles the integration of text and cell data using
        a feature encoder and decoder if they are provided.

        Parameters
        ----------
        input_ids: torch.Tensor
            Input token IDs of shape (batch_size, sequence_length) for the language model.
        input_counts: torch.Tensor, optional, default None
            Input counts of shape (num_cells, count_dim), which are passed through the feature encoder.
            If not provided, text-only generation is performed.
        return_conditions: bool, optional, default False
            Whether to return the condition embeddings (i.e. the last hidden states of signal tokens).
        attention_mask: torch.Tensor, optional, default None
            A mask that specifies which tokens should be attended to. If not provided, it is computed from ``input_ids``.
        prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], List[int]], optional, default None
            If provided, this function constraints the beam search to allowed tokens only at each step. If not provided,
            no constraint is applied. Please see
            https://huggingface.co/docs/transformers/v4.44.2/en/main_classes/text_generation#transformers.GenerationMixin.generate
            for more details.
        synced_gpus: bool, optional, default None
            Whether to continue running the while loop until max_length. Please see
            https://huggingface.co/docs/transformers/v4.44.2/en/main_classes/text_generation#transformers.GenerationMixin.generate
            for more details.
        assistant_model: PreTrainedModel, optional, default None
            An assistant model that can be used to accelerate generation. The assistant model must have the exact same tokenizer.
            The acceleration is achieved when forecasting candidate tokens with the assistent model is much faster than running
            generation with the model you are calling generate from. As such, the assistant model should be much smaller.
            For more details, please see https://arxiv.org/abs/2302.01318.
        streamer: BaseStreamer, optional, default None
            Streamer object that will be used to stream the generated sequences. Generated tokens are passed through
            ``streamer.put(token_ids)`` and the streamer is responsible for any further processing.
        **kwargs
            Additional keyword arguments passed to the ``generate`` method of the language model. In general,
            the language model is based on the transformers library. Thus, please see
            https://huggingface.co/docs/transformers/v4.44.2/en/main_classes/text_generation#transformers.GenerationMixin.generate
            for more details.

        Returns
        -------
        outputs: dict
            A dict containing the following keys:
            - "texts": list of str, the generated text sequences. Special tokens are removed.
            - "cells": list of cell data, where each entry is the decoded cell features. If no signal tokens are detected during
            generation, each entry is None.
            - "conditions": list of numpy array, where each entry is the condition embeddings. If no signal tokens are detected during
            generation, each entry is None. This key is only present if ``return_conditions`` is True.
        """
        # we don't need custom stopping criteria because the built-in stopping criteria is good enough
        # we will use EosTokenCriteria by setting generation configuration's eos token id
        # to the model's eos token id and the first signal token id
        eos_token_id = [self.base_model.config.eos_token_id]
        if self.feature_decoder is not None:
            eos_token_id.append(self.base_model.config.
                                special_tokens_index_dict["first_signal_token"])

        if attention_mask is None:
            attention_mask = (input_ids
                              != self.base_model.config.pad_token_id).to(
                                  input_ids.device)
            attention_mask = attention_mask.long()
        inputs_embeds = self.base_model.get_input_embeddings()(input_ids)
        if input_counts is not None:
            if self.normalize_total:
                input_counts = input_counts / input_counts.sum(
                    dim=-1, keepdim=True) * 1e4
            if self.log_variational:
                input_counts = torch.log1p(input_counts)
            assert self.feature_encoder is not None, "Feature encoder is not provided."
            features = self.feature_encoder(input_counts)
            merged_inputs = self._merge_input_ids_with_features(
                input_ids,
                inputs_embeds,
                features,
                attention_mask,
            )
            inputs_embeds = merged_inputs["inputs_embeds"]
            attention_mask = merged_inputs["attention_mask"]

        generated_ids = self.base_model.generate(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
            synced_gpus=synced_gpus,
            assistant_model=assistant_model,
            streamer=streamer,
            eos_token_id=eos_token_id,
            **kwargs,
        )
        outputs = {"cells": [None] * generated_ids.size(0)}
        if return_conditions:
            outputs["conditions"] = [None] * generated_ids.size(0)

        if self.feature_decoder is not None:
            # start to generate the cell
            num_signal_tokens = self.base_model.config.num_signal_tokens
            first_signal_token_index = self.base_model.config.special_tokens_index_dict[
                "first_signal_token"]
            batch_size, length = generated_ids.shape
            pad_token_id = self.base_model.config.pad_token_id
            cell_indices, signal_token_indices = torch.where(
                generated_ids == first_signal_token_index)
            if len(cell_indices) > 0:
                nb_pad = max(
                    (signal_token_indices + num_signal_tokens - 1).max().item()
                    - length + 1, 0)
                if nb_pad > 0:
                    final_generated_ids = torch.concat([
                        generated_ids,
                        torch.full((batch_size, nb_pad), pad_token_id).to(
                            device=generated_ids.device,
                            dtype=generated_ids.dtype)
                    ],
                                                       dim=1)
                else:
                    # deep copy because we will modify the previous generated tokens to remove the signal tokens
                    final_generated_ids = generated_ids.clone()
                # remove the signal tokens
                generated_ids[cell_indices, signal_token_indices] = pad_token_id
                shift = torch.arange(num_signal_tokens).to(
                    device=generated_ids.device, dtype=generated_ids.dtype)
                indexer = signal_token_indices[:, None] + shift[None]
                final_generated_ids[
                    cell_indices[:, None].expand_as(indexer),
                    indexer] = shift[None] + first_signal_token_index

                # after the all signal tokens are added, we get the corresponding hidden states
                if self.base_model.config.is_encoder_decoder:
                    decoder_input_ids = final_generated_ids
                    last_hidden_states = self.base_model(
                        inputs_embeds=inputs_embeds,
                        decoder_input_ids=decoder_input_ids,
                        return_dict=True,
                        output_hidden_states=True,
                    ).decoder_hidden_states[-1]
                    condition_embeds = self._get_condition_embeds(
                        last_hidden_states, decoder_input_ids)
                elif self.base_model.config.is_decoder:
                    # unlike encoder-decoder architecture, the original input_ids are concatenated with the generated_ids
                    # note that we should consider padding tokens in the original input_ids
                    input_length = input_ids.size(-1)
                    new_input_ids = input_ids.clone()
                    num_normal_tokens = torch.sum(input_ids[cell_indices]
                                                  != pad_token_id,
                                                  dim=-1).to(dtype=torch.long)
                    max_num_normal_tokens = num_normal_tokens.max()
                    max_num_generated_tokens = (signal_token_indices.max() +
                                                num_signal_tokens)
                    # note that real nb_pad may be less than the calculated nb_pad
                    # because for an input with maximum length, its generated tokens may not be long enough to reach the maximum length
                    # in this case, we do it for the convenience of parallel computing
                    # we will remove the redundant padding tokens later
                    nb_pad = max_num_normal_tokens.item(
                    ) + max_num_generated_tokens.item() - input_length
                    if nb_pad > 0:
                        new_input_ids = torch.concat([
                            new_input_ids,
                            torch.full((batch_size, nb_pad), pad_token_id).to(
                                device=new_input_ids.device,
                                dtype=new_input_ids.dtype)
                        ],
                                                     dim=1)
                    shift = torch.arange(0, max_num_generated_tokens.item()).to(
                        device=num_normal_tokens.device, dtype=torch.long)
                    src_col_indexer = num_normal_tokens[:, None] + shift[None]
                    # concatenate the original input_ids with the generated_ids including the signal tokens
                    new_input_ids[
                        cell_indices[:, None].expand_as(src_col_indexer),
                        src_col_indexer] = final_generated_ids[cell_indices][
                            0:max_num_generated_tokens.item()]
                    # remove the redundant padding tokens
                    new_input_ids = new_input_ids[:,
                                                  torch.any(new_input_ids !=
                                                            pad_token_id,
                                                            dim=0)]
                    new_inputs_embeds = self.base_model.get_input_embeddings()(
                        new_input_ids)
                    last_hidden_states = self.base_model(
                        inputs_embeds=new_inputs_embeds,
                        return_dict=True,
                        output_hidden_states=True,
                    ).hidden_states[-1]
                    condition_embeds = self._get_condition_embeds(
                        last_hidden_states, new_input_ids)
                else:
                    raise ValueError(
                        "The model must be either encoder-decoder or decoder-only architecture."
                    )

                # use downstream model to generate the cell
                cells = self.feature_decoder.generate(condition_embeds,
                                                      to_numpy=True)
                for cell_indice, cell in zip(cell_indices, cells):
                    outputs["cells"][cell_indice] = cell
                if return_conditions:
                    condition_embeds_ = condition_embeds
                    if condition_embeds_.device.type == "cuda":
                        condition_embeds_ = condition_embeds_.cpu()
                    condition_embeds_ = condition_embeds_.numpy()
                    for cell_indice, condition in zip(cell_indices,
                                                      condition_embeds_):
                        outputs["conditions"][cell_indice] = condition

        if generated_ids.device.type == "cuda":
            generated_ids = generated_ids.cpu()
        outputs["texts"] = self.tokenizer.batch_decode(generated_ids,
                                                       skip_special_tokens=True)

        return outputs

    def forward(self,
                input_ids: torch.Tensor,
                attention_mask: Optional[torch.Tensor] = None,
                input_counts: Optional[torch.Tensor] = None,
                output_counts: Optional[torch.Tensor] = None,
                labels: Optional[torch.Tensor] = None,
                output_attentions: Optional[bool] = None,
                output_hidden_states: Optional[bool] = None,
                generative_model_kwargs: Optional[Dict[str, Any]] = None,
                **base_model_kwargs) -> Seq2SeqLMOutput | CausalLMOutput:
        """
        Perform a forward pass for the CellTextLLM model, integrating cell feature generation
        and text generation in a multimodal setup.

        Parameters
        ----------
        input_ids: torch.Tensor
            The token IDs of the input text sequence of shape (batch_size, sequence_length).
        attention_mask: torch.Tensor, optional, default None
            A mask that indicates which tokens should be attended to, where 1 indicates a token should be attended
            to and 0 indicates it should not be. If None, it is computed from ``input_ids``.
        input_counts: torch.Tensor, optional, default None
            A tensor of shape (num_cells, count_dim) representing the input cell counts to be processed by the feature encoder.
        output_counts: torch.Tensor, optional, default None
            A tensor representing expected output cell counts for supervision in the cell generation process.
        labels: torch.Tensor, optional, default None
            Ground truth labels for text generation. Used for supervised training when provided.
        output_attentions: bool, optional, default None
            If set to True, the model returns attention weights in addition to the output.
        output_hidden_states: bool, optional, default None
            If set to True, the model returns hidden states in addition to the output.
        generative_model_kwargs: dict, optional, default None
            Additional keyword arguments passed to the ``forward`` method of cell generative model if output cell counts
            are provided.
        **base_model_kwargs: dict
            Additional keyword arguments passed to the ``forward`` method of base language model.

        Returns
        -------
        outputs: transformers.modeling_outputs.Seq2SeqLMOutput or transformers.modeling_outputs.CausalLMOutput
            The outputs of the model.
        """
        # get text embeddings
        inputs_embeds = self.base_model.get_input_embeddings()(input_ids)
        if input_counts is not None:
            if self.normalize_total:
                input_counts = input_counts / input_counts.sum(
                    dim=-1, keepdim=True) * 1e4
            if self.log_variational:
                input_counts = torch.log1p(input_counts)
        if attention_mask is None:
            attention_mask = (input_ids
                              != self.base_model.config.pad_token_id).to(
                                  input_ids.device)
            attention_mask = attention_mask.long()

        # get cell modality's features
        if input_counts is not None:
            assert self.feature_encoder is not None, "Feature encoder is not provided."
            features = self.feature_encoder(input_counts)
            merged_inputs = self._merge_input_ids_with_features(
                input_ids,
                inputs_embeds,
                features,
                attention_mask,
                labels=labels,
            )
            inputs_embeds = merged_inputs["inputs_embeds"]
            attention_mask = merged_inputs["attention_mask"]
            if merged_inputs["labels"] is not None:
                labels = merged_inputs["labels"]

        base_model_kwargs = {
            **base_model_kwargs,
            "return_dict": True,
        }
        if output_counts is None:
            outputs = self.base_model(
                attention_mask=attention_mask,
                inputs_embeds=inputs_embeds,
                labels=labels,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                **base_model_kwargs,
            )
            return outputs
        base_model_kwargs = {
            **base_model_kwargs,
            "output_hidden_states": True,
        }
        outputs = self.base_model(
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            labels=labels,
            output_attentions=output_attentions,
            **base_model_kwargs,
        )

        # get conditional embeddings
        if hasattr(outputs, "decoder_hidden_states"):
            if "decoder_input_ids" not in base_model_kwargs:
                raise ValueError(
                    "decoder_input_ids must be provided in base_model_kwargs when the encoder-decoder model "
                    " is used to generate cells.")
            last_hidden_states = outputs.decoder_hidden_states[-1]
            input_ids_ = base_model_kwargs["decoder_input_ids"]
        else:
            last_hidden_states = outputs.hidden_states[-1]
            input_ids_ = input_ids
        condition_embeds = self._get_condition_embeds(
            last_hidden_states,
            input_ids_,
            output_counts=output_counts,
        )
        if generative_model_kwargs is None:
            generative_model_kwargs = {}
        generative_outputs = self.feature_decoder(condition_embeds,
                                                  output_counts,
                                                  **generative_model_kwargs)
        batch_size = input_ids.shape[0]
        num_generative_samples = output_counts.shape[0]
        outputs.loss = outputs.loss + generative_outputs.loss * num_generative_samples / batch_size

        return outputs


class InstructCellConfig(PretrainedConfig):
    """
    A configuration class for the ``InstructCell`` model, which integrates a base language model
    configuration, a feature encoder configuration, and a feature decoder configuration for
    multimodal cell-language tasks.

    Parameters
    ----------
    base_model_config: transformers.PretrainedConfig, optional, default None
        The configuration of the base language model. If this argument is a dictionary, the
        corresponding configuration class will be instantiated based on the ``model_type`` key
        in the dictionary. If it is None, a default configuration is created (e.g., T5).
    feature_encoder_config: Dict[str, Any], default {
        "is_q_former_encoder": True,
        "count_dim": 18961,
        "cross_attention_frequency": 1,
        "num_hidden_layers": 4,
        "num_key_value_tokens": 6,
        "num_blocks": 3,
        "num_query_tokens": 8,
        "hidden_dropout_prob": 0.1,
    }
        The configuration dictionary for the feature encoder. This dictionary can specify
        whether a Q-Former module is used, the dimensionality of the input cell counts,
        dropout probabilities, and other hyperparameters for the encoder architecture.
    feature_decoder_config: Dict[str, Any], default {
        "use_layer_norm": "both",
        "use_batch_norm": "none",
        "n_latent": 256,
        "condition_input_dim": 256,
        "log_variational": True,
        "n_layers": 4,
        "n_hidden": 1024,
        "dropout_rate": 0.1,
        "adaptive_library": True,
    }
        The configuration dictionary for the feature decoder. This dictionary can specify
        normalization schemes, number of latent dimensions, dimensionality of condition
        inputs, number of hidden layers, dropout rate, and whether library size is adaptive.
    tokenizer_use_fast: bool, default False
        Whether to use the fast tokenizer implementation from ``transformers`` library.
    normalize_total: bool, default True
        Whether to normalize total read count for each single-cell input.
    log_variational: bool, default True
        Whether to apply the log1p-transformation to to each single-cell input.
    **kwargs
        Additional keyword arguments passed to ``transformers.PretrainedConfig``.
    """

    model_type = "instructcell"
    sub_configs = {"base_model_config": AutoConfig}

    def __init__(
        self,
        base_model_config: Optional[PretrainedConfig] = None,
        feature_encoder_config: Dict[str, Any] = {
            "is_q_former_encoder": True,
            "count_dim": 18961,
            "cross_attention_frequency": 1,
            "num_hidden_layers": 4,
            "num_key_value_tokens": 6,
            "num_blocks": 3,
            "num_query_tokens": 8,
            "hidden_dropout_prob": 0.1,
        },
        feature_decoder_config: Dict[str, Any] = {
            "use_layer_norm": "both",
            "use_batch_norm": "none",
            "n_latent": 256,
            "condition_input_dim": 256,
            "log_variational": True,
            "n_layers": 4,
            "n_hidden": 1024,
            "dropout_rate": 0.1,
            "adaptive_library": True,
        },
        tokenizer_use_fast: bool = False,
        normalize_total: bool = True,
        log_variational: bool = True,
        **kwargs,
    ) -> "InstructCellConfig":
        # initialize the configuration of base language model
        if isinstance(base_model_config, dict):
            base_model_config["model_type"] = (base_model_config["model_type"]
                                               if "model_type"
                                               in base_model_config else "t5")
            base_model_config = CONFIG_MAPPING[base_model_config["model_type"]](
                **base_model_config)
        elif base_model_config is None:
            base_model_config = CONFIG_MAPPING["t5"]()
        self.base_model_config = base_model_config
        for attr_names in [
                "special_tokens_index_dict",
                "pad_token_id",
                "eos_token_id",
                "ignore_index",
                "num_signal_tokens",
        ]:
            if not hasattr(self.base_model_config, attr_names):
                self.base_model_config.__setattr__(attr_names, None)

        # initialize the configuration of feature encoder
        self.feature_encoder_config = feature_encoder_config

        # initialize the configuration of feature decoder
        self.feature_decoder_config = feature_decoder_config

        # the configuration of tokenizer
        self.tokenizer_use_fast = tokenizer_use_fast

        # other parameters of InstructCell
        self.normalize_total = normalize_total
        self.log_variational = log_variational

        super().__init__(**kwargs)


class InstructCell(PreTrainedModel):
    """
    A model class that integrates a base language model, a feature encoder, and a feature
    decoder to support multimodal interaction between text and single-cell data.

    The ``InstructCell`` model contains:
    - A base language model (e.g., T5, GPT).
    - A feature encoder (e.g., SCQFormer) to encode the input cell counts.
    - A feature decoder (e.g., Generator) to reconstruct or generate cell data (gene expression)
      from sampled latent vectors and given biological conditions.

    This class leverages a ``CellTextLLM`` instance, referred to as the assistant, to combine
    text embeddings and cell embeddings in a single forward pass or generation pipeline.
    It provides methods for tokenizing text, performing multimodal forward passes, and
    generating text or cell data outputs.

    Parameters
    ----------
    config: InstructCellConfig
        The configuration object containing hyperparameters and settings for the base model,
        feature encoder, feature decoder, tokenizer usage, and other relevant attributes.

    Notes
    -----
    - The model relies on a placeholder or signal tokens to identify where the cell embeddings
      should be inserted. Therefore, the user must invoke ``initialize_special_tokens``
      to set up these tokens in the tokenizer and in the model configuration if cell-based
      functionalities are needed.
    - The feature encoder and decoder used here typically expect preprocessed cell data,
      such as counts that may be normalized or log1p-transformed (controlled by
      ``config.normalize_total`` and ``config.log_variational``).
    """

    config_class = InstructCellConfig

    def __init__(self, config: InstructCellConfig) -> "InstructCell":
        # load the pretrained language model
        if config.base_model_config.is_encoder_decoder:
            model = AutoModelForSeq2SeqLM.from_config(config.base_model_config)
        else:
            model = AutoModelForCausalLM.from_config(config.base_model_config)

        # load the Q-former module
        count_dim = config.feature_encoder_config["count_dim"]
        is_q_former_encoder = config.feature_encoder_config[
            "is_q_former_encoder"]
        if is_q_former_encoder:
            cross_attention_frequency = config.feature_encoder_config[
                "cross_attention_frequency"]
            num_hidden_layers = config.feature_encoder_config[
                "num_hidden_layers"]
            qformer_config = Blip2QFormerConfig(
                vocab_size=0,
                hidden_size=model.config.hidden_size,
                hidden_dropout_prob=config.
                feature_encoder_config["hidden_dropout_prob"],
                num_hidden_layers=num_hidden_layers,
                num_attention_heads=model.config.num_attention_heads,
                intermediate_size=model.config.hidden_size * 4,
                pad_token_id=model.config.pad_token_id,
                cross_attention_frequency=cross_attention_frequency,
                encoder_hidden_size=model.config.hidden_size,
            )
            num_key_value_tokens = config.feature_encoder_config[
                "num_key_value_tokens"]
            num_blocks = config.feature_encoder_config["num_blocks"]
            num_query_tokens = config.feature_encoder_config["num_query_tokens"]
            feature_encoder = SCQFormer(
                count_dim,
                num_query_tokens,
                num_key_value_tokens,
                qformer_config,
                num_hidden_layers=num_blocks,
            )
        else:
            feature_encoder = nn.Sequential(
                nn.Linear(count_dim,
                          (count_dim + model.config.hidden_size) // 2),
                nn.GELU(),
                nn.Linear((count_dim + model.config.hidden_size) // 2,
                          model.config.hidden_size),
                nn.Dropout(
                    config.feature_encoder_config["hidden_dropout_prob"]),
            )

        # load the cell reconstruction module
        condition_input_dim = config.feature_decoder_config[
            "condition_input_dim"]
        use_layer_norm = config.feature_decoder_config["use_layer_norm"]
        use_batch_norm = config.feature_decoder_config["use_batch_norm"]
        n_latent = config.feature_decoder_config["n_latent"]
        log_variational = config.feature_decoder_config["log_variational"]
        n_layers = config.feature_decoder_config["n_layers"]
        n_hidden = config.feature_decoder_config["n_hidden"]
        dropout_rate = config.feature_decoder_config["dropout_rate"]
        adaptive_library = config.feature_decoder_config["adaptive_library"]
        # create a dummy matrix to initialize library size
        dummy_counts_for_init = np.random.randint(0, 1000, (10, count_dim))
        library_log_means, library_log_vars = init_library_size(
            dummy_counts_for_init)
        feature_decoder = Generator(
            count_dim,
            condition_dim=model.config.hidden_size,
            condition_input_dim=condition_input_dim,
            n_layers=n_layers,
            n_hidden=n_hidden,
            n_latent=n_latent,
            dropout_rate=dropout_rate,
            use_layer_norm=use_layer_norm,
            use_batch_norm=use_batch_norm,
            encode_covariates=True,
            deeply_inject_covariates=False,
            log_variational=log_variational,
            adaptive_library=adaptive_library,
            use_observed_lib_size=False,
            library_log_means=library_log_means,
            library_log_vars=library_log_vars,
        )

        super().__init__(config)
        # the submodule of InstructCell
        self.assistant = CellTextLLM(
            model,
            None,
            feature_encoder=feature_encoder,
            feature_decoder=feature_decoder,
            normalize_total=config.normalize_total,
            log_variational=config.log_variational,
        )

    def _check_tokenizer(self) -> None:
        """Check whether the tokenizer is set."""
        if self.assistant.tokenizer is None:
            raise ValueError(
                "The tokenizer has not been set. Please configure the tokenizer first."
            )

    def save_pretrained(self,
                        save_directory: Optional[str],
                        save_tokenizer_kwargs: Dict[str, Any] = {},
                        **kwargs) -> None:
        """
        Save the ``InstructCell`` model and its tokenizer to a directory, so that it can be
        reloaded using the ``from_pretrained`` method.

        Parameters
        ----------
        save_directory: str, optional
            The directory where the model and tokenizer files will be saved.
        save_tokenizer_kwargs: dict, default {}
            Additional keyword arguments passed to the tokenizer's ``save_pretrained`` method.
        **kwargs
            Additional keyword arguments passed to the ``save_pretrained`` method of the parent
            class (``transformers.PreTrainedModel``).
        """
        self._check_tokenizer()
        self.assistant.tokenizer.save_pretrained(save_directory,
                                                 **save_tokenizer_kwargs)
        super().save_pretrained(save_directory, **kwargs)

    @classmethod
    def from_pretrained(cls,
                        pretrained_model_name_or_path: Optional[str],
                        load_tokenizer_kwargs: Dict[str, Any] = {},
                        **kwargs) -> "InstructCell":
        """
        Load a pre-trained ``InstructCell`` model from a local directory or remote model hub.

        This method loads both the model weights and its tokenizer from the specified
        directory or model repository.

        Parameters
        ----------
        pretrained_model_name_or_path: str, optional
            The path (local directory) or name of the remote repository to load the model from.
        load_tokenizer_kwargs: Dict[str, Any], default {}
            Additional keyword arguments passed to the ``from_pretrained`` method of the class
            ``transformers.PreTrainedTokenizer`` or ``transformers.PreTrainedTokenizerFast``.
        **kwargs
            Additional keyword arguments passed to the ``from_pretrained`` method of the parent
            class (``transformers.PreTrainedModel``).

        Returns
        -------
        model: InstructCell
            The InstructCell model with trained base language model, feature encoder, feature
            decoder, and tokenizer.
        """
        model = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
        load_tokenizer_kwargs = {
            "use_fast": model.config.tokenizer_use_fast,
            **load_tokenizer_kwargs
        }
        # load the corresponding tokenizer
        model.assistant.tokenizer = AutoTokenizer.from_pretrained(
            pretrained_model_name_or_path, **load_tokenizer_kwargs)
        return model

    def initialize_special_tokens(
            self,
            modality_tag: str = "CELL",
            num_signal_tokens: int = 1,
            pad_to_multiple_of: Optional[int] = 8) -> None:
        """
        Initialize and register special tokens for cell-language modeling, such as placeholder
        tokens and signal tokens for inserting or extracting cell embeddings.

        Parameters
        ----------
        modality_tag: str, default "CELL"
            The tag that represents the modality to be added to the tokenizer. It will be used to create
            start and end tags as well as signal tokens.
        num_signal_tokens: int, default 1
            The number of signal tokens to add. These tokens serve as signals enabling the model to generate
            data in the specified modality.
        pad_to_multiple_of: int, optional, default 8
            If set will pad the embedding matrix to a multiple of the provided value after resizing embedding
            layer. This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute
            capability ``>= 7.5`` (Volta), or on TPUs which benefit from having sequence lengths be a multiple
            of 128.
        """
        self._check_tokenizer()
        # ensure the tokenizer has a pad token
        if self.assistant.tokenizer.pad_token_id is None:
            self.assistant.tokenizer.pad_token = self.assistant.tokenizer.eos_token
        ignore_index = -100 if not hasattr(self.assistant.base_model.config,
                                           "ignore_index") else None
        pad_token_id = self.assistant.tokenizer.pad_token_id if not hasattr(
            self.assistant.base_model.config, "pad_token_id") else None

        # prepare the model for cell-language modelling
        prepare_cell_text_llm(
            self.assistant.base_model,
            self.assistant.tokenizer,
            modality_tag=modality_tag,
            num_signal_tokens=num_signal_tokens,
            ignore_index=ignore_index,
            pad_token_id=pad_token_id,
            pad_to_multiple_of=pad_to_multiple_of,
        )

    def prepare_model_inputs(
        self,
        prompt: str,
        gene_counts: Optional[np.ndarray] = None,
        sc_metadata: Dict[str, str] = {},
    ) -> Dict[str, torch.Tensor]:
        """
        Prepare the input dictionary for the model, which includes tokenized text (``input_ids``)
        and optionally gene counts (``input_counts``).

        This method formats the prompt with single-cell metadata placeholders, tokenizes the
        resulting string, and wraps the input tokens in a PyTorch tensor. If gene counts are
        provided, they are also wrapped in a tensor.

        Parameters
        ----------
        prompt: str
            The text prompt that might include placeholders for single-cell metadata (e.g.,
            `{input}`, `{species}`). For example: ``"Please give the cell type of {input} from {species}."``
        gene_counts: np.ndarray, optional, default None
            The matrix of shape (num_cells, count_dim) or (count_dim,) that represents the
            input cell gene counts.
        sc_metadata: dict, default {}
            Additional metadata used to format the prompt string. For instance,
            ``{"species": "human"}``, which can be used to substitute for ``{species}`` in the prompt.
            By default, ``{input}`` in the prompt is automatically replaced with the special placeholder
            token (see method ``initialize_special_tokens``).

        Returns
        -------
        inputs: dict
            A dictionary containing:
            - "input_ids": torch.Tensor
              The tokenized text prompt of shape (1, sequence_length).
            - "input_counts": torch.Tensor or None
              The cell count tensor of shape (num_cells, count_dim) if gene_counts is not None,
              otherwise None.
        """
        if not hasattr(self.config.base_model_config, "special_tokens_dict"):
            raise ValueError(
                "Cannot get special tokens. ",
                "Please invoke `initialize_special_tokens` for cell-language modelling."
            )
        self._check_tokenizer()

        if gene_counts is not None and len(gene_counts.shape) == 1:
            gene_counts = gene_counts.reshape(1, -1)

        placeholder = self.config.base_model_config.special_tokens_dict[
            "placeholder"]
        start_tag = self.config.base_model_config.special_tokens_dict[
            "start_tag"]
        end_tag = self.config.base_model_config.special_tokens_dict["end_tag"]

        # users can define the value of key `input` in `sc_metadata` by themselves
        sc_metadata = {
            "input": f"{start_tag}{placeholder}{end_tag}",
            **sc_metadata,
        }
        prompt = prompt.format(**sc_metadata)
        prompt = f"User:\n{prompt}\n\nAssistant:\n"

        input_ids = self.assistant.tokenizer(prompt,
                                             return_tensors="np").input_ids
        inputs = {"input_ids": torch.from_numpy(input_ids)}
        if gene_counts is not None:
            inputs["input_counts"] = torch.from_numpy(gene_counts)
        else:
            inputs["input_counts"] = None
        for key, value in inputs.items():
            if value is not None:
                inputs[key] = value.to(device=self.assistant.base_model.device)
        if inputs["input_counts"] is not None:
            inputs["input_counts"] = inputs["input_counts"].to(dtype=self.dtype)

        return inputs

    def forward(self,
                input_ids: torch.Tensor,
                attention_mask: Optional[torch.Tensor] = None,
                input_counts: Optional[torch.Tensor] = None,
                output_counts: Optional[torch.Tensor] = None,
                labels: Optional[torch.Tensor] = None,
                output_attentions: Optional[bool] = None,
                output_hidden_states: Optional[bool] = None,
                generative_model_kwargs: Optional[Dict[str, Any]] = None,
                **base_model_kwargs) -> Seq2SeqLMOutput | CausalLMOutput:
        """
        Perform a forward pass for the CellTextLLM model, integrating cell feature generation
        and text generation in a multimodal setup.

        Parameters
        ----------
        input_ids: torch.Tensor
            The token IDs of the input text sequence of shape (batch_size, sequence_length).
        attention_mask: torch.Tensor, optional, default None
            A mask that indicates which tokens should be attended to, where 1 indicates a token should be attended
            to and 0 indicates it should not be. If None, it is computed from ``input_ids``.
        input_counts: torch.Tensor, optional, default None
            A tensor of shape (num_cells, count_dim) representing the input cell counts to be processed by the feature encoder.
        output_counts: torch.Tensor, optional, default None
            A tensor representing expected output cell counts for supervision in the cell generation process.
        labels: torch.Tensor, optional, default None
            Ground truth labels for text generation. Used for supervised training when provided.
        output_attentions: bool, optional, default None
            If set to True, the model returns attention weights in addition to the output.
        output_hidden_states: bool, optional, default None
            If set to True, the model returns hidden states in addition to the output.
        generative_model_kwargs: dict, optional, default None
            Additional keyword arguments passed to the ``forward`` method of cell generative model if output cell counts
            are provided.
        **base_model_kwargs: dict
            Additional keyword arguments passed to the ``forward`` method of base language model.

        Returns
        -------
        outputs: transformers.modeling_outputs.Seq2SeqLMOutput or transformers.modeling_outputs.CausalLMOutput
            The outputs of the model.
        """
        return self.assistant(
            input_ids,
            attention_mask=attention_mask,
            input_counts=input_counts,
            output_counts=output_counts,
            labels=labels,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            generative_model_kwargs=generative_model_kwargs,
            **base_model_kwargs,
        )

    def generate(
        self,
        input_ids: torch.Tensor,
        input_counts: Optional[torch.Tensor] = None,
        return_conditions: bool = False,
        attention_mask: Optional[torch.Tensor] = None,
        prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor],
                                                    List[int]]] = None,
        synced_gpus: Optional[bool] = None,
        assistant_model: Optional["PreTrainedModel"] = None,
        streamer: Optional["BaseStreamer"] = None,
        **kwargs,
    ) -> Dict[str, List[str | None | np.ndarray]]:
        """
        Generate texts or cells from text sequences ``input_ids`` and input counts ``input_counts`` (optionally).

        This method combines tokenized text input with encoded cell features to generate both text
        sequences and optionally cell data. It handles the integration of text and cell data using
        a feature encoder and decoder if they are provided.

        Parameters
        ----------
        input_ids: torch.Tensor
            Input token IDs of shape (batch_size, sequence_length) for the language model.
        input_counts: torch.Tensor, optional, default None
            Input counts of shape (num_cells, count_dim), which are passed through the feature encoder.
            If not provided, text-only generation is performed.
        return_conditions: bool, optional, default False
            Whether to return the condition embeddings (i.e. the last hidden states of signal tokens).
        attention_mask: torch.Tensor, optional, default None
            A mask that specifies which tokens should be attended to. If not provided, it is computed from ``input_ids``.
        prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], List[int]], optional, default None
            If provided, this function constraints the beam search to allowed tokens only at each step. If not provided,
            no constraint is applied. Please see
            https://huggingface.co/docs/transformers/v4.44.2/en/main_classes/text_generation#transformers.GenerationMixin.generate
            for more details.
        synced_gpus: bool, optional, default None
            Whether to continue running the while loop until max_length. Please see
            https://huggingface.co/docs/transformers/v4.44.2/en/main_classes/text_generation#transformers.GenerationMixin.generate
            for more details.
        assistant_model: PreTrainedModel, optional, default None
            An assistant model that can be used to accelerate generation. The assistant model must have the exact same tokenizer.
            The acceleration is achieved when forecasting candidate tokens with the assistent model is much faster than running
            generation with the model you are calling generate from. As such, the assistant model should be much smaller.
            For more details, please see https://arxiv.org/abs/2302.01318.
        streamer: BaseStreamer, optional, default None
            Streamer object that will be used to stream the generated sequences. Generated tokens are passed through
            ``streamer.put(token_ids)`` and the streamer is responsible for any further processing.
        **kwargs
            Additional keyword arguments passed to the ``generate`` method of the language model. In general,
            the language model is based on the transformers library. Thus, please see
            https://huggingface.co/docs/transformers/v4.44.2/en/main_classes/text_generation#transformers.GenerationMixin.generate
            for more details.

        Returns
        -------
        outputs: dict
            A dict containing the following keys:
            - "texts": list of str, the generated text sequences. Special tokens are removed.
            - "cells": list of cell data, where each entry is the decoded cell features. If no signal tokens are detected during
            generation, each entry is None.
            - "conditions": list of numpy array, where each entry is the condition embeddings. If no signal tokens are detected during
            generation, each entry is None. This key is only present if ``return_conditions`` is True.
        """
        self._check_tokenizer()
        return self.assistant.generate(
            input_ids,
            input_counts=input_counts,
            return_conditions=return_conditions,
            attention_mask=attention_mask,
            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
            synced_gpus=synced_gpus,
            assistant_model=assistant_model,
            streamer=streamer,
            **kwargs,
        )

    def predict(self,
                prompt: str,
                gene_counts: Optional[np.ndarray] = None,
                sc_metadata: Dict[str, str] = {},
                **kwargs) -> Dict[str, str | np.ndarray]:
        """
        A high-level helper method for inference, combining text prompt processing and generation
        in a single call.

        This method:
        1. Prepares the model inputs using ``prepare_model_inputs`` (tokenizes the prompt and
           optionally includes the cell counts).
        2. Calls the model's ``generate`` method to produce text (and optionally cell data).
        3. Returns the generated text and cell data in a dictionary.

        Parameters
        ----------
        prompt: str
            The user prompt or instruction, which may contain placeholders for single-cell data.
        gene_counts: np.ndarray, optional, default None
            The cell count array of shape (num_cells, count_dim) or (count_dim,).
        sc_metadata: Dict[str, str], default {}
            Additional metadata to format the prompt. Keys can be placeholders used in ``prompt``.
        **kwargs
            Additional keyword arguments passed to the model's ``generate`` method
            (e.g., ``top_p``).

        Returns
        -------
        outputs: dict
            A dictionary containing:
            - "text": str, the generated text in response to the prompt.
            - "cell": np.ndarray or None, the generated cell data, if applicable.

        Notes
        -----
        - This method is particularly convenient for quick inference scenarios where the user
          has cell data (counts) and a text prompt, and wants a single call to retrieve both
          text and cell outputs.
        """
        inputs = self.prepare_model_inputs(prompt,
                                           gene_counts=gene_counts,
                                           sc_metadata=sc_metadata)
        outputs = self.generate(inputs["input_ids"],
                                input_counts=inputs["input_counts"],
                                **kwargs)
        return {"text": outputs["texts"][0], "cell": outputs["cells"][0]}
