

import logging
import os
from dataclasses import dataclass
from typing import List, Optional, Union

import tqdm
from filelock import FileLock

from transformers import (
    BartTokenizer,
    BartTokenizerFast,
    DataProcessor,
    PreTrainedTokenizer,
    RobertaTokenizer,
    RobertaTokenizerFast,
    XLMRobertaTokenizer,
    is_tf_available,
    is_torch_available,
)


UpperCamelCase__  :      Any			=    logging.getLogger(__name__)



@dataclass(frozen=lowerCAmelCase__       )
class       __snake_case  :
  __lowerCAmelCase     :	str
  __lowerCAmelCase     :	str
  __lowerCAmelCase     :	Optional[str]						 =  None
  __lowerCAmelCase     :	Optional[str]						 =  None
  __lowerCAmelCase     :	Optional[str]						 =  None






@dataclass(frozen=lowerCAmelCase__       )
class       __snake_case  :
  __lowerCAmelCase     :	List[int]
  __lowerCAmelCase     :	Optional[List[int]]						 =  None
  __lowerCAmelCase     :	Optional[List[int]]						 =  None
  __lowerCAmelCase     :	Optional[Union[int, float]]						 =  None
  __lowerCAmelCase     :	Optional[int]						 =  None


if is_torch_available():
       import torch
       from torch.utils.data import Dataset



       class       __snake_case  (    lowerCAmelCase__       ):
         __lowerCAmelCase     :	List[InputFeatures]

         def __init__(					self						,		_A						,		_A						,		_A						,		_A = None						,		_A=False						,		_A = False						,		):
             SCREAMING_SNAKE_CASE_									=						hans_processors[task]()

             SCREAMING_SNAKE_CASE_									=						os.path.join(
                 _A						,		'cached_{}_{}_{}_{}'.format(
                     'dev' if evaluate else 'train'						,		tokenizer.__class__.__name__						,		str(_A)						,		_A						,		)						,		)
             SCREAMING_SNAKE_CASE_									=						processor.get_labels()
             if tokenizer.__class__ in (
                 RobertaTokenizer,
                 RobertaTokenizerFast,
                 XLMRobertaTokenizer,
                 BartTokenizer,
                 BartTokenizerFast,
             ):
                 # HACK(label indices are swapped in RoBERTa pretrained model)
                 SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						label_list[2], label_list[1]
             SCREAMING_SNAKE_CASE_									=						label_list

             # Make sure only the first process in distributed training processes the dataset,
             # and the others will use the cache.
             SCREAMING_SNAKE_CASE_									=						cached_features_file + '.lock'
             with FileLock(_A):
                 if os.path.exists(_A) and not overwrite_cache:
                     logger.info(f"""Loading features from cached file {cached_features_file}""")
                     SCREAMING_SNAKE_CASE_									=						torch.load(_A)
                 else:
                     logger.info(f"""Creating features from dataset file at {data_dir}""")

                     SCREAMING_SNAKE_CASE_									=						(
                         processor.get_dev_examples(_A) if evaluate else processor.get_train_examples(_A)
                     )

                     logger.info('Training examples: %s'						,		len(_A))
                     SCREAMING_SNAKE_CASE_									=						hans_convert_examples_to_features(_A						,		_A						,		_A						,		_A)
                     logger.info('Saving features into cached file %s'						,		_A)
                     torch.save(self.features						,		_A)

         def __len__(					self):
             return len(self.features)

         def __getitem__(					self						,		_A):
             return self.features[i]


         def 						lowerCAmelCase__     (					self):
             return self.label_list


if is_tf_available():
       import tensorflow as tf



       class       __snake_case  :
         __lowerCAmelCase     :	List[InputFeatures]

         def __init__(					self						,		_A						,		_A						,		_A						,		_A = 128						,		_A=False						,		_A = False						,		):
             SCREAMING_SNAKE_CASE_									=						hans_processors[task]()
             SCREAMING_SNAKE_CASE_									=						processor.get_labels()
             if tokenizer.__class__ in (
                 RobertaTokenizer,
                 RobertaTokenizerFast,
                 XLMRobertaTokenizer,
                 BartTokenizer,
                 BartTokenizerFast,
             ):
                 # HACK(label indices are swapped in RoBERTa pretrained model)
                 SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						label_list[2], label_list[1]
             SCREAMING_SNAKE_CASE_									=						label_list

             SCREAMING_SNAKE_CASE_									=						processor.get_dev_examples(_A) if evaluate else processor.get_train_examples(_A)
             SCREAMING_SNAKE_CASE_									=						hans_convert_examples_to_features(_A						,		_A						,		_A						,		_A)

             def gen():
                 for ex_index, ex in tqdm.tqdm(enumerate(self.features)						,		desc='convert examples to features'):
                     if ex_index % 10000 == 0:
                         logger.info('Writing example %d of %d' % (ex_index, len(_A)))

                     yield (
                         {
                             "example_id": 0,
                             "input_ids": ex.input_ids,
                             "attention_mask": ex.attention_mask,
                             "token_type_ids": ex.token_type_ids,
                         },
                         ex.label,
                     )

             SCREAMING_SNAKE_CASE_									=						tf.data.Dataset.from_generator(
                 _A						,		(
                     {
                         'example_id': tf.intaa,
                         'input_ids': tf.intaa,
                         'attention_mask': tf.intaa,
                         'token_type_ids': tf.intaa,
                     },
                     tf.intaa,
                 )						,		(
                     {
                         'example_id': tf.TensorShape([]),
                         'input_ids': tf.TensorShape([None, None]),
                         'attention_mask': tf.TensorShape([None, None]),
                         'token_type_ids': tf.TensorShape([None, None]),
                     },
                     tf.TensorShape([]),
                 )						,		)

         def 						lowerCAmelCase__     (					self):
             return self.dataset

         def __len__(					self):
             return len(self.features)

         def __getitem__(					self						,		_A):
             return self.features[i]


         def 						lowerCAmelCase__     (					self):
             return self.label_list



class       __snake_case  (    lowerCAmelCase__       ):

  def 						lowerCAmelCase__     (					self						,		_A):
      return self._create_examples(self._read_tsv(os.path.join(_A						,		'heuristics_train_set.txt'))						,		'train')

  def 						lowerCAmelCase__     (					self						,		_A):
      return self._create_examples(self._read_tsv(os.path.join(_A						,		'heuristics_evaluation_set.txt'))						,		'dev')

  def 						lowerCAmelCase__     (					self):
      return ["contradiction", "entailment", "neutral"]


  def 						lowerCAmelCase__     (					self						,		_A						,		_A):
      SCREAMING_SNAKE_CASE_									=						[]
      for i, line in enumerate(_A):
          if i == 0:
              continue
          SCREAMING_SNAKE_CASE_									=						'%s-%s' % (set_type, line[0])
          SCREAMING_SNAKE_CASE_									=						line[5]
          SCREAMING_SNAKE_CASE_									=						line[6]
          SCREAMING_SNAKE_CASE_									=						line[7][2:] if line[7].startswith('ex') else line[7]
          SCREAMING_SNAKE_CASE_									=						line[0]
          examples.append(InputExample(guid=_A						,		text_a=_A						,		text_b=_A						,		label=_A						,		pairID=_A))
      return examples






def 					_UpperCAmelCase       (  _SCREAMING_SNAKE_CASE					:		List[InputExample] ,   _SCREAMING_SNAKE_CASE					:		List[str] ,   _SCREAMING_SNAKE_CASE					:		int ,   _SCREAMING_SNAKE_CASE					:		PreTrainedTokenizer ,   ):





    """simple docstring"""


    SCREAMING_SNAKE_CASE_									=						{label: i for i, label in enumerate(_SCREAMING_SNAKE_CASE       )}

    SCREAMING_SNAKE_CASE_									=						[]
    for ex_index, example in tqdm.tqdm(enumerate(_SCREAMING_SNAKE_CASE       ) ,   desc='convert examples to features'       ):
        if ex_index % 10_000 == 0:
            logger.info('Writing example %d' % (ex_index)       )

        SCREAMING_SNAKE_CASE_									=						tokenizer(
            example.text_a ,   example.text_b ,   add_special_tokens=_SCREAMING_SNAKE_CASE ,   max_length=_SCREAMING_SNAKE_CASE ,   padding='max_length' ,   truncation=_SCREAMING_SNAKE_CASE ,   return_overflowing_tokens=_SCREAMING_SNAKE_CASE ,   )

        SCREAMING_SNAKE_CASE_									=						label_map[example.label] if example.label in label_map else 0

        SCREAMING_SNAKE_CASE_									=						int(example.pairID       )

        features.append(InputFeatures(**_SCREAMING_SNAKE_CASE ,   label=_SCREAMING_SNAKE_CASE ,   pairID=_SCREAMING_SNAKE_CASE       )       )

    for i, example in enumerate(examples[:5]       ):
        logger.info('*** Example ***'       )
        logger.info(f"""guid: {example}"""       )
        logger.info(f"""features: {features[i]}"""       )

    return features


UpperCamelCase__  :      Any			=    {
    "hans": 3,
}





UpperCamelCase__  :      List[Any]			=    {
    "hans": HansProcessor,
}





import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM



@require_torch
class       __snake_case  :

  def __init__(					self						,		_A						,		_A=99						,		_A=13						,		_A=16						,		_A=7						,		_A=True						,		_A=True						,		_A=True						,		_A=False						,		_A=True						,		_A=2						,		_A=32						,		_A=4						,		_A=4						,		_A=30						,		_A=0						,		_A=1						,		_A=2						,		_A=None						,		):
      SCREAMING_SNAKE_CASE_									=						parent
      SCREAMING_SNAKE_CASE_									=						batch_size
      SCREAMING_SNAKE_CASE_									=						decoder_seq_length
      # For common tests
      SCREAMING_SNAKE_CASE_									=						self.decoder_seq_length
      SCREAMING_SNAKE_CASE_									=						is_training
      SCREAMING_SNAKE_CASE_									=						use_attention_mask
      SCREAMING_SNAKE_CASE_									=						use_labels

      SCREAMING_SNAKE_CASE_									=						vocab_size
      SCREAMING_SNAKE_CASE_									=						d_model
      SCREAMING_SNAKE_CASE_									=						d_model
      SCREAMING_SNAKE_CASE_									=						decoder_layers
      SCREAMING_SNAKE_CASE_									=						decoder_layers
      SCREAMING_SNAKE_CASE_									=						decoder_ffn_dim
      SCREAMING_SNAKE_CASE_									=						decoder_attention_heads
      SCREAMING_SNAKE_CASE_									=						decoder_attention_heads
      SCREAMING_SNAKE_CASE_									=						eos_token_id
      SCREAMING_SNAKE_CASE_									=						bos_token_id
      SCREAMING_SNAKE_CASE_									=						pad_token_id
      SCREAMING_SNAKE_CASE_									=						decoder_start_token_id
      SCREAMING_SNAKE_CASE_									=						use_cache
      SCREAMING_SNAKE_CASE_									=						max_position_embeddings

      SCREAMING_SNAKE_CASE_									=						None
      SCREAMING_SNAKE_CASE_									=						decoder_seq_length
      SCREAMING_SNAKE_CASE_									=						2
      SCREAMING_SNAKE_CASE_									=						1

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.decoder_seq_length]						,		self.vocab_size)

      SCREAMING_SNAKE_CASE_									=						None
      if self.use_attention_mask:
          SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.decoder_seq_length]						,		vocab_size=2)

      SCREAMING_SNAKE_CASE_									=						None
      if self.use_labels:
          SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.decoder_seq_length]						,		self.vocab_size)

      SCREAMING_SNAKE_CASE_									=						TrOCRConfig(
          vocab_size=self.vocab_size						,		d_model=self.d_model						,		decoder_layers=self.decoder_layers						,		decoder_ffn_dim=self.decoder_ffn_dim						,		decoder_attention_heads=self.decoder_attention_heads						,		eos_token_id=self.eos_token_id						,		bos_token_id=self.bos_token_id						,		use_cache=self.use_cache						,		pad_token_id=self.pad_token_id						,		decoder_start_token_id=self.decoder_start_token_id						,		max_position_embeddings=self.max_position_embeddings						,		)

      return (config, input_ids, attention_mask, lm_labels)

  def 						lowerCAmelCase__     (					self						,		_A						,		_A						,		_A						,		_A						,		):
      SCREAMING_SNAKE_CASE_									=						True
      SCREAMING_SNAKE_CASE_									=						TrOCRDecoder(config=_A).to(_A).eval()
      SCREAMING_SNAKE_CASE_									=						input_ids[:2]

      input_ids[input_ids == 0] += 1
      # first forward pass
      SCREAMING_SNAKE_CASE_									=						model(_A						,		use_cache=_A)
      SCREAMING_SNAKE_CASE_									=						model(_A)
      SCREAMING_SNAKE_CASE_									=						model(_A						,		use_cache=_A)

      self.parent.assertTrue(len(_A) == len(_A))
      self.parent.assertTrue(len(_A) == len(_A) + 1)

      SCREAMING_SNAKE_CASE_									=						outputs['past_key_values']

      # create hypothetical next token and extent to next_input_ids
      SCREAMING_SNAKE_CASE_									=						ids_tensor((2, 1)						,		config.vocab_size - 1) + 1

      # append to next input_ids and
      SCREAMING_SNAKE_CASE_									=						torch.cat([input_ids, next_tokens]						,		dim=-1)

      SCREAMING_SNAKE_CASE_									=						model(_A)['last_hidden_state']
      SCREAMING_SNAKE_CASE_									=						model(_A						,		past_key_values=_A)['last_hidden_state']

      # select random slice
      SCREAMING_SNAKE_CASE_									=						ids_tensor((1,)						,		output_from_past.shape[-1]).item()
      SCREAMING_SNAKE_CASE_									=						output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
      SCREAMING_SNAKE_CASE_									=						output_from_past[:, 0, random_slice_idx].detach()

      # test that outputs are equal for slice
      assert torch.allclose(_A						,		_A						,		atol=1E-3)


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.prepare_config_and_inputs()
      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						config_and_inputs

      SCREAMING_SNAKE_CASE_									=						{'input_ids': input_ids, 'attention_mask': attention_mask}
      return config, inputs_dict






@require_torch
class       __snake_case  (    lowerCAmelCase__     , lowerCAmelCase__     , lowerCAmelCase__     , unittest.TestCase       ):
  __lowerCAmelCase     :	Tuple						 =  (TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
  __lowerCAmelCase     :	Union[str, Any]						 =  (TrOCRForCausalLM,) if is_torch_available() else ()
  __lowerCAmelCase     :	str						 =  {'text-generation': TrOCRForCausalLM} if is_torch_available() else {}
  __lowerCAmelCase     :	Any						 =  True
  __lowerCAmelCase     :	str						 =  False

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						TrOCRStandaloneDecoderModelTester(self						,		is_training=_A)
      SCREAMING_SNAKE_CASE_									=						ConfigTester(self						,		config_class=_A)

  def 						lowerCAmelCase__     (					self):
      pass

  def 						lowerCAmelCase__     (					self):
      pass

  def 						lowerCAmelCase__     (					self):
      pass

  def 						lowerCAmelCase__     (					self):
      self.config_tester.run_common_tests()

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_decoder_model_past(*_A)

  def 						lowerCAmelCase__     (					self):
      return


  @unittest.skip('The model doesn\'t support left padding')  # and it's not used enough to be worth fixing :)
  def 						lowerCAmelCase__     (					self):
      pass


