

import random
import unittest

from torch.utils.data import BatchSampler, DataLoader, IterableDataset

from accelerate import Accelerator
from accelerate.data_loader import (
    BatchSamplerShard,
    DataLoaderDispatcher,
    DataLoaderShard,
    IterableDatasetShard,
    SkipBatchSampler,
    SkipDataLoader,
    skip_first_batches,
)



class        A__(_A							):
						"""simple docstring"""






						def __init__(    self  ,      _lowercase=0.0_1  ,      _lowercase=1_000    )     -> Optional[int]:
									a_						:						int 							= p_stop
									a_						:						List[str] 							= max_length





						def __iter__(    self    )     -> List[Any]:
									a_						:						int 							= 0
									a_						:						Any 							= False
									while not stop and count < self.max_length:
												yield count
												count += 1
												a_						:						int 							= random.random() < self.p_stop







class        A__(unittest.TestCase							):
						"""simple docstring"""






						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase  ,      _lowercase=False  ,      _lowercase=True    )     -> Optional[int]:
									a_						:						Tuple 							= [
									    BatchSamplerShard(a_  ,      2  ,      a_  ,      split_batches=a_  ,      even_batches=a_    )
									    for i in range(2    )
									]
									a_						:						Dict 							= [list(a_    ) for batch_sampler_shard in batch_sampler_shards]
									if not split_batches:
												self.assertListEqual([len(a_    ) for shard in batch_sampler_shards]  ,      [len(a_    ) for e in expected]    )
									self.assertListEqual(a_  ,      a_    )





						def 			UpperCamelCase__					(    self    )     -> int:
									# Check the shards when the dataset is a round multiple of total batch size.
									a_						:						Any 							= BatchSampler(range(24    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						str 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
									]
									self.check_batch_sampler_shards(a_  ,      a_    )

									a_						:						int 							= BatchSampler(range(24    )  ,      batch_size=3  ,      drop_last=a_    )
									# Expected shouldn't change
									self.check_batch_sampler_shards(a_  ,      a_    )

									# Check the shards when the dataset is a round multiple of batch size but not total batch size.
									a_						:						Union[str, Any] 							= BatchSampler(range(21    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Dict 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [0, 1, 2]],
									]
									self.check_batch_sampler_shards(a_  ,      a_    )

									a_						:						Union[str, Any] 							= BatchSampler(range(21    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						List[Any] 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(a_  ,      a_    )

									# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
									# num_processes batch.
									a_						:						List[str] 							= BatchSampler(range(22    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Union[str, Any] 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 0, 1]],
									]
									self.check_batch_sampler_shards(a_  ,      a_    )

									a_						:						Tuple 							= BatchSampler(range(22    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Union[str, Any] 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(a_  ,      a_    )

									# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
									# num_processes batch.
									a_						:						Union[str, Any] 							= BatchSampler(range(20    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Union[str, Any] 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 0]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [1, 2, 3]],
									]
									self.check_batch_sampler_shards(a_  ,      a_    )

									a_						:						str 							= BatchSampler(range(20    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						int 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(a_  ,      a_    )

									# Check the shards when the dataset is very small.
									a_						:						str 							= BatchSampler(range(2    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Optional[Any] 							= [[[0, 1, 0]], [[1, 0, 1]]]
									self.check_batch_sampler_shards(a_  ,      a_    )

									a_						:						Union[str, Any] 							= BatchSampler(range(2    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						int 							= [[], []]
									self.check_batch_sampler_shards(a_  ,      a_    )





						def 			UpperCamelCase__					(    self    )     -> List[str]:
									# Check the shards when the dataset is a round multiple of batch size.
									a_						:						Union[str, Any] 							= BatchSampler(range(24    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						Optional[int] 							= [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_    )

									a_						:						int 							= BatchSampler(range(24    )  ,      batch_size=4  ,      drop_last=a_    )
									# Expected shouldn't change
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_    )

									# Check the shards when the dataset is not a round multiple of batch size.
									a_						:						int 							= BatchSampler(range(22    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						Dict 							= [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [0, 1]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_    )

									a_						:						List[Any] 							= BatchSampler(range(22    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						int 							= [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_    )

									# Check the shards when the dataset is not a round multiple of batch size or num_processes.
									a_						:						str 							= BatchSampler(range(21    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						List[Any] 							= [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 0]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [1, 2]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_    )

									a_						:						Dict 							= BatchSampler(range(21    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						str 							= [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_    )

									# Check the shards when the dataset is very small.
									a_						:						Any 							= BatchSampler(range(2    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						Optional[int] 							= [[[0, 1]], [[0, 1]]]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_    )

									a_						:						List[Any] 							= BatchSampler(range(2    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						Dict 							= [[], []]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_    )





						def 			UpperCamelCase__					(    self    )     -> List[Any]:
									# Check the shards when the dataset is a round multiple of total batch size.
									a_						:						Union[str, Any] 							= BatchSampler(range(24    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Any 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      even_batches=a_    )

									a_						:						Tuple 							= BatchSampler(range(24    )  ,      batch_size=3  ,      drop_last=a_    )
									# Expected shouldn't change
									self.check_batch_sampler_shards(a_  ,      a_  ,      even_batches=a_    )

									# Check the shards when the dataset is a round multiple of batch size but not total batch size.
									a_						:						Tuple 							= BatchSampler(range(21    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Optional[Any] 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      even_batches=a_    )

									a_						:						int 							= BatchSampler(range(21    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Any 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      even_batches=a_    )

									# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
									# num_processes batch.
									a_						:						Union[str, Any] 							= BatchSampler(range(22    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Union[str, Any] 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      even_batches=a_    )

									a_						:						Optional[Any] 							= BatchSampler(range(22    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						str 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      even_batches=a_    )

									# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
									# num_processes batch.
									a_						:						Optional[Any] 							= BatchSampler(range(20    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Optional[Any] 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      even_batches=a_    )

									a_						:						Optional[int] 							= BatchSampler(range(20    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						List[str] 							= [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      even_batches=a_    )

									# Check the shards when the dataset is very small.
									a_						:						Union[str, Any] 							= BatchSampler(range(2    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Any 							= [[[0, 1]], []]
									self.check_batch_sampler_shards(a_  ,      a_  ,      even_batches=a_    )

									a_						:						Union[str, Any] 							= BatchSampler(range(2    )  ,      batch_size=3  ,      drop_last=a_    )
									a_						:						Union[str, Any] 							= [[], []]
									self.check_batch_sampler_shards(a_  ,      a_  ,      even_batches=a_    )





						def 			UpperCamelCase__					(    self    )     -> Union[str, Any]:
									# Check the shards when the dataset is a round multiple of batch size.
									a_						:						Optional[Any] 							= BatchSampler(range(24    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						Tuple 							= [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_  ,      even_batches=a_    )

									a_						:						List[str] 							= BatchSampler(range(24    )  ,      batch_size=4  ,      drop_last=a_    )
									# Expected shouldn't change
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_  ,      even_batches=a_    )

									# Check the shards when the dataset is not a round multiple of batch size.
									a_						:						Dict 							= BatchSampler(range(22    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						Optional[int] 							= [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_  ,      even_batches=a_    )

									a_						:						Any 							= BatchSampler(range(22    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						List[str] 							= [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_  ,      even_batches=a_    )

									# Check the shards when the dataset is not a round multiple of batch size or num_processes.
									a_						:						Dict 							= BatchSampler(range(21    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						Tuple 							= [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_  ,      even_batches=a_    )

									a_						:						str 							= BatchSampler(range(21    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						Tuple 							= [
									    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
									    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
									]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_  ,      even_batches=a_    )

									# Check the shards when the dataset is very small.
									a_						:						Tuple 							= BatchSampler(range(2    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						Tuple 							= [[[0, 1]], []]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_  ,      even_batches=a_    )

									a_						:						int 							= BatchSampler(range(2    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						Tuple 							= [[], []]
									self.check_batch_sampler_shards(a_  ,      a_  ,      split_batches=a_  ,      even_batches=a_    )





						def 			UpperCamelCase__					(    self    )     -> Tuple:
									a_						:						List[Any] 							= [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
									a_						:						Optional[Any] 							= [BatchSamplerShard(a_  ,      2  ,      a_  ,      even_batches=a_    ) for i in range(2    )]

									self.assertEqual(len(batch_sampler_shards[0]    )  ,      3    )
									self.assertEqual(len(batch_sampler_shards[1]    )  ,      2    )

									self.assertListEqual(list(batch_sampler_shards[0]    )  ,      [[0, 1, 2], [5, 6, 7, 8], [12, 13]]    )
									self.assertListEqual(list(batch_sampler_shards[1]    )  ,      [[3, 4], [9, 10, 11]]    )





						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase=False  ,      _lowercase=2  ,      _lowercase=False    )     -> str:
									random.seed(a_    )
									a_						:						Optional[Any] 							= list(a_    )

									a_						:						Dict 							= [
									    IterableDatasetShard(
									        a_  ,      batch_size=a_  ,      drop_last=a_  ,      num_processes=a_  ,      process_index=a_  ,      split_batches=a_  ,      )
									    for i in range(a_    )
									]
									a_						:						List[Any] 							= []
									for iterable_dataset_shard in iterable_dataset_shards:
												# Since our random iterable dataset will be... random... we need to use a seed to get reproducible results.
												random.seed(a_    )
												iterable_dataset_lists.append(list(a_    )    )

									a_						:						int 							= batch_size // num_processes if split_batches else batch_size
									# All iterable dataset shard should have the same length, a round multiple of shard_batch_size
									a_						:						Optional[Any] 							= iterable_dataset_lists[0]
									for l in iterable_dataset_lists[1:]:
												self.assertEqual(len(a_    )  ,      len(a_    )    )
												self.assertTrue(len(a_    ) % shard_batch_size == 0    )

									a_						:						Tuple 							= []
									for idx in range(0  ,      len(a_    )  ,      a_    ):
												for l in iterable_dataset_lists:
															observed += l[idx : idx + shard_batch_size]

									if not drop_last:
												while len(a_    ) < len(a_    ):
															reference += reference
									self.assertListEqual(a_  ,      reference[: len(a_    )]    )





						def 			UpperCamelCase__					(    self    )     -> str:
									a_						:						Tuple 							= 42
									a_						:						List[Any] 							= RandomIterableDataset()

									self.check_iterable_dataset_shards(a_  ,      a_  ,      batch_size=4  ,      drop_last=a_  ,      split_batches=a_    )
									self.check_iterable_dataset_shards(a_  ,      a_  ,      batch_size=4  ,      drop_last=a_  ,      split_batches=a_    )
									self.check_iterable_dataset_shards(a_  ,      a_  ,      batch_size=4  ,      drop_last=a_  ,      split_batches=a_    )
									self.check_iterable_dataset_shards(a_  ,      a_  ,      batch_size=4  ,      drop_last=a_  ,      split_batches=a_    )

									# Edge case with a very small dataset
									a_						:						int 							= RandomIterableDataset(max_length=2    )

									self.check_iterable_dataset_shards(a_  ,      a_  ,      batch_size=4  ,      drop_last=a_  ,      split_batches=a_    )
									self.check_iterable_dataset_shards(a_  ,      a_  ,      batch_size=4  ,      drop_last=a_  ,      split_batches=a_    )
									self.check_iterable_dataset_shards(a_  ,      a_  ,      batch_size=4  ,      drop_last=a_  ,      split_batches=a_    )
									self.check_iterable_dataset_shards(a_  ,      a_  ,      batch_size=4  ,      drop_last=a_  ,      split_batches=a_    )





						def 			UpperCamelCase__					(    self    )     -> int:
									a_						:						Tuple 							= BatchSampler(range(16    )  ,      batch_size=4  ,      drop_last=a_    )
									a_						:						int 							= SkipBatchSampler(a_  ,      2    )
									self.assertListEqual(list(a_    )  ,      [[8, 9, 10, 11], [12, 13, 14, 15]]    )





						def 			UpperCamelCase__					(    self    )     -> Union[str, Any]:
									a_						:						str 							= SkipDataLoader(list(range(16    )    )  ,      batch_size=4  ,      skip_batches=2    )
									self.assertListEqual([t.tolist() for t in dataloader]  ,      [[8, 9, 10, 11], [12, 13, 14, 15]]    )





						def 			UpperCamelCase__					(    self    )     -> str:
									a_						:						Optional[int] 							= DataLoader(list(range(16    )    )  ,      batch_size=4    )
									a_						:						Optional[int] 							= skip_first_batches(a_  ,      num_batches=2    )
									self.assertListEqual([t.tolist() for t in new_dataloader]  ,      [[8, 9, 10, 11], [12, 13, 14, 15]]    )





						def 			UpperCamelCase__					(    self    )     -> Union[str, Any]:
									a_						:						int 							= DataLoaderShard(list(range(16    )    )  ,      batch_size=4    )
									for idx, _ in enumerate(a_    ):
												self.assertEqual(dataloader.end_of_dataloader  ,      idx == 3    )

									# Test it also works on the second iteration
									for idx, _ in enumerate(a_    ):
												self.assertEqual(dataloader.end_of_dataloader  ,      idx == 3    )





						def 			UpperCamelCase__					(    self    )     -> Any:
									Accelerator()
									a_						:						str 							= DataLoaderDispatcher(range(16    )  ,      batch_size=4    )
									for idx, _ in enumerate(a_    ):
												self.assertEqual(dataloader.end_of_dataloader  ,      idx == 3    )

									# Test it also works on the second iteration
									for idx, _ in enumerate(a_    ):
												self.assertEqual(dataloader.end_of_dataloader  ,      idx == 3    )









'''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									=   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 		lowerCamelCase	(   _A							):



	def __init__( self		,   *a_		,   a_=None		,   a_=None		,   a_=None		,   **a_						):
			super().__init__(*a_		,   **a_						)
			lowerCAmelCase  :     Tuple       						= eval_examples
			lowerCAmelCase  :     List[Any]       						= post_process_function
			lowerCAmelCase  :     List[Any]       						= quant_trainer_args
			lowerCAmelCase  :     Dict       						= 128  # default number of calibration samples



	def 	_lowerCamelCase     ( self		,   a_=None						):
			if calib_dataset is None and self.calib_dataset is None:
					raise ValueError("Trainer: calibration requires an calib_dataset."						)
			lowerCAmelCase  :     List[str]       						= calib_dataset if calib_dataset is not None else self.calib_dataset

			lowerCAmelCase  :     Union[str, Any]       						= 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 	_lowerCamelCase     ( self		,   a_=None						):
			lowerCAmelCase  :     str       						= self.train_dataset if calib_dataset is None else calib_dataset
			lowerCAmelCase  :     str       						= self.get_calib_dataloader(a_						)

			lowerCAmelCase  :     Tuple       						= 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
					lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase  :     Tuple       						= 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						)
			lowerCAmelCase  :     Union[str, Any]       						= model



	def 	_lowerCamelCase     ( self		,   a_=None		,   a_=None		,   a_=None		,   a_ = "eval"						):
			lowerCAmelCase  :     Optional[Any]       						= self.eval_dataset if eval_dataset is None else eval_dataset
			lowerCAmelCase  :     Optional[Any]       						= self.get_eval_dataloader(a_						)
			lowerCAmelCase  :     Union[str, Any]       						= self.eval_examples if eval_examples is None else eval_examples

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

			if self.post_process_function is not None and self.compute_metrics is not None:
					lowerCAmelCase  :     Optional[Any]       						= self.post_process_function(a_		,   a_		,   output.predictions						)
					lowerCAmelCase  :     Optional[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}_'''						):
									lowerCAmelCase  :     Any       						= metrics.pop(a_						)

					self.log(a_						)
			else:
					lowerCAmelCase  :     str       						= {}

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

			lowerCAmelCase  :     Dict       						= self.callback_handler.on_evaluate(self.args		,   self.state		,   self.control		,   a_						)
			return metrics



	def 	_lowerCamelCase     ( self		,   a_		,   a_		,   a_=None		,   a_ = "test"						):
			lowerCAmelCase  :     int       						= self.get_test_dataloader(a_						)

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

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

			lowerCAmelCase  :     List[str]       						= self.post_process_function(a_		,   a_		,   output.predictions		,   "predict"						)
			lowerCAmelCase  :     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}_'''						):
							lowerCAmelCase  :     Optional[Any]       						= metrics.pop(a_						)

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





	def 	_lowerCamelCase     ( self		,   a_="./"						):
			lowerCAmelCase  :     Optional[Any]       						= self.eval_dataset
			lowerCAmelCase  :     Optional[int]       						= self.get_eval_dataloader(a_						)

			lowerCAmelCase  :     Union[str, Any]       						= next(iter(a_						)						)

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

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

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

			lowerCAmelCase  :     Any       						= True

			lowerCAmelCase  :     Dict       						= self.model.to(a_						)

			model.eval()
			model.float()

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

			lowerCAmelCase  :     Any       						= os.path.join(a_		,   "model.onnx"						)
			logger.info(F'''exporting model to {output_model_file}'''						)

			lowerCAmelCase  :     List[Any]       						= {0: "batch_size", 1: "seq_len"}

			torch.onnx.export(
			    a_		,   a_		,   a_		,   export_params=a_		,   opset_version=13		,   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"						)



