Aluode/PerceptionLabPortable

1# Copyright 2025 The HuggingFace 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.14from functools import partial15from typing import Optional16 17import torch18import torch.nn as nn19 20from .cache_utils import Cache21from .modeling_outputs import (22    BaseModelOutputWithPast,23    QuestionAnsweringModelOutput,24    SequenceClassifierOutputWithPast,25    TokenClassifierOutput,26)27from .models.auto import AutoModel28from .processing_utils import Unpack29from .utils import TransformersKwargs, auto_docstring, can_return_tuple, logging30 31 32logger = logging.get_logger(__name__)33 34 35class GradientCheckpointingLayer(nn.Module):36    """Base class for layers with gradient checkpointing.37 38    This class enables gradient checkpointing functionality for a layer. By default, gradient checkpointing is disabled39    (`gradient_checkpointing = False`). When `model.set_gradient_checkpointing()` is called, gradient checkpointing is40    enabled by setting `gradient_checkpointing = True` and assigning a checkpointing function to `_gradient_checkpointing_func`.41 42    Important:43 44        When using gradient checkpointing with `use_reentrant=True`, inputs that require gradients (e.g. hidden states)45        must be passed as positional arguments (`*args`) rather than keyword arguments to properly propagate gradients.46 47        Example:48 49            ```python50            >>> # Correct - hidden_states passed as positional arg51            >>> out = self.layer(hidden_states, attention_mask=attention_mask)52 53            >>> # Incorrect - hidden_states passed as keyword arg54            >>> out = self.layer(hidden_states=hidden_states, attention_mask=attention_mask)55            ```56    """57 58    gradient_checkpointing = False59 60    def __call__(self, *args, **kwargs):61        if self.gradient_checkpointing and self.training:62            do_warn = False63            layer_name = self.__class__.__name__64            message = f"Caching is incompatible with gradient checkpointing in {layer_name}. Setting"65 66            if "use_cache" in kwargs and kwargs["use_cache"]:67                kwargs["use_cache"] = False68                message += " `use_cache=False`,"69                do_warn = True70 71            # different names for the same thing in different layers72            # TODO cyril: this one without `S` can be removed after deprection cycle73            if "past_key_value" in kwargs and kwargs["past_key_value"] is not None:74                kwargs["past_key_value"] = None75                message += " `past_key_value=None`,"76                do_warn = True77 78            if "past_key_values" in kwargs and kwargs["past_key_values"] is not None:79                kwargs["past_key_values"] = None80                message += " `past_key_values=None`,"81                do_warn = True82 83            if "layer_past" in kwargs and kwargs["layer_past"] is not None:84                kwargs["layer_past"] = None85                message += " `layer_past=None`,"86                do_warn = True87 88            # warn if anything was changed89            if do_warn:90                message = message.rstrip(",") + "."91                logger.warning_once(message)92 93            return self._gradient_checkpointing_func(partial(super().__call__, **kwargs), *args)94        return super().__call__(*args, **kwargs)95 96 97@auto_docstring98class GenericForSequenceClassification:99    base_model_prefix = "model"100 101    def __init__(self, config):102        super().__init__(config)103        self.num_labels = config.num_labels104        # Similar to `self.model = AutoModel.from_config(config)` but allows to change the base model name if needed in the child class105        setattr(self, self.base_model_prefix, AutoModel.from_config(config))106        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)107 108        # Initialize weights and apply final processing109        self.post_init()110 111    @can_return_tuple112    @auto_docstring113    def forward(114        self,115        input_ids: Optional[torch.LongTensor] = None,116        attention_mask: Optional[torch.Tensor] = None,117        position_ids: Optional[torch.LongTensor] = None,118        past_key_values: Optional[Cache] = None,119        inputs_embeds: Optional[torch.FloatTensor] = None,120        labels: Optional[torch.LongTensor] = None,121        use_cache: Optional[bool] = None,122        **kwargs: Unpack[TransformersKwargs],123    ) -> SequenceClassifierOutputWithPast:124        transformer_outputs: BaseModelOutputWithPast = getattr(self, self.base_model_prefix)(125            input_ids,126            attention_mask=attention_mask,127            position_ids=position_ids,128            past_key_values=past_key_values,129            inputs_embeds=inputs_embeds,130            use_cache=use_cache,131            **kwargs,132        )133        hidden_states = transformer_outputs.last_hidden_state134        logits = self.score(hidden_states)135 136        if input_ids is not None:137            batch_size = input_ids.shape[0]138        else:139            batch_size = inputs_embeds.shape[0]140 141        if self.config.pad_token_id is None and batch_size != 1:142            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")143        if self.config.pad_token_id is None:144            last_non_pad_token = -1145        elif input_ids is not None:146            # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id147            non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)148            token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)149            last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)150        else:151            last_non_pad_token = -1152            logger.warning_once(153                f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "154                "unexpected if using padding tokens in conjunction with `inputs_embeds.`"155            )156 157        pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]158 159        loss = None160        if labels is not None:161            loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)162 163        return SequenceClassifierOutputWithPast(164            loss=loss,165            logits=pooled_logits,166            past_key_values=transformer_outputs.past_key_values,167            hidden_states=transformer_outputs.hidden_states,168            attentions=transformer_outputs.attentions,169        )170 171 172@auto_docstring173class GenericForQuestionAnswering:174    base_model_prefix = "model"175 176    def __init__(self, config):177        super().__init__(config)178        # Similar to `self.model = AutoModel.from_config(config)` but allows to change the base model name if needed in the child class179        setattr(self, self.base_model_prefix, AutoModel.from_config(config))180        self.qa_outputs = nn.Linear(config.hidden_size, 2)181 182        # Initialize weights and apply final processing183        self.post_init()184 185    def get_input_embeddings(self):186        return getattr(self, self.base_model_prefix).embed_tokens187 188    def set_input_embeddings(self, value):189        getattr(self, self.base_model_prefix).embed_tokens = value190 191    @can_return_tuple192    @auto_docstring193    def forward(194        self,195        input_ids: Optional[torch.LongTensor] = None,196        attention_mask: Optional[torch.Tensor] = None,197        position_ids: Optional[torch.LongTensor] = None,198        past_key_values: Optional[Cache] = None,199        inputs_embeds: Optional[torch.FloatTensor] = None,200        start_positions: Optional[torch.LongTensor] = None,201        end_positions: Optional[torch.LongTensor] = None,202        **kwargs: Unpack[TransformersKwargs],203    ) -> QuestionAnsweringModelOutput:204        outputs: BaseModelOutputWithPast = getattr(self, self.base_model_prefix)(205            input_ids,206            attention_mask=attention_mask,207            position_ids=position_ids,208            past_key_values=past_key_values,209            inputs_embeds=inputs_embeds,210            **kwargs,211        )212 213        sequence_output = outputs.last_hidden_state214 215        logits = self.qa_outputs(sequence_output)216        start_logits, end_logits = logits.split(1, dim=-1)217        start_logits = start_logits.squeeze(-1).contiguous()218        end_logits = end_logits.squeeze(-1).contiguous()219 220        loss = None221        if start_positions is not None and end_positions is not None:222            loss = self.loss_function(start_logits, end_logits, start_positions, end_positions, **kwargs)223 224        return QuestionAnsweringModelOutput(225            loss=loss,226            start_logits=start_logits,227            end_logits=end_logits,228            hidden_states=outputs.hidden_states,229            attentions=outputs.attentions,230        )231 232 233@auto_docstring234class GenericForTokenClassification:235    base_model_prefix = "model"236 237    def __init__(self, config):238        super().__init__(config)239        self.num_labels = config.num_labels240        # Similar to `self.model = AutoModel.from_config(config)` but allows to change the base model name if needed in the child class241        setattr(self, self.base_model_prefix, AutoModel.from_config(config))242        if getattr(config, "classifier_dropout", None) is not None:243            classifier_dropout = config.classifier_dropout244        elif getattr(config, "hidden_dropout", None) is not None:245            classifier_dropout = config.hidden_dropout246        else:247            classifier_dropout = 0.1248        self.dropout = nn.Dropout(classifier_dropout)249        self.score = nn.Linear(config.hidden_size, config.num_labels)250 251        # Initialize weights and apply final processing252        self.post_init()253 254    @can_return_tuple255    @auto_docstring256    def forward(257        self,258        input_ids: Optional[torch.LongTensor] = None,259        attention_mask: Optional[torch.Tensor] = None,260        position_ids: Optional[torch.LongTensor] = None,261        past_key_values: Optional[Cache] = None,262        inputs_embeds: Optional[torch.FloatTensor] = None,263        labels: Optional[torch.LongTensor] = None,264        use_cache: Optional[bool] = None,265        **kwargs: Unpack[TransformersKwargs],266    ) -> TokenClassifierOutput:267        outputs: BaseModelOutputWithPast = getattr(self, self.base_model_prefix)(268            input_ids,269            attention_mask=attention_mask,270            position_ids=position_ids,271            past_key_values=past_key_values,272            inputs_embeds=inputs_embeds,273            use_cache=use_cache,274            **kwargs,275        )276        sequence_output = outputs.last_hidden_state277        sequence_output = self.dropout(sequence_output)278        logits = self.score(sequence_output)279 280        loss = None281        if labels is not None:282            loss = self.loss_function(logits, labels, self.config)283 284        return TokenClassifierOutput(285            loss=loss,286            logits=logits,287            hidden_states=outputs.hidden_states,288            attentions=outputs.attentions,289        )290