






import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


A_						:						str						      =					logging.getLogger(__name__)




class 							_lowerCAmelCase:





       """simple docstring"""







       def __init__(   self				):
         UpperCamelCase_:   Optional[int]        =							False





       def      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         if not self.initialized:
           UpperCamelCase_:   Optional[Any]        =							RagRetriever(
               _lowerCamelCase    ,     question_encoder_tokenizer=_lowerCamelCase    ,     generator_tokenizer=_lowerCamelCase    ,     index=_lowerCamelCase    ,     init_retrieval=_lowerCamelCase    ,     )
           UpperCamelCase_:   str        =							True





       def      _a       (   self				):
         self.retriever.index.init_index()





       def      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         UpperCamelCase_ ,UpperCamelCase_:   Any        =							self.retriever._main_retrieve(_lowerCamelCase    ,     _lowerCamelCase				)
         return doc_ids, retrieved_doc_embeds







class 							_lowerCAmelCase(			UpperCAmelCase_   ):





       """simple docstring"""







       def __init__(   self    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase=None				):
         if index is not None and index.is_initialized() and len(_lowerCamelCase				) > 0:
           raise ValueError(
               'When using Ray for distributed fine-tuning, '
               'you\'ll need to provide the paths instead, '
               'as the dataset and the index are loaded '
               'separately. More info in examples/rag/use_own_knowledge_dataset.py '				)
         super().__init__(
             _lowerCamelCase    ,     question_encoder_tokenizer=_lowerCamelCase    ,     generator_tokenizer=_lowerCamelCase    ,     index=_lowerCamelCase    ,     init_retrieval=_lowerCamelCase    ,     )
         UpperCamelCase_:   List[str]        =							retrieval_workers
         if len(self.retrieval_workers				) > 0:
           ray.get(
               [
                   worker.create_rag_retriever.remote(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				)
                   for worker in self.retrieval_workers
               ]				)





       def      _a       (   self				):
         logger.info('initializing retrieval'				)

         if len(self.retrieval_workers				) > 0:
           ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]				)
         else:
           # Non-distributed training. Load index into this same process.
           self.index.init_index()





       def      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         if len(self.retrieval_workers				) > 0:
           # Select a random retrieval actor.
           UpperCamelCase_:   Union[str, Any]        =							self.retrieval_workers[random.randint(0    ,     len(self.retrieval_workers				) - 1				)]
           UpperCamelCase_ ,UpperCamelCase_:   str        =							ray.get(random_worker.retrieve.remote(_lowerCamelCase    ,     _lowerCamelCase				)				)
         else:
           UpperCamelCase_ ,UpperCamelCase_:   Dict        =							self._main_retrieve(_lowerCamelCase    ,     _lowerCamelCase				)
         return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(_lowerCamelCase				)





       @classmethod
       def      _a       (   cls    ,     _lowerCamelCase    ,     _lowerCamelCase=None    ,     **_lowerCamelCase				):
         return super(_lowerCamelCase    ,     cls				).get_tokenizers(_lowerCamelCase    ,     _lowerCamelCase    ,     **_lowerCamelCase				)





       @classmethod
       def      _a       (   cls    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase=None    ,     **_lowerCamelCase				):
         UpperCamelCase_:   List[str]        =							kwargs.pop('config'    ,     _lowerCamelCase				) or RagConfig.from_pretrained(_lowerCamelCase    ,     **_lowerCamelCase				)
         UpperCamelCase_:   Optional[int]        =							RagTokenizer.from_pretrained(_lowerCamelCase    ,     config=_lowerCamelCase				)
         UpperCamelCase_:   List[str]        =							rag_tokenizer.question_encoder
         UpperCamelCase_:   List[Any]        =							rag_tokenizer.generator
         if indexed_dataset is not None:
           UpperCamelCase_:   Union[str, Any]        =							'custom'
           UpperCamelCase_:   int        =							CustomHFIndex(config.retrieval_vector_size    ,     _lowerCamelCase				)
         else:
           UpperCamelCase_:   str        =							cls._build_index(_lowerCamelCase				)
         return cls(
             _lowerCamelCase    ,     question_encoder_tokenizer=_lowerCamelCase    ,     generator_tokenizer=_lowerCamelCase    ,     retrieval_workers=_lowerCamelCase    ,     index=_lowerCamelCase    ,     )






"""simple docstring"""
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, 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 import (
           OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
           OpenAIGPTConfig,
           OpenAIGPTDoubleHeadsModel,
           OpenAIGPTForSequenceClassification,
           OpenAIGPTLMHeadModel,
           OpenAIGPTModel,
       )







class       a							:



 def __init__(     self				,       _snake_case				,       _snake_case=13				,       _snake_case=7				,       _snake_case=True				,       _snake_case=True				,       _snake_case=True				,       _snake_case=99				,       _snake_case=32				,       _snake_case=5				,       _snake_case=4				,       _snake_case=37				,       _snake_case="gelu"				,       _snake_case=0.1				,       _snake_case=0.1				,       _snake_case=5_12				,       _snake_case=16				,       _snake_case=2				,       _snake_case=0.02				,       _snake_case=3				,       _snake_case=4				,       _snake_case=None				,       ):







  """simple docstring"""






  lowerCAmelCase    				= parent
  lowerCAmelCase    				= batch_size
  lowerCAmelCase    				= seq_length
  lowerCAmelCase    				= is_training
  lowerCAmelCase    				= use_token_type_ids
  lowerCAmelCase    				= use_labels
  lowerCAmelCase    				= vocab_size
  lowerCAmelCase    				= hidden_size
  lowerCAmelCase    				= num_hidden_layers
  lowerCAmelCase    				= num_attention_heads
  lowerCAmelCase    				= intermediate_size
  lowerCAmelCase    				= hidden_act
  lowerCAmelCase    				= hidden_dropout_prob
  lowerCAmelCase    				= attention_probs_dropout_prob
  lowerCAmelCase    				= max_position_embeddings
  lowerCAmelCase    				= type_vocab_size
  lowerCAmelCase    				= type_sequence_label_size
  lowerCAmelCase    				= initializer_range
  lowerCAmelCase    				= num_labels
  lowerCAmelCase    				= num_choices
  lowerCAmelCase    				= scope
  lowerCAmelCase    				= self.vocab_size - 1



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






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

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

  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    				= OpenAIGPTConfig(
      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,
      head_mask,
      token_type_ids,
      sequence_labels,
      token_labels,
      choice_labels,
  )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTModel(config=_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       head_mask=_snake_case     )
  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case     )
  lowerCAmelCase    				= model(_snake_case     )

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



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTLMHeadModel(_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.loss.shape				,       ()     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTDoubleHeadsModel(_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.loss.shape				,       ()     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= self.num_labels
  lowerCAmelCase    				= OpenAIGPTForSequenceClassification(_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= ids_tensor([self.batch_size]				,       self.type_sequence_label_size     )
  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.prepare_config_and_inputs()
  (
      (
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,
  )    				= config_and_inputs
  lowerCAmelCase    				= {
      'input_ids': input_ids,
      'token_type_ids': token_type_ids,
      'head_mask': head_mask,
  }

  return config, inputs_dict







@require_torch
class       a							(   a__   , a__   , a__   , unittest.TestCase  ):
 snake_case__             =       (
     (OpenAIGPTModel, OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel, OpenAIGPTForSequenceClassification)
     if is_torch_available()
     else ()
 )
 snake_case__             =       (
     (OpenAIGPTLMHeadModel,) if is_torch_available() else ()
 )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
 snake_case__             =       (
     {
         '''feature-extraction''': OpenAIGPTModel,
         '''text-classification''': OpenAIGPTForSequenceClassification,
         '''text-generation''': OpenAIGPTLMHeadModel,
         '''zero-shot''': OpenAIGPTForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
   # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
   # `OpenAIGPTConfig` 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    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case=False     ):







  """simple docstring"""






  lowerCAmelCase    				= super()._prepare_for_class(_snake_case				,       _snake_case				,       return_labels=_snake_case     )

  if return_labels:
   if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
    lowerCAmelCase    				= torch.zeros(
        (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length)				,       dtype=torch.long				,       device=_snake_case				,       )
    lowerCAmelCase    				= inputs_dict['labels']
    lowerCAmelCase    				= inputs_dict['labels']
    lowerCAmelCase    				= torch.zeros(
        (self.model_tester.batch_size, self.model_tester.num_choices)				,       dtype=torch.long				,       device=_snake_case				,       )
    lowerCAmelCase    				= torch.zeros(
        self.model_tester.batch_size				,       dtype=torch.long				,       device=_snake_case     )
  return inputs_dict



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTModelTester(self     )
  lowerCAmelCase    				= ConfigTester(self				,       config_class=_snake_case				,       n_embd=37     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  self.config_tester.run_common_tests()



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_openai_gpt_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_lm_head_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_double_lm_head_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*_snake_case     )



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
   lowerCAmelCase    				= OpenAIGPTModel.from_pretrained(_snake_case     )
   self.assertIsNotNone(_snake_case     )







@require_torch
class       a							(   unittest.TestCase  ):



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTLMHeadModel.from_pretrained('openai-gpt'     )
  model.to(_snake_case     )
  lowerCAmelCase    				= torch.tensor([[4_81, 47_35, 5_44]]				,       dtype=torch.long				,       device=_snake_case     )  # the president is
  lowerCAmelCase    				= [
      4_81,
      47_35,
      5_44,
      2_46,
      9_63,
      8_70,
      7_62,
      2_39,
      2_44,
      4_04_77,
      2_44,
      2_49,
      7_19,
      8_81,
      4_87,
      5_44,
      2_40,
      2_44,
      6_03,
      4_81,
  ]  # the president is a very good man. " \n " i\'m sure he is, " said the

  lowerCAmelCase    				= model.generate(_snake_case				,       do_sample=_snake_case     )
  self.assertListEqual(output_ids[0].tolist()				,       _snake_case     )





