






import unittest

from transformers import SpeechaTextaConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers.models.speech_to_text_a.modeling_speech_to_text_a import (
							    SpeechaTextaDecoder,
							    SpeechaTextaForCausalLM,
							)


@require_torch
class 					__SCREAMING_SNAKE_CASE:
			def __init__(       self: str    ,		UpperCamelCase: str    ,		UpperCamelCase: List[str]=99    ,		UpperCamelCase: Any=13    ,		UpperCamelCase: List[str]=16    ,		UpperCamelCase: Optional[Any]=7    ,		UpperCamelCase: str=True    ,		UpperCamelCase: Optional[Any]=True    ,		UpperCamelCase: Union[str, Any]=True    ,		UpperCamelCase: Optional[int]=False    ,		UpperCamelCase: List[str]=True    ,		UpperCamelCase: Any=2    ,		UpperCamelCase: Optional[int]=32    ,		UpperCamelCase: str=4    ,		UpperCamelCase: Any=4    ,		UpperCamelCase: Dict=30    ,		UpperCamelCase: Tuple=0    ,		UpperCamelCase: List[Any]=1    ,		UpperCamelCase: Optional[Any]=2    ,		UpperCamelCase: Optional[Any]=None    ,		)						->						List[Any]:
							snake_case__													=       parent
							snake_case__													=       batch_size
							snake_case__													=       decoder_seq_length
							# For common tests
							snake_case__													=       self.decoder_seq_length
							snake_case__													=       is_training
							snake_case__													=       use_attention_mask
							snake_case__													=       use_labels

							snake_case__													=       vocab_size
							snake_case__													=       d_model
							snake_case__													=       d_model
							snake_case__													=       decoder_layers
							snake_case__													=       decoder_layers
							snake_case__													=       decoder_ffn_dim
							snake_case__													=       decoder_attention_heads
							snake_case__													=       decoder_attention_heads
							snake_case__													=       eos_token_id
							snake_case__													=       bos_token_id
							snake_case__													=       pad_token_id
							snake_case__													=       decoder_start_token_id
							snake_case__													=       use_cache
							snake_case__													=       max_position_embeddings

							snake_case__													=       None
							snake_case__													=       decoder_seq_length
							snake_case__													=       2
							snake_case__													=       1
			def lowerCAmelCase_	(       self: List[str]  )						->						Optional[int]:
							snake_case__													=       ids_tensor([self.batch_size, self.decoder_seq_length]    ,		self.vocab_size  )

							snake_case__													=       None
							if self.use_attention_mask:
											snake_case__													=       ids_tensor([self.batch_size, self.decoder_seq_length]    ,		vocab_size=2  )

							snake_case__													=       None
							if self.use_labels:
											snake_case__													=       ids_tensor([self.batch_size, self.decoder_seq_length]    ,		self.vocab_size  )

							snake_case__													=       SpeechaTextaConfig(
							    vocab_size=self.vocab_size    ,		d_model=self.d_model    ,		decoder_layers=self.decoder_layers    ,		decoder_ffn_dim=self.decoder_ffn_dim    ,		decoder_attention_heads=self.decoder_attention_heads    ,		eos_token_id=self.eos_token_id    ,		bos_token_id=self.bos_token_id    ,		use_cache=self.use_cache    ,		pad_token_id=self.pad_token_id    ,		decoder_start_token_id=self.decoder_start_token_id    ,		max_position_embeddings=self.max_position_embeddings    ,		)

							return (
							    config,
							    input_ids,
							    attention_mask,
							    lm_labels,
							)
			def lowerCAmelCase_	(       self: List[Any]    ,		UpperCamelCase: Any    ,		UpperCamelCase: Any    ,		UpperCamelCase: Optional[Any]    ,		UpperCamelCase: Optional[int]    ,		)						->						Optional[Any]:
							snake_case__													=       True
							snake_case__													=       SpeechaTextaDecoder(config=UpperCamelCase  ).to(UpperCamelCase  ).eval()
							snake_case__													=       input_ids[:2]

							input_ids[input_ids == 0] += 1
							# first forward pass
							snake_case__													=       model(UpperCamelCase    ,		use_cache=UpperCamelCase  )
							snake_case__													=       model(UpperCamelCase  )
							snake_case__													=       model(UpperCamelCase    ,		use_cache=UpperCamelCase  )

							self.parent.assertTrue(len(UpperCamelCase  ) == len(UpperCamelCase  )  )
							self.parent.assertTrue(len(UpperCamelCase  ) == len(UpperCamelCase  ) + 1  )

							snake_case__													=       outputs["past_key_values"]

							# create hypothetical next token and extent to next_input_ids
							snake_case__													=       ids_tensor((2, 1)    ,		config.vocab_size - 1  ) + 1

							# append to next input_ids and
							snake_case__													=       torch.cat([input_ids, next_tokens]    ,		dim=-1  )

							print(UpperCamelCase  )

							snake_case__													=       model(UpperCamelCase  )["last_hidden_state"]
							snake_case__													=       model(UpperCamelCase    ,		past_key_values=UpperCamelCase  )["last_hidden_state"]

							# select random slice
							snake_case__													=       ids_tensor((1,)    ,		output_from_past.shape[-1]  ).item()
							snake_case__													=       output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
							snake_case__													=       output_from_past[:, 0, random_slice_idx].detach()

							# test that outputs are equal for slice
							assert torch.allclose(UpperCamelCase    ,		UpperCamelCase    ,		atol=1e-3  )


			def lowerCAmelCase_	(       self: str  )						->						Tuple:
							snake_case__													=       self.prepare_config_and_inputs()
							(
							    snake_case__       
							)													=       config_and_inputs

							snake_case__													=       {
							    "input_ids": input_ids,
							    "attention_mask": attention_mask,
							}
							return config, inputs_dict


@require_torch
class 					__SCREAMING_SNAKE_CASE(		snake_case_      ,					snake_case_      ,					snake_case_      ,					unittest.TestCase		):
			_UpperCAmelCase	      =    (SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
			_UpperCAmelCase	      =    (SpeechaTextaForCausalLM,) if is_torch_available() else ()
			_UpperCAmelCase	      =    {"""text-generation""": SpeechaTextaForCausalLM} if is_torch_available() else {}
			_UpperCAmelCase	      =    True
			_UpperCAmelCase	      =    False
			def lowerCAmelCase_	(       self: List[str]    ,		)						->						Tuple:
							snake_case__													=       SpeechaTextaStandaloneDecoderModelTester(self    ,		is_training=UpperCamelCase  )
							snake_case__													=       ConfigTester(self    ,		config_class=UpperCamelCase  )
			def lowerCAmelCase_	(       self: int  )						->						int:
							pass
			def lowerCAmelCase_	(       self: int  )						->						str:
							pass
			def lowerCAmelCase_	(       self: Any  )						->						str:
							pass
			def lowerCAmelCase_	(       self: str  )						->						str:
							self.config_tester.run_common_tests()
			def lowerCAmelCase_	(       self: Optional[Any]  )						->						Optional[int]:
							snake_case__													=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_decoder_model_past(*UpperCamelCase  )


			def lowerCAmelCase_	(       self: int  )						->						Union[str, Any]:
							return












import argparse
import logging
import os
import sys
from collections import namedtuple

import torch
from modeling_bertabs import BertAbs, build_predictor
from torch.utils.data import DataLoader, SequentialSampler
from tqdm import tqdm

from transformers import BertTokenizer

from .utils_summarization import (
    CNNDMDataset,
    build_mask,
    compute_token_type_ids,
    encode_for_summarization,
    truncate_or_pad,
)


__UpperCamelCase   :						List[str]							=  logging.getLogger(__name__)
logging.basicConfig(stream=sys.stdout, level=logging.INFO)


__UpperCamelCase   :						Tuple							=  namedtuple("""Batch""", ["""document_names""", """batch_size""", """src""", """segs""", """mask_src""", """tgt_str"""])
def        a_	( _A					)     ->     Union[str, Any]:







				"""simple docstring"""
				snake_case__													=       BertTokenizer.from_pretrained('bert-base-uncased' ,				do_lower_case=_A					)
				snake_case__													=       BertAbs.from_pretrained('remi/bertabs-finetuned-extractive-abstractive-summarization'					)
				model.to(args.device					)
				model.eval()

				snake_case__													=       {
				    'BOS': tokenizer.vocab['[unused0]'],
				    'EOS': tokenizer.vocab['[unused1]'],
				    'PAD': tokenizer.vocab['[PAD]'],
				}

				if args.compute_rouge:
								snake_case__													=       []
								snake_case__													=       []

								import nltk
								import rouge

								nltk.download('punkt'					)
								snake_case__													=       rouge.Rouge(
								    metrics=['rouge-n', 'rouge-l'] ,				max_n=2 ,				limit_length=_A ,				length_limit=args.beam_size ,				length_limit_type='words' ,				apply_avg=_A ,				apply_best=_A ,				alpha=0.5 ,				weight_factor=1.2 ,				stemming=_A ,				)

				# these (unused) arguments are defined to keep the compatibility
				# with the legacy code and will be deleted in a next iteration.
				snake_case__													=       ''
				snake_case__													=       ''

				snake_case__													=       build_data_iterator(_A ,				_A					)
				snake_case__													=       build_predictor(_A ,				_A ,				_A ,				_A					)

				logger.info('***** Running evaluation *****'					)
				logger.info('  Number examples = %d' ,				len(data_iterator.dataset					)					)
				logger.info('  Batch size = %d' ,				args.batch_size					)
				logger.info(''					)
				logger.info('***** Beam Search parameters *****'					)
				logger.info('  Beam size = %d' ,				args.beam_size					)
				logger.info('  Minimum length = %d' ,				args.min_length					)
				logger.info('  Maximum length = %d' ,				args.max_length					)
				logger.info('  Alpha (length penalty) = %.2f' ,				args.alpha					)
				logger.info('  Trigrams %s be blocked' ,				('will' if args.block_trigram else 'will NOT')					)

				for batch in tqdm(_A					):
								snake_case__													=       predictor.translate_batch(_A					)
								snake_case__													=       predictor.from_batch(_A					)
								snake_case__													=       [format_summary(_A					) for t in translations]
								save_summaries(_A ,				args.summaries_output_dir ,				batch.document_names					)

								if args.compute_rouge:
												reference_summaries += batch.tgt_str
												generated_summaries += summaries

				if args.compute_rouge:
								snake_case__													=       rouge_evaluator.get_scores(_A ,				_A					)
								snake_case__													=       format_rouge_scores(_A					)
								save_rouge_scores(_A					)
								print(_A					)
def        a_	( _A ,				_A ,				_A					)     ->     Optional[Any]:







				"""simple docstring"""
				for summary, document_name in zip(_A ,				_A					):
								# Prepare the summary file's name
								if "." in document_name:
												snake_case__													=       '.'.join(document_name.split('.'					)[:-1]					)
												snake_case__													=       document_name.split('.'					)[-1]
												snake_case__													=       bare_document_name + '_summary.' + extension
								else:
												snake_case__													=       document_name + '_summary'

								snake_case__													=       os.path.join(_A ,				_A					)
								with open(_A ,				'w'					) as output:
												output.write(_A					)
def        a_	( _A					)     ->     Union[str, Any]:







				"""simple docstring"""
				snake_case__							,  snake_case__							,  snake_case__													=       translation
				snake_case__													=       (
				    raw_summary.replace('[unused0]' ,				''					)
				    .replace('[unused3]' ,				''					)
				    .replace('[PAD]' ,				''					)
				    .replace('[unused1]' ,				''					)
				    .replace(R' +' ,				' '					)
				    .replace(' [unused2] ' ,				'. '					)
				    .replace('[unused2]' ,				''					)
				    .strip()
				)

				return summary
def        a_	( _A					)     ->     Any:







				"""simple docstring"""
				return """\n
****** ROUGE SCORES ******

** ROUGE 1
F1        >> {:.3f}
Precision >> {:.3f}
Recall    >> {:.3f}

** ROUGE 2
F1        >> {:.3f}
Precision >> {:.3f}
Recall    >> {:.3f}

** ROUGE L
F1        >> {:.3f}
Precision >> {:.3f}
Recall    >> {:.3f}""".format(
				    scores['rouge-1']['f'] ,				scores['rouge-1']['p'] ,				scores['rouge-1']['r'] ,				scores['rouge-2']['f'] ,				scores['rouge-2']['p'] ,				scores['rouge-2']['r'] ,				scores['rouge-l']['f'] ,				scores['rouge-l']['p'] ,				scores['rouge-l']['r'] ,				)
def        a_	( _A					)     ->     int:







				"""simple docstring"""
				with open('rouge_scores.txt' ,				'w'					) as output:
								output.write(_A					)
def        a_	( _A ,				_A					)     ->     Optional[Any]:







				"""simple docstring"""
				snake_case__													=       load_and_cache_examples(_A ,				_A					)
				snake_case__													=       SequentialSampler(_A					)

				def collate_fn(_A					):
								return collate(_A ,				_A ,				block_size=512 ,				device=args.device					)

				snake_case__													=       DataLoader(
				    _A ,				sampler=_A ,				batch_size=args.batch_size ,				collate_fn=_A ,				)

				return iterator
def        a_	( _A ,				_A					)     ->     Any:







				"""simple docstring"""
				snake_case__													=       CNNDMDataset(args.documents_dir					)
				return dataset
def        a_	( _A ,				_A ,				_A ,				_A					)     ->     str:







				"""simple docstring"""
				snake_case__													=       [x for x in data if not len(x[1]					) == 0]  # remove empty_files
				snake_case__													=       [name for name, _, _ in data]
				snake_case__													=       [' '.join(_A					) for _, _, summary_list in data]

				snake_case__													=       [encode_for_summarization(_A ,				_A ,				_A					) for _, story, summary in data]
				snake_case__													=       torch.tensor(
				    [truncate_or_pad(_A ,				_A ,				tokenizer.pad_token_id					) for story, _ in encoded_text]					)
				snake_case__													=       compute_token_type_ids(_A ,				tokenizer.cls_token_id					)
				snake_case__													=       build_mask(_A ,				tokenizer.pad_token_id					)

				snake_case__													=       Batch(
				    document_names=_A ,				batch_size=len(_A					) ,				src=encoded_stories.to(_A					) ,				segs=encoder_token_type_ids.to(_A					) ,				mask_src=encoder_mask.to(_A					) ,				tgt_str=_A ,				)

				return batch
def        a_	( _A ,				_A					)     ->     Optional[Any]:







				"""simple docstring"""
				snake_case__													=       summary_tokens.to('cpu'					).numpy()
				snake_case__													=       tokenizer.decode(_A					)
				snake_case__													=       summary.split('.'					)
				snake_case__													=       [s + '.' for s in sentences]
				return sentences
def        a_	( )     ->     int:







				"""simple docstring"""
				snake_case__													=       argparse.ArgumentParser()
				parser.add_argument(
				    '--documents_dir' ,				default=_A ,				type=_A ,				required=_A ,				help='The folder where the documents to summarize are located.' ,				)
				parser.add_argument(
				    '--summaries_output_dir' ,				default=_A ,				type=_A ,				required=_A ,				help='The folder in wich the summaries should be written. Defaults to the folder where the documents are' ,				)
				parser.add_argument(
				    '--compute_rouge' ,				default=_A ,				type=_A ,				required=_A ,				help='Compute the ROUGE metrics during evaluation. Only available for the CNN/DailyMail dataset.' ,				)
				# EVALUATION options
				parser.add_argument(
				    '--no_cuda' ,				default=_A ,				type=_A ,				help='Whether to force the execution on CPU.' ,				)
				parser.add_argument(
				    '--batch_size' ,				default=4 ,				type=_A ,				help='Batch size per GPU/CPU for training.' ,				)
				# BEAM SEARCH arguments
				parser.add_argument(
				    '--min_length' ,				default=50 ,				type=_A ,				help='Minimum number of tokens for the summaries.' ,				)
				parser.add_argument(
				    '--max_length' ,				default=200 ,				type=_A ,				help='Maixmum number of tokens for the summaries.' ,				)
				parser.add_argument(
				    '--beam_size' ,				default=5 ,				type=_A ,				help='The number of beams to start with for each example.' ,				)
				parser.add_argument(
				    '--alpha' ,				default=0.95 ,				type=_A ,				help='The value of alpha for the length penalty in the beam search.' ,				)
				parser.add_argument(
				    '--block_trigram' ,				default=_A ,				type=_A ,				help='Whether to block the existence of repeating trigrams in the text generated by beam search.' ,				)
				snake_case__													=       parser.parse_args()

				# Select device (distibuted not available)
				snake_case__													=       torch.device('cuda' if torch.cuda.is_available() and not args.no_cuda else 'cpu'					)

				# Check the existence of directories
				if not args.summaries_output_dir:
								snake_case__													=       args.documents_dir

				if not documents_dir_is_valid(args.documents_dir					):
								raise FileNotFoundError(
								    'We could not find the directory you specified for the documents to summarize, or it was empty. Please'
								    ' specify a valid path.'					)
				os.makedirs(args.summaries_output_dir ,				exist_ok=_A					)

				evaluate(_A					)
def        a_	( _A					)     ->     str:







				"""simple docstring"""
				if not os.path.exists(_A					):
								return False

				snake_case__													=       os.listdir(_A					)
				if len(_A					) == 0:
								return False

				return True


if __name__ == "__main__":
							main()




