




'''simple docstring'''




from __future__ import annotations

import unittest

from transformers import AutoTokenizer, GPTJConfig, is_tf_available
from transformers.testing_utils import require_tf, slow, tooslow

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


if is_tf_available():
							import tensorflow as tf

							from transformers.models.gptj.modeling_tf_gptj import (
							    TFGPTJForCausalLM,
							    TFGPTJForQuestionAnswering,
							    TFGPTJForSequenceClassification,
							    TFGPTJModel,
							    shape_list,
							)






class     UpperCamelCase_     :







						"""simple docstring"""


						def __init__(     self      :       Any							, _lowerCamelCase      :       Tuple    ):





											"""simple docstring"""







											A_ :     Optional[int] =							parent
											A_ :     Optional[Any] =							13
											A_ :     Dict =							7
											A_ :     int =							True
											A_ :     int =							True
											A_ :     Optional[Any] =							True
											A_ :     Tuple =							True
											A_ :     List[Any] =							True
											A_ :     Union[str, Any] =							99
											A_ :     Union[str, Any] =							32
											A_ :     Union[str, Any] =							4
											A_ :     List[str] =							2
											A_ :     Dict =							4
											A_ :     Tuple =							37
											A_ :     List[str] =							'''gelu'''
											A_ :     Tuple =							0.1
											A_ :     List[Any] =							0.1
											A_ :     Optional[Any] =							512
											A_ :     Union[str, Any] =							16
											A_ :     int =							2
											A_ :     Dict =							0.02
											A_ :     Union[str, Any] =							3
											A_ :     int =							4
											A_ :     Union[str, Any] =							None
											A_ :     Union[str, Any] =							self.vocab_size - 1
											A_ :     int =							self.vocab_size - 1
											A_ :     List[Any] =							self.vocab_size - 1







						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											A_ :     List[str] =							ids_tensor([self.batch_size, self.seq_length]							, self.vocab_size    )

											A_ :     Tuple =							None
											if self.use_input_mask:
																A_ :     Any =							random_attention_mask([self.batch_size, self.seq_length]    )

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

											A_ :     Any =							None
											if self.use_mc_token_ids:
																A_ :     Union[str, Any] =							ids_tensor([self.batch_size, self.num_choices]							, self.seq_length    )

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

											A_ :     List[str] =							GPTJConfig(
											    vocab_size=self.vocab_size							, n_embd=self.hidden_size							, n_layer=self.num_hidden_layers							, n_head=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							, n_positions=self.max_position_embeddings							, type_vocab_size=self.type_vocab_size							, initializer_range=self.initializer_range							, bos_token_id=self.bos_token_id							, eos_token_id=self.eos_token_id							, pad_token_id=self.pad_token_id							, rotary_dim=self.rotary_dim							, return_dict=_lowerCamelCase							, )

											A_ :     List[Any] =							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      :       List[Any]							, _lowerCamelCase      :       Union[str, Any]							, _lowerCamelCase      :       List[Any]							, _lowerCamelCase      :       Optional[int]							, _lowerCamelCase      :       List[str]							, _lowerCamelCase      :       Union[str, Any]							, *_lowerCamelCase      :       Union[str, Any]    ):





											"""simple docstring"""







											A_ :     Optional[int] =							TFGPTJModel(config=_lowerCamelCase    )
											A_ :     Optional[int] =							{
											    '''input_ids''': input_ids,
											    '''attention_mask''': input_mask,
											    '''token_type_ids''': token_type_ids,
											}
											A_ :     List[Any] =							model(_lowerCamelCase    )

											A_ :     Optional[int] =							[input_ids, None, input_mask]  # None is the input for 'past'
											A_ :     Dict =							model(_lowerCamelCase    )

											A_ :     Any =							model(_lowerCamelCase    )

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







						def       _a    (     self      :       int							, _lowerCamelCase      :       Dict							, _lowerCamelCase      :       List[str]							, _lowerCamelCase      :       int							, _lowerCamelCase      :       int							, _lowerCamelCase      :       List[Any]							, *_lowerCamelCase      :       Union[str, Any]    ):





											"""simple docstring"""







											A_ :     str =							TFGPTJModel(config=_lowerCamelCase    )

											# first forward pass
											A_ :     List[str] =							model(_lowerCamelCase							, token_type_ids=_lowerCamelCase							, use_cache=_lowerCamelCase    )
											A_ :     List[Any] =							model(_lowerCamelCase							, token_type_ids=_lowerCamelCase    )
											A_ :     Optional[Any] =							model(_lowerCamelCase							, token_type_ids=_lowerCamelCase							, use_cache=_lowerCamelCase    )

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

											A_ ,A_ :     Tuple =							outputs.to_tuple()

											# create hypothetical next token and extent to next_input_ids
											A_ :     Tuple =							ids_tensor((self.batch_size, 1)							, config.vocab_size    )
											A_ :     str =							ids_tensor([self.batch_size, 1]							, self.type_vocab_size    )

											# append to next input_ids and token_type_ids
											A_ :     Tuple =							tf.concat([input_ids, next_tokens]							, axis=-1    )
											A_ :     Optional[Any] =							tf.concat([token_type_ids, next_token_types]							, axis=-1    )

											A_ :     Optional[int] =							model(_lowerCamelCase							, token_type_ids=_lowerCamelCase    )['''last_hidden_state''']
											A_ :     str =							model(_lowerCamelCase							, token_type_ids=_lowerCamelCase							, past_key_values=_lowerCamelCase    )[
											    '''last_hidden_state'''
											]

											# select random slice
											A_ :     Any =							int(ids_tensor((1,)							, shape_list(_lowerCamelCase    )[-1]    )    )
											A_ :     Optional[Any] =							output_from_no_past[:, -1, random_slice_idx]
											A_ :     Optional[Any] =							output_from_past[:, 0, random_slice_idx]

											# test that outputs are equal for slice
											tf.debugging.assert_near(_lowerCamelCase							, _lowerCamelCase							, rtol=1E-6    )







						def       _a    (     self      :       Dict							, _lowerCamelCase      :       Optional[Any]							, _lowerCamelCase      :       Optional[Any]							, _lowerCamelCase      :       Optional[Any]							, _lowerCamelCase      :       Optional[int]							, _lowerCamelCase      :       List[Any]							, *_lowerCamelCase      :       int    ):





											"""simple docstring"""







											A_ :     List[str] =							TFGPTJModel(config=_lowerCamelCase    )

											# create attention mask
											A_ :     List[str] =							self.seq_length // 2
											A_ :     Union[str, Any] =							tf.ones((self.batch_size, half_seq_length)							, dtype=tf.intaa    )
											A_ :     str =							tf.zeros((self.batch_size, self.seq_length - half_seq_length)							, dtype=tf.intaa    )
											A_ :     List[Any] =							tf.concat([attn_mask_begin, attn_mask_end]							, axis=1    )

											# first forward pass
											A_ ,A_ :     Tuple =							model(_lowerCamelCase							, attention_mask=_lowerCamelCase    ).to_tuple()

											# create hypothetical next token and extent to next_input_ids
											A_ :     Union[str, Any] =							ids_tensor((self.batch_size, 1)							, config.vocab_size    )

											# change a random masked slice from input_ids
											A_ :     Union[str, Any] =							ids_tensor((1,)							, _lowerCamelCase    ).numpy() + 1
											A_ :     Union[str, Any] =							ids_tensor((self.batch_size, self.seq_length)							, config.vocab_size    )
											A_ :     Optional[int] =							tf.range(self.seq_length    ) == (self.seq_length - random_seq_idx_to_change)
											A_ :     Union[str, Any] =							tf.transpose(
											    tf.broadcast_to(tf.expand_dims(_lowerCamelCase							, -1    )							, (self.seq_length, self.batch_size)    )    )
											A_ :     Any =							tf.where(_lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase    )

											# append to next input_ids and attn_mask
											A_ :     List[Any] =							tf.concat([input_ids, next_tokens]							, axis=-1    )
											A_ :     Tuple =							tf.concat([attn_mask, tf.ones((shape_list(_lowerCamelCase    )[0], 1)							, dtype=tf.intaa    )]							, axis=1    )

											# get two different outputs
											A_ :     Tuple =							model(_lowerCamelCase							, attention_mask=_lowerCamelCase    )['''last_hidden_state''']
											A_ :     int =							model(_lowerCamelCase							, past_key_values=_lowerCamelCase							, attention_mask=_lowerCamelCase    )[
											    '''last_hidden_state'''
											]

											# select random slice
											A_ :     Tuple =							int(ids_tensor((1,)							, shape_list(_lowerCamelCase    )[-1]    )    )
											A_ :     Any =							output_from_no_past[:, -1, random_slice_idx]
											A_ :     str =							output_from_past[:, 0, random_slice_idx]

											# test that outputs are equal for slice
											tf.debugging.assert_near(_lowerCamelCase							, _lowerCamelCase							, rtol=1E-12    )







						def       _a    (     self      :       Optional[int]							, _lowerCamelCase      :       int							, _lowerCamelCase      :       Any							, _lowerCamelCase      :       int							, _lowerCamelCase      :       List[str]							, _lowerCamelCase      :       str							, *_lowerCamelCase      :       Tuple    ):





											"""simple docstring"""







											A_ :     int =							TFGPTJModel(config=_lowerCamelCase    )

											A_ :     Any =							input_ids[:1, :]
											A_ :     Optional[Any] =							input_mask[:1, :]
											A_ :     str =							token_type_ids[:1, :]
											A_ :     Any =							1

											# first forward pass
											A_ :     List[str] =							model(_lowerCamelCase							, attention_mask=_lowerCamelCase							, token_type_ids=_lowerCamelCase							, use_cache=_lowerCamelCase    )

											A_ ,A_ :     Union[str, Any] =							outputs.to_tuple()

											# create hypothetical next token and extent to next_input_ids
											A_ :     List[str] =							ids_tensor((self.batch_size, 3)							, config.vocab_size    )
											A_ :     Optional[int] =							ids_tensor((self.batch_size, 3)							, 2    )
											A_ :     Tuple =							ids_tensor((self.batch_size, 3)							, self.type_vocab_size    )

											# append to next input_ids and token_type_ids
											A_ :     int =							tf.concat([input_ids, next_tokens]							, axis=-1    )
											A_ :     Tuple =							tf.concat([input_mask, next_attn_mask]							, axis=-1    )
											A_ :     Dict =							tf.concat([token_type_ids, next_token_types]							, axis=-1    )

											A_ :     List[str] =							model(
											    _lowerCamelCase							, token_type_ids=_lowerCamelCase							, attention_mask=_lowerCamelCase    )['''last_hidden_state''']
											A_ :     Union[str, Any] =							model(
											    _lowerCamelCase							, token_type_ids=_lowerCamelCase							, attention_mask=_lowerCamelCase							, past_key_values=_lowerCamelCase							, )['''last_hidden_state''']
											self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]    )

											# select random slice
											A_ :     Optional[Any] =							int(ids_tensor((1,)							, shape_list(_lowerCamelCase    )[-1]    )    )
											A_ :     Optional[Any] =							output_from_no_past[:, -3:, random_slice_idx]
											A_ :     List[Any] =							output_from_past[:, :, random_slice_idx]

											# test that outputs are equal for slice
											tf.debugging.assert_near(_lowerCamelCase							, _lowerCamelCase							, rtol=1E-3    )







						def       _a    (     self      :       Tuple							, _lowerCamelCase      :       Any							, _lowerCamelCase      :       List[str]							, _lowerCamelCase      :       Any							, _lowerCamelCase      :       int							, _lowerCamelCase      :       Dict							, *_lowerCamelCase      :       List[str]    ):





											"""simple docstring"""







											A_ :     Any =							TFGPTJForCausalLM(config=_lowerCamelCase    )
											A_ :     Tuple =							{
											    '''input_ids''': input_ids,
											    '''attention_mask''': input_mask,
											    '''token_type_ids''': token_type_ids,
											}
											A_ :     Any =							model(_lowerCamelCase    )
											self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.seq_length, self.vocab_size)    )







						def       _a    (     self      :       List[Any]    ):





											"""simple docstring"""







											A_ :     List[Any] =							self.prepare_config_and_inputs()

											(
											    (
											    A_
											) ,(
											    A_
											) ,(
											    A_
											) ,(
											    A_
											) ,(
											    A_
											) ,(
											    A_
											) ,(
											    A_
											) ,(
											    A_
											) ,(
											    A_
											) ,
											) :     List[Any] =							config_and_inputs

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






@require_tf
class     UpperCamelCase_     (a__,      a__,      a__,      unittest.TestCase     ):







						"""simple docstring"""


						_lowerCAmelCase    						=					(
						    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
						    if is_tf_available()
						    else ()
						)

						_lowerCAmelCase    						=					(TFGPTJForCausalLM,) if is_tf_available() else ()
						_lowerCAmelCase    						=					(
						    {
						        'feature-extraction': TFGPTJModel,
						        'question-answering': TFGPTJForQuestionAnswering,
						        'text-classification': TFGPTJForSequenceClassification,
						        'text-generation': TFGPTJForCausalLM,
						        'zero-shot': TFGPTJForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						_lowerCAmelCase    						=					False
						_lowerCAmelCase    						=					False
						_lowerCAmelCase    						=					False
						_lowerCAmelCase    						=					False







						def       _a    (     self      :       List[Any]							, _lowerCamelCase      :       Optional[int]							, _lowerCamelCase      :       Optional[Any]							, _lowerCamelCase      :       List[Any]							, _lowerCamelCase      :       Union[str, Any]							, _lowerCamelCase      :       Tuple    ):





											"""simple docstring"""







											if (
											    pipeline_test_casse_name == "QAPipelineTests"
											    and tokenizer_name is not None
											    and not tokenizer_name.endswith('''Fast'''    )
											):
																# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
																# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
																# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
																return True

											return False







						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											A_ :     Tuple =							TFGPTJModelTester(self    )
											A_ :     str =							ConfigTester(self							, config_class=_lowerCamelCase							, n_embd=37    )







						def       _a    (     self      :       int    ):





											"""simple docstring"""







											self.config_tester.run_common_tests()







						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											A_ :     List[Any] =							self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_gptj_model(*_lowerCamelCase    )







						def       _a    (     self      :       Tuple    ):





											"""simple docstring"""







											A_ :     int =							self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_gptj_model_past(*_lowerCamelCase    )







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





											"""simple docstring"""







											A_ :     str =							self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_gptj_model_attention_mask_past(*_lowerCamelCase    )







						def       _a    (     self      :       str    ):





											"""simple docstring"""







											A_ :     int =							self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_gptj_model_past_large_inputs(*_lowerCamelCase    )







						def       _a    (     self      :       Any    ):





											"""simple docstring"""







											A_ :     List[str] =							self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_gptj_lm_head_model(*_lowerCamelCase    )







						@slow
						@unittest.skipIf(
						    not is_tf_available() or len(tf.config.list_physical_devices('''GPU'''    )    ) > 0							, '''skip testing on GPU for now to avoid GPU OOM.'''							, )
						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											A_ :     Union[str, Any] =							TFGPTJModel.from_pretrained('''EleutherAI/gpt-j-6B'''							, from_pt=_lowerCamelCase    )
											self.assertIsNotNone(_lowerCamelCase    )







						@unittest.skip(reason='''Currently, model embeddings are going to undergo a major refactor.'''    )
						def       _a    (     self      :       str    ):





											"""simple docstring"""







											super().test_resize_token_embeddings()






@require_tf
@tooslow
# Marked as @tooslow due to GPU OOM -- but still useful to run locally. Requires ~39GB of RAM.
class     UpperCamelCase_     (unittest.TestCase     ):







						"""simple docstring"""


						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											A_ :     Dict =							TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''							, from_pt=_lowerCamelCase    )
											A_ :     Tuple =							tf.convert_to_tensor([[464, 3290]]							, dtype=tf.intaa    )  # The dog
											# fmt: off
											# The dog is a man's best friend. It is a loyal companion, and it is a friend
											A_ :     List[Any] =							[464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 15185, 11, 290, 340, 318, 257, 1545]
											# fmt: on
											A_ :     List[Any] =							model.generate(_lowerCamelCase							, do_sample=_lowerCamelCase    )
											self.assertListEqual(output_ids[0].numpy().tolist()							, _lowerCamelCase    )







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





											"""simple docstring"""







											A_ :     List[str] =							AutoTokenizer.from_pretrained('''EleutherAI/gpt-j-6B'''							, revision='''float16'''    )
											A_ :     List[str] =							TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''							, revision='''float16'''							, from_pt=_lowerCamelCase    )

											A_ :     Union[str, Any] =							tokenizer('''Today is a nice day and'''							, return_tensors='''tf'''    )
											# forces the generation to happen on CPU, to avoid GPU-related quirks
											with tf.device(''':/CPU:0'''    ):
																A_ :     List[Any] =							model.generate(**_lowerCamelCase							, do_sample=_lowerCamelCase							, seed=[42, 0]    )
											A_ :     Tuple =							tokenizer.decode(output_ids[0]							, skip_special_tokens=_lowerCamelCase    )

											A_ :     int =							'''Today is a nice day and I’m going to go for a walk. I’'''
											self.assertEqual(_lowerCamelCase							, _lowerCamelCase    )







						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											A_ :     str =							TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''							, revision='''float16'''							, from_pt=_lowerCamelCase    )
											A_ :     Tuple =							AutoTokenizer.from_pretrained('''EleutherAI/gpt-j-6B'''							, revision='''float16'''    )

											A_ :     str =							'''left'''

											# Define PAD Token = EOS Token = 50256
											A_ :     Tuple =							tokenizer.eos_token
											A_ :     List[Any] =							model.config.eos_token_id

											# use different length sentences to test batching
											A_ :     Dict =							[
											    '''Hello, my dog is a little''',
											    '''Today, I''',
											]
											A_ :     Dict =							[
											    '''Hello, my dog is a little over a year old and has been diagnosed with hip dysplasia''',
											    '''Today, I’m going to be talking about a topic that’''',
											]
											return model, tokenizer, sentences, expected_output_sentences







						def       _a    (     self      :       int    ):





											"""simple docstring"""







											A_ ,A_ ,A_ ,A_ :     str =							self._get_beam_search_test_objects()

											A_ :     Optional[int] =							tokenizer(_lowerCamelCase							, return_tensors='''tf'''							, padding=_lowerCamelCase    )
											A_ :     Dict =							model.generate(**_lowerCamelCase							, do_sample=_lowerCamelCase							, num_beams=2    )
											A_ :     Optional[int] =							tokenizer.batch_decode(_lowerCamelCase							, skip_special_tokens=_lowerCamelCase    )
											self.assertListEqual(_lowerCamelCase							, _lowerCamelCase    )







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





											"""simple docstring"""







											A_ ,A_ ,A_ ,A_ :     Union[str, Any] =							self._get_beam_search_test_objects()

											A_ :     Tuple =							tokenizer(_lowerCamelCase							, return_tensors='''tf'''							, padding=_lowerCamelCase    )
											A_ :     str =							tokenizer(sentences[0]							, return_tensors='''tf'''    )
											A_ :     Optional[int] =							model.generate(**_lowerCamelCase							, do_sample=_lowerCamelCase							, num_beams=2    )
											A_ :     Optional[int] =							(
											    shape_list(inputs_non_padded['''input_ids''']    )[-1]
											    - tf.reduce_sum(tf.cast(inputs['''attention_mask'''][-1]							, tf.intaa    )    ).numpy()
											)
											A_ :     Any =							tokenizer(sentences[1]							, return_tensors='''tf'''    )
											A_ :     Optional[int] =							model.generate(
											    **_lowerCamelCase							, do_sample=_lowerCamelCase							, num_beams=2							, max_length=model.config.max_length - num_paddings    )
											A_ :     int =							tokenizer.decode(output_non_padded[0]							, skip_special_tokens=_lowerCamelCase    )
											A_ :     Tuple =							tokenizer.decode(output_padded[0]							, skip_special_tokens=_lowerCamelCase    )
											self.assertListEqual(_lowerCamelCase							, [non_padded_sentence, padded_sentence]    )







						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											A_ ,A_ ,A_ ,A_ :     Any =							self._get_beam_search_test_objects()

											A_ :     Any =							tokenizer(_lowerCamelCase							, return_tensors='''tf'''							, padding=_lowerCamelCase    )
											A_ :     int =							tf.function(model.generate							, jit_compile=_lowerCamelCase    )
											A_ :     Any =							xla_generate(**_lowerCamelCase							, do_sample=_lowerCamelCase							, num_beams=2    )
											A_ :     Any =							tokenizer.batch_decode(_lowerCamelCase							, skip_special_tokens=_lowerCamelCase    )
											self.assertListEqual(_lowerCamelCase							, _lowerCamelCase    )









'''simple docstring'''




import logging
import random

import ray

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


snake_case__          =						logging.getLogger(__name__)






class     UpperCamelCase_     :







						"""simple docstring"""


						def __init__(     self      :       Tuple    ):





											"""simple docstring"""







											A_ :     Optional[int] =							False







						def       _a    (     self      :       Any							, _lowerCamelCase      :       List[str]							, _lowerCamelCase      :       Any							, _lowerCamelCase      :       str							, _lowerCamelCase      :       int    ):





											"""simple docstring"""







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







						def       _a    (     self      :       int    ):





											"""simple docstring"""







											self.retriever.index.init_index()







						def       _a    (     self      :       Any    ):





											"""simple docstring"""







											del self.retriever
											A_ :     List[Any] =							False







						def       _a    (     self      :       Optional[int]							, _lowerCamelCase      :       Union[str, Any]							, _lowerCamelCase      :       Union[str, Any]    ):





											"""simple docstring"""







											A_ ,A_ :     Dict =							self.retriever._main_retrieve(_lowerCamelCase							, _lowerCamelCase    )
											A_ :     int =							self.retriever.index.get_doc_dicts(_lowerCamelCase    )
											return doc_ids, retrieved_doc_embeds, doc_dicts






class     UpperCamelCase_     (a__     ):







						"""simple docstring"""


						def __init__(     self      :       int							, _lowerCamelCase      :       Optional[Any]							, _lowerCamelCase      :       Dict							, _lowerCamelCase      :       Optional[int]							, _lowerCamelCase      :       Dict							, _lowerCamelCase      :       Tuple=None    ):





											"""simple docstring"""







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

											A_ :     Tuple =							retrieval_workers
											A_ :     Any =							question_encoder_tokenizer
											A_ :     Dict =							generator_tokenizer
											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      :       Dict    ):





											"""simple docstring"""







											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      :       Dict							, _lowerCamelCase      :       Tuple							, _lowerCamelCase      :       int    ):





											"""simple docstring"""







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







						@classmethod
						def       _a    (     cls      :       List[str]							, _lowerCamelCase      :       Union[str, Any]							, _lowerCamelCase      :       Optional[Any]=None							, **_lowerCamelCase      :       str    ):





											"""simple docstring"""







											return super(_lowerCamelCase							, cls    ).get_tokenizers(_lowerCamelCase							, _lowerCamelCase							, **_lowerCamelCase    )







						@classmethod
						def       _a    (     cls      :       Optional[Any]							, _lowerCamelCase      :       Any							, _lowerCamelCase      :       List[Any]							, _lowerCamelCase      :       str=None							, **_lowerCamelCase      :       Tuple    ):





											"""simple docstring"""







											A_ :     Tuple =							kwargs.pop('''config'''							, _lowerCamelCase    ) or RagConfig.from_pretrained(_lowerCamelCase							, **_lowerCamelCase    )
											A_ :     List[Any] =							RagTokenizer.from_pretrained(_lowerCamelCase							, config=_lowerCamelCase    )
											A_ :     Union[str, Any] =							rag_tokenizer.question_encoder
											A_ :     Dict =							rag_tokenizer.generator

											if indexed_dataset is not None:
																A_ :     Any =							'''custom'''
																A_ :     int =							CustomHFIndex(config.retrieval_vector_size							, _lowerCamelCase    )
											else:
																A_ :     Optional[int] =							cls._build_index(_lowerCamelCase    )

											return cls(
											    _lowerCamelCase							, question_encoder_tokenizer=_lowerCamelCase							, generator_tokenizer=_lowerCamelCase							, retrieval_workers=_lowerCamelCase							, index=_lowerCamelCase							, )







						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											logger.info('''re-loading the new dataset with embeddings'''    )
											# access from the training loop

											ray.get([worker.clear_object.remote() for worker in self.retrieval_workers]    )

											# build the index object again
											A_ :     Dict =							self._build_index(self.config    )

											ray.get(
											    [
											        worker.create_rag_retriever.remote(
											            self.config							, self.question_encoder_tokenizer							, self.generator_tokenizer							, _lowerCamelCase    )
											        for worker in self.retrieval_workers
											    ]    )



