# SPDX-License-Identifier: Apache-2.0
# Copyright 2024 The vLLM team.
"""Wrapper around `transformers` models for vLLM v0.6.2.

This module provides the Transformers modeling backend that wraps
any HuggingFace model with the vLLM interface, enabling support for custom
models that define their implementation via `auto_map` in config.json.

Architecture (following latest vLLM patterns):
- Base: Core functionality (meta init, PP/TP support, module replacement, attention, weight loading)
- CausalMixin: Causal LM specific (lm_head, compute_logits, sample)
- EmbeddingMixin: Embedding/pooling specific (pooler, pooling)
- SequenceClassificationMixin: Classification specific (classifier, pooling)

Composed model classes:
- TransformersForCausalLM = CausalMixin + Base
- TransformersForEmbedding = EmbeddingMixin + Base
- TransformersForSequenceClassification = SequenceClassificationMixin + Base

Key optimizations:
- Meta device initialization for memory efficiency
- Pipeline Parallel support (PPMissingLayer)
- Tensor Parallel support (tp_plan based module replacement)
- Module replacement (Linear, RMSNorm, Embedding) with vLLM optimized versions
- vLLM Attention instances for proper KV cache allocation
- AutoWeightsLoader for efficient weight loading with name mapping
"""

from vllm.model_executor.models.transformers.base import (
    Base,
    set_attention_context,
    clear_attention_context,
    get_attention_context,
    vllm_flash_attention_forward,
)
from vllm.model_executor.models.transformers.causal import CausalMixin
from vllm.model_executor.models.transformers.pooling import (
    EmbeddingMixin,
    SequenceClassificationMixin,
)
from vllm.model_executor.models.transformers.legacy import LegacyMixin
from vllm.model_executor.models.transformers.utils import (
    init_on_device_without_buffers,
    replace_linear_class,
    replace_rms_norm_class,
    log_replacement,
    maybe_prefix,
)


# ============================================================================
# Composed Model Classes (Mixin + Base pattern)
# ============================================================================

class TransformersForCausalLM(CausalMixin, Base):
    """
    Transformers backend wrapper for causal language models.
    
    Combines CausalMixin (lm_head, compute_logits, sample) with
    Base (meta init, PP/TP support, module replacement, attention, weight loading).
    
    Supports any HuggingFace model with auto_map in config.json.
    """
    pass


class TransformersForEmbedding(EmbeddingMixin, Base):
    """
    Transformers backend wrapper for embedding/sentence similarity models.
    
    Combines EmbeddingMixin (pooler, pooling) with
    Base (meta init, PP/TP support, module replacement, attention, weight loading).
    
    Supports embedding models like BERT, sentence-transformers, etc.
    """
    pass


class TransformersForSequenceClassification(SequenceClassificationMixin, Base):
    """
    Transformers backend wrapper for sequence classification models.
    
    Combines SequenceClassificationMixin (classifier, pooling) with
    Base (meta init, PP/TP support, module replacement, attention, weight loading).
    
    Supports cross-encoders and classification models.
    """
    pass


class TransformersForLegacy(LegacyMixin, EmbeddingMixin, Base):
    """
    Transformers backend wrapper for legacy/encoder models.
    
    Combines LegacyMixin (BERT/RoBERTa weight mapping, position handling) with
    EmbeddingMixin (pooler) and Base (core functionality).
    
    Supports BERT, RoBERTa, and similar encoder models.
    """
    pass


__all__ = [
    # Main wrapper classes
    "TransformersForCausalLM",
    "TransformersForEmbedding",
    "TransformersForSequenceClassification",
    "TransformersForLegacy",
    # Base class for extension
    "Base",
    # Mixin classes for custom combinations
    "CausalMixin",
    "EmbeddingMixin",
    "SequenceClassificationMixin",
    "LegacyMixin",
    # Attention context management
    "set_attention_context",
    "clear_attention_context",
    "get_attention_context",
    "vllm_flash_attention_forward",
    # Utility functions
    "init_on_device_without_buffers",
    "replace_linear_class",
    "replace_rms_norm_class",
    "log_replacement",
    "maybe_prefix",
]
