




"""simple docstring"""




from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase				:     str					=				logging.get_logger(__name__)

lowerCAmelCase				:     Tuple					=				{
    """google/realm-cc-news-pretrained-embedder""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-encoder""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-scorer""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-openqa""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/config.json"""
    ),
    """google/realm-orqa-nq-openqa""": """https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/config.json""",
    """google/realm-orqa-nq-reader""": """https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/config.json""",
    """google/realm-orqa-wq-openqa""": """https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/config.json""",
    """google/realm-orqa-wq-reader""": """https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/config.json""",
    # See all REALM models at https://huggingface.co/models?filter=realm
}






class        __magic_name__				(			UpperCAmelCase__							):






				'''simple docstring'''
				__UpperCamelCase         =		"realm"





				def __init__(			self						,					_a=30_522						,					_a=768						,					_a=128						,					_a=12						,					_a=12						,					_a=8						,					_a=3_072						,					_a="gelu_new"						,					_a=0.1						,					_a=0.1						,					_a=512						,					_a=2						,					_a=0.02						,					_a=1e-1_2						,					_a=256						,					_a=10						,					_a=1e-3						,					_a=5						,					_a=320						,					_a=13_353_718						,					_a=5_000						,					_a=1						,					_a=0						,					_a=2						,					**_a						,					):





							"""simple docstring"""
							super().__init__(pad_token_id=_a						,					bos_token_id=_a						,					eos_token_id=_a						,					**_a					)

							# Common config
							lowerCamelCase      =   vocab_size
							lowerCamelCase      =   max_position_embeddings
							lowerCamelCase      =   hidden_size
							lowerCamelCase      =   retriever_proj_size
							lowerCamelCase      =   num_hidden_layers
							lowerCamelCase      =   num_attention_heads
							lowerCamelCase      =   num_candidates
							lowerCamelCase      =   intermediate_size
							lowerCamelCase      =   hidden_act
							lowerCamelCase      =   hidden_dropout_prob
							lowerCamelCase      =   attention_probs_dropout_prob
							lowerCamelCase      =   initializer_range
							lowerCamelCase      =   type_vocab_size
							lowerCamelCase      =   layer_norm_eps

							# Reader config
							lowerCamelCase      =   span_hidden_size
							lowerCamelCase      =   max_span_width
							lowerCamelCase      =   reader_layer_norm_eps
							lowerCamelCase      =   reader_beam_size
							lowerCamelCase      =   reader_seq_len

							# Retrieval config
							lowerCamelCase      =   num_block_records
							lowerCamelCase      =   searcher_beam_size









"""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__							)							)



