





import datasets

from .nmt_bleu import compute_bleu  # From: https://github.com/tensorflow/nmt/blob/master/nmt/scripts/bleu.py


lowerCAmelCase_    =				'''\
@INPROCEEDINGS{Papineni02bleu:a,
    author = {Kishore Papineni and Salim Roukos and Todd Ward and Wei-jing Zhu},
    title = {BLEU: a Method for Automatic Evaluation of Machine Translation},
    booktitle = {},
    year = {2002},
    pages = {311--318}
}
@inproceedings{lin-och-2004-orange,
    title = "{ORANGE}: a Method for Evaluating Automatic Evaluation Metrics for Machine Translation",
    author = "Lin, Chin-Yew  and
      Och, Franz Josef",
    booktitle = "{COLING} 2004: Proceedings of the 20th International Conference on Computational Linguistics",
    month = "aug 23{--}aug 27",
    year = "2004",
    address = "Geneva, Switzerland",
    publisher = "COLING",
    url = "https://www.aclweb.org/anthology/C04-1072",
    pages = "501--507",
}
'''

lowerCAmelCase_    =				'''\
BLEU (bilingual evaluation understudy) is an algorithm for evaluating the quality of text which has been machine-translated from one natural language to another.
Quality is considered to be the correspondence between a machine\'s output and that of a human: "the closer a machine translation is to a professional human translation,
the better it is" – this is the central idea behind BLEU. BLEU was one of the first metrics to claim a high correlation with human judgements of quality, and
remains one of the most popular automated and inexpensive metrics.

Scores are calculated for individual translated segments—generally sentences—by comparing them with a set of good quality reference translations.
Those scores are then averaged over the whole corpus to reach an estimate of the translation\'s overall quality. Intelligibility or grammatical correctness
are not taken into account[citation needed].

BLEU\'s output is always a number between 0 and 1. This value indicates how similar the candidate text is to the reference texts, with values closer to 1
representing more similar texts. Few human translations will attain a score of 1, since this would indicate that the candidate is identical to one of the
reference translations. For this reason, it is not necessary to attain a score of 1. Because there are more opportunities to match, adding additional
reference translations will increase the BLEU score.
'''

lowerCAmelCase_    =				'''
Computes BLEU score of translated segments against one or more references.
Args:
    predictions: list of translations to score.
        Each translation should be tokenized into a list of tokens.
    references: list of lists of references for each translation.
        Each reference should be tokenized into a list of tokens.
    max_order: Maximum n-gram order to use when computing BLEU score.
    smooth: Whether or not to apply Lin et al. 2004 smoothing.
Returns:
    \'bleu\': bleu score,
    \'precisions\': geometric mean of n-gram precisions,
    \'brevity_penalty\': brevity penalty,
    \'length_ratio\': ratio of lengths,
    \'translation_length\': translation_length,
    \'reference_length\': reference_length
Examples:

    >>> predictions = [
    ...     ["hello", "there", "general", "kenobi"],                             # tokenized prediction of the first sample
    ...     ["foo", "bar", "foobar"]                                             # tokenized prediction of the second sample
    ... ]
    >>> references = [
    ...     [["hello", "there", "general", "kenobi"], ["hello", "there", "!"]],  # tokenized references for the first sample (2 references)
    ...     [["foo", "bar", "foobar"]]                                           # tokenized references for the second sample (1 reference)
    ... ]
    >>> bleu = datasets.load_metric("bleu")
    >>> results = bleu.compute(predictions=predictions, references=references)
    >>> print(results["bleu"])
    1.0
'''



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION,     _KWARGS_DESCRIPTION     )
class 					__lowerCAmelCase       (			datasets.Metric     ):






			def 			lowerCamelCase  (self    )     ->   Dict:


							'''simple docstring'''



							return datasets.MetricInfo(
							    description=_DESCRIPTION					,							citation=_CITATION					,							inputs_description=_KWARGS_DESCRIPTION					,							features=datasets.Features(
							        {
							            '''predictions''': datasets.Sequence(datasets.Value('''string'''					,							id='''token'''    )					,							id='''sequence'''    ),
							            '''references''': datasets.Sequence(
							                datasets.Sequence(datasets.Value('''string'''					,							id='''token'''    )					,							id='''sequence'''    )					,							id='''references'''    ),
							        }    )					,							codebase_urls=['''https://github.com/tensorflow/nmt/blob/master/nmt/scripts/bleu.py''']					,							reference_urls=[
							        '''https://en.wikipedia.org/wiki/BLEU''',
							        '''https://towardsdatascience.com/evaluating-text-output-in-nlp-bleu-at-your-own-risk-e8609665a213''',
							    ]					,							)






			def 			lowerCamelCase  (self					,							__magic_name__					,							__magic_name__					,							__magic_name__=4					,							__magic_name__=False    )     ->   Optional[Any]:


							'''simple docstring'''



							snake_case_						:    Any				=  compute_bleu(
							    reference_corpus=__magic_name__					,							translation_corpus=__magic_name__					,							max_order=__magic_name__					,							smooth=__magic_name__    )
							((snake_case_)			,							(snake_case_)			,							(snake_case_)			,							(snake_case_)			,							(snake_case_)			,							(snake_case_))						:    Any				=  score
							return {
							    "bleu": bleu,
							    "precisions": precisions,
							    "brevity_penalty": bp,
							    "length_ratio": ratio,
							    "translation_length": translation_length,
							    "reference_length": reference_length,
							}







import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


lowerCAmelCase_    =				logging.getLogger(__name__)






def      lowerCamelCase_						(				)		->						Optional[Any]:






				"""simple docstring"""


				snake_case_						:    List[str]				=  argparse.ArgumentParser(
				    description='''Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset.''' )
				parser.add_argument(
				    '''--dataset_name''' ,				type=_UpperCamelCase ,				default='''wikitext''' ,				help='''Name of the training. Explore datasets at: hf.co/datasets.''' ,				)
				parser.add_argument(
				    '''--dataset_config''' ,				type=_UpperCamelCase ,				default='''wikitext-103-raw-v1''' ,				help='''Configuration name of the dataset.''' )
				parser.add_argument(
				    '''--tokenizer_name_or_path''' ,				type=_UpperCamelCase ,				default='''sayakpaul/unigram-tokenizer-wikitext''' ,				help='''Tokenizer identifier. Can be a local filepath or a Hub identifier.''' ,				)
				parser.add_argument(
				    '''--shard_size''' ,				type=_UpperCamelCase ,				default=1_000 ,				help='''Number of entries to go in a single shard.''' ,				)
				parser.add_argument('''--split''' ,				type=_UpperCamelCase ,				default='''train''' ,				choices=['''train''', '''test''', '''validation'''] )
				parser.add_argument(
				    '''--limit''' ,				default=_UpperCamelCase ,				type=_UpperCamelCase ,				help='''Limit the number of shards (used for debugging).''' ,				)
				parser.add_argument(
				    '''--max_length''' ,				type=_UpperCamelCase ,				default=512 ,				help='''Maximum sequence length. For training on TPUs, it helps to have a maximum'''
				    ''' sequence length that is a multiple of 8.''' ,				)
				parser.add_argument(
				    '''--output_dir''' ,				default='''tf-tpu''' ,				type=_UpperCamelCase ,				help='''Output directory where the TFRecord shards will be saved. If the'''
				    ''' path is appended with `gs://` (\'gs://tf-tpu\', for example) then the TFRecord'''
				    ''' shards will be directly saved to a Google Cloud Storage bucket.''' ,				)

				snake_case_						:    List[Any]				=  parser.parse_args()
				return args






def      lowerCamelCase_						(				_UpperCamelCase )		->						Tuple:






				"""simple docstring"""


				def fn(_UpperCamelCase ):
								return tokenizer(examples['''text'''] )

				return fn






def      lowerCamelCase_						(				_UpperCamelCase )		->						Union[str, Any]:






				"""simple docstring"""


				snake_case_						:    Any				=  []
				for i in range(len(tokenized_data['''input_ids'''] ) ):
								snake_case_						:    Any				=  {
								    '''input_ids''': tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data['''input_ids'''][i] ) ),
								    '''attention_mask''': tf.train.Feature(
								        intaa_list=tf.train.IntaaList(value=tokenized_data['''attention_mask'''][i] ) ),
								}
								snake_case_						:    Optional[int]				=  tf.train.Features(feature=_UpperCamelCase )
								snake_case_						:    Optional[Any]				=  tf.train.Example(features=_UpperCamelCase )
								snake_case_						:    Optional[Any]				=  example.SerializeToString()
								records.append(_UpperCamelCase )
				return records






def      lowerCamelCase_						(				_UpperCamelCase )		->						Optional[int]:






				"""simple docstring"""


				snake_case_						:    int				=  datasets.load_dataset(args.dataset_name ,				args.dataset_config ,				split=args.split )

				if args.limit is not None:
								snake_case_						:    Union[str, Any]				=  min(len(_UpperCamelCase ) ,				args.limit )
								snake_case_						:    int				=  dataset.select(range(_UpperCamelCase ) )
								print(f'''Limiting the dataset to {args.limit} entries.''' )

				snake_case_						:    Dict				=  AutoTokenizer.from_pretrained(args.tokenizer_name_or_path )

				# Handle output directory creation.
				# For serializing into a Google Cloud Storage Bucket, one needs to first
				# create a bucket.
				if "gs" not in args.output_dir:
								if not os.path.exists(args.output_dir ):
												os.makedirs(args.output_dir )
								snake_case_						:    str				=  os.path.join(args.output_dir ,				args.split )
								if not os.path.exists(_UpperCamelCase ):
												os.makedirs(_UpperCamelCase )
				else:
								snake_case_						:    Optional[Any]				=  os.path.join(args.output_dir ,				args.split )

				# Tokenize the whole dataset at once.
				snake_case_						:    Optional[Any]				=  tokenize_function(_UpperCamelCase )
				snake_case_						:    List[Any]				=  dataset.map(_UpperCamelCase ,				batched=_UpperCamelCase ,				num_proc=4 ,				remove_columns=['''text'''] )

				# We need to concatenate all our texts together, and then split the result
				# into chunks of a fixed size, which we will call block_size. To do this, we
				# will use the map method again, with the option batched=True. When we use batched=True,
				# the function we pass to map() will be passed multiple inputs at once, allowing us
				# to group them into more or fewer examples than we had in the input.
				# This allows us to create our new fixed-length samples. The advantage of this
				# method is that we don't lose a whole lot of content from the dataset compared to the
				# case where we simply tokenize with a pre-defined max_length.

				def group_texts(_UpperCamelCase ):
								# Concatenate all texts.
								snake_case_						:    Tuple				=  {k: sum(examples[k] ,				[] ) for k in examples.keys()}
								snake_case_						:    List[str]				=  len(concatenated_examples[list(examples.keys() )[0]] )
								# We drop the small remainder, though you could add padding instead if the model supports it
								# In this, as in all things, we advise you to follow your heart 🫀
								snake_case_						:    int				=  (total_length // args.max_length) * args.max_length
								# Split by chunks of max_len.
								snake_case_						:    Union[str, Any]				=  {
								    k: [t[i : i + args.max_length] for i in range(0 ,				_UpperCamelCase ,				args.max_length )]
								    for k, t in concatenated_examples.items()
								}
								return result

				snake_case_						:    int				=  dataset_tokenized.map(_UpperCamelCase ,				batched=_UpperCamelCase ,				batch_size=1_000 ,				num_proc=4 )

				snake_case_						:    str				=  0
				snake_case_						:    Optional[Any]				=  0
				for shard in range(0 ,				len(_UpperCamelCase ) ,				args.shard_size ):
								snake_case_						:    Any				=  grouped_dataset[shard : shard + args.shard_size]
								snake_case_						:    str				=  len(dataset_snapshot['''input_ids'''] )
								snake_case_						:    Union[str, Any]				=  os.path.join(_UpperCamelCase ,				f'''dataset-{shard_count}-{records_containing}.tfrecord''' )
								snake_case_						:    Dict				=  get_serialized_examples(_UpperCamelCase )

								with tf.io.TFRecordWriter(_UpperCamelCase ) as out_file:
												for i in range(len(_UpperCamelCase ) ):
																snake_case_						:    List[str]				=  serialized_examples[i]
																out_file.write(_UpperCamelCase )
												print('''Wrote file {} containing {} records'''.format(_UpperCamelCase ,				_UpperCamelCase ) )

								shard_count += 1
								total_records += records_containing

				with open(f'''split-{args.split}-records-count.txt''' ,				'''w''' ) as f:
								print(f'''Total {args.split} records: {total_records}''' ,				file=_UpperCamelCase )


if __name__ == "__main__":
	lowerCAmelCase_    =				parse_args()
	main(args)
