


'''simple docstring'''





from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization import (
    Adafactor,
    AdamW,
    get_constant_schedule,
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
       from fairscale.optim import OSS


__UpperCamelCase				:  Optional[int]      					=							logging.get_logger(__name__)

__UpperCamelCase				:  Tuple      					=							{
    'linear': get_linear_schedule_with_warmup,
    'cosine': get_cosine_schedule_with_warmup,
    'cosine_w_restarts': get_cosine_with_hard_restarts_schedule_with_warmup,
    'polynomial': get_polynomial_decay_schedule_with_warmup,
    'constant': get_constant_schedule,
    'constant_w_warmup': get_constant_schedule_with_warmup,
}





class  lowerCAmelCase__(    __lowercase						):






    '''simple docstring'''





    def __init__(	self				:   Union[str, Any]       ,   __snake_case				:   List[Any]=None       ,   __snake_case				:   Any=None       ,   *__snake_case				:   Dict       ,   **__snake_case				:   int	):





       '''simple docstring'''
       super().__init__(*__snake_case       ,   **__snake_case	)

       if config is None:
          assert isinstance(self.model       ,   __snake_case	), (
              "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
              f''' {self.model.__class__}'''
          )
          UpperCAmelCase_							:      Optional[Any]						=   self.model.config
       else:
          UpperCAmelCase_							:      Tuple						=   config

       UpperCAmelCase_							:      int						=   data_args
       UpperCAmelCase_							:      Union[str, Any]						=   self.config.tgt_vocab_size if isinstance(self.config       ,   __snake_case	) else self.config.vocab_size

       if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
          assert self.config.pad_token_id is not None, (
              "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
              " calculation or doing label smoothing."
          )

       if self.config.pad_token_id is None and self.config.eos_token_id is not None:
          logger.warning(
              f'''The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for'''
              ''' padding..'''	)

       if self.args.label_smoothing == 0:
          UpperCAmelCase_							:      Dict						=   torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id	)
       else:
          # dynamically import label_smoothed_nll_loss
          from utils import label_smoothed_nll_loss

          UpperCAmelCase_							:      List[str]						=   label_smoothed_nll_loss





    def 					_lowerCamelCase (	self				:   List[str]       ,   __snake_case				:   int	):





       '''simple docstring'''
       if self.optimizer is None:
          UpperCAmelCase_							:      str						=   ["bias", "LayerNorm.weight"]
          UpperCAmelCase_							:      Any						=   [
              {
                  "params": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay	)],
                  "weight_decay": self.args.weight_decay,
              },
              {
                  "params": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay	)],
                  "weight_decay": 0.0,
              },
          ]
          UpperCAmelCase_							:      Dict						=   Adafactor if self.args.adafactor else AdamW
          if self.args.adafactor:
             UpperCAmelCase_							:      Dict						=   Adafactor
             UpperCAmelCase_							:      Tuple						=   {"scale_parameter": False, "relative_step": False}
          else:
             UpperCAmelCase_							:      Union[str, Any]						=   AdamW
             UpperCAmelCase_							:      List[Any]						=   {
                 "betas": (self.args.adam_betaa, self.args.adam_betaa),
                 "eps": self.args.adam_epsilon,
             }
          UpperCAmelCase_							:      Any						=   self.args.learning_rate
          if self.sharded_ddp:
             UpperCAmelCase_							:      Dict						=   OSS(
                 params=__snake_case       ,   optim=__snake_case       ,   **__snake_case       ,   )
          else:
             UpperCAmelCase_							:      Tuple						=   optimizer_cls(__snake_case       ,   **__snake_case	)

       if self.lr_scheduler is None:
          UpperCAmelCase_							:      int						=   self._get_lr_scheduler(__snake_case	)
       else:  # ignoring --lr_scheduler
          logger.warning('''scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'''	)





    def 					_lowerCamelCase (	self				:   Optional[Any]       ,   __snake_case				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      str						=   arg_to_scheduler[self.args.lr_scheduler]
       if self.args.lr_scheduler == "constant":
          UpperCAmelCase_							:      Dict						=   schedule_func(self.optimizer	)
       elif self.args.lr_scheduler == "constant_w_warmup":
          UpperCAmelCase_							:      Union[str, Any]						=   schedule_func(self.optimizer       ,   num_warmup_steps=self.args.warmup_steps	)
       else:
          UpperCAmelCase_							:      Tuple						=   schedule_func(
              self.optimizer       ,   num_warmup_steps=self.args.warmup_steps       ,   num_training_steps=__snake_case	)
       return scheduler





    def 					_lowerCamelCase (	self				:   Tuple	):





       '''simple docstring'''
       if isinstance(self.train_dataset       ,   torch.utils.data.IterableDataset	):
          return None
       elif is_torch_tpu_available():
          return get_tpu_sampler(self.train_dataset	)
       else:
          if self.args.sortish_sampler:
             self.train_dataset.make_sortish_sampler(
                 self.args.per_device_train_batch_size       ,   distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED)       ,   )

          return (
              RandomSampler(self.train_dataset	)
              if self.args.local_rank == -1
              else DistributedSampler(self.train_dataset	)
          )





    def 					_lowerCamelCase (	self				:   Dict       ,   __snake_case				:   Optional[int]       ,   __snake_case				:   int       ,   __snake_case				:   Dict	):





       '''simple docstring'''
       if self.args.label_smoothing == 0:
          if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
             # force training to ignore pad token
             UpperCAmelCase_							:      Any						=   model(**__snake_case       ,   use_cache=__snake_case	)[0]
             UpperCAmelCase_							:      Optional[Any]						=   self.loss_fn(logits.view(-1       ,   logits.shape[-1]	)       ,   labels.view(-1	)	)
          else:
             # compute usual loss via models
             UpperCAmelCase_							:      Optional[Any]						=   model(**__snake_case       ,   labels=__snake_case       ,   use_cache=__snake_case	)[:2]
       else:
          # compute label smoothed loss
          UpperCAmelCase_							:      List[Any]						=   model(**__snake_case       ,   use_cache=__snake_case	)[0]
          UpperCAmelCase_							:      List[Any]						=   torch.nn.functional.log_softmax(__snake_case       ,   dim=-1	)
          UpperCAmelCase_							:      str						=   self.loss_fn(__snake_case       ,   __snake_case       ,   self.args.label_smoothing       ,   ignore_index=self.config.pad_token_id	)
       return loss, logits





    def 					_lowerCamelCase (	self				:   Any       ,   __snake_case				:   str       ,   __snake_case				:   Optional[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Dict						=   inputs.pop('''labels'''	)
       UpperCAmelCase_							:      Any						=   self._compute_loss(__snake_case       ,   __snake_case       ,   __snake_case	)
       return loss





    def 					_lowerCamelCase (	self				:   Optional[Any]       ,   __snake_case				:   nn.Module       ,   __snake_case				:   Dict[str, Union[torch.Tensor, Any]]       ,   __snake_case				:   bool       ,   __snake_case				:   Optional[List[str]] = None       ,   ):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   self._prepare_inputs(__snake_case	)

       UpperCAmelCase_							:      str						=   {
           "max_length": self.data_args.val_max_target_length
           if self.data_args is not None
           else self.config.max_length,
           "num_beams": self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
       }

       if self.args.predict_with_generate and not self.args.prediction_loss_only:
          UpperCAmelCase_							:      str						=   self.model.generate(
              inputs['''input_ids''']       ,   attention_mask=inputs['''attention_mask''']       ,   **__snake_case       ,   )
          # in case the batch is shorter than max length, the output should be padded
          if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
             UpperCAmelCase_							:      Dict						=   self._pad_tensors_to_max_len(__snake_case       ,   gen_kwargs['''max_length''']	)

       UpperCAmelCase_							:      Any						=   inputs.pop('''labels'''	)
       with torch.no_grad():
          # compute loss on predict data
          UpperCAmelCase_							:      List[str]						=   self._compute_loss(__snake_case       ,   __snake_case       ,   __snake_case	)

       UpperCAmelCase_							:      Tuple						=   loss.mean().detach()
       if self.args.prediction_loss_only:
          return (loss, None, None)

       UpperCAmelCase_							:      Optional[int]						=   generated_tokens if self.args.predict_with_generate else logits

       if labels.shape[-1] < gen_kwargs["max_length"]:
          UpperCAmelCase_							:      str						=   self._pad_tensors_to_max_len(__snake_case       ,   gen_kwargs['''max_length''']	)

       return (loss, logits, labels)





    def 					_lowerCamelCase (	self				:   List[str]       ,   __snake_case				:   Tuple       ,   __snake_case				:   int	):





       '''simple docstring'''
       # If PAD token is not defined at least EOS token has to be defined
       UpperCAmelCase_							:      Optional[int]						=   self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

       if pad_token_id is None:
          raise ValueError(
              '''Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be'''
              f''' padded to `max_length`={max_length}'''	)

       UpperCAmelCase_							:      Union[str, Any]						=   pad_token_id * torch.ones(
           (tensor.shape[0], max_length)       ,   dtype=tensor.dtype       ,   device=tensor.device	)
       UpperCAmelCase_							:      int						=   tensor
       return padded_tensor



import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     str			= ['image_processor', 'tokenizer']
    A_  :     int			= 'LayoutLMv2ImageProcessor'
    A_  :     str			= ('LayoutXLMTokenizer', 'LayoutXLMTokenizerFast')





    def __init__(	self				:   Union[str, Any]       ,   __snake_case				:   List[Any]=None       ,   __snake_case				:   List[str]=None       ,   **__snake_case				:   Optional[int]	):





       '''simple docstring'''
       if "feature_extractor" in kwargs:
          warnings.warn(
              '''The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'''
              ''' instead.'''       ,   __snake_case       ,   )
          UpperCAmelCase_							:      List[Any]						=   kwargs.pop('''feature_extractor'''	)

       UpperCAmelCase_							:      Optional[int]						=   image_processor if image_processor is not None else feature_extractor
       if image_processor is None:
          raise ValueError('''You need to specify an `image_processor`.'''	)
       if tokenizer is None:
          raise ValueError('''You need to specify a `tokenizer`.'''	)

       super().__init__(__snake_case       ,   __snake_case	)





    def __call__(	self				:   List[str]       ,   __snake_case				:   Dict       ,   __snake_case				:   Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None       ,   __snake_case				:   Optional[Union[PreTokenizedInput, List[PreTokenizedInput]]] = None       ,   __snake_case				:   Union[List[List[int]], List[List[List[int]]]] = None       ,   __snake_case				:   Optional[Union[List[int], List[List[int]]]] = None       ,   __snake_case				:   bool = True       ,   __snake_case				:   Union[bool, str, PaddingStrategy] = False       ,   __snake_case				:   Union[bool, str, TruncationStrategy] = None       ,   __snake_case				:   Optional[int] = None       ,   __snake_case				:   int = 0       ,   __snake_case				:   Optional[int] = None       ,   __snake_case				:   Optional[bool] = None       ,   __snake_case				:   Optional[bool] = None       ,   __snake_case				:   bool = False       ,   __snake_case				:   bool = False       ,   __snake_case				:   bool = False       ,   __snake_case				:   bool = False       ,   __snake_case				:   bool = True       ,   __snake_case				:   Optional[Union[str, TensorType]] = None       ,   **__snake_case				:   Optional[int]       ,   ):





       '''simple docstring'''
       # verify input
       if self.image_processor.apply_ocr and (boxes is not None):
          raise ValueError(
              '''You cannot provide bounding boxes '''
              '''if you initialized the image processor with apply_ocr set to True.'''	)

       if self.image_processor.apply_ocr and (word_labels is not None):
          raise ValueError(
              '''You cannot provide word labels if you initialized the image processor with apply_ocr set to True.'''	)

       if return_overflowing_tokens is True and return_offsets_mapping is False:
          raise ValueError('''You cannot return overflowing tokens without returning the offsets mapping.'''	)

       # first, apply the image processor
       UpperCAmelCase_							:      Tuple						=   self.image_processor(images=__snake_case       ,   return_tensors=__snake_case	)

       # second, apply the tokenizer
       if text is not None and self.image_processor.apply_ocr and text_pair is None:
          if isinstance(__snake_case       ,   __snake_case	):
             UpperCAmelCase_							:      Optional[Any]						=   [text]  # add batch dimension (as the image processor always adds a batch dimension)
          UpperCAmelCase_							:      Any						=   features['''words''']

       UpperCAmelCase_							:      str						=   self.tokenizer(
           text=text if text is not None else features['''words''']       ,   text_pair=text_pair if text_pair is not None else None       ,   boxes=boxes if boxes is not None else features['''boxes''']       ,   word_labels=__snake_case       ,   add_special_tokens=__snake_case       ,   padding=__snake_case       ,   truncation=__snake_case       ,   max_length=__snake_case       ,   stride=__snake_case       ,   pad_to_multiple_of=__snake_case       ,   return_token_type_ids=__snake_case       ,   return_attention_mask=__snake_case       ,   return_overflowing_tokens=__snake_case       ,   return_special_tokens_mask=__snake_case       ,   return_offsets_mapping=__snake_case       ,   return_length=__snake_case       ,   verbose=__snake_case       ,   return_tensors=__snake_case       ,   **__snake_case       ,   )

       # add pixel values
       UpperCAmelCase_							:      List[str]						=   features.pop('''pixel_values'''	)
       if return_overflowing_tokens is True:
          UpperCAmelCase_							:      Optional[int]						=   self.get_overflowing_images(__snake_case       ,   encoded_inputs['''overflow_to_sample_mapping''']	)
       UpperCAmelCase_							:      List[Any]						=   images

       return encoded_inputs





    def 					_lowerCamelCase (	self				:   Union[str, Any]       ,   __snake_case				:   str       ,   __snake_case				:   List[Any]	):





       '''simple docstring'''
       # in case there's an overflow, ensure each `input_ids` sample is mapped to its corresponding image
       UpperCAmelCase_							:      List[str]						=   []
       for sample_idx in overflow_to_sample_mapping:
          images_with_overflow.append(images[sample_idx]	)

       if len(__snake_case	) != len(__snake_case	):
          raise ValueError(
              '''Expected length of images to be the same as the length of `overflow_to_sample_mapping`, but got'''
              f''' {len(__snake_case	)} and {len(__snake_case	)}'''	)

       return images_with_overflow





    def 					_lowerCamelCase (	self				:   List[Any]       ,   *__snake_case				:   Optional[int]       ,   **__snake_case				:   List[Any]	):





       '''simple docstring'''
       return self.tokenizer.batch_decode(*__snake_case       ,   **__snake_case	)





    def 					_lowerCamelCase (	self				:   str       ,   *__snake_case				:   Optional[Any]       ,   **__snake_case				:   Union[str, Any]	):





       '''simple docstring'''
       return self.tokenizer.decode(*__snake_case       ,   **__snake_case	)





    @property
    def 					_lowerCamelCase (	self				:   Optional[Any]	):





       '''simple docstring'''
       return ["input_ids", "bbox", "attention_mask", "image"]





    @property
    def 					_lowerCamelCase (	self				:   List[str]	):





       '''simple docstring'''
       warnings.warn(
           '''`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'''       ,   __snake_case       ,   )
       return self.image_processor_class





    @property
    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       warnings.warn(
           '''`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'''       ,   __snake_case       ,   )
       return self.image_processor