



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 	lowerCamelCase__ (		unittest.TestCase   ):
						_lowerCAmelCase								= MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
						_lowerCAmelCase								= TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING


						def   _lowerCamelCase			( self						:  Optional[int]      ,       _a						:  List[str]      ,       _a						:  Union[str, Any]      ,       _a						:  int			):
													a__:       str											=SummarizationPipeline(model=__UpperCamelCase      ,       tokenizer=__UpperCamelCase			)
													return summarizer, ["(CNN)The Palestinian Authority officially became", "Some other text"]


						def   _lowerCamelCase			( self						:  Union[str, Any]      ,       _a						:  Optional[Any]      ,       _a						:  int			):
													a__:       Tuple											=summarizer.model

													a__:       Optional[Any]											=summarizer("(CNN)The Palestinian Authority officially became"			)
													self.assertEqual(__UpperCamelCase      ,       [{"summary_text": ANY(__UpperCamelCase			)}]			)

													a__:       Optional[int]											=summarizer(
													    "(CNN)The Palestinian Authority officially became "      ,       num_beams=2      ,       min_length=2      ,       max_length=5      ,       )
													self.assertEqual(__UpperCamelCase      ,       [{"summary_text": ANY(__UpperCamelCase			)}]			)

													# Some models (Switch Transformers, LED, T5, LongT5, etc) can handle long sequences.
													a__:       List[str]											=[
													    "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(__UpperCamelCase      ,       __UpperCamelCase			)
																				    and get_gpu_count() > 0
																				    and len(summarizer.model.trainable_weights			) > 0
																				):
																											with self.assertRaises(__UpperCamelCase			):
																																		a__:       Tuple											=summarizer("This " * 1_0_0_0			)
													a__:       List[str]											=summarizer("This " * 1_0_0_0      ,       truncation=TruncationStrategy.ONLY_FIRST			)


						@require_torch
						def   _lowerCamelCase			( self						:  Union[str, Any]			):
													a__:       Optional[int]											=pipeline(task="summarization"      ,       model="sshleifer/tiny-mbart"      ,       framework="pt"			)
													a__:       Dict											=summarizer("This is a small test"			)
													self.assertEqual(
													    __UpperCamelCase      ,       [
													        {
													            "summary_text": "เข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไป"
													        }
													    ]      ,       )


						@require_tf
						def   _lowerCamelCase			( self						:  List[str]			):
													a__:       Optional[int]											=pipeline(task="summarization"      ,       model="sshleifer/tiny-mbart"      ,       framework="tf"			)
													a__:       Optional[Any]											=summarizer("This is a small test"			)
													self.assertEqual(
													    __UpperCamelCase      ,       [
													        {
													            "summary_text": "เข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไป"
													        }
													    ]      ,       )







						@require_torch
						@slow
						def   _lowerCamelCase			( self						:  Union[str, Any]			):
													a__:       Tuple											=pipeline(task="summarization"      ,       device=__UpperCamelCase			)
													a__:       List[str]											=(
													    " (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."
													)
													a__:       Optional[Any]											=(
													    " 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 ."
													)
													a__:       List[str]											=summarizer(__UpperCamelCase			)
													self.assertEqual(result[0]["summary_text"]      ,       __UpperCamelCase			)










# This test is meant to be run in on an instance with TPUs like this:
#
#   python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
#
# Replace 8 with the number of TPU cores you have.
#

import sys
from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.utils import logging


__UpperCAmelCase						  =     logging.get_logger(__name__)


if is_torch_available():
		import torch
		from torch import nn
		from torch.utils.data import Dataset

		from transformers import Trainer


		class 	lowerCamelCase__ (		_a   ):


								def __init__( self						:  Any      ,       _a						:  int = 1_0_1			):
															a__:       Dict											=length


								def __len__( self						:  Optional[Any]			):
															return self.length







								def __getitem__( self						:  Dict      ,       _a						:  int			):
															return i


		class 	lowerCamelCase__ :


								def __call__( self						:  Any      ,       _a						:  List[str]			):
															return {"input_ids": torch.tensor(_a			), "labels": torch.tensor(_a			)}


		class 	lowerCamelCase__ (		nn.Module   ):


								def __init__( self						:  Tuple			):
															super().__init__()
															# Add some (unused) params otherwise DDP will complain.
															a__:       int											=nn.Linear(1_2_0      ,       8_0			)







								def   _lowerCamelCase			( self						:  str      ,       _a						:  List[Any]      ,       _a						:  List[str]=None			):
															if labels is not None:
																						return torch.tensor(0.0      ,       device=input_ids.device			), input_ids
															else:
																						return input_ids





def 						__lowerCamelCase       (      ):
							a__:       List[str]											=HfArgumentParser((TrainingArguments,)						)
							sys.argv += ["--output_dir", "./examples"]
							a__:       int											=parser.parse_args_into_dataclasses()[0]

							logger.warning(
							    F"Process rank: {training_args.local_rank}, device: {training_args.device}, "
							    F"tpu_num_cores: {training_args.tpu_num_cores}"     , )

							# Essentially, what we want to verify in the distributed case is
							# that we get all samples back, in the right order.
							# (this is crucial for prediction for instance)
							for dataset_length in [1_001, 256, 15]:
														a__:       Optional[Any]											=DummyDataset(__magic_name__						)

														def compute_metrics(__magic_name__	:  EvalPrediction						) -> Dict:
																					a__:       List[str]											=list(range(len(__magic_name__						)						)						)
																					a__:       Union[str, Any]											=p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
																					return {"success": success}

														a__:       Union[str, Any]											=Trainer(
														    model=DummyModel()     , args=__magic_name__     , data_collator=DummyDataCollator()     , eval_dataset=__magic_name__     , compute_metrics=__magic_name__     , )
														a__:       str											=trainer.evaluate()
														logger.info(__magic_name__						)
														if metrics["eval_success"] is not True:
																					logger.error(__magic_name__						)
																					exit(1						)

														a__:       List[str]											=trainer.predict(__magic_name__						)
														logger.info(p.metrics						)
														if p.metrics["test_success"] is not True:
																					logger.error(p.metrics						)
																					exit(1						)

														a__:       Tuple											=2

														a__:       Optional[int]											=trainer.evaluate()
														logger.info(__magic_name__						)
														if metrics["eval_success"] is not True:
																					logger.error(__magic_name__						)
																					exit(1						)

														a__:       List[str]											=trainer.predict(__magic_name__						)
														logger.info(p.metrics						)
														if p.metrics["test_success"] is not True:
																					logger.error(p.metrics						)
																					exit(1						)

														a__:       List[str]											=None

							logger.info("🔥 All distributed tests successful"						)





def 						__lowerCamelCase       (      __magic_name__	:  Optional[Any]						):
							# For xla_spawn (TPUs)
							main()


if __name__ == "__main__":
		main()





