src/hf_model.py

"""
Hugging Face Transformers Integration for Q-TensorFormer.

Provides first-class Hugging Face ecosystem compatibility:
  - QTensorFormerConfig (inherits PretrainedConfig)
  - QTensorFormerForCausalLM (inherits PreTrainedModel, GenerationMixin)
  - AutoConfig & AutoModelForCausalLM registration
  - Generation support with past_key_values and Adaptive KV Cache
  - save_pretrained and from_pretrained serialization
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Tuple, Dict, List, Union
from dataclasses import dataclass

from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers import AutoConfig, AutoModelForCausalLM

from .blocks import HybridBlock
from .kv_cache import AdaptiveKVCache, KVPrecision
from .resource_allocator import AllocationBudget


class QTensorFormerConfig(PretrainedConfig):
    """
    Configuration class for Q-TensorFormer.
    """
    model_type = "qtensorformer"

    def __init__(
        self,
        vocab_size: int = 10000,
        d_model: int = 128,
        n_layers: int = 2,
        n_heads: int = 4,
        n_kv_heads: Optional[int] = None,
        ff_multiplier: int = 4,
        max_seq_len: int = 128,
        dropout: float = 0.1,
        tt_rank: int = 8,
        tt_min_rank: int = 2,
        use_quantum: bool = True,
        n_qubits: int = 4,
        backend_type: str = "classical_surrogate",
        default_preset: str = "balanced",
        hysteresis_tau: float = 0.15,
        tie_word_embeddings: bool = True,
        enable_early_exit: bool = False,
        early_exit_threshold: float = 0.20,
        **kwargs,
    ):
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads
        self.ff_multiplier = ff_multiplier
        self.max_seq_len = max_seq_len
        self.dropout = dropout
        self.tt_rank = tt_rank
        self.tt_min_rank = tt_min_rank
        self.use_quantum = use_quantum
        self.n_qubits = n_qubits
        self.backend_type = backend_type
        self.default_preset = default_preset
        self.hysteresis_tau = hysteresis_tau
        self.enable_early_exit = enable_early_exit
        self.early_exit_threshold = early_exit_threshold


class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 512, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len).float().unsqueeze(1)
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x: torch.Tensor, start_pos: int = 0) -> torch.Tensor:
        seq_len = x.size(1)
        return self.dropout(x + self.pe[:, start_pos:start_pos + seq_len, :])


class QTensorFormerForCausalLM(PreTrainedModel, GenerationMixin):
    """
    Q-TensorFormer Causal Language Model.
    Compatible with Hugging Face generate() and AutoModelForCausalLM.
    """
    config_class = QTensorFormerConfig
    base_model_prefix = "qtensorformer"
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: QTensorFormerConfig):
        super().__init__(config)
        self.config = config

        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.pos_encoding = PositionalEncoding(config.d_model, config.max_seq_len, config.dropout)

        self.blocks = nn.ModuleList([
            HybridBlock(
                d_model=config.d_model,
                n_heads=config.n_heads,
                n_kv_heads=config.n_kv_heads,
                ff_multiplier=config.ff_multiplier,
                tt_rank=config.tt_rank,
                tt_min_rank=config.tt_min_rank,
                use_quantum=config.use_quantum,
                n_qubits=config.n_qubits,
                backend_type=config.backend_type,
                dropout=config.dropout,
                max_seq_len=config.max_seq_len,
                hysteresis_tau=config.hysteresis_tau,
                default_preset=config.default_preset,
                enable_early_exit=config.enable_early_exit,
                early_exit_threshold=config.early_exit_threshold,
                vocab_size=config.vocab_size,
            )
            for _ in range(config.n_layers)
        ])

        self.ln_f = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.lm_head.weight = self.embedding.weight  # Weight tying

        self.post_init()

    def get_input_embeddings(self):
        return self.embedding

    def set_input_embeddings(self, value):
        self.embedding = value

    def get_output_embeddings(self):
        return self.lm_head

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

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.Tensor,
        past_key_values: Optional[List[AdaptiveKVCache]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> Dict:
        # If past_key_values are present, only pass the latest single token
        if past_key_values is not None and past_key_values[0].seq_len > 0:
            input_ids = input_ids[:, -1:]

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

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[AdaptiveKVCache]] = None,
        labels: Optional[torch.Tensor] = None,
        use_cache: bool = True,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        preset: Optional[str] = None,
        budget: Optional[AllocationBudget] = None,
        force_classical: bool = False,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        B, T = input_ids.shape
        start_pos = 0
        if past_key_values is not None and past_key_values[0].seq_len > 0:
            start_pos = past_key_values[0].seq_len

        # Initialize KV caches if requested and not provided
        if use_cache and past_key_values is None:
            past_key_values = [
                AdaptiveKVCache(max_capacity=self.config.max_seq_len)
                for _ in range(self.config.n_layers)
            ]

        x = self.embedding(input_ids)
        x = self.pos_encoding(x, start_pos=start_pos)

        all_stats = []
        early_exit_logits = None
        for i, block in enumerate(self.blocks):
            layer_kv = past_key_values[i] if past_key_values is not None else None
            x, stats = block(
                x,
                mask=attention_mask,
                kv_cache=layer_kv,
                budget=budget,
                preset=preset or self.config.default_preset,
                force_classical=force_classical,
            )
            all_stats.append(stats)
            if stats.get("early_exit_triggered", False) and "early_exit_logits" in stats:
                early_exit_logits = stats["early_exit_logits"]
                break

        if early_exit_logits is not None:
            logits = early_exit_logits
        else:
            x = self.ln_f(x)
            logits = self.lm_head(x)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        if not return_dict:
            output = (logits, past_key_values)
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values,
            hidden_states=None,
            attentions=None,
        )

    def set_preset(self, preset_name: str):
        """Set deployment preset: 'full', 'balanced', 'latency', 'memory', 'energy', 'edge', 'classical_only'."""
        self.config.default_preset = preset_name.lower()


# Register with Hugging Face Auto classes
try:
    AutoConfig.register("qtensorformer", QTensorFormerConfig)
    AutoModelForCausalLM.register(QTensorFormerConfig, QTensorFormerForCausalLM)
except Exception:
    pass