




"""simple docstring"""




# Lint as: python3
import os
import re
import urllib.parse
from pathlib import Path
from typing import Callable, List, Optional, Union
from zipfile import ZipFile

from ..utils.file_utils import cached_path, hf_github_url
from ..utils.logging import get_logger
from ..utils.version import Version


lowerCAmelCase				:     str					=				get_logger(__name__)






class        __magic_name__				:






				'''simple docstring'''
				__UpperCamelCase         =		"dummy_data"
				__UpperCamelCase         =		"datasets"
				__UpperCamelCase         =		False





				def __init__(			self						,					_a						,					_a						,					_a						,					_a = None						,					_a = False						,					_a = True						,					_a = None						,					):





							"""simple docstring"""
							lowerCamelCase      =   0
							lowerCamelCase      =   dataset_name
							lowerCamelCase      =   cache_dir
							lowerCamelCase      =   use_local_dummy_data
							lowerCamelCase      =   config
							# download_callbacks take a single url as input
							lowerCamelCase      =   download_callbacks or []
							# if False, it doesn't load existing files and it returns the paths of the dummy files relative
							# to the dummy_data zip file root
							lowerCamelCase      =   load_existing_dummy_data

							# TODO(PVP, QL) might need to make this more general
							lowerCamelCase      =   str(_a					)
							# to be downloaded
							lowerCamelCase      =   None
							lowerCamelCase      =   None





				@property
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							if self._dummy_file is None:
										lowerCamelCase      =   self.download_dummy_data()
							return self._dummy_file





				@property
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							if self.config is not None:
										# structure is dummy / config_name / version_name
										return os.path.join("""dummy"""						,					self.config.name						,					self.version_name					)
							# structure is dummy / version_name
							return os.path.join("""dummy"""						,					self.version_name					)





				@property
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							return os.path.join(self.dummy_data_folder						,					"""dummy_data.zip"""					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   (
							    self.local_path_to_dummy_data if self.use_local_dummy_data is True else self.github_path_to_dummy_data
							)

							lowerCamelCase      =   cached_path(
							    _a						,					cache_dir=self.cache_dir						,					extract_compressed_file=_a						,					force_extract=_a					)

							return os.path.join(_a						,					self.dummy_file_name					)





				@property
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							return os.path.join(self.datasets_scripts_dir						,					self.dataset_name						,					self.dummy_zip_file					)





				@property
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							if self._bucket_url is None:
										lowerCamelCase      =   hf_github_url(self.dataset_name						,					self.dummy_zip_file.replace(os.sep						,					"""/"""					)					)
							return self._bucket_url





				@property
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							# return full path if its a dir
							if os.path.isdir(self.dummy_file					):
										return self.dummy_file
							# else cut off path to file -> example `xsum`.
							return "/".join(self.dummy_file.replace(os.sep						,					"""/"""					).split("""/"""					)[:-1]					)





				def   _lowerCAmelCase							(			self						,					_a						,					*_a					):





							"""simple docstring"""
							if self.load_existing_dummy_data:
										# dummy data is downloaded and tested
										lowerCamelCase      =   self.dummy_file
							else:
										# dummy data cannot be downloaded and only the path to dummy file is returned
										lowerCamelCase      =   self.dummy_file_name

							# special case when data_url is a dict
							if isinstance(_a						,					_a					):
										return self.create_dummy_data_dict(_a						,					_a					)
							elif isinstance(_a						,					(list, tuple)					):
										return self.create_dummy_data_list(_a						,					_a					)
							else:
										return self.create_dummy_data_single(_a						,					_a					)





				def   _lowerCAmelCase							(			self						,					_a						,					*_a					):





							"""simple docstring"""
							return self.download_and_extract(_a					)





				def   _lowerCAmelCase							(			self						,					_a						,					_a					):





							"""simple docstring"""
							return self.download_and_extract(_a					)





				def   _lowerCAmelCase							(			self						,					_a						,					*_a						,					**_a					):





							"""simple docstring"""
							return path





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							return {}





				def   _lowerCAmelCase							(			self						,					_a						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   {}
							for key, single_urls in data_url.items():
										for download_callback in self.download_callbacks:
													if isinstance(_a						,					_a					):
																for single_url in single_urls:
																			download_callback(_a					)
													else:
																lowerCamelCase      =   single_urls
																download_callback(_a					)
            # we force the name of each key to be the last file / folder name of the url path
            # if the url has arguments, we need to encode them with urllib.parse.quote_plus
										if isinstance(_a						,					_a					):
													lowerCamelCase      =   [os.path.join(_a						,					urllib.parse.quote_plus(Path(_a					).name					)					) for x in single_urls]
										else:
													lowerCamelCase      =   single_urls
													lowerCamelCase      =   os.path.join(_a						,					urllib.parse.quote_plus(Path(_a					).name					)					)
										lowerCamelCase      =   value

							# make sure that values are unique
							if all(isinstance(_a						,					_a					) for i in dummy_data_dict.values()					) and len(set(dummy_data_dict.values()					)					) < len(
							    dummy_data_dict.values()					):
										# append key to value to make its name unique
										lowerCamelCase      =   {key: value + key for key, value in dummy_data_dict.items()}

							return dummy_data_dict





				def   _lowerCAmelCase							(			self						,					_a						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   []
							# trick: if there are many shards named like `data.txt-000001-of-00300`, only use the first one
							lowerCamelCase      =   all(bool(re.findall("""[0-9]{3,}-of-[0-9]{3,}"""						,					_a					)					) for url in data_url					)
							lowerCamelCase      =   all(
							    url.startswith("""https://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed"""					) for url in data_url					)
							if data_url and (is_tf_records or is_pubmed_records):
										lowerCamelCase      =   [data_url[0]] * len(_a					)
							for single_url in data_url:
										for download_callback in self.download_callbacks:
													download_callback(_a					)
										# we force the name of each key to be the last file / folder name of the url path
										# if the url has arguments, we need to encode them with urllib.parse.quote_plus
										lowerCamelCase      =   os.path.join(_a						,					urllib.parse.quote_plus(single_url.split("""/"""					)[-1]					)					)
										dummy_data_list.append(_a					)
							return dummy_data_list





				def   _lowerCAmelCase							(			self						,					_a						,					_a					):





							"""simple docstring"""
							for download_callback in self.download_callbacks:
										download_callback(_a					)
							# we force the name of each key to be the last file / folder name of the url path
							# if the url has arguments, we need to encode them with urllib.parse.quote_plus
							lowerCamelCase      =   os.path.join(_a						,					urllib.parse.quote_plus(data_url.split("""/"""					)[-1]					)					)
							if os.path.exists(_a					) or not self.load_existing_dummy_data:
										return value
							else:
										# Backward compatibility, maybe deprecate at one point.
										# For many datasets with single url calls to dl_manager.download_and_extract,
										# the dummy_data.zip file is actually the zipped downloaded file
										# while now we expected the dummy_data.zip file to be a directory containing
										# the downloaded file.
										return path_to_dummy_data





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							pass





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							pass





				def   _lowerCAmelCase							(			self						,					_a					):





							"""simple docstring"""
							def _iter_archive_members(_a					):
										# this preserves the order of the members inside the ZIP archive
										lowerCamelCase      =   Path(self.dummy_file					).parent
										lowerCamelCase      =   path.relative_to(_a					)
										with ZipFile(self.local_path_to_dummy_data					) as zip_file:
													lowerCamelCase      =   zip_file.namelist()
										for member in members:
													if member.startswith(relative_path.as_posix()					):
																yield dummy_parent_path.joinpath(_a					)

							lowerCamelCase      =   Path(_a					)
							lowerCamelCase      =   _iter_archive_members(_a					) if self.use_local_dummy_data else path.rglob("""*"""					)
							for file_path in file_paths:
										if file_path.is_file() and not file_path.name.startswith((""".""", """__""")					):
													yield file_path.relative_to(_a					).as_posix(), file_path.open("""rb"""					)





				def   _lowerCAmelCase							(			self						,					_a					):





							"""simple docstring"""
							if not isinstance(_a						,					_a					):
										lowerCamelCase      =   [paths]
							for path in paths:
										if os.path.isfile(_a					):
													if os.path.basename(_a					).startswith((""".""", """__""")					):
																return
													yield path
										else:
													for dirpath, dirnames, filenames in os.walk(_a					):
																if os.path.basename(_a					).startswith((""".""", """__""")					):
																			continue
																dirnames.sort()
																for filename in sorted(_a					):
																			if filename.startswith((""".""", """__""")					):
																						continue
																			yield os.path.join(_a						,					_a					)









"""simple docstring"""




import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
lowerCAmelCase				:     List[str]					=				object()

# For specifying empty leaf dict `{}`
lowerCAmelCase				:     Any					=				object()



def        a__					(							snake_case__			,					snake_case__							)		->   List[Any]:
			lowerCamelCase      =   tuple((re.compile(x + """$"""							) for x in qs)							)
			for i in range(len(snake_case__							) - len(snake_case__							) + 1							):
						lowerCamelCase      =   [x.match(snake_case__							) for x, y in zip(snake_case__			,					ks[i:]							)]
						if matches and all(snake_case__							):
									return True
			return False



def        a__					(							snake_case__							)		->   str:
			def replace(snake_case__			,					snake_case__							):
						for rule, replacement in rules:
									if _match(snake_case__			,					snake_case__							):
												return replacement
						return val

			return replace



def        a__					(							)		->   Union[str, Any]:
			return [
			    # embeddings
			    (("transformer", "wpe", "embedding"), P("""mp"""			,					snake_case__							)),
			    (("transformer", "wte", "embedding"), P("""mp"""			,					snake_case__							)),
			    # atention
			    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(snake_case__			,					"""mp"""							)),
			    (("attention", "out_proj", "kernel"), P("""mp"""			,					snake_case__							)),
			    (("attention", "out_proj", "bias"), None),
			    # mlp
			    (("mlp", "c_fc", "kernel"), P(snake_case__			,					"""mp"""							)),
			    (("mlp", "c_fc", "bias"), P("""mp"""							)),
			    (("mlp", "c_proj", "kernel"), P("""mp"""			,					snake_case__							)),
			    (("mlp", "c_proj", "bias"), None),
			    # layer norms
			    ((r"ln_\d+", "bias"), None),
			    ((r"\d+", r"ln_\d+", "scale"), None),
			    (("ln_f", "bias"), None),
			    (("ln_f", "scale"), None),
			]



def        a__					(							snake_case__							)		->   Optional[Any]:
			lowerCamelCase      =   _get_partition_rules()
			lowerCamelCase      =   _replacement_rules(snake_case__							)
			lowerCamelCase      =   {k: _unmatched for k in flatten_dict(snake_case__							)}
			lowerCamelCase      =   {k: replace(snake_case__			,					snake_case__							) for k, v in initd.items()}
			assert _unmatched not in result.values(), "Incomplete partition spec."
			return freeze(unflatten_dict(snake_case__							)							)



