





from __future__ import annotations

import unittest

from transformers import CTRLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers.models.ctrl.modeling_tf_ctrl import (
          TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFCTRLForSequenceClassification,
          TFCTRLLMHeadModel,
          TFCTRLModel,
      )

class 	__snake_case					(      a_ ):






    '''simple docstring'''





    def __init__( self    :		Dict      ,    snake_case    :		List[str]      ,    ):

      """simple docstring"""





      __lowerCAmelCase								=					parent
      __lowerCAmelCase								=					13
      __lowerCAmelCase								=					7
      __lowerCAmelCase								=					True
      __lowerCAmelCase								=					True
      __lowerCAmelCase								=					True
      __lowerCAmelCase								=					True
      __lowerCAmelCase								=					True
      __lowerCAmelCase								=					99
      __lowerCAmelCase								=					32
      __lowerCAmelCase								=					2
      __lowerCAmelCase								=					4
      __lowerCAmelCase								=					37
      __lowerCAmelCase								=					'gelu'
      __lowerCAmelCase								=					0.1
      __lowerCAmelCase								=					0.1
      __lowerCAmelCase								=					512
      __lowerCAmelCase								=					16
      __lowerCAmelCase								=					2
      __lowerCAmelCase								=					0.02
      __lowerCAmelCase								=					3
      __lowerCAmelCase								=					4
      __lowerCAmelCase								=					None
      __lowerCAmelCase								=					self.vocab_size - 1





    def 						A		( self    :		int				):

      """simple docstring"""





      __lowerCAmelCase								=					ids_tensor([self.batch_size, self.seq_length]      ,    self.vocab_size				)

      __lowerCAmelCase								=					None
      if self.use_input_mask:
        __lowerCAmelCase								=					random_attention_mask([self.batch_size, self.seq_length]				)

      __lowerCAmelCase								=					None
      if self.use_token_type_ids:
        __lowerCAmelCase								=					ids_tensor([self.batch_size, self.seq_length]      ,    self.type_vocab_size				)

      __lowerCAmelCase								=					None
      if self.use_mc_token_ids:
        __lowerCAmelCase								=					ids_tensor([self.batch_size, self.num_choices]      ,    self.seq_length				)

      __lowerCAmelCase								=					None
      __lowerCAmelCase								=					None
      __lowerCAmelCase								=					None
      if self.use_labels:
        __lowerCAmelCase								=					ids_tensor([self.batch_size]      ,    self.type_sequence_label_size				)
        __lowerCAmelCase								=					ids_tensor([self.batch_size, self.seq_length]      ,    self.num_labels				)
        __lowerCAmelCase								=					ids_tensor([self.batch_size]      ,    self.num_choices				)

      __lowerCAmelCase								=					CTRLConfig(
          vocab_size=self.vocab_size      ,    n_embd=self.hidden_size      ,    n_layer=self.num_hidden_layers      ,    n_head=self.num_attention_heads      ,    n_positions=self.max_position_embeddings      ,    pad_token_id=self.pad_token_id      ,    )

      __lowerCAmelCase								=					ids_tensor([self.num_hidden_layers, self.num_attention_heads]      ,    2				)

      return (
          config,
          input_ids,
          input_mask,
          head_mask,
          token_type_ids,
          mc_token_ids,
          sequence_labels,
          token_labels,
          choice_labels,
      )





    def 						A		( self    :		Union[str, Any]      ,    snake_case    :		Optional[Any]      ,    snake_case    :		Optional[Any]      ,    snake_case    :		Dict      ,    snake_case    :		List[Any]      ,    snake_case    :		str      ,    *snake_case    :		List[str]				):

      """simple docstring"""





      __lowerCAmelCase								=					TFCTRLModel(config=UpperCAmelCase_				)
      __lowerCAmelCase								=					{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
      __lowerCAmelCase								=					model(UpperCAmelCase_				)

      __lowerCAmelCase								=					[input_ids, None, input_mask]  # None is the input for 'past'
      __lowerCAmelCase								=					model(UpperCAmelCase_				)

      __lowerCAmelCase								=					model(UpperCAmelCase_				)

      self.parent.assertEqual(result.last_hidden_state.shape      ,    (self.batch_size, self.seq_length, self.hidden_size)				)





    def 						A		( self    :		List[str]      ,    snake_case    :		Optional[Any]      ,    snake_case    :		List[str]      ,    snake_case    :		Dict      ,    snake_case    :		Tuple      ,    snake_case    :		Dict      ,    *snake_case    :		str				):

      """simple docstring"""





      __lowerCAmelCase								=					TFCTRLLMHeadModel(config=UpperCAmelCase_				)
      __lowerCAmelCase								=					{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
      __lowerCAmelCase								=					model(UpperCAmelCase_				)
      self.parent.assertEqual(result.logits.shape      ,    (self.batch_size, self.seq_length, self.vocab_size)				)





    def 						A		( self    :		Tuple      ,    snake_case    :		str      ,    snake_case    :		Optional[Any]      ,    snake_case    :		int      ,    snake_case    :		Tuple      ,    snake_case    :		Optional[Any]      ,    *snake_case    :		List[str]				):

      """simple docstring"""





      __lowerCAmelCase								=					self.num_labels
      __lowerCAmelCase								=					ids_tensor([self.batch_size]      ,    self.type_sequence_label_size				)
      __lowerCAmelCase								=					{
          'input_ids': input_ids,
          'token_type_ids': token_type_ids,
          'labels': sequence_labels,
      }
      __lowerCAmelCase								=					TFCTRLForSequenceClassification(UpperCAmelCase_				)
      __lowerCAmelCase								=					model(UpperCAmelCase_				)
      self.parent.assertEqual(result.logits.shape      ,    (self.batch_size, self.num_labels)				)





    def 						A		( self    :		int				):

      """simple docstring"""





      __lowerCAmelCase								=					self.prepare_config_and_inputs()

      (
          __lowerCAmelCase			
      )								=					config_and_inputs

      __lowerCAmelCase								=					{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
      return config, inputs_dict


@require_tf
class 	__snake_case					(      a_							,					a_							,					unittest.TestCase ):






    '''simple docstring'''
    _lowerCAmelCase  =						(TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
    _lowerCAmelCase  =						(TFCTRLLMHeadModel,) if is_tf_available() else ()
    _lowerCAmelCase  =						(
        {
            "feature-extraction": TFCTRLModel,
            "text-classification": TFCTRLForSequenceClassification,
            "text-generation": TFCTRLLMHeadModel,
            "zero-shot": TFCTRLForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    _lowerCAmelCase  =						False
    _lowerCAmelCase  =						False





    def 						A		( self    :		Union[str, Any]      ,    snake_case    :		Dict      ,    snake_case    :		Any      ,    snake_case    :		Optional[Any]      ,    snake_case    :		Tuple      ,    snake_case    :		Optional[Any]				):

      """simple docstring"""





      if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
        # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
        # `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
        # config could not be created.
        return True

      return False





    def 						A		( self    :		Optional[Any]				):

      """simple docstring"""





      __lowerCAmelCase								=					TFCTRLModelTester(self				)
      __lowerCAmelCase								=					ConfigTester(self      ,    config_class=UpperCAmelCase_      ,    n_embd=37				)





    def 						A		( self    :		List[Any]				):

      """simple docstring"""





      self.config_tester.run_common_tests()





    def 						A		( self    :		str				):

      """simple docstring"""





      __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_ctrl_model(*UpperCAmelCase_				)





    def 						A		( self    :		Dict				):

      """simple docstring"""





      __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_ctrl_lm_head(*UpperCAmelCase_				)





    def 						A		( self    :		List[str]				):

      """simple docstring"""





      __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_ctrl_for_sequence_classification(*UpperCAmelCase_				)





    def 						A		( self    :		Optional[Any]				):

      """simple docstring"""





      __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs_for_common()
      __lowerCAmelCase								=					[TFCTRLLMHeadModel]
      __lowerCAmelCase								=					[TFCTRLForSequenceClassification]

      for model_class in self.all_model_classes:
        __lowerCAmelCase								=					model_class(UpperCAmelCase_				)
        model.build()  # may be needed for the get_bias() call below
        assert isinstance(model.get_input_embeddings()      ,    tf.keras.layers.Layer				)

        if model_class in list_lm_models:
          __lowerCAmelCase								=					model.get_output_embeddings()
          assert isinstance(UpperCAmelCase_      ,    tf.keras.layers.Layer				)
          __lowerCAmelCase								=					model.get_bias()
          assert isinstance(UpperCAmelCase_      ,    UpperCAmelCase_				)
          for k, v in name.items():
            assert isinstance(UpperCAmelCase_      ,    tf.Variable				)
        elif model_class in list_other_models_with_output_ebd:
          __lowerCAmelCase								=					model.get_output_embeddings()
          assert isinstance(UpperCAmelCase_      ,    tf.keras.layers.Layer				)
          __lowerCAmelCase								=					model.get_bias()
          assert name is None
        else:
          __lowerCAmelCase								=					model.get_output_embeddings()
          assert x is None
          __lowerCAmelCase								=					model.get_bias()
          assert name is None





    @slow
    def 						A		( self    :		Dict				):

      """simple docstring"""





      for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __lowerCAmelCase								=					TFCTRLModel.from_pretrained(UpperCAmelCase_				)
        self.assertIsNotNone(UpperCAmelCase_				)


@require_tf
class 	__snake_case					(      unittest.TestCase ):






    '''simple docstring'''





    @slow
    def 						A		( self    :		Tuple				):

      """simple docstring"""





      __lowerCAmelCase								=					TFCTRLLMHeadModel.from_pretrained('''ctrl'''				)
      __lowerCAmelCase								=					tf.convert_to_tensor([[11_859, 0, 1_611, 8]]      ,    dtype=tf.intaa				)  # Legal the president is
      __lowerCAmelCase								=					[
          11_859,
          0,
          1_611,
          8,
          5,
          150,
          26_449,
          2,
          19,
          348,
          469,
          3,
          2_595,
          48,
          20_740,
          246_533,
          246_533,
          19,
          30,
          5,
      ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

      __lowerCAmelCase								=					model.generate(UpperCAmelCase_      ,    do_sample=UpperCAmelCase_				)
      self.assertListEqual(output_ids[0].numpy().tolist()      ,    UpperCAmelCase_				)

'''simple docstring'''





import os
from functools import partial
from glob import glob

import faiss
from datasets import Features, Sequence, Value, concatenate_datasets, load_dataset, load_from_disk

from transformers import DPRContextEncoder, DPRContextEncoderTokenizerFast






def SCREAMING_SNAKE_CASE__  (				snake_case       :				List[Any]				,	snake_case       :				List[str]=100				,	snake_case       :				Any=" "  )   ->				Any:







     """simple docstring"""
     a					:							Union[str, Any]							= text.split(snake_case  )
     return [character.join(text[i : i + n]  ).strip() for i in range(0				,	len(snake_case  )				,	snake_case  )]






def SCREAMING_SNAKE_CASE__  (				snake_case       :				Optional[int]  )   ->				Union[str, Any]:







     """simple docstring"""
     a							,			a					:							int							= [], []
     for title, text in zip(documents['title']				,	documents['text']  ):
          if text is not None:
               for passage in split_text(snake_case  ):
                    titles.append(title if title is not None else ''  )
                    texts.append(snake_case  )
     return {"title": titles, "text": texts}






def SCREAMING_SNAKE_CASE__  (				snake_case       :				int				,	snake_case       :				str				,	snake_case       :				List[Any]				,	snake_case       :				Dict				,	snake_case       :				Union[str, Any]				,	snake_case       :				int  )   ->				Dict:







     """simple docstring"""
     a					:							Tuple							= load_dataset(
         'csv'				,	data_files=[csv_path]				,	split='train'				,	delimiter='\t'				,	column_names=['title', 'text']  )
     a					:							Union[str, Any]							= kb_dataset.map(
         snake_case				,	batched=snake_case				,	num_proc=1  )  # if you want you can load already splitted csv.
     a					:							Optional[Any]							= [kb_dataset.shard(snake_case				,	snake_case				,	contiguous=snake_case  ) for i in range(snake_case  )]
     a					:							List[Any]							= kb_list[process_num]

     a					:							Dict							= 'data_' + str(snake_case  )
     a					:							int							= os.path.join(snake_case				,	snake_case  )

     a					:							Any							= DPRContextEncoderTokenizerFast.from_pretrained('facebook/dpr-ctx_encoder-multiset-base'  )
     a					:							List[str]							= ctx_encoder.to(device=snake_case  )

     def embed(
         snake_case       :				dict				,	snake_case       :				DPRContextEncoder				,	snake_case       :				DPRContextEncoderTokenizerFast				,	snake_case       :				Any  ) -> dict:
          a					:							Tuple							= ctx_tokenizer(
              documents['title']				,	documents['text']				,	truncation=snake_case				,	padding='longest'				,	return_tensors='pt'  )['input_ids']
          a					:							Any							= ctx_encoder(input_ids.to(device=snake_case  )				,	return_dict=snake_case  ).pooler_output
          return {"embeddings": embeddings.detach().cpu().numpy()}

     a					:							Union[str, Any]							= Features(
         {'text': Value('string'  ), 'title': Value('string'  ), 'embeddings': Sequence(Value('float32'  )  )}  )  # optional, save as float32 instead of float64 to save space

     a					:							List[str]							= data_shrad.map(
         partial(snake_case				,	ctx_encoder=snake_case				,	ctx_tokenizer=snake_case				,	device=snake_case  )				,	batched=snake_case				,	batch_size=16				,	features=snake_case				,	)
     dataset.save_to_disk(snake_case  )






def SCREAMING_SNAKE_CASE__  (				snake_case       :				Any				,	snake_case       :				Any  )   ->				List[Any]:







     """simple docstring"""
     a					:							int							= []

     for shard_address in glob(str(snake_case  ) + '/*/'  ):
          data_shard_list.append(load_from_disk(snake_case  )  )

     a					:							Tuple							= concatenate_datasets(snake_case  )
     faiss.omp_set_num_threads(96  )

     a					:							List[str]							= faiss.IndexHNSWFlat(768				,	128				,	faiss.METRIC_INNER_PRODUCT  )
     concat.add_faiss_index('embeddings'				,	custom_index=snake_case  )
     concat.get_index('embeddings'  ).save(
         snake_case  )  # since we load the index in to memory,we can directly update the index in the disk

