





'''simple docstring'''







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,
)


snake_case_  :							Any            =	logging.getLogger(__name__)
logging.basicConfig(stream=sys.stdout, level=logging.INFO)


snake_case_  :							Dict            =	namedtuple("Batch", ["document_names", "batch_size", "src", "segs", "mask_src", "tgt_str"])





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Tuple )			->						Dict:
						UpperCAmelCase_      :   Any      =	BertTokenizer.from_pretrained('''bert-base-uncased''',      do_lower_case=SCREAMING_SNAKE_CASE__ )
						UpperCAmelCase_      :   str      =	BertAbs.from_pretrained('''remi/bertabs-finetuned-extractive-abstractive-summarization''' )
						model.to(args.device )
						model.eval()

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

						if args.compute_rouge:
												UpperCAmelCase_      :   Union[str, Any]      =	[]
												UpperCAmelCase_      :   Optional[int]      =	[]

												import nltk
												import rouge

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

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

						UpperCAmelCase_      :   Any      =	build_data_iterator(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ )
						UpperCAmelCase_      :   Any      =	build_predictor(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ )

						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(SCREAMING_SNAKE_CASE__ ):
												UpperCAmelCase_      :   List[Any]      =	predictor.translate_batch(SCREAMING_SNAKE_CASE__ )
												UpperCAmelCase_      :   Union[str, Any]      =	predictor.from_batch(SCREAMING_SNAKE_CASE__ )
												UpperCAmelCase_      :   Tuple      =	[format_summary(SCREAMING_SNAKE_CASE__ ) for t in translations]
												save_summaries(SCREAMING_SNAKE_CASE__,      args.summaries_output_dir,      batch.document_names )

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

						if args.compute_rouge:
												UpperCAmelCase_      :   Union[str, Any]      =	rouge_evaluator.get_scores(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ )
												UpperCAmelCase_      :   Any      =	format_rouge_scores(SCREAMING_SNAKE_CASE__ )
												save_rouge_scores(SCREAMING_SNAKE_CASE__ )
												print(SCREAMING_SNAKE_CASE__ )





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Optional[Any],      SCREAMING_SNAKE_CASE__   :      Dict,      SCREAMING_SNAKE_CASE__   :      Union[str, Any] )			->						Optional[Any]:
						for summary, document_name in zip(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ ):
												# Prepare the summary file's name
												if "." in document_name:
																		UpperCAmelCase_      :   List[Any]      =	'''.'''.join(document_name.split('''.''' )[:-1] )
																		UpperCAmelCase_      :   Dict      =	document_name.split('''.''' )[-1]
																		UpperCAmelCase_      :   Optional[int]      =	bare_document_name + '''_summary.''' + extension
												else:
																		UpperCAmelCase_      :   List[str]      =	document_name + '''_summary'''

												UpperCAmelCase_      :   Any      =	os.path.join(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ )
												with open(SCREAMING_SNAKE_CASE__,      '''w''' ) as output:
																		output.write(SCREAMING_SNAKE_CASE__ )





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Optional[Any] )			->						Optional[int]:
						UpperCAmelCase_    ,							UpperCAmelCase_    ,							UpperCAmelCase_      :   int      =	translation
						UpperCAmelCase_      :   Any      =	(
						    raw_summary.replace('''[unused0]''',      '''''' )
						    .replace('''[unused3]''',      '''''' )
						    .replace('''[PAD]''',      '''''' )
						    .replace('''[unused1]''',      '''''' )
						    .replace(R''' +''',      ''' ''' )
						    .replace(''' [unused2] ''',      '''. ''' )
						    .replace('''[unused2]''',      '''''' )
						    .strip()
						)

						return summary





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      str )			->						Any:
						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 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Any )			->						Tuple:
						with open('''rouge_scores.txt''',      '''w''' ) as output:
												output.write(SCREAMING_SNAKE_CASE__ )





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Dict,      SCREAMING_SNAKE_CASE__   :      Dict )			->						Optional[int]:
						UpperCAmelCase_      :   Dict      =	load_and_cache_examples(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ )
						UpperCAmelCase_      :   Any      =	SequentialSampler(SCREAMING_SNAKE_CASE__ )

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

						UpperCAmelCase_      :   Any      =	DataLoader(
						    SCREAMING_SNAKE_CASE__,      sampler=SCREAMING_SNAKE_CASE__,      batch_size=args.batch_size,      collate_fn=SCREAMING_SNAKE_CASE__,      )

						return iterator





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      List[Any],      SCREAMING_SNAKE_CASE__   :      Union[str, Any] )			->						str:
						UpperCAmelCase_      :   List[Any]      =	CNNDMDataset(args.documents_dir )
						return dataset





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Any,      SCREAMING_SNAKE_CASE__   :      int,      SCREAMING_SNAKE_CASE__   :      str,      SCREAMING_SNAKE_CASE__   :      int )			->						Optional[int]:
						UpperCAmelCase_      :   Optional[int]      =	[x for x in data if not len(x[1] ) == 0]  # remove empty_files
						UpperCAmelCase_      :   Optional[int]      =	[name for name, _, _ in data]
						UpperCAmelCase_      :   List[Any]      =	[''' '''.join(SCREAMING_SNAKE_CASE__ ) for _, _, summary_list in data]

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

						UpperCAmelCase_      :   Optional[int]      =	Batch(
						    document_names=SCREAMING_SNAKE_CASE__,      batch_size=len(SCREAMING_SNAKE_CASE__ ),      src=encoded_stories.to(SCREAMING_SNAKE_CASE__ ),      segs=encoder_token_type_ids.to(SCREAMING_SNAKE_CASE__ ),      mask_src=encoder_mask.to(SCREAMING_SNAKE_CASE__ ),      tgt_str=SCREAMING_SNAKE_CASE__,      )

						return batch





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      List[Any],      SCREAMING_SNAKE_CASE__   :      Optional[Any] )			->						Union[str, Any]:
						UpperCAmelCase_      :   Dict      =	summary_tokens.to('''cpu''' ).numpy()
						UpperCAmelCase_      :   Optional[int]      =	tokenizer.decode(SCREAMING_SNAKE_CASE__ )
						UpperCAmelCase_      :   str      =	summary.split('''.''' )
						UpperCAmelCase_      :   Dict      =	[s + '''.''' for s in sentences]
						return sentences





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

						# Select device (distibuted not available)
						UpperCAmelCase_      :   int      =	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:
												UpperCAmelCase_      :   Optional[Any]      =	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=SCREAMING_SNAKE_CASE__ )

						evaluate(SCREAMING_SNAKE_CASE__ )





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Union[str, Any] )			->						Optional[Any]:
						if not os.path.exists(SCREAMING_SNAKE_CASE__ ):
												return False

						UpperCAmelCase_      :   str      =	os.listdir(SCREAMING_SNAKE_CASE__ )
						if len(SCREAMING_SNAKE_CASE__ ) == 0:
												return False

						return True


if __name__ == "__main__":
					main()













"""simple docstring"""







import jax
import jax.numpy as jnp
from bigbird_flax import FlaxBigBirdForNaturalQuestions
from datasets import load_from_disk

from transformers import BigBirdTokenizerFast


__a							:						Optional[int]   							=			{0: 'null', 1: 'short', 2: 'long', 3: 'yes', 4: 'no'}
__a							:						Dict   							=			set(''.join(['‘', '’', '´', '`', '.', ',', '-', '"']))





def 							SCREAMING_SNAKE_CASE							(   lowerCamelCase_			,  lowerCamelCase_=0			,  lowerCamelCase_=None):
					return [" ".join(x.split(''' ''')[begin:end]) for x in answers if len(x.split(''' ''')) > 1]





def 							SCREAMING_SNAKE_CASE							(   lowerCamelCase_			,  lowerCamelCase_=False):
					if make_sub_answers:
										# if answers are longer than one word, make sure a predictions is correct if it coresponds to the complete 1: or :-1 sub word
										# *e.g.* if the correct answer contains a prefix such as "the", or "a"
										a__ 							=	(
										    given_answers + get_sub_answers(lowerCamelCase_			,  begin=1) + get_sub_answers(lowerCamelCase_			,  end=-1)
										)
					a__ 							=	[]
					for answer in given_answers:
										a__ 							=	answer.replace('''_'''			,  ''' ''').lower()
										a__ 							=	''''''.join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else ''' ''' for c in alias)
										answers.append(''' '''.join(alias.split()).strip())
					return set(lowerCamelCase_)





def 							SCREAMING_SNAKE_CASE							(   lowerCamelCase_			,  lowerCamelCase_			,  lowerCamelCase_=1			,  lowerCamelCase_=100):
					a__							,a__ 							=	jax.lax.top_k(lowerCamelCase_			,  lowerCamelCase_)
					a__							,a__ 							=	jax.lax.top_k(lowerCamelCase_			,  lowerCamelCase_)

					a__ 							=	best_end_idx[:, None] - best_start_idx[None, :]
					a__ 							=	jnp.logical_or(widths < 0			,  widths > max_size)
					a__ 							=	(best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
					a__ 							=	jnp.argmax(lowerCamelCase_).item()

					return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]





def 							SCREAMING_SNAKE_CASE							(   lowerCamelCase_):
					a__ 							=	sample['''question''']['''text''']
					a__ 							=	sample['''document''']['''tokens''']['''token''']
					a__ 							=	sample['''document''']['''tokens''']['''is_html''']
					a__ 							=	sample['''annotations''']['''long_answer''']
					a__ 							=	sample['''annotations''']['''short_answers''']

					a__ 							=	''' '''.join([context[i] for i in range(len(lowerCamelCase_)) if not is_html[i]])

					# 0 - No ; 1 - Yes
					for answer in sample["annotations"]["yes_no_answer"]:
										if answer == 0 or answer == 1:
															return {
															    "question": question,
															    "context": context_string,
															    "short": [],
															    "long": [],
															    "category": "no" if answer == 0 else "yes",
															}

					a__ 							=	[]
					for s in short_answers:
										short_targets.extend(s['''text'''])
					a__ 							=	list(set(lowerCamelCase_))

					a__ 							=	[]
					for s in long_answers:
										if s["start_token"] == -1:
															continue
										a__ 							=	context[s['''start_token'''] : s['''end_token''']]
										a__ 							=	is_html[s['''start_token'''] : s['''end_token''']]
										a__ 							=	''' '''.join([answer[i] for i in range(len(lowerCamelCase_)) if not html[i]])
										if new_answer not in long_targets:
															long_targets.append(lowerCamelCase_)

					a__ 							=	'''long_short''' if len(short_targets + long_targets) > 0 else '''null'''

					return {
					    "question": question,
					    "context": context_string,
					    "short": short_targets,
					    "long": long_targets,
					    "category": category,
					}





def 							SCREAMING_SNAKE_CASE							(   ):
					a__ 							=	load_from_disk('''natural-questions-validation''')
					a__ 							=	dataset.map(lowerCamelCase_).remove_columns(['''annotations''', '''document''', '''id'''])
					print(lowerCamelCase_)

					a__ 							=	dataset.filter(lambda lowerCamelCase_: (len(x['''question''']) + len(x['''context'''])) < 4 * 4096)
					a__ 							=	short_validation_dataset.filter(lambda lowerCamelCase_: x["category"] != "null")
					short_validation_dataset

					a__ 							=	'''vasudevgupta/flax-bigbird-natural-questions'''
					a__ 							=	FlaxBigBirdForNaturalQuestions.from_pretrained(lowerCamelCase_)
					a__ 							=	BigBirdTokenizerFast.from_pretrained(lowerCamelCase_)

					@jax.jit
					def forward(*lowerCamelCase_			,  **lowerCamelCase_):
										a__							,a__							,a__ 							=	model(*lowerCamelCase_			,  **lowerCamelCase_)
										return start_logits, end_logits, jnp.argmax(lowerCamelCase_			,  axis=-1)

					def evaluate(lowerCamelCase_):
										# encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
										a__ 							=	tokenizer(
										    example['''question''']			,  example['''context''']			,  return_tensors='''np'''			,  max_length=4096			,  padding='''max_length'''			,  truncation=lowerCamelCase_			,  )

										a__							,a__							,a__ 							=	forward(**lowerCamelCase_)

										a__ 							=	CATEGORY_MAPPING[category.item()]

										a__ 							=	example['''long'''] + example['''short''']
										if example["category"] in ["yes", "no", "null"]:
															a__ 							=	[example['''category''']]
										a__ 							=	example['''category'''] != '''null'''
										# Now target can be: "yes", "no", "null", "list of long & short answers"

										if predicted_category in ["yes", "no", "null"]:
															a__ 							=	[predicted_category]
															a__ 							=	example['''output'''] == example['''targets''']
															a__ 							=	predicted_category != '''null'''
															return example

										a__ 							=	38 if predicted_category == '''short''' else 1024
										a__							,a__ 							=	get_best_valid_start_end_idx(
										    start_scores[0]			,  end_scores[0]			,  top_k=8			,  max_size=lowerCamelCase_)

										a__ 							=	inputs['''input_ids'''][0].tolist()
										a__ 							=	[tokenizer.decode(input_ids[start_score : end_score + 1])]

										a__ 							=	expand_to_aliases(example['''targets''']			,  make_sub_answers=lowerCamelCase_)
										a__ 							=	expand_to_aliases(example['''output'''])

										# some preprocessing to both prediction and answer
										a__ 							=	{''''''.join(a.split()) for a in answers}
										a__ 							=	{''''''.join(p.split()) for p in predictions}
										a__ 							=	{s for s in predictions if s not in ['''``''', '''\'\'''', '''`''', '''\'''']}

										# if there is a common element, it's a exact match
										a__ 							=	len(list(answers & predictions)) > 0
										a__ 							=	predicted_category != '''null''' and len(lowerCamelCase_) > 0

										return example

					a__ 							=	short_validation_dataset.map(lowerCamelCase_)

					a__ 							=	len(lowerCamelCase_)
					a__ 							=	len(short_validation_dataset.filter(lambda lowerCamelCase_: x["match"] == 1))
					print('''EM score:'''			,  (matched / total) * 100			,  '''%''')


if __name__ == "__main__":
	main()


