


'''simple docstring'''
import os
import zipfile

import pytest

from datasets.utils.extract import (
    BzipaExtractor,
    Extractor,
    GzipExtractor,
    LzaExtractor,
    SevenZipExtractor,
    TarExtractor,
    XzExtractor,
    ZipExtractor,
    ZstdExtractor,
)

from .utils import require_lza, require_pyazr, require_zstandard





@pytest.mark.parametrize(
    '''compression_format, is_archive'''    ,   [
        ('''7z''', True),
        ('''bz2''', False),
        ('''gzip''', False),
        ('''lz4''', False),
        ('''tar''', True),
        ('''xz''', False),
        ('''zip''', True),
        ('''zstd''', False),
    ]    ,   )
def    a__			(   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   )       ->      Optional[int]:
    UpperCAmelCase__     :						Optional[int]						=							{
        '''7z''': (seven_zip_file, SevenZipExtractor),
        '''bz2''': (bza_file, BzipaExtractor),
        '''gzip''': (gz_file, GzipExtractor),
        '''lz4''': (lza_file, LzaExtractor),
        '''tar''': (tar_file, TarExtractor),
        '''xz''': (xz_file, XzExtractor),
        '''zip''': (zip_file, ZipExtractor),
        '''zstd''': (zstd_file, ZstdExtractor),
    }
    UpperCAmelCase__		,		UpperCAmelCase__     :						int						=							input_paths_and_base_extractors[compression_format]
    if input_path is None:
        UpperCAmelCase__     :						Tuple						=							F"""for '{compression_format}' compression_format, """
        if compression_format == "7z":
            reason += require_pyazr.kwargs["reason"]
        elif compression_format == "lz4":
            reason += require_lza.kwargs["reason"]
        elif compression_format == "zstd":
            reason += require_zstandard.kwargs["reason"]
        pytest.skip(lowerCAmelCase__		)
    assert base_extractor.is_extractable(lowerCAmelCase__		)
    UpperCAmelCase__     :						List[str]						=							tmp_path / ('''extracted''' if is_archive else '''extracted.txt''')
    base_extractor.extract(lowerCAmelCase__    ,   lowerCAmelCase__		)
    if is_archive:
        assert output_path.is_dir()
        for file_path in output_path.iterdir():
            assert file_path.name == text_file.name
            UpperCAmelCase__     :						Dict						=							file_path.read_text(encoding='''utf-8'''		)
    else:
        UpperCAmelCase__     :						Union[str, Any]						=							output_path.read_text(encoding='''utf-8'''		)
    UpperCAmelCase__     :						Dict						=							text_file.read_text(encoding='''utf-8'''		)
    assert extracted_file_content == expected_file_content





@pytest.mark.parametrize(
    '''compression_format, is_archive'''    ,   [
        ('''7z''', True),
        ('''bz2''', False),
        ('''gzip''', False),
        ('''lz4''', False),
        ('''tar''', True),
        ('''xz''', False),
        ('''zip''', True),
        ('''zstd''', False),
    ]    ,   )
def    a__			(   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   )       ->      Optional[Any]:
    UpperCAmelCase__     :						Union[str, Any]						=							{
        '''7z''': seven_zip_file,
        '''bz2''': bza_file,
        '''gzip''': gz_file,
        '''lz4''': lza_file,
        '''tar''': tar_file,
        '''xz''': xz_file,
        '''zip''': zip_file,
        '''zstd''': zstd_file,
    }
    UpperCAmelCase__     :						Dict						=							input_paths[compression_format]
    if input_path is None:
        UpperCAmelCase__     :						int						=							F"""for '{compression_format}' compression_format, """
        if compression_format == "7z":
            reason += require_pyazr.kwargs["reason"]
        elif compression_format == "lz4":
            reason += require_lza.kwargs["reason"]
        elif compression_format == "zstd":
            reason += require_zstandard.kwargs["reason"]
        pytest.skip(lowerCAmelCase__		)
    UpperCAmelCase__     :						Any						=							Extractor.infer_extractor_format(lowerCAmelCase__		)
    assert extractor_format is not None
    UpperCAmelCase__     :						Any						=							tmp_path / ('''extracted''' if is_archive else '''extracted.txt''')
    Extractor.extract(lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__		)
    if is_archive:
        assert output_path.is_dir()
        for file_path in output_path.iterdir():
            assert file_path.name == text_file.name
            UpperCAmelCase__     :						Optional[int]						=							file_path.read_text(encoding='''utf-8'''		)
    else:
        UpperCAmelCase__     :						Optional[int]						=							output_path.read_text(encoding='''utf-8'''		)
    UpperCAmelCase__     :						Optional[int]						=							text_file.read_text(encoding='''utf-8'''		)
    assert extracted_file_content == expected_file_content





@pytest.fixture
def    a__			(   lowerCAmelCase__    ,   lowerCAmelCase__		)       ->      str:
    import tarfile

    UpperCAmelCase__     :						int						=							tmp_path / '''data_dot_dot'''
    directory.mkdir()
    UpperCAmelCase__     :						Tuple						=							directory / '''tar_file_with_dot_dot.tar'''
    with tarfile.TarFile(lowerCAmelCase__    ,   '''w'''		) as f:
        f.add(lowerCAmelCase__    ,   arcname=os.path.join('''..'''    ,   text_file.name		)		)
    return path





@pytest.fixture
def    a__			(   lowerCAmelCase__		)       ->      Any:
    import tarfile

    UpperCAmelCase__     :						str						=							tmp_path / '''data_sym_link'''
    directory.mkdir()
    UpperCAmelCase__     :						Optional[Any]						=							directory / '''tar_file_with_sym_link.tar'''
    os.symlink('''..'''    ,   directory / '''subdir'''    ,   target_is_directory=lowerCAmelCase__		)
    with tarfile.TarFile(lowerCAmelCase__    ,   '''w'''		) as f:
        f.add(str(directory / '''subdir'''		)    ,   arcname='''subdir'''		)  # str required by os.readlink on Windows and Python < 3.8
    return path





@pytest.mark.parametrize(
    '''insecure_tar_file, error_log'''    ,   [('''tar_file_with_dot_dot''', '''illegal path'''), ('''tar_file_with_sym_link''', '''Symlink''')]    ,   )
def    a__			(   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__		)       ->      int:
    UpperCAmelCase__     :						Dict						=							{
        '''tar_file_with_dot_dot''': tar_file_with_dot_dot,
        '''tar_file_with_sym_link''': tar_file_with_sym_link,
    }
    UpperCAmelCase__     :						Any						=							insecure_tar_files[insecure_tar_file]
    UpperCAmelCase__     :						Dict						=							tmp_path / '''extracted'''
    TarExtractor.extract(lowerCAmelCase__    ,   lowerCAmelCase__		)
    assert caplog.text
    for record in caplog.records:
        assert record.levelname == "ERROR"
        assert error_log in record.msg







def    a__			(   lowerCAmelCase__		)       ->      Optional[Any]:
    # We should have less false positives than zipfile.is_zipfile
    # We do that by checking only the magic number
    UpperCAmelCase__     :						Dict						=							tmpdir / '''not_a_zip_file'''
    # From: https://github.com/python/cpython/pull/5053
    UpperCAmelCase__     :						List[str]						=							(
        b'''\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00'''
        b'''\x00\x02\x08\x06\x00\x00\x00\x99\x81\xb6\'\x00\x00\x00\x15I'''
        b'''DATx\x01\x01\n\x00\xf5\xff\x00PK\x05\x06\x00PK\x06\x06\x07'''
        b'''\xac\x01N\xc6|a\r\x00\x00\x00\x00IEND\xaeB`\x82'''
    )
    with not_a_zip_file.open('''wb'''		) as f:
        f.write(lowerCAmelCase__		)
    assert zipfile.is_zipfile(str(lowerCAmelCase__		)		)  # is a false positive for `zipfile`
    assert not ZipExtractor.is_extractable(lowerCAmelCase__		)  # but we're right




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 	__lowercase    (__SCREAMING_SNAKE_CASE  ):





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					tempfile.mkdtemp()
										SCREAMING_SNAKE_CASE_		:     Dict					=					5

										# Realm tok
										SCREAMING_SNAKE_CASE_		:     Dict					=					[
										    '[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',
										]
										SCREAMING_SNAKE_CASE_		:     Tuple					=					os.path.join(self.tmpdirname			,						'realm_tokenizer'       )
										os.makedirs(lowerCAmelCase__			,						exist_ok=lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					os.path.join(lowerCAmelCase__			,						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]       )       )

										SCREAMING_SNAKE_CASE_		:     str					=					os.path.join(self.tmpdirname			,						'realm_block_records'       )
										os.makedirs(lowerCAmelCase__			,						exist_ok=lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										shutil.rmtree(self.tmpdirname       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					RealmConfig(num_block_records=self.num_block_records       )
										return config






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					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=lowerCAmelCase__			,						)
										return block_records






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					RealmRetriever(
										    block_records=self.get_dummy_block_records()			,						tokenizer=self.get_tokenizer()			,						)
										return retriever






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					self.get_config()
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					self.get_dummy_retriever()
										SCREAMING_SNAKE_CASE_		:     List[str]					=					retriever.tokenizer

										SCREAMING_SNAKE_CASE_		:     Any					=					np.array([0, 3]			,						dtype='long'       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					tokenizer(['Test question']       ).input_ids
										SCREAMING_SNAKE_CASE_		:     int					=					tokenizer(
										    ['the fourth']			,						add_special_tokens=lowerCAmelCase__			,						return_token_type_ids=lowerCAmelCase__			,						return_attention_mask=lowerCAmelCase__			,						).input_ids
										SCREAMING_SNAKE_CASE_		:     Any					=					config.reader_seq_len

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     str					=					retriever(
										    lowerCAmelCase__			,						lowerCAmelCase__			,						answer_ids=lowerCAmelCase__			,						max_length=lowerCAmelCase__			,						return_tensors='np'       )

										self.assertEqual(len(lowerCAmelCase__       )			,						2       )
										self.assertEqual(len(lowerCAmelCase__       )			,						2       )
										self.assertEqual(len(lowerCAmelCase__       )			,						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 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					self.get_config()
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.get_dummy_retriever()
										SCREAMING_SNAKE_CASE_		:     Tuple					=					retriever.tokenizer

										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					np.array([0, 3, 5]			,						dtype='long'       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					tokenizer(['Test question']       ).input_ids
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					tokenizer(
										    ['the fourth', 'longer longer']			,						add_special_tokens=lowerCAmelCase__			,						return_token_type_ids=lowerCAmelCase__			,						return_attention_mask=lowerCAmelCase__			,						).input_ids
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					config.reader_seq_len

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Tuple					=					retriever(
										    lowerCAmelCase__			,						lowerCAmelCase__			,						answer_ids=lowerCAmelCase__			,						max_length=lowerCAmelCase__			,						return_tensors='np'       )

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






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					self.get_dummy_retriever()
										retriever.save_pretrained(os.path.join(self.tmpdirname			,						'realm_block_records'       )       )

										# Test local path
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					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:
													SCREAMING_SNAKE_CASE_		:     Optional[int]					=					os.path.join(
													    os.path.join(self.tmpdirname			,						'realm_block_records'       )			,						_REALM_BLOCK_RECORDS_FILENAME       )
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					RealmRetriever.from_pretrained('google/realm-cc-news-pretrained-openqa'       )

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





