"""simple docstring"""


from copy import deepcopy
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import Dataset

from .deepspeed import is_deepspeed_zeroa_enabled
from .generation.configuration_utils import GenerationConfig
from .trainer import Trainer
from .utils import logging


if TYPE_CHECKING:
    from .data.data_collator import DataCollator
    from .modeling_utils import PreTrainedModel
    from .tokenization_utils_base import PreTrainedTokenizerBase
    from .trainer_callback import TrainerCallback
    from .trainer_utils import EvalPrediction, PredictionOutput
    from .training_args import TrainingArguments


__UpperCamelCase								=    logging.get_logger(__name__)




class 				UpperCamelCase      (       lowerCAmelCase__     ):







       def __init__(				self,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = (None, None),							lowerCAmelCase__ = None,							)   ->   Dict:
            super().__init__(
                model=lowerCAmelCase__,							args=lowerCAmelCase__,							data_collator=lowerCAmelCase__,							train_dataset=lowerCAmelCase__,							eval_dataset=lowerCAmelCase__,							tokenizer=lowerCAmelCase__,							model_init=lowerCAmelCase__,							compute_metrics=lowerCAmelCase__,							callbacks=lowerCAmelCase__,							optimizers=lowerCAmelCase__,							preprocess_logits_for_metrics=lowerCAmelCase__,							)

            # Override self.model.generation_config if a GenerationConfig is specified in args.
            # Priority: args.generation_config > model.generation_config > default GenerationConfig.
            if self.args.generation_config is not None:
                 snake_case_       =      self.load_generation_config(self.args.generation_config)
                 snake_case_       =      gen_config







       @staticmethod
       def 	a_		(				lowerCAmelCase__)   ->   GenerationConfig:

            # GenerationConfig provided, nothing to do
            if isinstance(lowerCAmelCase__,							lowerCAmelCase__):
                 return deepcopy(lowerCAmelCase__)

            # str or Path
            snake_case_       =      Path(lowerCAmelCase__) if isinstance(lowerCAmelCase__,							lowerCAmelCase__) else gen_config_arg
            snake_case_       =      None

            # Figuring if it is path pointing to a file, pointing to a directory or else a model id or URL
            # This step is required in order to determine config_file_name
            if pretrained_model_name.is_file():
                 snake_case_       =      pretrained_model_name.name
                 snake_case_       =      pretrained_model_name.parent
            # dir path
            elif pretrained_model_name.is_dir():
                 pass
            # model id or URL
            else:
                 snake_case_       =      gen_config_arg

            snake_case_       =      GenerationConfig.from_pretrained(lowerCAmelCase__,							lowerCAmelCase__)
            return gen_config







       def 	a_		(				self,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = "eval",							**lowerCAmelCase__,							)   ->   Dict[str, float]:
            snake_case_       =      gen_kwargs.copy()
            if gen_kwargs.get('max_length') is None and gen_kwargs.get('max_new_tokens') is None:
                 snake_case_       =      self.args.generation_max_length
            snake_case_       =      (
                gen_kwargs['num_beams'] if gen_kwargs.get('num_beams') is not None else self.args.generation_num_beams
            )
            snake_case_       =      gen_kwargs

            return super().evaluate(lowerCAmelCase__,							ignore_keys=lowerCAmelCase__,							metric_key_prefix=lowerCAmelCase__)







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__ = None,							lowerCAmelCase__ = "test",							**lowerCAmelCase__,							)   ->   "PredictionOutput":
            snake_case_       =      gen_kwargs.copy()
            if gen_kwargs.get('max_length') is None and gen_kwargs.get('max_new_tokens') is None:
                 snake_case_       =      self.args.generation_max_length
            snake_case_       =      (
                gen_kwargs['num_beams'] if gen_kwargs.get('num_beams') is not None else self.args.generation_num_beams
            )
            snake_case_       =      gen_kwargs

            return super().predict(lowerCAmelCase__,							ignore_keys=lowerCAmelCase__,							metric_key_prefix=lowerCAmelCase__)







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__ = None,							**lowerCAmelCase__,							)   ->   Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:

            if not self.args.predict_with_generate or prediction_loss_only:
                 return super().prediction_step(
                     lowerCAmelCase__,							lowerCAmelCase__,							prediction_loss_only=lowerCAmelCase__,							ignore_keys=lowerCAmelCase__)

            snake_case_       =      'labels' in inputs
            snake_case_       =      self._prepare_inputs(lowerCAmelCase__)

            # XXX: adapt synced_gpus for fairscale as well
            # Priority (handled in generate):
            # gen_kwargs > model.generation_config > default GenerationConfig()

            if len(lowerCAmelCase__) == 0 and hasattr(self,							'_gen_kwargs'):
                 snake_case_       =      self._gen_kwargs.copy()

            if gen_kwargs.get('max_length') is None and gen_kwargs.get('max_new_tokens') is None:
                 snake_case_       =      self.model.config.max_length
            snake_case_       =      (
                gen_kwargs['num_beams'] if gen_kwargs.get('num_beams') is not None else self.model.config.num_beams
            )
            snake_case_       =      True if is_deepspeed_zeroa_enabled() else False
            snake_case_       =      (
                gen_kwargs['synced_gpus'] if gen_kwargs.get('synced_gpus') is not None else default_synced_gpus
            )

            # If the `decoder_input_ids` was created from `labels`, evict the former, so that the model can freely generate
            # (otherwise, it would continue generating from the padded `decoder_input_ids`)
            if (
                "labels" in inputs
                and "decoder_input_ids" in inputs
                and inputs["labels"].shape == inputs["decoder_input_ids"].shape
            ):
                 snake_case_       =      {k: v for k, v in inputs.items() if k != 'decoder_input_ids'}
            snake_case_       =      self.model.generate(**lowerCAmelCase__,							**lowerCAmelCase__)

            # Temporary hack to ensure the generation config is not initialized for each iteration of the evaluation loop
            # TODO: remove this hack when the legacy code that initializes generation_config from a model config is
            # removed in https://github.com/huggingface/transformers/blob/98d88b23f54e5a23e741833f1e973fdf600cc2c5/src/transformers/generation/utils.py#L1183
            if self.model.generation_config._from_model_config:
                 snake_case_       =      False

            # Retrieves GenerationConfig from model.generation_config
            snake_case_       =      self.model.generation_config
            # in case the batch is shorter than max length, the output should be padded
            if generated_tokens.shape[-1] < gen_config.max_length:
                 snake_case_       =      self._pad_tensors_to_max_len(lowerCAmelCase__,							gen_config.max_length)
            elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1:
                 snake_case_       =      self._pad_tensors_to_max_len(lowerCAmelCase__,							gen_config.max_new_tokens + 1)

            with torch.no_grad():
                 if has_labels:
                      with self.compute_loss_context_manager():
                           snake_case_       =      model(**lowerCAmelCase__)
                      if self.label_smoother is not None:
                           snake_case_       =      self.label_smoother(lowerCAmelCase__,							inputs['labels']).mean().detach()
                      else:
                           snake_case_       =      (outputs['loss'] if isinstance(lowerCAmelCase__,							lowerCAmelCase__) else outputs[0]).mean().detach()
                 else:
                      snake_case_       =      None

            if self.args.prediction_loss_only:
                 return loss, None, None

            if has_labels:
                 snake_case_       =      inputs['labels']
                 if labels.shape[-1] < gen_config.max_length:
                      snake_case_       =      self._pad_tensors_to_max_len(lowerCAmelCase__,							gen_config.max_length)
                 elif gen_config.max_new_tokens is not None and labels.shape[-1] < gen_config.max_new_tokens + 1:
                      snake_case_       =      self._pad_tensors_to_max_len(lowerCAmelCase__,							gen_config.max_new_tokens + 1)
            else:
                 snake_case_       =      None

            return loss, generated_tokens, labels







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__)   ->   Any:
            if self.tokenizer is not None and hasattr(self.tokenizer,							'pad_token_id'):
                 # If PAD token is not defined at least EOS token has to be defined
                 snake_case_       =      (
                     self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id
                 )
            else:
                 if self.model.config.pad_token_id is not None:
                      snake_case_       =      self.model.config.pad_token_id
                 else:
                      raise ValueError('Pad_token_id must be set in the configuration of the model, in order to pad tensors')

            snake_case_       =      pad_token_id * torch.ones(
                (tensor.shape[0], max_length),							dtype=tensor.dtype,							device=tensor.device)
            snake_case_       =      tensor
            return padded_tensor


"""simple docstring"""


import unittest

from transformers import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    SummarizationPipeline,
    TFPreTrainedModel,
    pipeline,
)
from transformers.testing_utils import get_gpu_count, is_pipeline_test, require_tf, require_torch, slow, torch_device
from transformers.tokenization_utils import TruncationStrategy

from .test_pipelines_common import ANY


__UpperCamelCase								=    -1 if torch_device == '''cpu''' else 0




@is_pipeline_test
class 				UpperCamelCase      (       unittest.TestCase     ):
       SCREAMING_SNAKE_CASE_				=   MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
       SCREAMING_SNAKE_CASE_				=   TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__)   ->   int:
            snake_case_       =      SummarizationPipeline(model=lowerCAmelCase__,							tokenizer=lowerCAmelCase__)
            return summarizer, ["(CNN)The Palestinian Authority officially became", "Some other text"]







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__)   ->   Union[str, Any]:
            snake_case_       =      summarizer.model

            snake_case_       =      summarizer('(CNN)The Palestinian Authority officially became')
            self.assertEqual(lowerCAmelCase__,							[{'summary_text': ANY(lowerCAmelCase__)}])

            snake_case_       =      summarizer(
                '(CNN)The Palestinian Authority officially became ',							num_beams=2,							min_length=2,							max_length=5,							)
            self.assertEqual(lowerCAmelCase__,							[{'summary_text': ANY(lowerCAmelCase__)}])

            # Some models (Switch Transformers, LED, T5, LongT5, etc) can handle long sequences.
            snake_case_       =      [
                'SwitchTransformersConfig',
                'T5Config',
                'LongT5Config',
                'LEDConfig',
                'PegasusXConfig',
                'FSMTConfig',
                'M2M100Config',
                'ProphetNetConfig',  # positional embeddings up to a fixed maximum size (otherwise clamping the values)
            ]
            if model.config.__class__.__name__ not in model_can_handle_longer_seq:
                 # Too long and exception is expected.
                 # For TF models, if the weights are initialized in GPU context, we won't get expected index error from
                 # the embedding layer.
                 if not (
                     isinstance(lowerCAmelCase__,							lowerCAmelCase__)
                     and get_gpu_count() > 0
                     and len(summarizer.model.trainable_weights) > 0
                 ):
                      with self.assertRaises(lowerCAmelCase__):
                           snake_case_       =      summarizer('This ' * 1000)
            snake_case_       =      summarizer('This ' * 1000,							truncation=TruncationStrategy.ONLY_FIRST)







       @require_torch
       def 	a_		(				self)   ->   Union[str, Any]:
            snake_case_       =      pipeline(task='summarization',							model='sshleifer/tiny-mbart',							framework='pt')
            snake_case_       =      summarizer('This is a small test')
            self.assertEqual(
                lowerCAmelCase__,							[
                    {
                        'summary_text': 'เข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไป'
                    }
                ],							)







       @require_tf
       def 	a_		(				self)   ->   Any:
            snake_case_       =      pipeline(task='summarization',							model='sshleifer/tiny-mbart',							framework='tf')
            snake_case_       =      summarizer('This is a small test')
            self.assertEqual(
                lowerCAmelCase__,							[
                    {
                        'summary_text': 'เข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไป'
                    }
                ],							)







       @require_torch
       @slow
       def 	a_		(				self)   ->   List[str]:
            snake_case_       =      pipeline(task='summarization',							device=lowerCAmelCase__)
            snake_case_       =      (
                ' (CNN)The Palestinian Authority officially became the 123rd member of the International Criminal Court on'
                ' Wednesday, a step that gives the court jurisdiction over alleged crimes in Palestinian territories. The'
                ' formal accession was marked with a ceremony at The Hague, in the Netherlands, where the court is based.'
                ' The Palestinians signed the ICC\'s founding Rome Statute in January, when they also accepted its'
                ' jurisdiction over alleged crimes committed "in the occupied Palestinian territory, including East'
                ' Jerusalem, since June 13, 2014." Later that month, the ICC opened a preliminary examination into the'
                ' situation in Palestinian territories, paving the way for possible war crimes investigations against'
                ' Israelis. As members of the court, Palestinians may be subject to counter-charges as well. Israel and'
                ' the United States, neither of which is an ICC member, opposed the Palestinians\' efforts to join the'
                ' body. But Palestinian Foreign Minister Riad al-Malki, speaking at Wednesday\'s ceremony, said it was a'
                ' move toward greater justice. "As Palestine formally becomes a State Party to the Rome Statute today, the'
                ' world is also a step closer to ending a long era of impunity and injustice," he said, according to an'
                ' ICC news release. "Indeed, today brings us closer to our shared goals of justice and peace." Judge'
                ' Kuniko Ozaki, a vice president of the ICC, said acceding to the treaty was just the first step for the'
                ' Palestinians. "As the Rome Statute today enters into force for the State of Palestine, Palestine'
                ' acquires all the rights as well as responsibilities that come with being a State Party to the Statute.'
                ' These are substantive commitments, which cannot be taken lightly," she said. Rights group Human Rights'
                ' Watch welcomed the development. "Governments seeking to penalize Palestine for joining the ICC should'
                ' immediately end their pressure, and countries that support universal acceptance of the court\'s treaty'
                ' should speak out to welcome its membership," said Balkees Jarrah, international justice counsel for the'
                ' group. \"What\'s objectionable is the attempts to undermine international justice, not Palestine\'s'
                ' decision to join a treaty to which over 100 countries around the world are members." In January, when'
                ' the preliminary ICC examination was opened, Israeli Prime Minister Benjamin Netanyahu described it as an'
                ' outrage, saying the court was overstepping its boundaries. The United States also said it "strongly"'
                ' disagreed with the court\'s decision. \"As we have said repeatedly, we do not believe that Palestine is a'
                ' state and therefore we do not believe that it is eligible to join the ICC," the State Department said in'
                ' a statement. It urged the warring sides to resolve their differences through direct negotiations. "We'
                ' will continue to oppose actions against Israel at the ICC as counterproductive to the cause of peace,"'
                ' it said. But the ICC begs to differ with the definition of a state for its purposes and refers to the'
                ' territories as "Palestine." While a preliminary examination is not a formal investigation, it allows the'
                ' court to review evidence and determine whether to investigate suspects on both sides. Prosecutor Fatou'
                ' Bensouda said her office would "conduct its analysis in full independence and impartiality." The war'
                ' between Israel and Hamas militants in Gaza last summer left more than 2,000 people dead. The inquiry'
                ' will include alleged war crimes committed since June. The International Criminal Court was set up in'
                ' 2002 to prosecute genocide, crimes against humanity and war crimes. CNN\'s Vasco Cotovio, Kareem Khadder'
                ' and Faith Karimi contributed to this report.'
            )
            snake_case_       =      (
                ' The Palestinian Authority becomes the 123rd member of the International Criminal Court . The move gives'
                ' the court jurisdiction over alleged crimes in Palestinian territories . Israel and the United States'
                ' opposed the Palestinians\' efforts to join the court . Rights group Human Rights Watch welcomes the move,'
                ' says governments seeking to penalize Palestine should end pressure .'
            )
            snake_case_       =      summarizer(lowerCAmelCase__)
            self.assertEqual(result[0]['summary_text'],							lowerCAmelCase__)

