Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2025 the HuggingFace Team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16from typing import Optional17 18import torch19 20from ...cache_utils import Cache21from ..gemma2.configuration_gemma2 import Gemma2Config22from ..gemma2.modeling_gemma2 import Gemma2DecoderLayer, Gemma2ForCausalLM23 24 25class VaultGemmaConfig(Gemma2Config):26    pass27 28 29class VaultGemmaDecoderLayer(Gemma2DecoderLayer):30    def __init__(self, **super_kwargs):31        super().__init__(**super_kwargs)32        del self.post_attention_layernorm33        del self.post_feedforward_layernorm34 35    def forward(36        self,37        hidden_states: torch.Tensor,38        position_embeddings: tuple[torch.Tensor, torch.Tensor],39        attention_mask: Optional[torch.Tensor] = None,40        position_ids: Optional[torch.LongTensor] = None,41        past_key_values: Optional[Cache] = None,42        output_attentions: Optional[bool] = False,43        use_cache: Optional[bool] = False,44        cache_position: Optional[torch.LongTensor] = None,45        **kwargs,46    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:47        residual = hidden_states48        hidden_states = self.input_layernorm(hidden_states)49        # Self Attention50        hidden_states, self_attn_weights = self.self_attn(51            hidden_states=hidden_states,52            position_embeddings=position_embeddings,53            attention_mask=attention_mask,54            position_ids=position_ids,55            past_key_values=past_key_values,56            output_attentions=output_attentions,57            use_cache=use_cache,58            cache_position=cache_position,59            **kwargs,60        )61        hidden_states = residual + hidden_states62 63        residual = hidden_states64        hidden_states = self.pre_feedforward_layernorm(hidden_states)65        hidden_states = self.mlp(hidden_states)66        hidden_states = residual + hidden_states67 68        outputs = (hidden_states,)69        if output_attentions:70            outputs += (self_attn_weights,)71 72        return outputs73 74 75class VaultGemmaForCausalLM(Gemma2ForCausalLM):76    pass77 78 79__all__ = [80    "VaultGemmaConfig",81    "VaultGemmaForCausalLM",82    "VaultGemmaModel",  # noqa: F82283    "VaultGemmaPreTrainedModel",  # noqa: F82284]85