#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/pp_ocrv6_tiny_rec/modular_pp_ocrv6_tiny_rec.py.
#               Do NOT edit this file manually as any edits will be overwritten by the generation of
#             the file from the modular. If any change should be done, please apply the change to the
#                          modular_pp_ocrv6_tiny_rec.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2026 The PaddlePaddle Team and The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import torch.nn as nn
import torch.nn.functional as F

from ...backbone_utils import load_backbone
from ...modeling_outputs import BaseModelOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...processing_utils import Unpack
from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
from .configuration_pp_ocrv6_tiny_rec import PPOCRV6TinyRecConfig


@auto_docstring
class PPOCRV6TinyRecPreTrainedModel(PreTrainedModel):
    config: PPOCRV6TinyRecConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = []
    _skip_keys_device_placement = ["past_key_values"]
    _supports_flash_attn = True
    _supports_sdpa = True
    _supports_flex_attn = True

    _can_compile_fullgraph = True
    _supports_attention_backend = True
    _can_record_outputs = {}
    main_input_name = "pixel_values"
    input_modalities = ("image",)


class PPOCRV6TinyRecHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        in_channels = config.backbone_config.block_configs[-1][-1][2]
        mid_channels = config.hidden_size

        self.conv1 = nn.Conv1d(
            in_channels,
            in_channels,
            kernel_size=5,
            padding=2,
            groups=in_channels,
            bias=False,
        )
        self.norm1 = nn.BatchNorm1d(in_channels)
        self.conv2 = nn.Conv1d(
            in_channels,
            in_channels,
            kernel_size=1,
            bias=False,
        )
        self.norm2 = nn.BatchNorm1d(in_channels)
        self.act_fn = nn.Hardswish()
        self.fc1 = nn.Linear(in_channels, mid_channels)
        self.fc2 = nn.Linear(mid_channels, config.head_out_channels)

    def forward(self, hidden_states: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]):
        hidden_states = hidden_states.squeeze(2)
        hidden_states = self.act_fn(self.norm1(self.conv1(hidden_states)))
        hidden_states = self.act_fn(self.norm2(self.conv2(hidden_states)))

        hidden_states = hidden_states.transpose(1, 2)
        hidden_states = self.fc1(hidden_states)
        hidden_states = self.fc2(hidden_states)

        hidden_states = F.softmax(hidden_states, dim=2, dtype=torch.float32).to(hidden_states.dtype)

        return hidden_states


@auto_docstring(custom_intro="PPOCRV6TinyRec model, consisting of Backbone and Head networks.")
class PPOCRV6TinyRecModel(PPOCRV6TinyRecPreTrainedModel):
    def __init__(self, config: PPOCRV6TinyRecConfig):
        super().__init__(config)
        self.backbone = load_backbone(config)

        self.gradient_checkpointing = False
        self.post_init()

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.FloatTensor,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.FloatTensor] | BaseModelOutputWithNoAttention:
        outputs = self.backbone(pixel_values, **kwargs)
        hidden_state = outputs.feature_maps[-1]
        hidden_state = F.avg_pool2d(hidden_state, (3, 2))

        return BaseModelOutputWithNoAttention(
            last_hidden_state=hidden_state,
            hidden_states=outputs.hidden_states,
        )


@auto_docstring(custom_intro="PPOCR6TinyRec model for text recognition tasks.")
class PPOCRV6TinyRecForTextRecognition(PPOCRV6TinyRecPreTrainedModel):
    _keys_to_ignore_on_load_missing = ["num_batches_tracked"]

    def __init__(self, config: PPOCRV6TinyRecConfig):
        super().__init__(config)
        self.model = PPOCRV6TinyRecModel(config)
        self.head = PPOCRV6TinyRecHead(config)

        self.post_init()

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.FloatTensor,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.FloatTensor] | BaseModelOutputWithNoAttention:
        outputs = self.model(pixel_values, **kwargs)
        head_outputs = self.head(outputs.last_hidden_state, **kwargs)

        return BaseModelOutputWithNoAttention(
            last_hidden_state=head_outputs,
            hidden_states=outputs.hidden_states,
        )


__all__ = ["PPOCRV6TinyRecForTextRecognition", "PPOCRV6TinyRecModel", "PPOCRV6TinyRecPreTrainedModel"]
