



import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

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


if is_torch_available():
   import torch

   from transformers import OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel






class 	__snake_case	:




  def __init__(			self					:						Dict						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						Tuple=1_3						, __lowerCAmelCase					:						str=7						, __lowerCAmelCase					:						List[str]=True						, __lowerCAmelCase					:						Optional[int]=True						, __lowerCAmelCase					:						int=False						, __lowerCAmelCase					:						Optional[Any]=True						, __lowerCAmelCase					:						int=9_9						, __lowerCAmelCase					:						Any=3_2						, __lowerCAmelCase					:						Any=5						, __lowerCAmelCase					:						Tuple=4						, __lowerCAmelCase					:						Any=3_7						, __lowerCAmelCase					:						Any="gelu"						, __lowerCAmelCase					:						Optional[Any]=0.1						, __lowerCAmelCase					:						int=0.1						, __lowerCAmelCase					:						Dict=5_1_2						, __lowerCAmelCase					:						Tuple=1_6						, __lowerCAmelCase					:						Optional[Any]=2						, __lowerCAmelCase					:						Dict=0.02						, __lowerCAmelCase					:						str=3						, __lowerCAmelCase					:						Optional[Any]=4						, __lowerCAmelCase					:						Optional[int]=None						, ):






         """simple docstring"""






         _lowerCamelCase						:    str         =			parent
         _lowerCamelCase						:    Dict         =			batch_size
         _lowerCamelCase						:    Optional[int]         =			seq_length
         _lowerCamelCase						:    int         =			is_training
         _lowerCamelCase						:    str         =			use_input_mask
         _lowerCamelCase						:    List[str]         =			use_token_type_ids
         _lowerCamelCase						:    Optional[int]         =			use_labels
         _lowerCamelCase						:    Tuple         =			vocab_size
         _lowerCamelCase						:    int         =			hidden_size
         _lowerCamelCase						:    Optional[Any]         =			num_hidden_layers
         _lowerCamelCase						:    Tuple         =			num_attention_heads
         _lowerCamelCase						:    List[Any]         =			intermediate_size
         _lowerCamelCase						:    Union[str, Any]         =			hidden_act
         _lowerCamelCase						:    List[str]         =			hidden_dropout_prob
         _lowerCamelCase						:    Union[str, Any]         =			attention_probs_dropout_prob
         _lowerCamelCase						:    int         =			max_position_embeddings
         _lowerCamelCase						:    int         =			type_vocab_size
         _lowerCamelCase						:    Any         =			type_sequence_label_size
         _lowerCamelCase						:    str         =			initializer_range
         _lowerCamelCase						:    Optional[Any]         =			num_labels
         _lowerCamelCase						:    List[str]         =			num_choices
         _lowerCamelCase						:    int         =			scope




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size		)

         _lowerCamelCase						:    Union[str, Any]         =			None
         if self.use_input_mask:
                _lowerCamelCase						:    Optional[Any]         =			random_attention_mask([self.batch_size, self.seq_length]		)

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

         _lowerCamelCase						:    Optional[int]         =			None
         _lowerCamelCase						:    Union[str, Any]         =			None
         _lowerCamelCase						:    List[Any]         =			None
         if self.use_labels:
                _lowerCamelCase						:    Optional[int]         =			ids_tensor([self.batch_size]						, self.type_sequence_label_size		)
                _lowerCamelCase						:    List[str]         =			ids_tensor([self.batch_size, self.seq_length]						, self.num_labels		)
                _lowerCamelCase						:    str         =			ids_tensor([self.batch_size]						, self.num_choices		)

         _lowerCamelCase						:    Any         =			self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         return OpenLlamaConfig(
             vocab_size=self.vocab_size						, hidden_size=self.hidden_size						, num_hidden_layers=self.num_hidden_layers						, num_attention_heads=self.num_attention_heads						, intermediate_size=self.intermediate_size						, hidden_act=self.hidden_act						, hidden_dropout_prob=self.hidden_dropout_prob						, attention_probs_dropout_prob=self.attention_probs_dropout_prob						, max_position_embeddings=self.max_position_embeddings						, type_vocab_size=self.type_vocab_size						, is_decoder=lowercase_						, initializer_range=self.initializer_range						, use_stable_embedding=lowercase_						, )




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			OpenLlamaModel(config=lowercase_		)
         model.to(lowercase_		)
         model.eval()
         _lowerCamelCase						:    Tuple         =			model(lowercase_						, attention_mask=lowercase_		)
         _lowerCamelCase						:    int         =			model(lowercase_		)
         self.parent.assertEqual(result.last_hidden_state.shape						, (self.batch_size, self.seq_length, self.hidden_size)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						List[str]						, ):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			True
         _lowerCamelCase						:    Optional[int]         =			OpenLlamaModel(lowercase_		)
         model.to(lowercase_		)
         model.eval()
         _lowerCamelCase						:    Optional[Any]         =			model(
             lowercase_						, attention_mask=lowercase_						, encoder_hidden_states=lowercase_						, encoder_attention_mask=lowercase_						, )
         _lowerCamelCase						:    Optional[int]         =			model(
             lowercase_						, attention_mask=lowercase_						, encoder_hidden_states=lowercase_						, )
         _lowerCamelCase						:    Tuple         =			model(lowercase_						, attention_mask=lowercase_		)
         self.parent.assertEqual(result.last_hidden_state.shape						, (self.batch_size, self.seq_length, self.hidden_size)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						List[str]						, ):






         """simple docstring"""






         _lowerCamelCase						:    str         =			OpenLlamaForCausalLM(config=lowercase_		)
         model.to(lowercase_		)
         model.eval()
         _lowerCamelCase						:    List[Any]         =			model(lowercase_						, attention_mask=lowercase_						, labels=lowercase_		)
         self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.vocab_size)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						Dict						, ):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			True
         _lowerCamelCase						:    Dict         =			True
         _lowerCamelCase						:    Optional[int]         =			OpenLlamaForCausalLM(config=lowercase_		)
         model.to(lowercase_		)
         model.eval()

         # first forward pass
         _lowerCamelCase						:    int         =			model(
             lowercase_						, attention_mask=lowercase_						, encoder_hidden_states=lowercase_						, encoder_attention_mask=lowercase_						, use_cache=lowercase_						, )
         _lowerCamelCase						:    int         =			outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         _lowerCamelCase						:    Dict         =			ids_tensor((self.batch_size, 3)						, config.vocab_size		)
         _lowerCamelCase						:    Tuple         =			ids_tensor((self.batch_size, 3)						, vocab_size=2		)

         # append to next input_ids and
         _lowerCamelCase						:    Any         =			torch.cat([input_ids, next_tokens]						, dim=-1		)
         _lowerCamelCase						:    Tuple         =			torch.cat([input_mask, next_mask]						, dim=-1		)

         _lowerCamelCase						:    Union[str, Any]         =			model(
             lowercase_						, attention_mask=lowercase_						, encoder_hidden_states=lowercase_						, encoder_attention_mask=lowercase_						, output_hidden_states=lowercase_						, )["hidden_states"][0]
         _lowerCamelCase						:    List[Any]         =			model(
             lowercase_						, attention_mask=lowercase_						, encoder_hidden_states=lowercase_						, encoder_attention_mask=lowercase_						, past_key_values=lowercase_						, output_hidden_states=lowercase_						, )["hidden_states"][0]

         # select random slice
         _lowerCamelCase						:    Tuple         =			ids_tensor((1,)						, output_from_past.shape[-1]		).item()
         _lowerCamelCase						:    int         =			output_from_no_past[:, -3:, random_slice_idx].detach()
         _lowerCamelCase						:    Optional[int]         =			output_from_past[:, :, random_slice_idx].detach()

         self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]		)

         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(lowercase_						, lowercase_						, atol=1E-3		)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			self.prepare_config_and_inputs()
         (
             _lowerCamelCase       
         )						:    Optional[Any]         =			config_and_inputs
         _lowerCamelCase						:    Optional[Any]         =			{"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict






@require_torch
class 	__snake_case	(     lowercase__  ,					lowercase__  ,					lowercase__  ,					unittest.TestCase):
  snake_case__    :       List[Any]        =						(
      (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
  )
  snake_case__    :       Optional[int]        =						(OpenLlamaForCausalLM,) if is_torch_available() else ()
  snake_case__    :       Any        =						(
      {
          """feature-extraction""": OpenLlamaModel,
          """text-classification""": OpenLlamaForSequenceClassification,
          """text-generation""": OpenLlamaForCausalLM,
          """zero-shot""": OpenLlamaForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  snake_case__    :       Tuple        =						False
  snake_case__    :       List[Any]        =						False




  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			OpenLlamaModelTester(self		)
         _lowerCamelCase						:    Union[str, Any]         =			ConfigTester(self						, config_class=lowercase_						, hidden_size=3_7		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         self.config_tester.run_common_tests()




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*lowercase_		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
                _lowerCamelCase						:    Optional[Any]         =			type
                self.model_tester.create_and_check_model(*lowercase_		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			self.model_tester.prepare_config_and_inputs_for_common()
         _lowerCamelCase						:    Optional[int]         =			3
         _lowerCamelCase						:    Optional[Any]         =			input_dict["input_ids"]
         _lowerCamelCase						:    Union[str, Any]         =			input_ids.ne(1		).to(lowercase_		)
         _lowerCamelCase						:    Union[str, Any]         =			ids_tensor([self.model_tester.batch_size]						, self.model_tester.type_sequence_label_size		)
         _lowerCamelCase						:    str         =			OpenLlamaForSequenceClassification(lowercase_		)
         model.to(lowercase_		)
         model.eval()
         _lowerCamelCase						:    int         =			model(lowercase_						, attention_mask=lowercase_						, labels=lowercase_		)
         self.assertEqual(result.logits.shape						, (self.model_tester.batch_size, self.model_tester.num_labels)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    str         =			self.model_tester.prepare_config_and_inputs_for_common()
         _lowerCamelCase						:    str         =			3
         _lowerCamelCase						:    str         =			"single_label_classification"
         _lowerCamelCase						:    Any         =			input_dict["input_ids"]
         _lowerCamelCase						:    int         =			input_ids.ne(1		).to(lowercase_		)
         _lowerCamelCase						:    Union[str, Any]         =			ids_tensor([self.model_tester.batch_size]						, self.model_tester.type_sequence_label_size		)
         _lowerCamelCase						:    Optional[Any]         =			OpenLlamaForSequenceClassification(lowercase_		)
         model.to(lowercase_		)
         model.eval()
         _lowerCamelCase						:    List[Any]         =			model(lowercase_						, attention_mask=lowercase_						, labels=lowercase_		)
         self.assertEqual(result.logits.shape						, (self.model_tester.batch_size, self.model_tester.num_labels)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			self.model_tester.prepare_config_and_inputs_for_common()
         _lowerCamelCase						:    int         =			3
         _lowerCamelCase						:    Union[str, Any]         =			"multi_label_classification"
         _lowerCamelCase						:    Dict         =			input_dict["input_ids"]
         _lowerCamelCase						:    Tuple         =			input_ids.ne(1		).to(lowercase_		)
         _lowerCamelCase						:    Tuple         =			ids_tensor(
             [self.model_tester.batch_size, config.num_labels]						, self.model_tester.type_sequence_label_size		).to(torch.float		)
         _lowerCamelCase						:    List[str]         =			OpenLlamaForSequenceClassification(lowercase_		)
         model.to(lowercase_		)
         model.eval()
         _lowerCamelCase						:    List[Any]         =			model(lowercase_						, attention_mask=lowercase_						, labels=lowercase_		)
         self.assertEqual(result.logits.shape						, (self.model_tester.batch_size, self.model_tester.num_labels)		)




  @unittest.skip('''Open-Llama buffers include complex numbers, which breaks this test'''		)
  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         pass




  @parameterized.expand([('''linear''',), ('''dynamic''',)]		)
  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			self.model_tester.prepare_config_and_inputs_for_common()
         _lowerCamelCase						:    Dict         =			ids_tensor([1, 1_0]						, config.vocab_size		)
         _lowerCamelCase						:    List[Any]         =			ids_tensor([1, int(config.max_position_embeddings * 1.5		)]						, config.vocab_size		)

         set_seed(4_2		)  # Fixed seed at init time so the two models get the same random weights
         _lowerCamelCase						:    Union[str, Any]         =			OpenLlamaModel(lowercase_		)
         original_model.to(lowercase_		)
         original_model.eval()
         _lowerCamelCase						:    Optional[int]         =			original_model(lowercase_		).last_hidden_state
         _lowerCamelCase						:    Optional[int]         =			original_model(lowercase_		).last_hidden_state

         set_seed(4_2		)  # Fixed seed at init time so the two models get the same random weights
         _lowerCamelCase						:    Optional[int]         =			{"type": scaling_type, "factor": 10.0}
         _lowerCamelCase						:    List[Any]         =			OpenLlamaModel(lowercase_		)
         scaled_model.to(lowercase_		)
         scaled_model.eval()
         _lowerCamelCase						:    Tuple         =			scaled_model(lowercase_		).last_hidden_state
         _lowerCamelCase						:    Any         =			scaled_model(lowercase_		).last_hidden_state

         # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
         # maximum sequence length, so the outputs for the short input should match.
         if scaling_type == "dynamic":
                self.assertTrue(torch.allclose(lowercase_						, lowercase_						, atol=1E-5		)		)
         else:
                self.assertFalse(torch.allclose(lowercase_						, lowercase_						, atol=1E-5		)		)

         # The output should be different for long inputs
         self.assertFalse(torch.allclose(lowercase_						, lowercase_						, atol=1E-5		)		)







"""simple docstring"""






import logging
import os
from typing import List, Tuple

import numpy as np
import psutil
import torch
import torch.distributed as dist

from transformers import RagRetriever


lowerCAmelCase__      =    logging.getLogger(__name__)






class 	__snake_case	(     _lowercase):




  def __init__(			self					:						Union[str, Any]						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						int=None		):






         """simple docstring"""






         super().__init__(
             __lowerCAmelCase						, question_encoder_tokenizer=__lowerCAmelCase						, generator_tokenizer=__lowerCAmelCase						, index=__lowerCAmelCase						, init_retrieval=__lowerCAmelCase						, )
         _lowerCamelCase						:    Dict         =			None




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict						, __lowerCAmelCase					:						int		):






         """simple docstring"""






         logger.info('''initializing retrieval'''		)

         # initializing a separate process group for retrieval as the default
         # nccl backend doesn't support gather/scatter operations while gloo
         # is too slow to replace nccl for the core gpu communication
         if dist.is_initialized():
                logger.info('''dist initialized'''		)
                # needs to be set manually
                _lowerCamelCase						:    List[str]         =			self._infer_socket_ifname()
                # avoid clash with the NCCL port
                _lowerCamelCase						:    Dict         =			str(distributed_port + 1		)
                _lowerCamelCase						:    str         =			dist.new_group(ranks=__lowerCAmelCase						, backend='''gloo'''		)

         # initialize retriever only on the main worker
         if not dist.is_initialized() or self._is_main():
                logger.info('''dist not initialized / main'''		)
                self.index.init_index()

         # all processes wait untill the retriever is initialized by the main process
         if dist.is_initialized():
                torch.distributed.barrier(group=self.process_group		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         return dist.get_rank(group=self.process_group		) == 0




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Tuple=torch.floataa		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			torch.empty(__lowerCAmelCase						, dtype=__lowerCAmelCase		)
         dist.scatter(__lowerCAmelCase						, src=0						, scatter_list=__lowerCAmelCase						, group=self.process_group		)
         return target_tensor




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			psutil.net_if_addrs()
         # a hacky way to deal with varying network interface names
         _lowerCamelCase						:    str         =			next((addr for addr in addrs if addr.startswith('''e'''		))						, __lowerCAmelCase		)
         return ifname




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						np.ndarray						, __lowerCAmelCase					:						int		):






         """simple docstring"""






         if not dist.is_initialized():
                _lowerCamelCase      ,   _lowerCamelCase						:    Any         =			self._main_retrieve(__lowerCAmelCase						, __lowerCAmelCase		)
                return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(__lowerCAmelCase		)

         # distributed training
         _lowerCamelCase						:    Dict         =			dist.get_world_size(group=self.process_group		)

         # gather logic
         _lowerCamelCase						:    str         =			None
         if self._is_main():
                _lowerCamelCase						:    List[Any]         =			[torch.empty(question_hidden_states.shape						, dtype=torch.floataa		) for _ in range(__lowerCAmelCase		)]
         dist.gather(torch.tensor(__lowerCAmelCase		)						, dst=0						, gather_list=__lowerCAmelCase						, group=self.process_group		)

         # scatter logic
         _lowerCamelCase						:    int         =			question_hidden_states.shape[0]
         _lowerCamelCase						:    str         =			[]
         _lowerCamelCase						:    Optional[int]         =			[]
         if self._is_main():
                assert len(__lowerCAmelCase		) == world_size
                _lowerCamelCase      ,   _lowerCamelCase						:    Tuple         =			self._main_retrieve(torch.cat(__lowerCAmelCase		).numpy()						, __lowerCAmelCase		)
                _lowerCamelCase      ,   _lowerCamelCase						:    List[str]         =			torch.tensor(__lowerCAmelCase		), torch.tensor(__lowerCAmelCase		)
                _lowerCamelCase						:    List[Any]         =			self._chunk_tensor(__lowerCAmelCase						, __lowerCAmelCase		)
                _lowerCamelCase						:    Optional[Any]         =			self._chunk_tensor(__lowerCAmelCase						, __lowerCAmelCase		)
         _lowerCamelCase						:    int         =			self._scattered(__lowerCAmelCase						, [n_queries, n_docs]						, target_type=torch.intaa		)
         _lowerCamelCase						:    str         =			self._scattered(__lowerCAmelCase						, [n_queries, n_docs, question_hidden_states.shape[1]]		)

         return retrieved_doc_embeds.numpy(), doc_ids.numpy(), self.index.get_doc_dicts(__lowerCAmelCase		)


