


import json
import os
from typing import Dict, List, Optional, Tuple

import regex as re

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowerCAmelCase__			:      str								=logging.get_logger(__name__)


lowerCAmelCase__			:      Any								={
    'vocab_file': 'vocab.json',
    'merges_file': 'merges.txt',
    'tokenizer_config_file': 'tokenizer_config.json',
}

lowerCAmelCase__			:      Union[str, Any]								={
    'vocab_file': {
        'facebook/blenderbot_small-90M': 'https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/vocab.json'
    },
    'merges_file': {
        'facebook/blenderbot_small-90M': 'https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/merges.txt'
    },
    'tokenizer_config_file': {
        'facebook/blenderbot_small-90M': (
            'https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/tokenizer_config.json'
        )
    },
}

lowerCAmelCase__			:      Tuple								={'facebook/blenderbot_small-90M': 5_12}






def 						a__			(				A__				):
			SCREAMING_SNAKE_CASE_		:     List[Any]					=					set()
			SCREAMING_SNAKE_CASE_		:     int					=					word[0]
			for char in word[1:]:
						pairs.add((prev_char, char)				)
						SCREAMING_SNAKE_CASE_		:     List[Any]					=					char

			SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					set(A__				)
			return pairs







class 	__lowercase    (__SCREAMING_SNAKE_CASE  ):





							"""simple docstring"""







							_UpperCAmelCase      =						VOCAB_FILES_NAMES
							_UpperCAmelCase      =						PRETRAINED_VOCAB_FILES_MAP
							_UpperCAmelCase      =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							_UpperCAmelCase      =						["""input_ids""", """attention_mask"""]






							def __init__(      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__="__start__"			,						lowerCAmelCase__="__end__"			,						lowerCAmelCase__="__unk__"			,						lowerCAmelCase__="__null__"			,						**lowerCAmelCase__			,						):





										"""simple docstring"""




										super().__init__(unk_token=lowerCAmelCase__			,						bos_token=lowerCAmelCase__			,						eos_token=lowerCAmelCase__			,						pad_token=lowerCAmelCase__			,						**lowerCAmelCase__       )

										with open(lowerCAmelCase__			,						encoding='utf-8'       ) as vocab_handle:
													SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					json.load(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					{v: k for k, v in self.encoder.items()}
										with open(lowerCAmelCase__			,						encoding='utf-8'       ) as merges_handle:
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					merges_handle.read().split('\n'       )[1:-1]
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					[tuple(merge.split()       ) for merge in merges]
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					dict(zip(lowerCAmelCase__			,						range(len(lowerCAmelCase__       )       )       )       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					{}






							@property
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										return len(self.encoder       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										return dict(self.encoder			,						**self.added_tokens_encoder       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__       ):





										"""simple docstring"""




										if token in self.cache:
													return self.cache[token]
										SCREAMING_SNAKE_CASE_		:     Tuple					=					re.sub('([.,!?()])'			,						r' \1'			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     List[str]					=					re.sub('(\')'			,						r' \1 '			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					re.sub(r'\s{2,}'			,						' '			,						lowerCAmelCase__       )
										if "\n" in token:
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					token.replace('\n'			,						' __newln__'       )

										SCREAMING_SNAKE_CASE_		:     Dict					=					token.split(' '       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					[]
										for token in tokens:
													if not len(lowerCAmelCase__       ):
																continue

													SCREAMING_SNAKE_CASE_		:     Any					=					token.lower()
													SCREAMING_SNAKE_CASE_		:     Any					=					tuple(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     str					=					tuple(list(word[:-1]       ) + [word[-1] + '</w>']       )
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					get_pairs(lowerCAmelCase__       )

													if not pairs:
																words.append(lowerCAmelCase__       )
																continue

													while True:
																SCREAMING_SNAKE_CASE_		:     List[str]					=					min(lowerCAmelCase__			,						key=lambda lowerCAmelCase__       : self.bpe_ranks.get(lowerCAmelCase__			,						float('inf'       )       )       )
																if bigram not in self.bpe_ranks:
																			break
																SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     int					=					bigram
																SCREAMING_SNAKE_CASE_		:     str					=					[]
																SCREAMING_SNAKE_CASE_		:     Optional[int]					=					0

																while i < len(lowerCAmelCase__       ):
																			try:
																						SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					word.index(lowerCAmelCase__			,						lowerCAmelCase__       )
																						new_word.extend(word[i:j]       )
																						SCREAMING_SNAKE_CASE_		:     List[Any]					=					j
																			except ValueError:
																						new_word.extend(word[i:]       )
																						break

																			if word[i] == first and i < len(lowerCAmelCase__       ) - 1 and word[i + 1] == second:
																						new_word.append(first + second       )
																						i += 2
																			else:
																						new_word.append(word[i]       )
																						i += 1
																SCREAMING_SNAKE_CASE_		:     Dict					=					tuple(lowerCAmelCase__       )
																SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					new_word
																if len(lowerCAmelCase__       ) == 1:
																			break
																else:
																			SCREAMING_SNAKE_CASE_		:     str					=					get_pairs(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     str					=					'@@ '.join(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     Dict					=					word[:-4]

													SCREAMING_SNAKE_CASE_		:     Tuple					=					word
													words.append(lowerCAmelCase__       )
										return " ".join(lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[Any]					=					[]

										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					re.findall(r'\S+\n?'			,						lowerCAmelCase__       )

										for token in words:
													split_tokens.extend(list(self.bpe(lowerCAmelCase__       ).split(' '       )       )       )
										return split_tokens






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					token.lower()
										return self.encoder.get(lowerCAmelCase__			,						self.encoder.get(self.unk_token       )       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__       ):





										"""simple docstring"""




										return self.decoder.get(lowerCAmelCase__			,						self.unk_token       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					' '.join(lowerCAmelCase__       ).replace('@@ '			,						''       ).strip()
										return out_string






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__ = None       ):





										"""simple docstring"""




										if not os.path.isdir(lowerCAmelCase__       ):
													logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''       )
													return
										SCREAMING_SNAKE_CASE_		:     int					=					os.path.join(
										    lowerCAmelCase__			,						(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']       )
										SCREAMING_SNAKE_CASE_		:     List[str]					=					os.path.join(
										    lowerCAmelCase__			,						(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['merges_file']       )

										with open(lowerCAmelCase__			,						'w'			,						encoding='utf-8'       ) as f:
													f.write(json.dumps(self.encoder			,						indent=2			,						sort_keys=lowerCAmelCase__			,						ensure_ascii=lowerCAmelCase__       ) + '\n'       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					0
										with open(lowerCAmelCase__			,						'w'			,						encoding='utf-8'       ) as writer:
													writer.write('#version: 0.2\n'       )
													for bpe_tokens, token_index in sorted(self.bpe_ranks.items()			,						key=lambda lowerCAmelCase__       : kv[1]       ):
																if index != token_index:
																			logger.warning(
																			    F'''Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'''
																			    ' Please check that the tokenizer is not corrupted!'       )
																			SCREAMING_SNAKE_CASE_		:     Tuple					=					token_index
																writer.write(' '.join(lowerCAmelCase__       ) + '\n'       )
																index += 1

										return vocab_file, merge_file









'''simple docstring'''
import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


lowerCAmelCase   : Dict             =				logging.getLogger(__name__)

if is_torch_tpu_available(check_device=False):
					import torch_xla.core.xla_model as xm
					import torch_xla.debug.metrics as met




class        _UpperCamelCase	(    SCREAMING_SNAKE_CASE):


				'''simple docstring'''


				def __init__(     self							,      *a_							,      a_=None							,      a_=None							,      a_=None							,      **a_							)    ->  Any:
					super().__init__(*a_							,      **a_							)
					lowercase						:		Optional[Any]   					=		eval_examples
					lowercase						:		Any   					=		post_process_function
					lowercase						:		Dict   					=		quant_trainer_args
					lowercase						:		List[Any]   					=		1_2_8  # default number of calibration samples

				def 	a__		(     self							,      a_=None							)    ->  List[str]:
					if calib_dataset is None and self.calib_dataset is None:
						raise ValueError("Trainer: calibration requires an calib_dataset."							)
					lowercase						:		Dict   					=		calib_dataset if calib_dataset is not None else self.calib_dataset

					lowercase						:		Optional[int]   					=		self._remove_unused_columns(a_							,      description="Calibration"							)

					return DataLoader(
					    a_							,      batch_size=self.args.eval_batch_size							,      collate_fn=self.data_collator							,      drop_last=self.args.dataloader_drop_last							,      num_workers=self.args.dataloader_num_workers							,      pin_memory=self.args.dataloader_pin_memory							,      shuffle=a_							,      )

				def 	a__		(     self							,      a_=None							)    ->  str:
					lowercase						:		Optional[Any]   					=		self.train_dataset if calib_dataset is None else calib_dataset
					lowercase						:		Optional[Any]   					=		self.get_calib_dataloader(a_							)

					lowercase						:		Union[str, Any]   					=		self.model
					quant_trainer.configure_model(a_							,      self.quant_trainer_args							,      calib=a_							)
					model.eval()
					quant_trainer.enable_calibration(a_							)

					logger.info("***** Running calibration *****"							)
					logger.info(F'''  Num examples = {self.calib_num}'''							)
					logger.info(F'''  Batch size = {calib_dataloader.batch_size}'''							)

					for step, inputs in enumerate(a_							):
						# Prediction step
						lowercase					,		lowercase					,		lowercase						:		Union[str, Any]   					=		self.prediction_step(a_							,      a_							,      prediction_loss_only=a_							)
						if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
							break

					quant_trainer.finish_calibration(a_							,      self.quant_trainer_args							)
					lowercase						:		List[Any]   					=		model

				def 	a__		(     self							,      a_=None							,      a_=None							,      a_=None							,      a_ = "eval"							)    ->  int:
					lowercase						:		List[str]   					=		self.eval_dataset if eval_dataset is None else eval_dataset
					lowercase						:		Union[str, Any]   					=		self.get_eval_dataloader(a_							)
					lowercase						:		Tuple   					=		self.eval_examples if eval_examples is None else eval_examples

					# Temporarily disable metric computation, we will do it in the loop here.
					lowercase						:		List[str]   					=		self.compute_metrics
					lowercase						:		Dict   					=		None
					lowercase						:		int   					=		self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
					try:
						lowercase						:		Tuple   					=		eval_loop(
						    a_							,      description="Evaluation"							,      prediction_loss_only=True if compute_metrics is None else None							,      ignore_keys=a_							,      )
					finally:
						lowercase						:		List[Any]   					=		compute_metrics

					if self.post_process_function is not None and self.compute_metrics is not None:
						lowercase						:		Union[str, Any]   					=		self.post_process_function(a_							,      a_							,      output.predictions							)
						lowercase						:		Union[str, Any]   					=		self.compute_metrics(a_							)

						# Prefix all keys with metric_key_prefix + '_'
						for key in list(metrics.keys()							):
							if not key.startswith(F'''{metric_key_prefix}_'''							):
								lowercase						:		Dict   					=		metrics.pop(a_							)

						self.log(a_							)
					else:
						lowercase						:		Tuple   					=		{}

					if self.args.tpu_metrics_debug or self.args.debug:
						# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
						xm.master_print(met.metrics_report()							)

					lowercase						:		int   					=		self.callback_handler.on_evaluate(self.args							,      self.state							,      self.control							,      a_							)
					return metrics

				def 	a__		(     self							,      a_							,      a_							,      a_=None							,      a_ = "test"							)    ->  Optional[int]:
					lowercase						:		Dict   					=		self.get_test_dataloader(a_							)

					# Temporarily disable metric computation, we will do it in the loop here.
					lowercase						:		List[str]   					=		self.compute_metrics
					lowercase						:		List[str]   					=		None
					lowercase						:		int   					=		self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
					try:
						lowercase						:		Optional[Any]   					=		eval_loop(
						    a_							,      description="Prediction"							,      prediction_loss_only=True if compute_metrics is None else None							,      ignore_keys=a_							,      )
					finally:
						lowercase						:		int   					=		compute_metrics

					if self.post_process_function is None or self.compute_metrics is None:
						return output

					lowercase						:		List[str]   					=		self.post_process_function(a_							,      a_							,      output.predictions							,      "predict"							)
					lowercase						:		Dict   					=		self.compute_metrics(a_							)

					# Prefix all keys with metric_key_prefix + '_'
					for key in list(metrics.keys()							):
						if not key.startswith(F'''{metric_key_prefix}_'''							):
							lowercase						:		Union[str, Any]   					=		metrics.pop(a_							)

					return PredictionOutput(predictions=predictions.predictions							,      label_ids=predictions.label_ids							,      metrics=a_							)


				def 	a__		(     self							,      a_="./"							)    ->  Optional[Any]:
					lowercase						:		str   					=		self.eval_dataset
					lowercase						:		Dict   					=		self.get_eval_dataloader(a_							)

					lowercase						:		Any   					=		next(iter(a_							)							)

					# saving device - to make it consistent
					lowercase						:		List[Any]   					=		torch.device("cuda" if torch.cuda.is_available() else "cpu"							)

					# convert to tuple
					lowercase						:		Tuple   					=		tuple(v.to(a_							) for k, v in batch.items()							)

					logger.info("Converting model to be onnx compatible"							)
					from pytorch_quantization.nn import TensorQuantizer

					lowercase						:		Optional[Any]   					=		True

					lowercase						:		int   					=		self.model.to(a_							)

					model.eval()
					model.float()

					lowercase						:		Union[str, Any]   					=		model.module if hasattr(a_							,      "module"							) else model
					quant_trainer.configure_model(a_							,      self.quant_trainer_args							)

					lowercase						:		List[str]   					=		os.path.join(a_							,      "model.onnx"							)
					logger.info(F'''exporting model to {output_model_file}'''							)

					lowercase						:		Any   					=		{0: "batch_size", 1: "seq_len"}

					torch.onnx.export(
					    a_							,      a_							,      a_							,      export_params=a_							,      opset_version=1_3							,      do_constant_folding=a_							,      input_names=["input_ids", "attention_mask", "token_type_ids"]							,      output_names=["output_start_logits", "output_end_logits"]							,      dynamic_axes={
					        "input_ids": axes,
					        "attention_mask": axes,
					        "token_type_ids": axes,
					        "output_start_logits": axes,
					        "output_end_logits": axes,
					    }							,      verbose=a_							,      )
					logger.info("onnx export finished"							)


