


"""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        __A       (snake_case_):





							'''simple docstring'''

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



														"""simple docstring"""



														snake_case_          =						tempfile.mkdtemp()
														snake_case_          =						5

														# Realm tok
														snake_case_          =						[
														    """[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""",
														]
														snake_case_          =						os.path.join(self.tmpdirname						,    """realm_tokenizer"""    )
														os.makedirs(UpperCAmelCase_						,    exist_ok=UpperCAmelCase_    )
														snake_case_          =						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]    )    )

														snake_case_          =						os.path.join(self.tmpdirname						,    """realm_block_records"""    )
														os.makedirs(UpperCAmelCase_						,    exist_ok=UpperCAmelCase_    )



							def 						lowerCAmelCase    (			self							: int    )   ->RealmTokenizer:



														"""simple docstring"""



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



							def 						lowerCAmelCase    (			self							: int    )   ->Union[str, Any]:



														"""simple docstring"""



														shutil.rmtree(self.tmpdirname    )



							def 						lowerCAmelCase    (			self							: Union[str, Any]    )   ->Dict:



														"""simple docstring"""



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



							def 						lowerCAmelCase    (			self							: Union[str, Any]    )   ->Optional[int]:



														"""simple docstring"""



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



							def 						lowerCAmelCase    (			self							: Dict    )   ->Any:



														"""simple docstring"""



														snake_case_          =						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							: Union[str, Any]    )   ->List[str]:



														"""simple docstring"""



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



							def 						lowerCAmelCase    (			self							: Dict    )   ->Dict:



														"""simple docstring"""



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

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

														snake_case_    ,  snake_case_    ,  snake_case_    ,  snake_case_          =						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, 10)    )
														self.assertEqual(concat_inputs.attention_mask.shape						,    (2, 10)    )
														self.assertEqual(concat_inputs.token_type_ids.shape						,    (2, 10)    )
														self.assertEqual(concat_inputs.special_tokens_mask.shape						,    (2, 10)    )
														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							: Any    )   ->List[str]:



														"""simple docstring"""



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

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

														snake_case_    ,  snake_case_    ,  snake_case_    ,  snake_case_          =						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							: List[str]    )   ->Union[str, Any]:



														"""simple docstring"""



														snake_case_          =						self.get_dummy_retriever()
														retriever.save_pretrained(os.path.join(self.tmpdirname						,    """realm_block_records"""    )    )

														# Test local path
														snake_case_          =						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:
																					snake_case_          =						os.path.join(
																					    os.path.join(self.tmpdirname						,    """realm_block_records"""    )						,    _REALM_BLOCK_RECORDS_FILENAME    )
																					snake_case_          =						RealmRetriever.from_pretrained("""google/realm-cc-news-pretrained-openqa"""    )

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







"""simple docstring"""
from dataclasses import dataclass, field
from typing import Optional





@dataclass
class        __A       :





							'''simple docstring'''

							__lowercase:						Optional[str]							 =      field(
							    default="""codeparrot/codeparrot"""  ,						metadata={"""help""": """Model name or path of model to be trained."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""./"""  ,						metadata={"""help""": """Save dir where model repo is cloned and models updates are saved to."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""codeparrot/codeparrot-clean-train"""  ,						metadata={"""help""": """Name or path of training dataset."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""codeparrot/codeparrot-clean-valid"""  ,						metadata={"""help""": """Name or path of validation dataset."""})
							__lowercase:						Optional[int]							 =      field(default=2  ,						metadata={"""help""": """Batch size for training."""})
							__lowercase:						Optional[int]							 =      field(default=2  ,						metadata={"""help""": """Batch size for evaluation."""})
							__lowercase:						Optional[float]							 =      field(default=0.1  ,						metadata={"""help""": """Value of weight decay."""})
							__lowercase:						Optional[int]							 =      field(
							    default=1_00_00  ,						metadata={"""help""": """Size of buffer used to shuffle streaming dataset."""})
							__lowercase:						Optional[float]							 =      field(default=2E-4  ,						metadata={"""help""": """Learning rate fo training."""})
							__lowercase:						Optional[str]							 =      field(default="""cosine"""  ,						metadata={"""help""": """Learning rate."""})
							__lowercase:						Optional[int]							 =      field(
							    default=7_50  ,						metadata={"""help""": """Number of warmup steps in the learning rate schedule."""})
							__lowercase:						Optional[int]							 =      field(
							    default=16  ,						metadata={"""help""": """Number of gradient accumulation steps."""})
							__lowercase:						Optional[bool]							 =      field(
							    default=snake_case__  ,						metadata={"""help""": """Use gradient checkpointing to reduce memory footprint."""})
							__lowercase:						Optional[int]							 =      field(default=5_00_00  ,						metadata={"""help""": """Maximum number of training steps."""})
							__lowercase:						Optional[int]							 =      field(
							    default=-1  ,						metadata={"""help""": """Maximum number of evaluation steps. If -1 the full dataset is evaluated."""})
							__lowercase:						Optional[int]							 =      field(default=10_24  ,						metadata={"""help""": """Sequence lengths used for training."""})
							__lowercase:						Optional[int]							 =      field(default=1  ,						metadata={"""help""": """Training seed."""})
							__lowercase:						Optional[int]							 =      field(
							    default=10_24  ,						metadata={"""help""": """Interval to save checkpoints. Measured as number of forward passes not training steps."""}  ,						)
							__lowercase:						Optional[str]							 =      field(
							    default=snake_case__  ,						metadata={"""help""": """States path if the training should continue from a checkpoint folder."""})
							__lowercase:						Optional[bool]							 =      field(default=snake_case__  ,						metadata={"""help""": """If True the data is pretokenized."""})





@dataclass
class        __A       :





							'''simple docstring'''

							__lowercase:						Optional[str]							 =      field(
							    default="""codeparrot/codeparrot"""  ,						metadata={"""help""": """Model name or path of model to be evaluated."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""codeparrot/codeparrot-clean-valid"""  ,						metadata={"""help""": """Name or path of validation dataset."""})
							__lowercase:						Optional[int]							 =      field(default=2  ,						metadata={"""help""": """Batch size used for evaluation."""})
							__lowercase:						Optional[int]							 =      field(
							    default=-1  ,						metadata={"""help""": """Maximum number of evaluation steps. If -1 the full dataset is evaluated."""})
							__lowercase:						Optional[int]							 =      field(default=10_24  ,						metadata={"""help""": """Length of sequences to be evaluated."""})
							__lowercase:						Optional[int]							 =      field(default=1  ,						metadata={"""help""": """Random seed used for evaluation."""})





@dataclass
class        __A       :





							'''simple docstring'''

							__lowercase:						Optional[str]							 =      field(
							    default="""codeparrot/codeparrot"""  ,						metadata={"""help""": """Model name or path of model to be evaluated."""})
							__lowercase:						Optional[int]							 =      field(default=snake_case__  ,						metadata={"""help""": """Number of workers used for code evaluation."""})
							__lowercase:						Optional[int]							 =      field(
							    default=snake_case__  ,						metadata={"""help""": """The number of human-eval tasks to run. If not included all tasks are evaluated."""}  ,						)
							__lowercase:						Optional[bool]							 =      field(
							    default=snake_case__  ,						metadata={"""help""": """Sample from the language model's output distribution."""})
							__lowercase:						Optional[float]							 =      field(default=0.2  ,						metadata={"""help""": """Sampling temperature used for generation."""})
							__lowercase:						Optional[int]							 =      field(default=2_56  ,						metadata={"""help""": """Maximum number of newly generated tokens."""})
							__lowercase:						Optional[int]							 =      field(default=0  ,						metadata={"""help""": """Top-k parameter used for generation."""})
							__lowercase:						Optional[float]							 =      field(default=0.9_5  ,						metadata={"""help""": """Top-p parameter used for nucleus sampling."""})
							__lowercase:						Optional[int]							 =      field(default=10  ,						metadata={"""help""": """Number of generations to run in parallel."""})
							__lowercase:						Optional[int]							 =      field(
							    default=2_00  ,						metadata={"""help""": """Number of completions to generate for each sample."""})
							__lowercase:						Optional[int]							 =      field(default=1  ,						metadata={"""help""": """Random seed used for evaluation."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""eval_results.json"""  ,						metadata={"""help""": """Random seed used for evaluation."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""0"""  ,						metadata={"""help""": """Allow `code_eval` to execute Python code on machine"""})
							__lowercase:						Optional[int]							 =      field(
							    default=-1  ,						metadata={
							        """help""": (
							            """Determine which device to run the `text-generation` Pipeline on. -1 is CPU and any zero or positive"""
							            """ number corresponds to which GPU device id to run on."""
							        )
							    }  ,						)





@dataclass
class        __A       :





							'''simple docstring'''

							__lowercase:						Optional[int]							 =      field(
							    default=snake_case__  ,						metadata={
							        """help""": """The number of CPU cores to use for parallel preprocessing. Default uses the maximum available."""
							    }  ,						)
							__lowercase:						Optional[str]							 =      field(
							    default="""transformersbook/codeparrot"""  ,						metadata={"""help""": """Folder or name of dataset to process."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""codeparrot-clean"""  ,						metadata={"""help""": """Folder to save processed processed dataset."""})
							__lowercase:						Optional[int]							 =      field(
							    default=10_00_00  ,						metadata={"""help""": """Number of files to save per JSON output file."""})
							__lowercase:						Optional[str]							 =      field(default="""content"""  ,						metadata={"""help""": """Column containing text data to process."""})
							__lowercase:						Optional[float]							 =      field(
							    default=10_00  ,						metadata={"""help""": """Maximum line length in file, otherwise file is filtered."""})
							__lowercase:						Optional[float]							 =      field(
							    default=1_00  ,						metadata={"""help""": """Maximum mean line length in file, otherwise file is filtered."""})
							__lowercase:						Optional[float]							 =      field(
							    default=0.2_5  ,						metadata={"""help""": """Maximum fraction of non-alphanumeric characters, otherwise file is filtered."""})
							__lowercase:						Optional[float]							 =      field(
							    default=1.5  ,						metadata={"""help""": """Minimum character token ratio for the file, otherwise file is filtered."""})
							__lowercase:						Optional[float]							 =      field(
							    default=0.7  ,						metadata={"""help""": """Probability for filtering config, test and uncommon files."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""codeparrot/codeparrot"""  ,						metadata={"""help""": """Name or path to the tokenizer."""}  ,						)
							__lowercase:						Optional[bool]							 =      field(
							    default=snake_case__  ,						metadata={"""help""": """If True, near-duplicate samples are removed."""})
							__lowercase:						Optional[float]							 =      field(
							    default=0.8_5  ,						metadata={"""help""": """Jaccard threshold for near-duplicate samples."""})





@dataclass
class        __A       :





							'''simple docstring'''

							__lowercase:						Optional[str]							 =      field(
							    default="""gpt2"""  ,						metadata={"""help""": """Base tokenizer to build new tokenizer from."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""transformersbook/codeparrot-train"""  ,						metadata={"""help""": """Dataset to train tokenizer on."""})
							__lowercase:						Optional[str]							 =      field(default="""content"""  ,						metadata={"""help""": """Column containing text data to process."""})
							__lowercase:						Optional[int]							 =      field(default=20_00_00  ,						metadata={"""help""": """Number of examples to train tokenizer on."""})
							__lowercase:						Optional[int]							 =      field(
							    default=3_27_68  ,						metadata={"""help""": """Number of examples to train the tokenizer on."""})
							__lowercase:						Optional[str]							 =      field(default="""codeparrot"""  ,						metadata={"""help""": """Name of new tokenizer."""})
							__lowercase:						Optional[bool]							 =      field(default=snake_case__  ,						metadata={"""help""": """Push saved tokenizer to the hub."""})





@dataclass
class        __A       :





							'''simple docstring'''

							__lowercase:						Optional[str]							 =      field(
							    default="""codeparrot/codeparrot"""  ,						metadata={"""help""": """Name or path to the tokenizer."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""codeparrot/codeparrot-clean-train"""  ,						metadata={"""help""": """Name or path to the dataset to pretokenize."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""tokenized-codeparrot-train"""  ,						metadata={"""help""": """Repo name of the pretokenized data."""})
							__lowercase:						Optional[int]							 =      field(default=snake_case__  ,						metadata={"""help""": """Number of workers used for code evaluation."""})





@dataclass
class        __A       :





							'''simple docstring'''

							__lowercase:						Optional[str]							 =      field(
							    default="""gpt2-large"""  ,						metadata={"""help""": """Configuration to use for model initialization."""})
							__lowercase:						Optional[str]							 =      field(
							    default="""codeparrot/codeparrot"""  ,						metadata={"""help""": """Tokenizer attached to model."""})
							__lowercase:						Optional[str]							 =      field(default="""codeparrot"""  ,						metadata={"""help""": """Name of the created model."""})
							__lowercase:						Optional[bool]							 =      field(default=snake_case__  ,						metadata={"""help""": """Push saved tokenizer to the hub."""})



