"""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"]