modeling_logos.py

"""HuggingFace causal-LM wrapper for Logos."""

from __future__ import annotations

from copy import deepcopy
import importlib
from typing import Any, Dict, Optional

import torch
from torch import nn
from transformers import GenerationConfig, PreTrainedModel
try:
    from transformers.generation import GenerationMixin
except Exception:  # pragma: no cover - older/newer transformers layout guard
    GenerationMixin = object
from transformers.modeling_outputs import CausalLMOutputWithPast

try:
    from .configuration_logos import LogosConfig
except ImportError:
    from configuration_logos import LogosConfig

try:
    from .baseline import BaselineConfig as _BaselineConfig
    from .hybrid import HybridConfig as _HybridConfig
    from .linear import LinearConfig as _LinearConfig
    from .lm_loss import lm_cross_entropy_from_logits as _lm_cross_entropy_from_logits
    from .logos import LogosTransformer
    from .recursive import RecursiveConfig as _RecursiveConfig
    from .residual import ResidualConfig as _ResidualConfig
except ImportError:
    LogosTransformer = importlib.import_module("models.logos").LogosTransformer


_GENERATION_CONFIG_KEYS = set(GenerationConfig().to_dict())
_SAMPLING_ONLY_KEYS = (
    "temperature",
    "top_k",
    "top_p",
    "min_p",
    "typical_p",
    "epsilon_cutoff",
    "eta_cutoff",
)


def _reorder_cache_value(value: Any, beam_idx: torch.LongTensor) -> Any:
    if isinstance(value, torch.Tensor):
        if value.dim() > 0 and value.size(0) == beam_idx.size(0):
            return value.index_select(0, beam_idx.to(value.device))
        return value
    if isinstance(value, dict):
        return {k: _reorder_cache_value(v, beam_idx) for k, v in value.items()}
    if isinstance(value, list):
        return [_reorder_cache_value(v, beam_idx) for v in value]
    if isinstance(value, tuple):
        return tuple(_reorder_cache_value(v, beam_idx) for v in value)
    return value


class LogosForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = LogosConfig
    base_model_prefix = "model"
    _tied_weights_keys = {"model.lm_head.weight": "model.token_emb.weight"}
    supports_gradient_checkpointing = False
    _supports_cache_class = False
    _supports_static_cache = False
    _no_split_modules = [
        "LogosTransformerBlock",
        "HybridAttentionLayer",
        "KimiDeltaAttention",
        "LocalAttention",
        "CompressedGlobalAttention",
        "BlockAttentionResidual",
        "MoELayer",
        "SwiGLU",
    ]

    @classmethod
    def _supports_default_dynamic_cache(cls) -> bool:
        return False

    def __init__(self, config: LogosConfig):
        arch_top_k = config.__dict__.pop("top_k", None)
        try:
            super().__init__(config)
        finally:
            if arch_top_k is not None:
                config.top_k = arch_top_k
        native_config = config.to_native_config()
        self.model = LogosTransformer(native_config)
        self.all_tied_weights_keys = self.get_expanded_tied_weights_keys()

    def get_input_embeddings(self) -> nn.Module:
        return self.model.token_emb

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.model.token_emb = value
        self.model.lm_head.weight = self.model.token_emb.weight

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

    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
        self.model.lm_head = new_embeddings

    def tie_weights(self, *args: Any, **kwargs: Any) -> None:
        self.model.lm_head.weight = self.model.token_emb.weight
        self._reset_rotary_buffers()

    def generate(self, *args: Any, **kwargs: Any):
        generation_config = kwargs.get("generation_config")
        if generation_config is not None:
            explicit_keys = _GENERATION_CONFIG_KEYS.intersection(kwargs)
            if explicit_keys:
                merged = deepcopy(generation_config)
                for key in sorted(explicit_keys):
                    setattr(merged, key, kwargs.pop(key))
                if "max_new_tokens" in explicit_keys and "max_length" not in explicit_keys:
                    merged.max_length = None
                if getattr(merged, "do_sample", None) is False:
                    for key in _SAMPLING_ONLY_KEYS:
                        if key not in explicit_keys and hasattr(merged, key):
                            setattr(merged, key, None)
                kwargs["generation_config"] = merged

        return super().generate(*args, **kwargs)

    def _reset_rotary_buffers(self) -> None:
        for module in self.modules():
            if module.__class__.__name__ != "RotaryEmbedding":
                continue
            device = module.cos.device
            if device.type == "meta":
                device = self.model.token_emb.weight.device
            inv_freq = 1.0 / (
                float(getattr(module, "base", 10000.0)) ** (
                    torch.arange(0, module.head_dim, 2, device=device).float()
                    / module.head_dim
                )
            )
            t = torch.arange(module.max_seq_len, device=device, dtype=torch.float32)
            freqs = torch.einsum("i,j->ij", t, inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1)
            module.cos = emb.cos().unsqueeze(0).unsqueeze(0)
            module.sin = emb.sin().unsqueeze(0).unsqueeze(0)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Dict[str, Any]] = None,
        use_cache: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.Tensor] = None,
        **kwargs: Any,
    ):
        if input_ids is None:
            raise ValueError("LogosForCausalLM requires input_ids")
        if use_cache is None:
            use_cache = bool(getattr(self.config, "use_cache", True))
        if return_dict is None:
            return_dict = bool(getattr(self.config, "return_dict", True))

        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
            is_causal=True,
            past_key_values=past_key_values,
            use_cache=bool(use_cache),
        )

        if not return_dict:
            result = (outputs["logits"],)
            if use_cache:
                result = result + (outputs.get("past_key_values"),)
            if outputs.get("loss") is not None:
                result = (outputs["loss"],) + result
            return result

        return CausalLMOutputWithPast(
            loss=outputs.get("loss"),
            logits=outputs["logits"],
            past_key_values=outputs.get("past_key_values"),
        )

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        past_key_values: Optional[Dict[str, Any]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        cache_position: Optional[torch.Tensor] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        if past_key_values is not None:
            seen_tokens = int(past_key_values.get("seen_tokens", 0) or 0)
            if input_ids.size(1) > seen_tokens:
                input_ids = input_ids[:, seen_tokens:]
            else:
                input_ids = input_ids[:, -1:]

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
            "use_cache": kwargs.get("use_cache", True),
            "cache_position": cache_position,
        }

    @staticmethod
    def _reorder_cache(
        past_key_values: Optional[Dict[str, Any]],
        beam_idx: torch.LongTensor,
    ) -> Optional[Dict[str, Any]]:
        if past_key_values is None:
            return None
        return _reorder_cache_value(past_key_values, beam_idx)


__all__ = ["LogosForCausalLM"]