Aluode/PerceptionLabPortable

1# Copyright (c) 2025 Baidu, Inc. and HuggingFace Inc. team. All Rights Reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14"""PyTorch Ernie 4.5 model"""15 16import torch17from torch import nn18 19from ...modeling_rope_utils import dynamic_rope_update20from ...utils import auto_docstring, can_return_tuple21from ..glm.modeling_glm import rotate_half22from ..llama.modeling_llama import (23    LlamaAttention,24    LlamaForCausalLM,25    LlamaMLP,26    LlamaRotaryEmbedding,27)28from .configuration_ernie4_5 import Ernie4_5Config29 30 31class Ernie4_5RotaryEmbedding(LlamaRotaryEmbedding):32    @torch.no_grad()33    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)34    def forward(self, x, position_ids):35        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)36        position_ids_expanded = position_ids[:, None, :].float()37 38        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"39        with torch.autocast(device_type=device_type, enabled=False):  # Force float3240            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)41            emb = torch.cat((freqs, freqs), dim=-1)42            cos = emb.cos() * self.attention_scaling43            sin = emb.sin() * self.attention_scaling44 45        # keeping it in full precision46        return cos, sin47 48 49def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):50    """Applies Rotary Position Embedding to the query and key tensors.51 52    Args:53        q (`torch.Tensor`): The query tensor.54        k (`torch.Tensor`): The key tensor.55        cos (`torch.Tensor`): The cosine part of the rotary embedding.56        sin (`torch.Tensor`): The sine part of the rotary embedding.57        position_ids (`torch.Tensor`, *optional*):58            Deprecated and unused.59        unsqueeze_dim (`int`, *optional*, defaults to 1):60            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and61            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note62            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and63            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes64            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have65            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.66    Returns:67        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.68    """69    # glm rope style (with full dim) and full precision70    original_dtype = q.dtype71 72    cos = cos.unsqueeze(unsqueeze_dim)73    sin = sin.unsqueeze(unsqueeze_dim)74 75    # Interleave them instead of usual shape76    cos = cos[..., : cos.shape[-1] // 2].repeat_interleave(2, dim=-1)77    sin = sin[..., : sin.shape[-1] // 2].repeat_interleave(2, dim=-1)78 79    q_embed = (q.float() * cos) + (rotate_half(q).float() * sin)80    k_embed = (k.float() * cos) + (rotate_half(k).float() * sin)81 82    return q_embed.to(original_dtype), k_embed.to(original_dtype)83 84 85class Ernie4_5MLP(LlamaMLP):86    def __init__(self, config: Ernie4_5Config):87        super().__init__(config)88 89        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.use_bias)90        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.use_bias)91        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)92 93 94class Ernie4_5Attention(LlamaAttention):95    def __init__(self, config: Ernie4_5Config, layer_idx: int):96        super().__init__(config, layer_idx)97 98        self.attention_dropout = 0.099 100        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.use_bias)101        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.use_bias)102        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.use_bias)103        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.use_bias)104 105 106class Ernie4_5ForCausalLM(LlamaForCausalLM):107    @can_return_tuple108    @auto_docstring109    def forward(self, **super_kwargs):110        r"""111        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):112            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,113            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored114            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.115        """116        super().forward(**super_kwargs)117 118 119__all__ = [120    "Ernie4_5ForCausalLM",121    "Ernie4_5Model",  # noqa: F822122    "Ernie4_5PreTrainedModel",  # noqa: F822123]124