





import json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_gpta import GPTaTokenizer


if TYPE_CHECKING:
      from transformers.pipelines.conversational import Conversation


SCREAMING_SNAKE_CASE_						       =	logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_						       =	{"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt""", """tokenizer_file""": """tokenizer.json"""}

SCREAMING_SNAKE_CASE_						       =	{
    """vocab_file""": {
        """gpt2""": """https://huggingface.co/gpt2/resolve/main/vocab.json""",
        """gpt2-medium""": """https://huggingface.co/gpt2-medium/resolve/main/vocab.json""",
        """gpt2-large""": """https://huggingface.co/gpt2-large/resolve/main/vocab.json""",
        """gpt2-xl""": """https://huggingface.co/gpt2-xl/resolve/main/vocab.json""",
        """distilgpt2""": """https://huggingface.co/distilgpt2/resolve/main/vocab.json""",
    },
    """merges_file""": {
        """gpt2""": """https://huggingface.co/gpt2/resolve/main/merges.txt""",
        """gpt2-medium""": """https://huggingface.co/gpt2-medium/resolve/main/merges.txt""",
        """gpt2-large""": """https://huggingface.co/gpt2-large/resolve/main/merges.txt""",
        """gpt2-xl""": """https://huggingface.co/gpt2-xl/resolve/main/merges.txt""",
        """distilgpt2""": """https://huggingface.co/distilgpt2/resolve/main/merges.txt""",
    },
    """tokenizer_file""": {
        """gpt2""": """https://huggingface.co/gpt2/resolve/main/tokenizer.json""",
        """gpt2-medium""": """https://huggingface.co/gpt2-medium/resolve/main/tokenizer.json""",
        """gpt2-large""": """https://huggingface.co/gpt2-large/resolve/main/tokenizer.json""",
        """gpt2-xl""": """https://huggingface.co/gpt2-xl/resolve/main/tokenizer.json""",
        """distilgpt2""": """https://huggingface.co/distilgpt2/resolve/main/tokenizer.json""",
    },
}

SCREAMING_SNAKE_CASE_						       =	{
    """gpt2""": 1_0_2_4,
    """gpt2-medium""": 1_0_2_4,
    """gpt2-large""": 1_0_2_4,
    """gpt2-xl""": 1_0_2_4,
    """distilgpt2""": 1_0_2_4,
}

class 		UpperCamelCase__			(				a__       ):



     '''simple docstring'''

     __snake_case   :     str		     =      VOCAB_FILES_NAMES
     __snake_case   :     List[str]		     =      PRETRAINED_VOCAB_FILES_MAP
     __snake_case   :     Dict		     =      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __snake_case   :     Optional[Any]		     =      ["input_ids", "attention_mask"]
     __snake_case   :     Dict		     =      GPTaTokenizer






     def __init__(							self						:   Any						,lowerCamelCase__						:   Optional[Any]=None						,lowerCamelCase__						:   Optional[int]=None						,lowerCamelCase__						:   List[Any]=None						,lowerCamelCase__						:   List[Any]="<|endoftext|>"						,lowerCamelCase__						:   Any="<|endoftext|>"						,lowerCamelCase__						:   Tuple="<|endoftext|>"						,lowerCamelCase__						:   Optional[int]=False						,**lowerCamelCase__						:   Optional[int]						,)					->		int:




           '''simple docstring'''


           super().__init__(
               lowerCamelCase__						,lowerCamelCase__						,tokenizer_file=lowerCamelCase__						,unk_token=lowerCamelCase__						,bos_token=lowerCamelCase__						,eos_token=lowerCamelCase__						,add_prefix_space=lowerCamelCase__						,**lowerCamelCase__						,)

           SCREAMING_SNAKE_CASE		=						kwargs.pop("""add_bos_token"""						,lowerCamelCase__	)

           SCREAMING_SNAKE_CASE		=						json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()	)
           if pre_tok_state.get("""add_prefix_space"""						,lowerCamelCase__	) != add_prefix_space:
                 SCREAMING_SNAKE_CASE		=						getattr(lowerCamelCase__						,pre_tok_state.pop("""type"""	)	)
                 SCREAMING_SNAKE_CASE		=						add_prefix_space
                 SCREAMING_SNAKE_CASE		=						pre_tok_class(**lowerCamelCase__	)

           SCREAMING_SNAKE_CASE		=						add_prefix_space






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[str]						,*lowerCamelCase__						:   Tuple						,**lowerCamelCase__						:   Union[str, Any]	)					->		Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						kwargs.get("""is_split_into_words"""						,lowerCamelCase__	)
           assert self.add_prefix_space or not is_split_into_words, (
               F"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
               "to use it with pretokenized inputs."
           )

           return super()._batch_encode_plus(*lowerCamelCase__						,**lowerCamelCase__	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any						,*lowerCamelCase__						:   Tuple						,**lowerCamelCase__						:   Any	)					->		List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						kwargs.get("""is_split_into_words"""						,lowerCamelCase__	)

           assert self.add_prefix_space or not is_split_into_words, (
               F"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
               "to use it with pretokenized inputs."
           )

           return super()._encode_plus(*lowerCamelCase__						,**lowerCamelCase__	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[int]						,lowerCamelCase__						:   str						,lowerCamelCase__						:   Optional[str] = None	)					->		Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self._tokenizer.model.save(lowerCamelCase__						,name=lowerCamelCase__	)
           return tuple(lowerCamelCase__	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[int]						,lowerCamelCase__						:   "Conversation"	)					->		Union[str, Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						[]
           for is_user, text in conversation.iter_texts():
                 input_ids.extend(self.encode(lowerCamelCase__						,add_special_tokens=lowerCamelCase__	) + [self.eos_token_id]	)

           if len(lowerCamelCase__	) > self.model_max_length:
                 SCREAMING_SNAKE_CASE		=						input_ids[-self.model_max_length :]
           return input_ids









import logging
import random

import ray

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


SCREAMING_SNAKE_CASE_						       =	logging.getLogger(__name__)

class 		UpperCamelCase__			:



     '''simple docstring'''

     def __init__(							self						:   str	)					->		List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						False






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any						,lowerCamelCase__						:   int						,lowerCamelCase__						:   Dict						,lowerCamelCase__						:   List[Any]						,lowerCamelCase__						:   Tuple	)					->		Optional[Any]:




           '''simple docstring'''


           if not self.initialized:
                 SCREAMING_SNAKE_CASE		=						RagRetriever(
                     lowerCamelCase__						,question_encoder_tokenizer=lowerCamelCase__						,generator_tokenizer=lowerCamelCase__						,index=lowerCamelCase__						,init_retrieval=lowerCamelCase__						,)
                 SCREAMING_SNAKE_CASE		=						True






     def     SCREAMING_SNAKE_CASE__  (							self						:   Dict	)					->		Any:




           '''simple docstring'''


           self.retriever.index.init_index()






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any						,lowerCamelCase__						:   List[Any]						,lowerCamelCase__						:   List[str]	)					->		List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE, SCREAMING_SNAKE_CASE		=						self.retriever._main_retrieve(lowerCamelCase__						,lowerCamelCase__	)
           return doc_ids, retrieved_doc_embeds

class 		UpperCamelCase__			(				lowerCAmelCase_       ):



     '''simple docstring'''

     def __init__(							self						:   int						,lowerCamelCase__						:   Dict						,lowerCamelCase__						:   Optional[Any]						,lowerCamelCase__						:   List[str]						,lowerCamelCase__						:   Optional[Any]						,lowerCamelCase__						:   Dict=None	)					->		Any:




           '''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__						,)
           SCREAMING_SNAKE_CASE		=						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     SCREAMING_SNAKE_CASE__  (							self						:   List[str]	)					->		Union[str, Any]:




           '''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     SCREAMING_SNAKE_CASE__  (							self						:   Any						,lowerCamelCase__						:   Any						,lowerCamelCase__						:   int	)					->		Dict:




           '''simple docstring'''


           if len(self.retrieval_workers	) > 0:
                 # Select a random retrieval actor.
                 SCREAMING_SNAKE_CASE		=						self.retrieval_workers[random.randint(0						,len(self.retrieval_workers	) - 1	)]
                 SCREAMING_SNAKE_CASE, SCREAMING_SNAKE_CASE		=						ray.get(random_worker.retrieve.remote(lowerCamelCase__						,lowerCamelCase__	)	)
           else:
                 SCREAMING_SNAKE_CASE, SCREAMING_SNAKE_CASE		=						self._main_retrieve(lowerCamelCase__						,lowerCamelCase__	)
           return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(lowerCamelCase__	)






     @classmethod
     def     SCREAMING_SNAKE_CASE__  (							cls						:   Optional[Any]						,lowerCamelCase__						:   List[Any]						,lowerCamelCase__						:   Union[str, Any]=None						,**lowerCamelCase__						:   Optional[Any]	)					->		Any:




           '''simple docstring'''


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






     @classmethod
     def     SCREAMING_SNAKE_CASE__  (							cls						:   Any						,lowerCamelCase__						:   Tuple						,lowerCamelCase__						:   List[Any]						,lowerCamelCase__						:   List[Any]=None						,**lowerCamelCase__						:   Any	)					->		Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						kwargs.pop("""config"""						,lowerCamelCase__	) or RagConfig.from_pretrained(lowerCamelCase__						,**lowerCamelCase__	)
           SCREAMING_SNAKE_CASE		=						RagTokenizer.from_pretrained(lowerCamelCase__						,config=lowerCamelCase__	)
           SCREAMING_SNAKE_CASE		=						rag_tokenizer.question_encoder
           SCREAMING_SNAKE_CASE		=						rag_tokenizer.generator
           if indexed_dataset is not None:
                 SCREAMING_SNAKE_CASE		=						"""custom"""
                 SCREAMING_SNAKE_CASE		=						CustomHFIndex(config.retrieval_vector_size						,lowerCamelCase__	)
           else:
                 SCREAMING_SNAKE_CASE		=						cls._build_index(lowerCamelCase__	)
           return cls(
               lowerCamelCase__						,question_encoder_tokenizer=lowerCamelCase__						,generator_tokenizer=lowerCamelCase__						,retrieval_workers=lowerCamelCase__						,index=lowerCamelCase__						,)


