modeling_ministu.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput
from .modules import STU, Attention, MLP
from .utils import nearest_power_of_two
from .layers import STULayer, AttentionLayer
from .configuration_ministu import MiniSTUConfig
from .filters import get_spectral_filters

try:
    from liger_kernel.transformers.rms_norm import LigerRMSNorm as TritonNorm
    triton_norm = True
except ImportError as e:
    print(
        f"Unable to import Triton-based RMSNorm: {e}. Falling back to PyTorch implementation."
    )
    from torch.nn import RMSNorm
    triton_norm = False
# Load the tokenizer
#from transformers import AutoModelForCausalLM, AutoTokenizer
#model_name = "Hazan-Lab/STU-426M"
#tokenizer = AutoTokenizer.from_pretrained(
#    model_name,
#    trust_remote_code=True
#)

class MiniSTU(PreTrainedModel):
    config_class = MiniSTUConfig

    def __init__(self, config) -> None:
        super(MiniSTU, self).__init__(config)
        self.n_layers = config.n_layers
        self.n = nearest_power_of_two(config.seq_len * 2 - 1, round_up=True)
        
        if isinstance(config.torch_dtype, torch.dtype):
            torch_dtype = config.torch_dtype
        else:
            torch_dtype = getattr(torch, config.torch_dtype)

        device = torch.device(config.device)

        self.phi = get_spectral_filters(
            config.seq_len,
            config.num_eigh,
            config.use_hankel_L,
            device=device,
            dtype=torch_dtype,
        )
        
        self.use_approx = config.use_approx
        self.use_hankel_L = config.use_hankel_L

        self.tok_emb = nn.Embedding(
            config.vocab_size, config.n_embd, dtype=torch_dtype, device=device
        )
        self.dropout = nn.Dropout(config.dropout)

        self.layers = nn.ModuleList()
        for layer_idx in range(self.n_layers):
            if layer_idx % 2 == 0:
                self.layers.append(STULayer(config, self.phi, self.n))
            else:
                self.layers.append(
                    AttentionLayer(config)
                    if config.use_attn
                    else STULayer(config, self.phi, self.n)
                )

        self.norm = TritonNorm(config.n_embd) if triton_norm else RMSNorm(config.n_embd)

        self.lm_head = nn.Linear(
            config.n_embd, config.vocab_size, bias=config.bias, dtype=torch_dtype, device=device
        )
        self.tok_emb.weight = self.lm_head.weight

        self.std = (config.n_embd) ** -0.5
        self.apply(self._init_weights)
        print("Model Parameter Count: %.2fM\n" % (self._get_num_params() / 1e6,))

    def forward(
        self,
        input_ids: torch.Tensor,
        labels: torch.Tensor = None,
        **kwargs
    ) -> CausalLMOutput:
        # Compute embeddings
        tok_emb = self.tok_emb(input_ids)
        x = self.dropout(tok_emb)

        # Pass through layers
        for layer in self.layers:
            x = layer(x)

        # Normalize and project to vocabulary
        x = self.norm(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            # Shift so that tokens predict the next token
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1)
            )

        return CausalLMOutput(
            loss=loss,
            logits=logits,
        )

    def _get_num_params(self):
        n_params = sum(p.numel() for p in self.parameters())
        if hasattr(self, "pos_emb") and self.pos_emb is not None:
            n_params -= self.pos_emb.weight.numel()
        if self.tok_emb.weight is not self.lm_head.weight:
            n_params -= self.tok_emb.weight.numel()
        return n_params

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            if hasattr(module, "SCALE_INIT"):
                self.std *= (2 * self.n_layers) ** -0.5
            torch.nn.init.normal_(module.weight, mean=0.0, std=self.std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=self.std)
        elif isinstance(module, STU):
            if self.use_approx:
                torch.nn.init.xavier_normal_(module.M_inputs)
                torch.nn.init.xavier_normal_(module.M_filters)
            else:
                torch.nn.init.xavier_normal_(module.M_phi_plus)
                if not self.use_hankel_L:
                    torch.nn.init.xavier_normal_(module.M_phi_minus)
        elif isinstance(module, Attention):
            torch.nn.init.xavier_normal_(module.c_attn.weight)
            torch.nn.init.xavier_normal_(module.c_proj.weight)
            if module.c_attn.bias is not None:
                torch.nn.init.zeros_(module.c_attn.bias)
            if module.c_proj.bias is not None:
                torch.nn.init.zeros_(module.c_proj.bias)
    @staticmethod
    def top_k_top_p_filtering(
        logits: torch.Tensor,
        top_k: int = 50,
        top_p: float = 0.95,
        filter_value: float = float("-inf"),
    ):
        """
        Filters a distribution of logits using top-k and/or nucleus (top-p) filtering.
        """
        # top_k
        if top_k > 0:
            top_k = min(top_k, logits.size(-1))
            # Remove all logits that are not in the top k
            indices_to_remove = logits < torch.topk(logits, top_k, dim=-1).values[:, -1, None]
            logits[indices_to_remove] = filter_value

        # top_p (nucleus)
        if 0 < top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

            # Remove tokens with cumulative probability above the threshold
            sorted_indices_to_remove = cumulative_probs > top_p
            # Shift the indices to the right to keep also the first token above the threshold
            sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
            sorted_indices_to_remove[:, 0] = False

            indices_to_remove = sorted_indices_to_remove.scatter(
                dim=1, index=sorted_indices, src=sorted_indices_to_remove
            )
            logits[indices_to_remove] = filter_value

        return logits

    def generate(
        self,
        input_ids: torch.LongTensor,
        max_new_tokens: int = 50,
        temperature: float = 0.5,
        top_k: int = 50,
        top_p: float = 0.95,
        eos_token_id: int = None,
        pad_token_id: int = 0,
        **kwargs
    ):
        """
        Naive token-by-token generation loop that uses top-k/top-p filtering and optional temperature.

        Args:
            input_ids (torch.LongTensor): shape (batch_size, sequence_length).
            max_new_tokens (int): max number of tokens to generate (beyond input_ids length).
            temperature (float): sampling temperature (>=0).
            top_k (int): Top-K sampling cutoff.
            top_p (float): Nucleus sampling cutoff.
            eos_token_id (int): If set, stop generation when this token is produced.
            pad_token_id (int): If set, can be used to pad sequences. (Not fully used here.)
            kwargs: Unused arguments (like num_beams) for compatibility.

        Returns:
            torch.LongTensor: shape (batch_size, sequence_length + generated_tokens).
        """
        device = input_ids.device
        #print("1=====================")
        #print(tokenizer.decode(input_ids[0], skip_special_tokens=True))
        #print("1=====================")

        # We'll accumulate new tokens into generated_ids
        generated_ids = input_ids.clone()

        for _ in range(max_new_tokens):
            # Forward pass to get logits for the last token
            outputs = self.forward(generated_ids)
            logits = outputs.logits[:, -1, :]  # shape: (batch_size, vocab_size)

            # Scale logits by temperature
            if temperature != 1.0:
                logits = logits / temperature

            # Filter logits using top-k and/or top-p
            logits = self.top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)

            # Convert to probabilities
            probabilities = F.softmax(logits, dim=-1)

            # Sample from the distribution
            next_token = torch.multinomial(probabilities, num_samples=1)  # (batch_size, 1)

            # Append next token
            generated_ids = torch.cat([generated_ids, next_token], dim=1)

            # If eos_token_id is set and any sample produced it, we optionally could break early
            if eos_token_id is not None:
                # Check if all sequences in the batch ended
                # or if you want to do a more fine-grained approach
                if (next_token == eos_token_id).all():
                    break
        #print("2=====================")
        #print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))
        #print("2=====================")
        return generated_ids