codingutf8 2022 the huggingface inc team 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 orglicenseslicense2 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 pytorch donut swin transformer model this implementation is identical to a regular swin transformer without final layer norm on top of the final hidden states import collections abc import math from dataclasses import dataclass from typing import optional tuple union import torch import torch utils checkpoint from torch import nn from activations import act2fn from modelingutils import pretrainedmodel from pytorchutils import findpruneableheadsandindices meshgrid prunelinearlayer from utils import modeloutput addcodesampledocstrings addstartdocstrings addstartdocstringstomodelforward logging from configurationdonutswin import donutswinconfig logger logging getloggername general docstring configfordoc donutswinconfig base docstring checkpointfordoc https huggingface conaverclovaixdonutbase expectedoutputshape 1 49 768 donutswinpretrainedmodelarchivelist naverclovaixdonutbase see all donut swin models at https huggingface comodels filterdonut dataclass copied from transformers models swin modelingswin swinencoderoutput with swindonutswin class donutswinencoderoutputmodeloutput lasthiddenstate torch floattensor none hiddenstates optionaltupletorch floattensor none attentions optionaltupletorch floattensor none reshapedhiddenstates optionaltupletorch floattensor none dataclass copied from transformers models swin modelingswin swinmodeloutput with swindonutswin class donutswinmodeloutputmodeloutput lasthiddenstate torch floattensor none pooleroutput optionaltorch floattensor none hiddenstates optionaltupletorch floattensor none attentions optionaltupletorch floattensor none reshapedhiddenstates optionaltupletorch floattensor none copied from transformers models swin modelingswin windowpartition def windowpartitioninputfeature windowsize batchsize height width numchannels inputfeature shape inputfeature inputfeature view batchsize height windowsize windowsize width windowsize windowsize numchannels windows inputfeature permute0 1 3 2 4 5 contiguous view1 windowsize windowsize numchannels return windows copied from transformers models swin modelingswin windowreverse def windowreversewindows windowsize height width numchannels windows shape1 windows windows view1 height windowsize width windowsize windowsize windowsize numchannels windows windows permute0 1 3 2 4 5 contiguous view1 height width numchannels return windows copied from transformers models swin modelingswin swinembeddings with swindonutswin class donutswinembeddingsnn module def initself config usemasktokenfalse super init self patchembeddings donutswinpatchembeddingsconfig numpatches self patchembeddings numpatches self patchgrid self patchembeddings gridsize self masktoken nn parametertorch zeros1 1 config embeddim if usemasktoken else none if config useabsoluteembeddings self positionembeddings nn parametertorch zeros1 numpatches 1 config embeddim else self positionembeddings none self norm nn layernormconfig embeddim self dropout nn dropoutconfig hiddendropoutprob def forward self pixelvalues optionaltorch floattensor boolmaskedpos optionaltorch booltensor none tupletorch tensor embeddings outputdimensions self patchembeddingspixelvalues embeddings self normembeddings batchsize seqlen embeddings size if boolmaskedpos is not none masktokens self masktoken expandbatchsize seqlen 1 replace the masked visual tokens by masktokens mask boolmaskedpos unsqueeze1 typeasmasktokens embeddings embeddings 1 0 mask masktokens mask if self positionembeddings is not none embeddings embeddings self positionembeddings embeddings self dropoutembeddings return embeddings outputdimensions copied from transformers models swin modelingswin swinpatchembeddings class donutswinpatchembeddingsnn module def initself config super init imagesize patchsize config imagesize config patchsize numchannels hiddensize config numchannels config embeddim imagesize imagesize if isinstanceimagesize collections abc iterable else imagesize imagesize patchsize patchsize if isinstancepatchsize collections abc iterable else patchsize patchsize numpatches imagesize1 patchsize1 imagesize0 patchsize0 self imagesize imagesize self patchsize patchsize self numchannels numchannels self numpatches numpatches self gridsize imagesize0 patchsize0 imagesize1 patchsize1 self projection nn conv2dnumchannels hiddensize kernelsizepatchsize stridepatchsize def maybepadself pixelvalues height width if width self patchsize1 0 padvalues 0 self patchsize1 width self patchsize1 pixelvalues nn functional padpixelvalues padvalues if height self patchsize0 0 padvalues 0 0 0 self patchsize0 height self patchsize0 pixelvalues nn functional padpixelvalues padvalues return pixelvalues def forwardself pixelvalues optionaltorch floattensor tupletorch tensor tupleint numchannels height width pixelvalues shape if numchannels self numchannels raise valueerror make sure that the channel dimension of the pixel values match with the one set in the configuration pad the input to be divisible by self patchsize if needed pixelvalues self maybepadpixelvalues height width embeddings self projectionpixelvalues height width embeddings shape outputdimensions height width embeddings embeddings flatten2 transpose1 2 return embeddings outputdimensions copied from transformers models swin modelingswin swinpatchmerging class donutswinpatchmergingnn module def initself inputresolution tupleint dim int normlayer nn module nn layernorm none super init self inputresolution inputresolution self dim dim self reduction nn linear4 dim 2 dim biasfalse self norm normlayer4 dim def maybepadself inputfeature height width shouldpad height 2 1 or width 2 1 if shouldpad padvalues 0 0 0 width 2 0 height 2 inputfeature nn functional padinputfeature padvalues return inputfeature def forwardself inputfeature torch tensor inputdimensions tupleint int torch tensor height width inputdimensions dim is height width batchsize dim numchannels inputfeature shape inputfeature inputfeature viewbatchsize height width numchannels pad input to be disible by width and height if needed inputfeature self maybepadinputfeature height width batchsize height2 width2 numchannels inputfeature0 inputfeature 0 2 0 2 batchsize height2 width2 numchannels inputfeature1 inputfeature 1 2 0 2 batchsize height2 width2 numchannels inputfeature2 inputfeature 0 2 1 2 batchsize height2 width2 numchannels inputfeature3 inputfeature 1 2 1 2 batchsize height2 width2 4numchannels inputfeature torch catinputfeature0 inputfeature1 inputfeature2 inputfeature3 1 inputfeature inputfeature viewbatchsize 1 4 numchannels batchsize height2width2 4c inputfeature self norminputfeature inputfeature self reductioninputfeature return inputfeature copied from transformers models beit modelingbeit droppath def droppathinput torch tensor dropprob float 0 0 training bool false torch tensor if dropprob 0 0 or not training return input keepprob 1 dropprob shape input shape0 1 input ndim 1 work with diff dim tensors not just 2d convnets randomtensor keepprob torch randshape dtypeinput dtype deviceinput device randomtensor floor binarize output input divkeepprob randomtensor return output copied from transformers models swin modelingswin swindroppath class donutswindroppathnn module copied from transformers models swin modelingswin swinselfattention with swindonutswin get pairwise relative position index for each token inside the window take the dot product between query and key to get the raw attention scores apply the attention mask is precomputed for all layers in donutswinmodel forward function normalize the attention scores to probabilities this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper mask heads if we want to copied from transformers models swin modelingswin swinselfoutput copied from transformers models swin modelingswin swinattention with swindonutswin prune linear layers update hyper params and store pruned heads copied from transformers models swin modelingswin swinintermediate copied from transformers models swin modelingswin swinoutput copied from transformers models swin modelingswin swinlayer with swindonutswin if window size is larger than input resolution we don t partition windows calculate attention mask for swmsa pad hiddenstates to multiples of window size cyclic shift partition windows reverse cyclic shift copied from transformers models swin modelingswin swinstage with swindonutswin patch merging layer copied from transformers models swin modelingswin swinencoder with swindonutswin rearrange b h w c b c h w rearrange b h w c b c h w here we use the original not downsampled height and width rearrange b h w c b c h w copied from transformers models swin modelingswin swinpretrainedmodel with swindonutswin an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights if isinstancemodule nn linear nn conv2d slightly different from the tf version which uses truncatednormal for initialization cf https github compytorchpytorchpull5617 module weight data normalmean0 0 stdself config initializerrange if module bias is not none module bias data zero elif isinstancemodule nn layernorm module bias data zero module weight data fill1 0 swinstartdocstring r this model is a pytorch torch nn modulehttps pytorch orgdocsstablenn htmltorch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config donutswinconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel frompretrained method to load the model weights addstartdocstrings the bare donut swin model transformer outputting raw hiddenstates without any specific head on top swinstartdocstring class donutswinmodeldonutswinpretrainedmodel def initself config addpoolinglayertrue usemasktokenfalse super initconfig self config config self numlayers lenconfig depths self numfeatures intconfig embeddim 2 self numlayers 1 self embeddings donutswinembeddingsconfig usemasktokenusemasktoken self encoder donutswinencoderconfig self embeddings patchgrid self pooler nn adaptiveavgpool1d1 if addpoolinglayer else none initialize weights and apply final processing self postinit def getinputembeddingsself return self embeddings patchembeddings def pruneheadsself headstoprune for layer heads in headstoprune items self encoder layerlayer attention pruneheadsheads addstartdocstringstomodelforwardswininputsdocstring addcodesampledocstrings checkpointcheckpointfordoc outputtypedonutswinmodeloutput configclassconfigfordoc modalityvision expectedoutputexpectedoutputshape def forward self pixelvalues optionaltorch floattensor none boolmaskedpos optionaltorch booltensor none headmask optionaltorch floattensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple donutswinmodeloutput r boolmaskedpos torch booltensor of shape batchsize numpatches boolean masked positions indicates which patches are masked 1 and which aren t 0 prepare head mask if needed 1 0 in headmask indicate we keep the head attentionprobs has shape bsz x nheads x n x n input headmask has shape numheads or numhiddenlayers x numheads and headmask is converted to shape numhiddenlayers x batch x numheads x seqlength x seqlength coding utf 8 2022 the huggingface inc team 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 pytorch donut swin transformer model this implementation is identical to a regular swin transformer without final layer norm on top of the final hidden states general docstring base docstring see all donut swin models at https huggingface co models filter donut copied from transformers models swin modeling_swin swinencoderoutput with swin donutswin donutswin encoder s outputs with potential hidden states and attentions args last_hidden_state torch floattensor of shape batch_size sequence_length hidden_size sequence of hidden states at the output of the last layer of the model hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each stage of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads reshaped_hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size hidden_size height width hidden states of the model at the output of each layer plus the initial embedding outputs reshaped to include the spatial dimensions copied from transformers models swin modeling_swin swinmodeloutput with swin donutswin donutswin model s outputs that also contains a pooling of the last hidden states args last_hidden_state torch floattensor of shape batch_size sequence_length hidden_size sequence of hidden states at the output of the last layer of the model pooler_output torch floattensor of shape batch_size hidden_size optional returned when add_pooling_layer true is passed average pooling of the last layer hidden state hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each stage of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads reshaped_hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size hidden_size height width hidden states of the model at the output of each layer plus the initial embedding outputs reshaped to include the spatial dimensions copied from transformers models swin modeling_swin window_partition partitions the given input into windows copied from transformers models swin modeling_swin window_reverse merges windows to produce higher resolution features copied from transformers models swin modeling_swin swinembeddings with swin donutswin construct the patch and position embeddings optionally also the mask token replace the masked visual tokens by mask_tokens copied from transformers models swin modeling_swin swinpatchembeddings this class turns pixel_values of shape batch_size num_channels height width into the initial hidden_states patch embeddings of shape batch_size seq_length hidden_size to be consumed by a transformer pad the input to be divisible by self patch_size if needed copied from transformers models swin modeling_swin swinpatchmerging patch merging layer args input_resolution tuple int resolution of input feature dim int number of input channels norm_layer nn module optional defaults to nn layernorm normalization layer class dim is height width pad input to be disible by width and height if needed batch_size height 2 width 2 num_channels batch_size height 2 width 2 num_channels batch_size height 2 width 2 num_channels batch_size height 2 width 2 num_channels batch_size height 2 width 2 4 num_channels batch_size height 2 width 2 4 c copied from transformers models beit modeling_beit drop_path drop paths stochastic depth per sample when applied in main path of residual blocks comment by ross wightman this is the same as the dropconnect impl i created for efficientnet etc networks however the original name is misleading as drop connect is a different form of dropout in a separate paper see discussion https github com tensorflow tpu issues 494 issuecomment 532968956 i ve opted for changing the layer and argument names to drop path rather than mix dropconnect as a layer name and use survival rate as the argument work with diff dim tensors not just 2d convnets binarize copied from transformers models swin modeling_swin swindroppath drop paths stochastic depth per sample when applied in main path of residual blocks copied from transformers models swin modeling_swin swinselfattention with swin donutswin get pair wise relative position index for each token inside the window take the dot product between query and key to get the raw attention scores apply the attention mask is precomputed for all layers in donutswinmodel forward function normalize the attention scores to probabilities this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper mask heads if we want to copied from transformers models swin modeling_swin swinselfoutput copied from transformers models swin modeling_swin swinattention with swin donutswin prune linear layers update hyper params and store pruned heads add attentions if we output them copied from transformers models swin modeling_swin swinintermediate copied from transformers models swin modeling_swin swinoutput copied from transformers models swin modeling_swin swinlayer with swin donutswin if window size is larger than input resolution we don t partition windows calculate attention mask for sw msa pad hidden_states to multiples of window size cyclic shift partition windows reverse cyclic shift copied from transformers models swin modeling_swin swinstage with swin donutswin patch merging layer copied from transformers models swin modeling_swin swinencoder with swin donutswin rearrange b h w c b c h w rearrange b h w c b c h w here we use the original not downsampled height and width rearrange b h w c b c h w copied from transformers models swin modeling_swin swinpretrainedmodel with swin donutswin an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights slightly different from the tf version which uses truncated_normal for initialization cf https github com pytorch pytorch pull 5617 this model is a pytorch torch nn module https pytorch org docs stable nn html torch nn module sub class use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config donutswinconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel from_pretrained method to load the model weights args pixel_values torch floattensor of shape batch_size num_channels height width pixel values pixel values can be obtained using autoimageprocessor see donutimageprocessor __call__ for details head_mask torch floattensor of shape num_heads or num_layers num_heads optional mask to nullify selected heads of the self attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple initialize weights and apply final processing prunes heads of the model heads_to_prune dict of layer_num list of heads to prune in this layer see base class pretrainedmodel bool_masked_pos torch booltensor of shape batch_size num_patches boolean masked positions indicates which patches are masked 1 and which aren t 0 prepare head mask if needed 1 0 in head_mask indicate we keep the head attention_probs has shape bsz x n_heads x n x n input head_mask has shape num_heads or num_hidden_layers x num_heads and head_mask is converted to shape num_hidden_layers x batch x num_heads x seq_length x seq_length
import collections.abc
import math
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn

from ...activations import ACT2FN
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import find_pruneable_heads_and_indices, meshgrid, prune_linear_layer
from ...utils import (
    ModelOutput,
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
)
from .configuration_donut_swin import DonutSwinConfig


logger = logging.get_logger(__name__)


_CONFIG_FOR_DOC = "DonutSwinConfig"


_CHECKPOINT_FOR_DOC = "https://huggingface.co/naver-clova-ix/donut-base"
_EXPECTED_OUTPUT_SHAPE = [1, 49, 768]

DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "naver-clova-ix/donut-base",
    
]


@dataclass

class DonutSwinEncoderOutput(ModelOutput):
    

    last_hidden_state: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    reshaped_hidden_states: Optional[Tuple[torch.FloatTensor]] = None


@dataclass

class DonutSwinModelOutput(ModelOutput):
    

    last_hidden_state: torch.FloatTensor = None
    pooler_output: Optional[torch.FloatTensor] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    reshaped_hidden_states: Optional[Tuple[torch.FloatTensor]] = None



def window_partition(input_feature, window_size):
    
    batch_size, height, width, num_channels = input_feature.shape
    input_feature = input_feature.view(
        batch_size, height // window_size, window_size, width // window_size, window_size, num_channels
    )
    windows = input_feature.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels)
    return windows



def window_reverse(windows, window_size, height, width):
    
    num_channels = windows.shape[-1]
    windows = windows.view(-1, height // window_size, width // window_size, window_size, window_size, num_channels)
    windows = windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, height, width, num_channels)
    return windows



class DonutSwinEmbeddings(nn.Module):
    

    def __init__(self, config, use_mask_token=False):
        super().__init__()

        self.patch_embeddings = DonutSwinPatchEmbeddings(config)
        num_patches = self.patch_embeddings.num_patches
        self.patch_grid = self.patch_embeddings.grid_size
        self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim)) if use_mask_token else None

        if config.use_absolute_embeddings:
            self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.embed_dim))
        else:
            self.position_embeddings = None

        self.norm = nn.LayerNorm(config.embed_dim)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(
        self, pixel_values: Optional[torch.FloatTensor], bool_masked_pos: Optional[torch.BoolTensor] = None
    ) -> Tuple[torch.Tensor]:
        embeddings, output_dimensions = self.patch_embeddings(pixel_values)
        embeddings = self.norm(embeddings)
        batch_size, seq_len, _ = embeddings.size()

        if bool_masked_pos is not None:
            mask_tokens = self.mask_token.expand(batch_size, seq_len, -1)
            
            mask = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)
            embeddings = embeddings * (1.0 - mask) + mask_tokens * mask

        if self.position_embeddings is not None:
            embeddings = embeddings + self.position_embeddings

        embeddings = self.dropout(embeddings)

        return embeddings, output_dimensions



class DonutSwinPatchEmbeddings(nn.Module):
    

    def __init__(self, config):
        super().__init__()
        image_size, patch_size = config.image_size, config.patch_size
        num_channels, hidden_size = config.num_channels, config.embed_dim
        image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
        patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
        num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_channels = num_channels
        self.num_patches = num_patches
        self.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])

        self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)

    def maybe_pad(self, pixel_values, height, width):
        if width % self.patch_size[1] != 0:
            pad_values = (0, self.patch_size[1] - width % self.patch_size[1])
            pixel_values = nn.functional.pad(pixel_values, pad_values)
        if height % self.patch_size[0] != 0:
            pad_values = (0, 0, 0, self.patch_size[0] - height % self.patch_size[0])
            pixel_values = nn.functional.pad(pixel_values, pad_values)
        return pixel_values

    def forward(self, pixel_values: Optional[torch.FloatTensor]) -> Tuple[torch.Tensor, Tuple[int]]:
        _, num_channels, height, width = pixel_values.shape
        if num_channels != self.num_channels:
            raise ValueError(
                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
            )
        
        pixel_values = self.maybe_pad(pixel_values, height, width)
        embeddings = self.projection(pixel_values)
        _, _, height, width = embeddings.shape
        output_dimensions = (height, width)
        embeddings = embeddings.flatten(2).transpose(1, 2)

        return embeddings, output_dimensions



class DonutSwinPatchMerging(nn.Module):
    

    def __init__(self, input_resolution: Tuple[int], dim: int, norm_layer: nn.Module = nn.LayerNorm) -> None:
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(4 * dim)

    def maybe_pad(self, input_feature, height, width):
        should_pad = (height % 2 == 1) or (width % 2 == 1)
        if should_pad:
            pad_values = (0, 0, 0, width % 2, 0, height % 2)
            input_feature = nn.functional.pad(input_feature, pad_values)

        return input_feature

    def forward(self, input_feature: torch.Tensor, input_dimensions: Tuple[int, int]) -> torch.Tensor:
        height, width = input_dimensions
        
        batch_size, dim, num_channels = input_feature.shape

        input_feature = input_feature.view(batch_size, height, width, num_channels)
        
        input_feature = self.maybe_pad(input_feature, height, width)
        
        input_feature_0 = input_feature[:, 0::2, 0::2, :]
        
        input_feature_1 = input_feature[:, 1::2, 0::2, :]
        
        input_feature_2 = input_feature[:, 0::2, 1::2, :]
        
        input_feature_3 = input_feature[:, 1::2, 1::2, :]
        
        input_feature = torch.cat([input_feature_0, input_feature_1, input_feature_2, input_feature_3], -1)
        input_feature = input_feature.view(batch_size, -1, 4 * num_channels)  

        input_feature = self.norm(input_feature)
        input_feature = self.reduction(input_feature)

        return input_feature



def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
    
    if drop_prob == 0.0 or not training:
        return input
    keep_prob = 1 - drop_prob
    shape = (input.shape[0],) + (1,) * (input.ndim - 1)  
    random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
    random_tensor.floor_()  
    output = input.div(keep_prob) * random_tensor
    return output



class DonutSwinDropPath(nn.Module):
    

    def __init__(self, drop_prob: Optional[float] = None) -> None:
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return drop_path(hidden_states, self.drop_prob, self.training)

    def extra_repr(self) -> str:
        return "p={}".format(self.drop_prob)



class DonutSwinSelfAttention(nn.Module):
    def __init__(self, config, dim, num_heads, window_size):
        super().__init__()
        if dim % num_heads != 0:
            raise ValueError(
                f"The hidden size ({dim}) is not a multiple of the number of attention heads ({num_heads})"
            )

        self.num_attention_heads = num_heads
        self.attention_head_size = int(dim / num_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size
        self.window_size = (
            window_size if isinstance(window_size, collections.abc.Iterable) else (window_size, window_size)
        )

        self.relative_position_bias_table = nn.Parameter(
            torch.zeros((2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1), num_heads)
        )

        
        coords_h = torch.arange(self.window_size[0])
        coords_w = torch.arange(self.window_size[1])
        coords = torch.stack(meshgrid([coords_h, coords_w], indexing="ij"))
        coords_flatten = torch.flatten(coords, 1)
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
        relative_coords[:, :, 0] += self.window_size[0] - 1
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
        relative_position_index = relative_coords.sum(-1)
        self.register_buffer("relative_position_index", relative_position_index)

        self.query = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
        self.key = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
        self.value = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        batch_size, dim, num_channels = hidden_states.shape
        mixed_query_layer = self.query(hidden_states)

        key_layer = self.transpose_for_scores(self.key(hidden_states))
        value_layer = self.transpose_for_scores(self.value(hidden_states))
        query_layer = self.transpose_for_scores(mixed_query_layer)

        
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)

        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)]
        relative_position_bias = relative_position_bias.view(
            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1
        )

        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
        attention_scores = attention_scores + relative_position_bias.unsqueeze(0)

        if attention_mask is not None:
            
            mask_shape = attention_mask.shape[0]
            attention_scores = attention_scores.view(
                batch_size // mask_shape, mask_shape, self.num_attention_heads, dim, dim
            )
            attention_scores = attention_scores + attention_mask.unsqueeze(1).unsqueeze(0)
            attention_scores = attention_scores.view(-1, self.num_attention_heads, dim, dim)

        
        attention_probs = nn.functional.softmax(attention_scores, dim=-1)

        
        
        attention_probs = self.dropout(attention_probs)

        
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(new_context_layer_shape)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        return outputs



class DonutSwinSelfOutput(nn.Module):
    def __init__(self, config, dim):
        super().__init__()
        self.dense = nn.Linear(dim, dim)
        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)

        return hidden_states



class DonutSwinAttention(nn.Module):
    def __init__(self, config, dim, num_heads, window_size):
        super().__init__()
        self.self = DonutSwinSelfAttention(config, dim, num_heads, window_size)
        self.output = DonutSwinSelfOutput(config, dim)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
        )

        
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        self_outputs = self.self(hidden_states, attention_mask, head_mask, output_attentions)
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[1:]  
        return outputs



class DonutSwinIntermediate(nn.Module):
    def __init__(self, config, dim):
        super().__init__()
        self.dense = nn.Linear(dim, int(config.mlp_ratio * dim))
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states



class DonutSwinOutput(nn.Module):
    def __init__(self, config, dim):
        super().__init__()
        self.dense = nn.Linear(int(config.mlp_ratio * dim), dim)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        return hidden_states



class DonutSwinLayer(nn.Module):
    def __init__(self, config, dim, input_resolution, num_heads, shift_size=0):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.shift_size = shift_size
        self.window_size = config.window_size
        self.input_resolution = input_resolution
        self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps)
        self.attention = DonutSwinAttention(config, dim, num_heads, window_size=self.window_size)
        self.drop_path = DonutSwinDropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity()
        self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps)
        self.intermediate = DonutSwinIntermediate(config, dim)
        self.output = DonutSwinOutput(config, dim)

    def set_shift_and_window_size(self, input_resolution):
        if min(input_resolution) <= self.window_size:
            
            self.shift_size = 0
            self.window_size = min(input_resolution)

    def get_attn_mask(self, height, width, dtype):
        if self.shift_size > 0:
            
            img_mask = torch.zeros((1, height, width, 1), dtype=dtype)
            height_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            width_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            count = 0
            for height_slice in height_slices:
                for width_slice in width_slices:
                    img_mask[:, height_slice, width_slice, :] = count
                    count += 1

            mask_windows = window_partition(img_mask, self.window_size)
            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
        else:
            attn_mask = None
        return attn_mask

    def maybe_pad(self, hidden_states, height, width):
        pad_right = (self.window_size - width % self.window_size) % self.window_size
        pad_bottom = (self.window_size - height % self.window_size) % self.window_size
        pad_values = (0, 0, 0, pad_right, 0, pad_bottom)
        hidden_states = nn.functional.pad(hidden_states, pad_values)
        return hidden_states, pad_values

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_dimensions: Tuple[int, int],
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
        always_partition: Optional[bool] = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if not always_partition:
            self.set_shift_and_window_size(input_dimensions)
        else:
            pass
        height, width = input_dimensions
        batch_size, _, channels = hidden_states.size()
        shortcut = hidden_states

        hidden_states = self.layernorm_before(hidden_states)

        hidden_states = hidden_states.view(batch_size, height, width, channels)

        
        hidden_states, pad_values = self.maybe_pad(hidden_states, height, width)

        _, height_pad, width_pad, _ = hidden_states.shape
        
        if self.shift_size > 0:
            shifted_hidden_states = torch.roll(hidden_states, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
        else:
            shifted_hidden_states = hidden_states

        
        hidden_states_windows = window_partition(shifted_hidden_states, self.window_size)
        hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels)
        attn_mask = self.get_attn_mask(height_pad, width_pad, dtype=hidden_states.dtype)
        if attn_mask is not None:
            attn_mask = attn_mask.to(hidden_states_windows.device)

        attention_outputs = self.attention(
            hidden_states_windows, attn_mask, head_mask, output_attentions=output_attentions
        )

        attention_output = attention_outputs[0]

        attention_windows = attention_output.view(-1, self.window_size, self.window_size, channels)
        shifted_windows = window_reverse(attention_windows, self.window_size, height_pad, width_pad)

        
        if self.shift_size > 0:
            attention_windows = torch.roll(shifted_windows, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
        else:
            attention_windows = shifted_windows

        was_padded = pad_values[3] > 0 or pad_values[5] > 0
        if was_padded:
            attention_windows = attention_windows[:, :height, :width, :].contiguous()

        attention_windows = attention_windows.view(batch_size, height * width, channels)

        hidden_states = shortcut + self.drop_path(attention_windows)

        layer_output = self.layernorm_after(hidden_states)
        layer_output = self.intermediate(layer_output)
        layer_output = hidden_states + self.output(layer_output)

        layer_outputs = (layer_output, attention_outputs[1]) if output_attentions else (layer_output,)
        return layer_outputs



class DonutSwinStage(nn.Module):
    def __init__(self, config, dim, input_resolution, depth, num_heads, drop_path, downsample):
        super().__init__()
        self.config = config
        self.dim = dim
        self.blocks = nn.ModuleList(
            [
                DonutSwinLayer(
                    config=config,
                    dim=dim,
                    input_resolution=input_resolution,
                    num_heads=num_heads,
                    shift_size=0 if (i % 2 == 0) else config.window_size // 2,
                )
                for i in range(depth)
            ]
        )

        
        if downsample is not None:
            self.downsample = downsample(input_resolution, dim=dim, norm_layer=nn.LayerNorm)
        else:
            self.downsample = None

        self.pointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_dimensions: Tuple[int, int],
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
        always_partition: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        height, width = input_dimensions
        for i, layer_module in enumerate(self.blocks):
            layer_head_mask = head_mask[i] if head_mask is not None else None

            layer_outputs = layer_module(
                hidden_states, input_dimensions, layer_head_mask, output_attentions, always_partition
            )

            hidden_states = layer_outputs[0]

        hidden_states_before_downsampling = hidden_states
        if self.downsample is not None:
            height_downsampled, width_downsampled = (height + 1) // 2, (width + 1) // 2
            output_dimensions = (height, width, height_downsampled, width_downsampled)
            hidden_states = self.downsample(hidden_states_before_downsampling, input_dimensions)
        else:
            output_dimensions = (height, width, height, width)

        stage_outputs = (hidden_states, hidden_states_before_downsampling, output_dimensions)

        if output_attentions:
            stage_outputs += layer_outputs[1:]
        return stage_outputs



class DonutSwinEncoder(nn.Module):
    def __init__(self, config, grid_size):
        super().__init__()
        self.num_layers = len(config.depths)
        self.config = config
        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths))]
        self.layers = nn.ModuleList(
            [
                DonutSwinStage(
                    config=config,
                    dim=int(config.embed_dim * 2**i_layer),
                    input_resolution=(grid_size[0] // (2**i_layer), grid_size[1] // (2**i_layer)),
                    depth=config.depths[i_layer],
                    num_heads=config.num_heads[i_layer],
                    drop_path=dpr[sum(config.depths[:i_layer]) : sum(config.depths[: i_layer + 1])],
                    downsample=DonutSwinPatchMerging if (i_layer < self.num_layers - 1) else None,
                )
                for i_layer in range(self.num_layers)
            ]
        )

        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_dimensions: Tuple[int, int],
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
        output_hidden_states: Optional[bool] = False,
        output_hidden_states_before_downsampling: Optional[bool] = False,
        always_partition: Optional[bool] = False,
        return_dict: Optional[bool] = True,
    ) -> Union[Tuple, DonutSwinEncoderOutput]:
        all_hidden_states = () if output_hidden_states else None
        all_reshaped_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None

        if output_hidden_states:
            batch_size, _, hidden_size = hidden_states.shape
            
            reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)
            reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
            all_hidden_states += (hidden_states,)
            all_reshaped_hidden_states += (reshaped_hidden_state,)

        for i, layer_module in enumerate(self.layers):
            layer_head_mask = head_mask[i] if head_mask is not None else None

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    layer_module.__call__, hidden_states, input_dimensions, layer_head_mask, output_attentions
                )
            else:
                layer_outputs = layer_module(
                    hidden_states, input_dimensions, layer_head_mask, output_attentions, always_partition
                )

            hidden_states = layer_outputs[0]
            hidden_states_before_downsampling = layer_outputs[1]
            output_dimensions = layer_outputs[2]

            input_dimensions = (output_dimensions[-2], output_dimensions[-1])

            if output_hidden_states and output_hidden_states_before_downsampling:
                batch_size, _, hidden_size = hidden_states_before_downsampling.shape
                
                
                reshaped_hidden_state = hidden_states_before_downsampling.view(
                    batch_size, *(output_dimensions[0], output_dimensions[1]), hidden_size
                )
                reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
                all_hidden_states += (hidden_states_before_downsampling,)
                all_reshaped_hidden_states += (reshaped_hidden_state,)
            elif output_hidden_states and not output_hidden_states_before_downsampling:
                batch_size, _, hidden_size = hidden_states.shape
                
                reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)
                reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
                all_hidden_states += (hidden_states,)
                all_reshaped_hidden_states += (reshaped_hidden_state,)

            if output_attentions:
                all_self_attentions += layer_outputs[3:]

        if not return_dict:
            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)

        return DonutSwinEncoderOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            reshaped_hidden_states=all_reshaped_hidden_states,
        )



class DonutSwinPreTrainedModel(PreTrainedModel):
    

    config_class = DonutSwinConfig
    base_model_prefix = "swin"
    main_input_name = "pixel_values"
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            
            
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


SWIN_START_DOCSTRING = r

SWIN_INPUTS_DOCSTRING = r


@add_start_docstrings(
    "The bare Donut Swin Model transformer outputting raw hidden-states without any specific head on top.",
    SWIN_START_DOCSTRING,
)
class DonutSwinModel(DonutSwinPreTrainedModel):
    def __init__(self, config, add_pooling_layer=True, use_mask_token=False):
        super().__init__(config)
        self.config = config
        self.num_layers = len(config.depths)
        self.num_features = int(config.embed_dim * 2 ** (self.num_layers - 1))

        self.embeddings = DonutSwinEmbeddings(config, use_mask_token=use_mask_token)
        self.encoder = DonutSwinEncoder(config, self.embeddings.patch_grid)

        self.pooler = nn.AdaptiveAvgPool1d(1) if add_pooling_layer else None

        
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.patch_embeddings

    def _prune_heads(self, heads_to_prune):
        
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    @add_start_docstrings_to_model_forward(SWIN_INPUTS_DOCSTRING)
    @add_code_sample_docstrings(
        checkpoint=_CHECKPOINT_FOR_DOC,
        output_type=DonutSwinModelOutput,
        config_class=_CONFIG_FOR_DOC,
        modality="vision",
        expected_output=_EXPECTED_OUTPUT_SHAPE,
    )
    def forward(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        bool_masked_pos: Optional[torch.BoolTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, DonutSwinModelOutput]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

        
        
        
        
        
        head_mask = self.get_head_mask(head_mask, len(self.config.depths))

        embedding_output, input_dimensions = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)

        encoder_outputs = self.encoder(
            embedding_output,
            input_dimensions,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output = encoder_outputs[0]

        pooled_output = None
        if self.pooler is not None:
            pooled_output = self.pooler(sequence_output.transpose(1, 2))
            pooled_output = torch.flatten(pooled_output, 1)

        if not return_dict:
            output = (sequence_output, pooled_output) + encoder_outputs[1:]

            return output

        return DonutSwinModelOutput(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            reshaped_hidden_states=encoder_outputs.reshaped_hidden_states,
        )