
'''simple docstring'''



import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer





class _UpperCamelCase		(     A		):





				'''simple docstring'''





				def 				__lowerCamelCase     (							self  :  List[Any]):







											'''simple docstring'''

											__lowercase       	=tempfile.mkdtemp()
											__lowercase       	=5

											# Realm tok
											__lowercase       	=[
											    '''[UNK]''',
											    '''[CLS]''',
											    '''[SEP]''',
											    '''[PAD]''',
											    '''[MASK]''',
											    '''test''',
											    '''question''',
											    '''this''',
											    '''is''',
											    '''the''',
											    '''first''',
											    '''second''',
											    '''third''',
											    '''fourth''',
											    '''fifth''',
											    '''record''',
											    '''want''',
											    '''##want''',
											    '''##ed''',
											    '''wa''',
											    '''un''',
											    '''runn''',
											    '''##ing''',
											    ''',''',
											    '''low''',
											    '''lowest''',
											]
											__lowercase       	=os.path.join(self.tmpdirname  ,       'realm_tokenizer')
											os.makedirs(UpperCAmelCase__  ,       exist_ok=UpperCAmelCase__)
											__lowercase       	=os.path.join(UpperCAmelCase__  ,       VOCAB_FILES_NAMES['vocab_file'])
											with open(self.vocab_file  ,       'w'  ,       encoding='utf-8') as vocab_writer:
																		vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]))

											__lowercase       	=os.path.join(self.tmpdirname  ,       'realm_block_records')
											os.makedirs(UpperCAmelCase__  ,       exist_ok=UpperCAmelCase__)

				def 				__lowerCamelCase     (							self  :  List[Any]):







											'''simple docstring'''

											return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname  ,       'realm_tokenizer'))

				def 				__lowerCamelCase     (							self  :  Any):







											'''simple docstring'''

											shutil.rmtree(self.tmpdirname)

				def 				__lowerCamelCase     (							self  :  Optional[int]):







											'''simple docstring'''

											__lowercase       	=RealmConfig(num_block_records=self.num_block_records)
											return config

				def 				__lowerCamelCase     (							self  :  int):







											'''simple docstring'''

											__lowercase       	=Dataset.from_dict(
											    {
											        'id': ['0', '1'],
											        'question': ['foo', 'bar'],
											        'answers': [['Foo', 'Bar'], ['Bar']],
											    })
											return dataset

				def 				__lowerCamelCase     (							self  :  str):







											'''simple docstring'''

											__lowercase       	=np.array(
											    [
											        B'This is the first record',
											        B'This is the second record',
											        B'This is the third record',
											        B'This is the fourth record',
											        B'This is the fifth record',
											        B'This is a longer longer longer record',
											    ]  ,       dtype=UpperCAmelCase__  ,       )
											return block_records

				def 				__lowerCamelCase     (							self  :  Optional[Any]):







											'''simple docstring'''

											__lowercase       	=RealmRetriever(
											    block_records=self.get_dummy_block_records()  ,       tokenizer=self.get_tokenizer()  ,       )
											return retriever

				def 				__lowerCamelCase     (							self  :  Optional[Any]):







											'''simple docstring'''

											__lowercase       	=self.get_config()
											__lowercase       	=self.get_dummy_retriever()
											__lowercase       	=retriever.tokenizer

											__lowercase       	=np.array([0, 3]  ,       dtype='long')
											__lowercase       	=tokenizer(['Test question']).input_ids
											__lowercase       	=tokenizer(
											    ['the fourth']  ,       add_special_tokens=UpperCAmelCase__  ,       return_token_type_ids=UpperCAmelCase__  ,       return_attention_mask=UpperCAmelCase__  ,       ).input_ids
											__lowercase       	=config.reader_seq_len

											__lowercase       	=retriever(
											    UpperCAmelCase__  ,       UpperCAmelCase__  ,       answer_ids=UpperCAmelCase__  ,       max_length=UpperCAmelCase__  ,       return_tensors='np')

											self.assertEqual(len(UpperCAmelCase__)  ,       2)
											self.assertEqual(len(UpperCAmelCase__)  ,       2)
											self.assertEqual(len(UpperCAmelCase__)  ,       2)
											self.assertEqual(concat_inputs.input_ids.shape  ,       (2, 1_0))
											self.assertEqual(concat_inputs.attention_mask.shape  ,       (2, 1_0))
											self.assertEqual(concat_inputs.token_type_ids.shape  ,       (2, 1_0))
											self.assertEqual(concat_inputs.special_tokens_mask.shape  ,       (2, 1_0))
											self.assertEqual(
											    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0])  ,       ['[CLS]', 'test', 'question', '[SEP]', 'this', 'is', 'the', 'first', 'record', '[SEP]']  ,       )
											self.assertEqual(
											    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1])  ,       ['[CLS]', 'test', 'question', '[SEP]', 'this', 'is', 'the', 'fourth', 'record', '[SEP]']  ,       )

				def 				__lowerCamelCase     (							self  :  Optional[Any]):







											'''simple docstring'''

											__lowercase       	=self.get_config()
											__lowercase       	=self.get_dummy_retriever()
											__lowercase       	=retriever.tokenizer

											__lowercase       	=np.array([0, 3, 5]  ,       dtype='long')
											__lowercase       	=tokenizer(['Test question']).input_ids
											__lowercase       	=tokenizer(
											    ['the fourth', 'longer longer']  ,       add_special_tokens=UpperCAmelCase__  ,       return_token_type_ids=UpperCAmelCase__  ,       return_attention_mask=UpperCAmelCase__  ,       ).input_ids
											__lowercase       	=config.reader_seq_len

											__lowercase       	=retriever(
											    UpperCAmelCase__  ,       UpperCAmelCase__  ,       answer_ids=UpperCAmelCase__  ,       max_length=UpperCAmelCase__  ,       return_tensors='np')

											self.assertEqual([False, True, True]  ,       UpperCAmelCase__)
											self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]  ,       UpperCAmelCase__)
											self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]  ,       UpperCAmelCase__)

				def 				__lowerCamelCase     (							self  :  Optional[int]):







											'''simple docstring'''

											__lowercase       	=self.get_dummy_retriever()
											retriever.save_pretrained(os.path.join(self.tmpdirname  ,       'realm_block_records'))

											# Test local path
											__lowercase       	=retriever.from_pretrained(os.path.join(self.tmpdirname  ,       'realm_block_records'))
											self.assertEqual(retriever.block_records[0]  ,       B'This is the first record')

											# Test mocked remote path
											with patch('transformers.models.realm.retrieval_realm.hf_hub_download') as mock_hf_hub_download:
																		__lowercase       	=os.path.join(
																		    os.path.join(self.tmpdirname  ,       'realm_block_records')  ,       _REALM_BLOCK_RECORDS_FILENAME)
																		__lowercase       	=RealmRetriever.from_pretrained('google/realm-cc-news-pretrained-openqa')

											self.assertEqual(retriever.block_records[0]  ,       B'This is the first record')








'''simple docstring'''



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

import numpy as np

from ...utils import is_tf_available, is_torch_available, logging


if TYPE_CHECKING:
		if is_torch_available():
				import torch
		if is_tf_available():
				import tensorflow as tf

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from .tokenization_codegen import CodeGenTokenizer


lowerCamelCase	 =       logging.get_logger(__name__)

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

lowerCamelCase	 =       {
    """vocab_file""": {
        """Salesforce/codegen-350M-mono""": """https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/vocab.json""",
    },
    """merges_file""": {
        """Salesforce/codegen-350M-mono""": """https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/merges.txt""",
    },
    """tokenizer_file""": {
        """Salesforce/codegen-350M-mono""": (
            """https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/tokenizer.json"""
        ),
    },
}

lowerCamelCase	 =       {
    """Salesforce/codegen-350M-mono""": 2048,
}





class _UpperCamelCase		(     A		):





				'''simple docstring'''





				lowerCAmelCase__					    =    VOCAB_FILES_NAMES
				lowerCAmelCase__					    =    PRETRAINED_VOCAB_FILES_MAP
				lowerCAmelCase__					    =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				lowerCAmelCase__					    =    ["""input_ids""", """attention_mask"""]
				lowerCAmelCase__					    =    CodeGenTokenizer

				def __init__(							self  :  int  ,       _lowerCAmelCase  :  str=None  ,       _lowerCAmelCase  :  List[Any]=None  ,       _lowerCAmelCase  :  Optional[Any]=None  ,       _lowerCAmelCase  :  List[str]="<|endoftext|>"  ,       _lowerCAmelCase  :  Optional[int]="<|endoftext|>"  ,       _lowerCAmelCase  :  Optional[Any]="<|endoftext|>"  ,       _lowerCAmelCase  :  Optional[int]=False  ,       **_lowerCAmelCase  :  Union[str, Any]  ,       ):







											'''simple docstring'''

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

											if kwargs.pop('add_bos_token'  ,       _lowerCAmelCase):
																		__lowercase       	=kwargs.pop('name_or_path'  ,       '')
																		raise ValueError(
																		    'Currenty GPT2\'s fast tokenizer does NOT support adding a BOS token.'
																		    'Instead you should use GPT2\'s slow tokenizer class `CodeGenTokenizer` as follows: \n'
																		    f"""`CodeGenTokenizer.from_pretrained('{model_id}')`\nor\n"""
																		    f"""`AutoTokenizer.from_pretrained('{model_id}', use_fast=False)`\n"""
																		    'This issue will be fixed soon, see: https://github.com/huggingface/tokenizers/pull/1005.'
																		    ' so that the fast tokenizer works correctly.')

											__lowercase       	=json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__())
											if pre_tok_state.get('add_prefix_space'  ,       _lowerCAmelCase) != add_prefix_space:
																		__lowercase       	=getattr(_lowerCAmelCase  ,       pre_tok_state.pop('type'))
																		__lowercase       	=add_prefix_space
																		__lowercase       	=pre_tok_class(**_lowerCAmelCase)

											__lowercase       	=add_prefix_space

				def 				__lowerCamelCase     (							self  :  Any  ,       *_lowerCAmelCase  :  Tuple  ,       **_lowerCAmelCase  :  Union[str, Any]):







											'''simple docstring'''

											__lowercase       	=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 				__lowerCamelCase     (							self  :  List[str]  ,       *_lowerCAmelCase  :  List[Any]  ,       **_lowerCAmelCase  :  List[str]):







											'''simple docstring'''

											__lowercase       	=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 				__lowerCamelCase     (							self  :  str  ,       _lowerCAmelCase  :  str  ,       _lowerCAmelCase  :  Optional[str] = None):







											'''simple docstring'''

											__lowercase       	=self._tokenizer.model.save(_lowerCAmelCase  ,       name=_lowerCAmelCase)
											return tuple(_lowerCAmelCase)

				def 				__lowerCamelCase     (							self  :  str  ,       _lowerCAmelCase  :  Union[int, List[int], "np.ndarray", "torch.Tensor", "tf.Tensor"]  ,       _lowerCAmelCase  :  bool = False  ,       _lowerCAmelCase  :  bool = None  ,       _lowerCAmelCase  :  Optional[List[str]] = None  ,       **_lowerCAmelCase  :  Any  ,       ):







											'''simple docstring'''

											__lowercase       	=super().decode(
											    token_ids=_lowerCAmelCase  ,       skip_special_tokens=_lowerCAmelCase  ,       clean_up_tokenization_spaces=_lowerCAmelCase  ,       **_lowerCAmelCase  ,       )

											if truncate_before_pattern is not None and len(_lowerCAmelCase) > 0:
																		__lowercase       	=self.truncate(_lowerCAmelCase  ,       _lowerCAmelCase)

											return decoded_text

				def 				__lowerCamelCase     (							self  :  Any  ,       _lowerCAmelCase  :  Union[str, Any]  ,       _lowerCAmelCase  :  Optional[int]):







											'''simple docstring'''

											def find_re(_lowerCAmelCase  :  List[Any]  ,       _lowerCAmelCase  :  Union[str, Any]  ,       _lowerCAmelCase  :  str):
																		__lowercase       	=pattern.search(_lowerCAmelCase  ,       _lowerCAmelCase)
																		return m.start() if m else -1

											__lowercase       	=[re.compile(_lowerCAmelCase  ,       re.MULTILINE) for pattern in truncate_before_pattern]

											__lowercase       	=list(re.finditer('^print'  ,       _lowerCAmelCase  ,       re.MULTILINE))

											if len(_lowerCAmelCase) > 1:
																		__lowercase       	=completion[: prints[1].start()]

											__lowercase       	=list(re.finditer('^def'  ,       _lowerCAmelCase  ,       re.MULTILINE))

											if len(_lowerCAmelCase) > 1:
																		__lowercase       	=completion[: defs[1].start()]

											__lowercase       	=0

											__lowercase       	=[
											    pos for pos in [find_re(_lowerCAmelCase  ,       _lowerCAmelCase  ,       _lowerCAmelCase) for terminal in terminals] if pos != -1
											]

											if len(_lowerCAmelCase) > 0:
																		return completion[: min(_lowerCAmelCase)]
											else:
																		return completion



