



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_			      =  logging.getLogger(__name__)
logging.basicConfig(stream=sys.stdout, level=logging.INFO)


UpperCAmelCase_			      =  namedtuple('''Batch''', ['''document_names''', '''batch_size''', '''src''', '''segs''', '''mask_src''', '''tgt_str'''])
def   UpperCAmelCase							(		A__     )							->					Tuple:
	_snake_case					:							str					     =   BertTokenizer.from_pretrained("""bert-base-uncased"""				,  do_lower_case=A__     )
	_snake_case					:							List[str]					     =   BertAbs.from_pretrained("""remi/bertabs-finetuned-extractive-abstractive-summarization"""     )
	model.to(args.device     )
	model.eval()

	_snake_case					:							int					     =   {
	    """BOS""": tokenizer.vocab["""[unused0]"""],
	    """EOS""": tokenizer.vocab["""[unused1]"""],
	    """PAD""": tokenizer.vocab["""[PAD]"""],
	}

	if args.compute_rouge:
		_snake_case					:							Any					     =   []
		_snake_case					:							List[str]					     =   []

		import nltk
		import rouge

		nltk.download("""punkt"""     )
		_snake_case					:							Any					     =   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					:							List[Any]					     =   """"""
	_snake_case					:							int					     =   """"""

	_snake_case					:							Any					     =   build_data_iterator(A__				,  A__     )
	_snake_case					:							str					     =   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					:							Tuple					     =   predictor.translate_batch(A__     )
		_snake_case					:							Optional[int]					     =   predictor.from_batch(A__     )
		_snake_case					:							int					     =   [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					:							Dict					     =   rouge_evaluator.get_scores(A__				,  A__     )
		_snake_case					:							Optional[int]					     =   format_rouge_scores(A__     )
		save_rouge_scores(A__     )
		print(A__     )
def   UpperCAmelCase							(		A__				,  A__				,  A__     )							->					Optional[int]:
	for summary, document_name in zip(A__				,  A__     ):
		# Prepare the summary file's name
		if "." in document_name:
			_snake_case					:							List[Any]					     =   """.""".join(document_name.split("""."""     )[:-1]     )
			_snake_case					:							List[Any]					     =   document_name.split("""."""     )[-1]
			_snake_case					:							int					     =   bare_document_name + """_summary.""" + extension
		else:
			_snake_case					:							Dict					     =   document_name + """_summary"""

		_snake_case					:							List[str]					     =   os.path.join(A__				,  A__     )
		with open(A__				,  """w"""     ) as output:
			output.write(A__     )
def   UpperCAmelCase							(		A__     )							->					List[Any]:
	_snake_case     , _snake_case     , _snake_case					:							Optional[Any]					     =   translation
	_snake_case					:							Optional[Any]					     =   (
	    raw_summary.replace("""[unused0]"""				,  """"""     )
	    .replace("""[unused3]"""				,  """"""     )
	    .replace("""[PAD]"""				,  """"""     )
	    .replace("""[unused1]"""				,  """"""     )
	    .replace(r""" +"""				,  """ """     )
	    .replace(""" [unused2] """				,  """. """     )
	    .replace("""[unused2]"""				,  """"""     )
	    .strip()
	)

	return summary
def   UpperCAmelCase							(		A__     )							->					Dict:
	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   UpperCAmelCase							(		A__     )							->					Union[str, Any]:
	with open("""rouge_scores.txt"""				,  """w"""     ) as output:
		output.write(A__     )
def   UpperCAmelCase							(		A__				,  A__     )							->					Optional[int]:
	_snake_case					:							int					     =   load_and_cache_examples(A__				,  A__     )
	_snake_case					:							Optional[int]					     =   SequentialSampler(A__     )

	def collate_fn(A__     ):
		return collate(A__				,  A__				,  block_size=5_12				,  device=args.device     )

	_snake_case					:							Optional[Any]					     =   DataLoader(
	    A__				,  sampler=A__				,  batch_size=args.batch_size				,  collate_fn=A__				,  )

	return iterator
def   UpperCAmelCase							(		A__				,  A__     )							->					str:
	_snake_case					:							str					     =   CNNDMDataset(args.documents_dir     )
	return dataset
def   UpperCAmelCase							(		A__				,  A__				,  A__				,  A__     )							->					List[Any]:
	_snake_case					:							Union[str, Any]					     =   [x for x in data if not len(x[1]     ) == 0]  # remove empty_files
	_snake_case					:							Dict					     =   [name for name, _, _ in data]
	_snake_case					:							Tuple					     =   [""" """.join(A__     ) for _, _, summary_list in data]

	_snake_case					:							str					     =   [encode_for_summarization(A__				,  A__				,  A__     ) for _, story, summary in data]
	_snake_case					:							Any					     =   torch.tensor(
	    [truncate_or_pad(A__				,  A__				,  tokenizer.pad_token_id     ) for story, _ in encoded_text]     )
	_snake_case					:							str					     =   compute_token_type_ids(A__				,  tokenizer.cls_token_id     )
	_snake_case					:							Any					     =   build_mask(A__				,  tokenizer.pad_token_id     )

	_snake_case					:							Tuple					     =   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   UpperCAmelCase							(		A__				,  A__     )							->					List[str]:
	_snake_case					:							List[str]					     =   summary_tokens.to("""cpu"""     ).numpy()
	_snake_case					:							Dict					     =   tokenizer.decode(A__     )
	_snake_case					:							int					     =   summary.split("""."""     )
	_snake_case					:							Tuple					     =   [s + """.""" for s in sentences]
	return sentences
def   UpperCAmelCase							(		)							->					Dict:
	_snake_case					:							Tuple					     =   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=2_00				,  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					:							Dict					     =   parser.parse_args()

	# Select device (distibuted not available)
	_snake_case					:							Union[str, Any]					     =   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					:							Optional[int]					     =   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   UpperCAmelCase							(		A__     )							->					Dict:
	if not os.path.exists(A__     ):
		return False

	_snake_case					:							int					     =   os.listdir(A__     )
	if len(A__     ) == 0:
		return False

	return True


if __name__ == "__main__":
						main()










# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# 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.
import os
from typing import List
def   UpperCAmelCase							(		A__				,  A__     )							->					Any:
	if name != "training_function" and name != "main":
		raise ValueError(f'''Incorrect function name passed: {name}, choose either \'main\' or \'training_function\''''     )
	_snake_case     , _snake_case					:							Union[str, Any]					     =   [], False
	for line in lines:
		if not found_start and f'''def {name}''' in line:
			_snake_case					:							Tuple					     =   True
			good_lines.append(A__     )
			continue
		if found_start:
			if name == "training_function" and "def main" in line:
				return good_lines
			if name == "main" and "if __name__" in line:
				return good_lines
			good_lines.append(A__     )
def   UpperCAmelCase							(		A__     )							->					Dict:
	return [line for line in lines if not line.lstrip().startswith("""#"""     ) and line != "\n"]


def   UpperCAmelCase							(		A__				,  A__				,  A__				,  A__ = None     )							->					str:
	with open(A__				,  """r"""     ) as f:
		_snake_case					:							Tuple					     =   f.readlines()
	with open(os.path.abspath(os.path.join("""examples"""				,  """nlp_example.py"""     )     )				,  """r"""     ) as f:
		_snake_case					:							Optional[Any]					     =   f.readlines()
	with open(A__				,  """r"""     ) as f:
		_snake_case					:							Optional[Any]					     =   f.readlines()
	if secondary_filename is not None:
		with open(A__				,  """r"""     ) as f:
			_snake_case					:							List[str]					     =   f.readlines()

    # This is our base, we remove all the code from here in our `full_filename` and `feature_filename` to find the new content
	if parser_only:
		_snake_case					:							Dict					     =   clean_lines(get_function_contents_by_name(A__				,  """main"""     )     )
		_snake_case					:							int					     =   clean_lines(get_function_contents_by_name(A__				,  """main"""     )     )
		_snake_case					:							Optional[Any]					     =   clean_lines(get_function_contents_by_name(A__				,  """main"""     )     )
		if secondary_filename is not None:
			_snake_case					:							Optional[Any]					     =   clean_lines(get_function_contents_by_name(A__				,  """main"""     )     )
	else:
		_snake_case					:							str					     =   clean_lines(get_function_contents_by_name(A__				,  """training_function"""     )     )
		_snake_case					:							List[Any]					     =   clean_lines(get_function_contents_by_name(A__				,  """training_function"""     )     )
		_snake_case					:							Union[str, Any]					     =   clean_lines(get_function_contents_by_name(A__				,  """training_function"""     )     )
		if secondary_filename is not None:
			_snake_case					:							List[Any]					     =   clean_lines(
			    get_function_contents_by_name(A__				,  """training_function"""     )     )

	_snake_case					:							str					     =   """train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n"""

	# Specific code in our script that differs from the full version, aka what is new
	_snake_case					:							Dict					     =   []
	_snake_case					:							str					     =   []  # We keep track of the idxs just in case it's a repeated statement
	_snake_case					:							Optional[Any]					     =   iter(A__     )
	for i in range(len(A__     ) - 1     ):
		if i not in passed_idxs:
			_snake_case					:							Any					     =   next(A__     )
			if (line not in full_file_func) and (line.lstrip() != _dl_line):
				if "TESTING_MOCKED_DATALOADERS" not in line:
					new_feature_code.append(A__     )
					passed_idxs.append(A__     )
				else:
					# Skip over the `config['num_epochs'] = 2` statement
					_snake_case					:							Optional[int]					     =   next(A__     )

    # Extract out just the new parts from the full_file_training_func
	_snake_case					:							Any					     =   []
	_snake_case					:							Tuple					     =   []  # We keep track of the idxs just in case it's a repeated statement
	for i, line in enumerate(A__     ):
		if i not in passed_idxs:
			if (line not in full_file_func) and (line.lstrip() != _dl_line):
				if "TESTING_MOCKED_DATALOADERS" not in line:
					new_full_example_parts.append(A__     )
					passed_idxs.append(A__     )

    # Finally, get the overall diff
	_snake_case					:							Optional[int]					     =   [line for line in new_feature_code if line not in new_full_example_parts]
	if secondary_filename is not None:
		_snake_case					:							Optional[int]					     =   [line for line in full_file_contents if line not in secondary_file_func]
		_snake_case					:							int					     =   [line for line in diff_from_example if line not in diff_from_two]

	return diff_from_example





