
import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu
class 					lowercase__(							nn.Module				):



	"""simple docstring"""




	def __init__(			self    :     Optional[int]			)			->	Optional[Any]:
					super().__init__()
					lowercase_	    =   nn.Linear(3			,				4			)
					lowercase_	    =   nn.BatchNormad(4			)
					lowercase_	    =   nn.Linear(4			,				5			)




	def 		_lowercase    (			self    :     Union[str, Any]			,				SCREAMING_SNAKE_CASE_    :     Optional[int]			)			->	Tuple:
					return self.lineara(self.batchnorm(self.lineara(SCREAMING_SNAKE_CASE_			)			)			)

class 					lowercase__(							UpperCAmelCase				):



	"""simple docstring"""




	def 		_lowercase    (			self    :     str			,				SCREAMING_SNAKE_CASE_    :     int			,				*SCREAMING_SNAKE_CASE_    :     Dict			,				**SCREAMING_SNAKE_CASE_    :     str			)			->	Optional[Any]:
					return (args[0] + 1,) + args[1:], kwargs

class 					lowercase__(							UpperCAmelCase				):



	"""simple docstring"""




	def 		_lowercase    (			self    :     int			,				SCREAMING_SNAKE_CASE_    :     Dict			,				SCREAMING_SNAKE_CASE_    :     List[Any]			)			->	List[str]:
					return output + 1

class 					lowercase__(							unittest.TestCase				):



	"""simple docstring"""




	def 		_lowercase    (			self    :     List[Any]			)			->	Dict:
					lowercase_	    =   ModelForTest()
					lowercase_	    =   ModelHook()

					add_hook_to_module(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					self.assertEqual(test_model._hf_hook			,				SCREAMING_SNAKE_CASE_			)
					self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_			,				'''_old_forward'''			)			)

					# Check adding the hook did not change the name or the signature
					self.assertEqual(test_model.forward.__name__			,				'''forward'''			)
					self.assertListEqual(list(inspect.signature(test_model.forward			).parameters			)			,				['''x''']			)

					remove_hook_from_module(SCREAMING_SNAKE_CASE_			)
					self.assertFalse(hasattr(SCREAMING_SNAKE_CASE_			,				'''_hf_hook'''			)			)
					self.assertFalse(hasattr(SCREAMING_SNAKE_CASE_			,				'''_old_forward'''			)			)
	def 		_lowercase    (			self    :     int			)			->	Tuple:
					lowercase_	    =   ModelForTest()
					lowercase_	    =   ModelHook()

					add_hook_to_module(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					add_hook_to_module(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				append=SCREAMING_SNAKE_CASE_			)

					self.assertEqual(isinstance(test_model._hf_hook			,				SCREAMING_SNAKE_CASE_			)			,				SCREAMING_SNAKE_CASE_			)
					self.assertEqual(len(test_model._hf_hook.hooks			)			,				2			)
					self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_			,				'''_old_forward'''			)			)

					# Check adding the hook did not change the name or the signature
					self.assertEqual(test_model.forward.__name__			,				'''forward'''			)
					self.assertListEqual(list(inspect.signature(test_model.forward			).parameters			)			,				['''x''']			)

					remove_hook_from_module(SCREAMING_SNAKE_CASE_			)
					self.assertFalse(hasattr(SCREAMING_SNAKE_CASE_			,				'''_hf_hook'''			)			)
					self.assertFalse(hasattr(SCREAMING_SNAKE_CASE_			,				'''_old_forward'''			)			)
	def 		_lowercase    (			self    :     List[Any]			)			->	int:
					lowercase_	    =   ModelForTest()
					lowercase_	    =   torch.randn(2			,				3			)
					lowercase_	    =   test_model(x + 1			)
					lowercase_	    =   test_model(x + 2			)

					lowercase_	    =   PreForwardHook()
					add_hook_to_module(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   test_model(SCREAMING_SNAKE_CASE_			)
					self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				atol=1e-5			)			)

					# Attaching a hook to a model when it already has one replaces, does not chain
					lowercase_	    =   PreForwardHook()
					add_hook_to_module(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   test_model(SCREAMING_SNAKE_CASE_			)
					self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				atol=1e-5			)			)

					# You need to use the sequential hook to chain two or more hooks
					lowercase_	    =   SequentialHook(PreForwardHook()			,				PreForwardHook()			)
					add_hook_to_module(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   test_model(SCREAMING_SNAKE_CASE_			)
					assert torch.allclose(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				atol=1e-5			)
	def 		_lowercase    (			self    :     Optional[Any]			)			->	int:
					lowercase_	    =   ModelForTest()
					lowercase_	    =   torch.randn(2			,				3			)
					lowercase_	    =   test_model(SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   PostForwardHook()
					add_hook_to_module(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   test_model(SCREAMING_SNAKE_CASE_			)
					self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_			,				output + 1			,				atol=1e-5			)			)

					# Attaching a hook to a model when it already has one replaces, does not chain
					lowercase_	    =   PostForwardHook()
					add_hook_to_module(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   test_model(SCREAMING_SNAKE_CASE_			)
					self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_			,				output + 1			,				atol=1e-5			)			)

					# You need to use the sequential hook to chain two or more hooks
					lowercase_	    =   SequentialHook(PostForwardHook()			,				PostForwardHook()			)
					add_hook_to_module(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   test_model(SCREAMING_SNAKE_CASE_			)
					assert torch.allclose(SCREAMING_SNAKE_CASE_			,				output + 2			,				atol=1e-5			)
	def 		_lowercase    (			self    :     List[str]			)			->	Optional[Any]:
					lowercase_	    =   ModelForTest()
					lowercase_	    =   torch.randn(2			,				3			)
					lowercase_	    =   test_model(SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   PostForwardHook()
					add_hook_to_module(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   test_model(SCREAMING_SNAKE_CASE_			)
					self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_			,				output + 1			)			)
					self.assertTrue(outputa.requires_grad			)

					lowercase_	    =   True
					lowercase_	    =   test_model(SCREAMING_SNAKE_CASE_			)
					self.assertFalse(outputa.requires_grad			)
	@require_multi_gpu
	def 		_lowercase    (			self    :     Optional[Any]			)			->	int:
					lowercase_	    =   ModelForTest()
					# Everything is on CPU
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)

					# This will move each submodule on different devices
					add_hook_to_module(model.lineara			,				AlignDevicesHook(execution_device=0			)			)
					add_hook_to_module(model.batchnorm			,				AlignDevicesHook(execution_device=0			)			)
					add_hook_to_module(model.lineara			,				AlignDevicesHook(execution_device=1			)			)

					self.assertEqual(model.lineara.weight.device			,				torch.device(0			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device(0			)			)
					self.assertEqual(model.batchnorm.running_mean.device			,				torch.device(0			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device(1			)			)

					# We can still make a forward pass. The input does not need to be on any particular device
					lowercase_	    =   torch.randn(2			,				3			)
					lowercase_	    =   model(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(output.device			,				torch.device(1			)			)

					# We can add a general hook to put back output on same device as input.
					add_hook_to_module(SCREAMING_SNAKE_CASE_			,				AlignDevicesHook(io_same_device=SCREAMING_SNAKE_CASE_			)			)
					lowercase_	    =   torch.randn(2			,				3			).to(0			)
					lowercase_	    =   model(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(output.device			,				torch.device(0			)			)
	def 		_lowercase    (			self    :     List[str]			)			->	Dict:
					lowercase_	    =   ModelForTest()

					# Everything is on CPU
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)

					# This will move each submodule on different devices
					lowercase_	    =   {'''execution_device''': 0 if torch.cuda.is_available() else '''cpu''', '''offload''': True}

					add_hook_to_module(model.lineara			,				AlignDevicesHook(**SCREAMING_SNAKE_CASE_			)			)
					add_hook_to_module(model.batchnorm			,				AlignDevicesHook(**SCREAMING_SNAKE_CASE_			)			)
					add_hook_to_module(model.lineara			,				AlignDevicesHook(**SCREAMING_SNAKE_CASE_			)			)

					# Parameters have been offloaded, so on the meta device
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					# Buffers are not included in the offload by default, so are on the execution device
					lowercase_	    =   torch.device(hook_kwargs['''execution_device''']			)
					self.assertEqual(model.batchnorm.running_mean.device			,				SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   torch.randn(2			,				3			)
					lowercase_	    =   model(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(output.device			,				SCREAMING_SNAKE_CASE_			)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_module(model.lineara			)
					remove_hook_from_module(model.batchnorm			)
					remove_hook_from_module(model.lineara			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)

					# Now test with buffers included in the offload
					lowercase_	    =   {
					    '''execution_device''': 0 if torch.cuda.is_available() else '''cpu''',
					    '''offload''': True,
					    '''offload_buffers''': True,
					}

					add_hook_to_module(model.lineara			,				AlignDevicesHook(**SCREAMING_SNAKE_CASE_			)			)
					add_hook_to_module(model.batchnorm			,				AlignDevicesHook(**SCREAMING_SNAKE_CASE_			)			)
					add_hook_to_module(model.lineara			,				AlignDevicesHook(**SCREAMING_SNAKE_CASE_			)			)

					# Parameters have been offloaded, so on the meta device, buffers included
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.batchnorm.running_mean.device			,				torch.device('''meta'''			)			)

					lowercase_	    =   torch.randn(2			,				3			)
					lowercase_	    =   model(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(output.device			,				SCREAMING_SNAKE_CASE_			)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_module(model.lineara			)
					remove_hook_from_module(model.batchnorm			)
					remove_hook_from_module(model.lineara			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)
	def 		_lowercase    (			self    :     Optional[Any]			)			->	List[str]:
					lowercase_	    =   ModelForTest()

					# Everything is on CPU
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)

					# This will move each submodule on different devices
					lowercase_	    =   0 if torch.cuda.is_available() else '''cpu'''
					attach_align_device_hook(SCREAMING_SNAKE_CASE_			,				execution_device=SCREAMING_SNAKE_CASE_			,				offload=SCREAMING_SNAKE_CASE_			)

					# Parameters have been offloaded, so on the meta device
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					# Buffers are not included in the offload by default, so are on the execution device
					lowercase_	    =   torch.device(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(model.batchnorm.running_mean.device			,				SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   torch.randn(2			,				3			)
					lowercase_	    =   model(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(output.device			,				SCREAMING_SNAKE_CASE_			)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_submodules(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)

					# Now test with buffers included in the offload
					attach_align_device_hook(SCREAMING_SNAKE_CASE_			,				execution_device=SCREAMING_SNAKE_CASE_			,				offload=SCREAMING_SNAKE_CASE_			,				offload_buffers=SCREAMING_SNAKE_CASE_			)

					# Parameters have been offloaded, so on the meta device, buffers included
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.batchnorm.running_mean.device			,				torch.device('''meta'''			)			)

					lowercase_	    =   torch.randn(2			,				3			)
					lowercase_	    =   model(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(output.device			,				SCREAMING_SNAKE_CASE_			)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_submodules(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)




	def 		_lowercase    (			self    :     Dict			)			->	Any:
					lowercase_	    =   ModelForTest()

					# Everything is on CPU
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)

					# This will move each submodule on different devices
					lowercase_	    =   0 if torch.cuda.is_available() else '''cpu'''
					attach_align_device_hook(
					    SCREAMING_SNAKE_CASE_			,				execution_device=SCREAMING_SNAKE_CASE_			,				offload=SCREAMING_SNAKE_CASE_			,				weights_map=model.state_dict()			)

					# Parameters have been offloaded, so on the meta device
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					# Buffers are not included in the offload by default, so are on the execution device
					lowercase_	    =   torch.device(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(model.batchnorm.running_mean.device			,				SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   torch.randn(2			,				3			)
					lowercase_	    =   model(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(output.device			,				SCREAMING_SNAKE_CASE_			)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_submodules(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)

					# Now test with buffers included in the offload
					attach_align_device_hook(
					    SCREAMING_SNAKE_CASE_			,				execution_device=SCREAMING_SNAKE_CASE_			,				offload=SCREAMING_SNAKE_CASE_			,				weights_map=model.state_dict()			,				offload_buffers=SCREAMING_SNAKE_CASE_			,				)

					# Parameters have been offloaded, so on the meta device, buffers included
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''meta'''			)			)
					self.assertEqual(model.batchnorm.running_mean.device			,				torch.device('''meta'''			)			)

					lowercase_	    =   torch.randn(2			,				3			)
					lowercase_	    =   model(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(output.device			,				SCREAMING_SNAKE_CASE_			)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_submodules(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.batchnorm.weight.device			,				torch.device('''cpu'''			)			)
					self.assertEqual(model.lineara.weight.device			,				torch.device('''cpu'''			)			)








"""simple docstring"""




import json
import os
from pathlib import Path

import pytest

from datasets.download.download_config import DownloadConfig
from datasets.download.download_manager import DownloadManager
from datasets.utils.file_utils import hash_url_to_filename


__magic_name__						: Optional[Any]    =				'http://www.mocksite.com/file1.txt'
__magic_name__						: Tuple    =				'"text": ["foo", "foo"]'
__magic_name__						: str    =				'6d8ce9aa78a471c7477201efbeabd3bb01ac2e7d100a6dc024ba1608361f90a8'




class 			__snake_case   :
							__a					 =       200
							__a					 =       {'''Content-Length''': '''100'''}
							__a					 =       {}




							def   __a   (     self:    List[str]     ,							**A_:    List[Any]   ):
										return [bytes(A_     ,							"""utf-8"""   )]





def       a_ (	*lowercase__ :List[Any],			**lowercase__ :Any     ):
			return MockResponse()





@pytest.mark.parametrize("""urls_type""",			[str, list, dict]     )
def       a_ (	lowercase__ :Optional[int],			lowercase__ :Any,			lowercase__ :Optional[int]     ):
			import requests

			monkeypatch.setattr(lowercase__,			"""request""",			lowercase__     )

			__lowerCamelCase          =      URL
			if issubclass(lowercase__,			lowercase__     ):
						__lowerCamelCase          =      url
			elif issubclass(lowercase__,			lowercase__     ):
						__lowerCamelCase          =      [url]
			elif issubclass(lowercase__,			lowercase__     ):
						__lowerCamelCase          =      {"""train""": url}
			__lowerCamelCase          =      """dummy"""
			__lowerCamelCase          =      """downloads"""
			__lowerCamelCase          =      tmp_path
			__lowerCamelCase          =      DownloadConfig(
			    cache_dir=os.path.join(lowercase__,			lowercase__     ),			use_etag=lowercase__,			)
			__lowerCamelCase          =      DownloadManager(dataset_name=lowercase__,			download_config=lowercase__     )
			__lowerCamelCase          =      dl_manager.download(lowercase__     )
			__lowerCamelCase          =      urls
			for downloaded_paths in [downloaded_paths]:
						if isinstance(lowercase__,			lowercase__     ):
									__lowerCamelCase          =      [downloaded_paths]
									__lowerCamelCase          =      [urls]
						elif isinstance(lowercase__,			lowercase__     ):
									assert "train" in downloaded_paths.keys()
									__lowerCamelCase          =      downloaded_paths.values()
									__lowerCamelCase          =      urls.values()
						assert downloaded_paths
						for downloaded_path, input_url in zip(lowercase__,			lowercase__     ):
									assert downloaded_path == dl_manager.downloaded_paths[input_url]
									__lowerCamelCase          =      Path(lowercase__     )
									__lowerCamelCase          =      downloaded_path.parts
									assert parts[-1] == HASH
									assert parts[-2] == cache_subdir
									assert downloaded_path.exists()
									__lowerCamelCase          =      downloaded_path.read_text()
									assert content == CONTENT
									__lowerCamelCase          =      downloaded_path.with_suffix(""".json"""     )
									assert metadata_downloaded_path.exists()
									__lowerCamelCase          =      json.loads(metadata_downloaded_path.read_text()     )
									assert metadata_content == {"url": URL, "etag": None}





@pytest.mark.parametrize("""paths_type""",			[str, list, dict]     )
def       a_ (	lowercase__ :Dict,			lowercase__ :Optional[Any],			lowercase__ :Dict     ):
			__lowerCamelCase          =      str(lowercase__     )
			if issubclass(lowercase__,			lowercase__     ):
						__lowerCamelCase          =      filename
			elif issubclass(lowercase__,			lowercase__     ):
						__lowerCamelCase          =      [filename]
			elif issubclass(lowercase__,			lowercase__     ):
						__lowerCamelCase          =      {"""train""": filename}
			__lowerCamelCase          =      """dummy"""
			__lowerCamelCase          =      xz_file.parent
			__lowerCamelCase          =      """extracted"""
			__lowerCamelCase          =      DownloadConfig(
			    cache_dir=lowercase__,			use_etag=lowercase__,			)
			__lowerCamelCase          =      DownloadManager(dataset_name=lowercase__,			download_config=lowercase__     )
			__lowerCamelCase          =      dl_manager.extract(lowercase__     )
			__lowerCamelCase          =      paths
			for extracted_paths in [extracted_paths]:
						if isinstance(lowercase__,			lowercase__     ):
									__lowerCamelCase          =      [extracted_paths]
									__lowerCamelCase          =      [paths]
						elif isinstance(lowercase__,			lowercase__     ):
									assert "train" in extracted_paths.keys()
									__lowerCamelCase          =      extracted_paths.values()
									__lowerCamelCase          =      paths.values()
						assert extracted_paths
						for extracted_path, input_path in zip(lowercase__,			lowercase__     ):
									assert extracted_path == dl_manager.extracted_paths[input_path]
									__lowerCamelCase          =      Path(lowercase__     )
									__lowerCamelCase          =      extracted_path.parts
									assert parts[-1] == hash_url_to_filename(lowercase__,			etag=lowercase__     )
									assert parts[-2] == extracted_subdir
									assert extracted_path.exists()
									__lowerCamelCase          =      extracted_path.read_text()
									__lowerCamelCase          =      text_file.read_text()
									assert extracted_file_content == expected_file_content





def       a_ (	lowercase__ :List[str],			lowercase__ :int     ):
			assert path.endswith(""".jsonl"""     )
			for num_items, line in enumerate(lowercase__,			start=1     ):
						__lowerCamelCase          =      json.loads(line.decode("""utf-8"""     )     )
						assert item.keys() == {"col_1", "col_2", "col_3"}
			assert num_items == 4





@pytest.mark.parametrize("""archive_jsonl""",			["""tar_jsonl_path""", """zip_jsonl_path"""]     )
def       a_ (	lowercase__ :Optional[int],			lowercase__ :Union[str, Any]     ):
			__lowerCamelCase          =      request.getfixturevalue(lowercase__     )
			__lowerCamelCase          =      DownloadManager()
			for num_jsonl, (path, file) in enumerate(dl_manager.iter_archive(lowercase__     ),			start=1     ):
						_test_jsonl(lowercase__,			lowercase__     )
			assert num_jsonl == 2





@pytest.mark.parametrize("""archive_nested_jsonl""",			["""tar_nested_jsonl_path""", """zip_nested_jsonl_path"""]     )
def       a_ (	lowercase__ :Optional[int],			lowercase__ :List[Any]     ):
			__lowerCamelCase          =      request.getfixturevalue(lowercase__     )
			__lowerCamelCase          =      DownloadManager()
			for num_tar, (path, file) in enumerate(dl_manager.iter_archive(lowercase__     ),			start=1     ):
						for num_jsonl, (subpath, subfile) in enumerate(dl_manager.iter_archive(lowercase__     ),			start=1     ):
									_test_jsonl(lowercase__,			lowercase__     )
			assert num_tar == 1
			assert num_jsonl == 2





def       a_ (	lowercase__ :Tuple     ):
			__lowerCamelCase          =      DownloadManager()
			for num_file, file in enumerate(dl_manager.iter_files(lowercase__     ),			start=1     ):
						assert os.path.basename(lowercase__     ) == ("test.txt" if num_file == 1 else "train.txt")
			assert num_file == 2




