




import dataclasses
import re
import string
from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple

import numpy as np

from . import residue_constants


UpperCamelCase__					    =    Mapping[str, np.ndarray]
UpperCamelCase__					    =    Mapping[str, Any]  # Is a nested dict.
UpperCamelCase__					    =    0.01
@dataclasses.dataclass(frozen=snake_case__   )
class     a__					:
						_a							:					np.ndarray  # [num_res, num_atom_type, 3]

						# Amino-acid type for each residue represented as an integer between 0 and
						# 20, where 20 is 'X'.
						_a							:					np.ndarray  # [num_res]

						# Binary float mask to indicate presence of a particular atom. 1.0 if an atom
						# is present and 0.0 if not. This should be used for loss masking.
						_a							:					np.ndarray  # [num_res, num_atom_type]

						# Residue index as used in PDB. It is not necessarily continuous or 0-indexed.
						_a							:					np.ndarray  # [num_res]

						# B-factors, or temperature factors, of each residue (in sq. angstroms units),
						# representing the displacement of the residue from its ground truth mean
						# value.
						_a							:					np.ndarray  # [num_res, num_atom_type]

						# Chain indices for multi-chain predictions
						_a							:					Optional[np.ndarray]      =					None

						# Optional remark about the protein. Included as a comment in output PDB
						# files
						_a							:					Optional[str]      =					None

						# Templates used to generate this protein (prediction-only)
						_a							:					Optional[Sequence[str]]      =					None

						# Chain corresponding to each parent
						_a							:					Optional[Sequence[int]]      =					None







def 			_a   (		SCREAMING_SNAKE_CASE_							:       str  ):
					__lowerCAmelCase            = R"(\[[A-Z]+\]\n)"
					__lowerCAmelCase            = [tag.strip() for tag in re.split(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  ) if len(SCREAMING_SNAKE_CASE_  ) > 0]
					__lowerCAmelCase            = zip(tags[0::2]   ,   [l.split("\n"  ) for l in tags[1::2]]  )

					__lowerCAmelCase            = ["N", "CA", "C"]
					__lowerCAmelCase            = None
					__lowerCAmelCase            = None
					__lowerCAmelCase            = None
					for g in groups:
										if "[PRIMARY]" == g[0]:
															__lowerCAmelCase            = g[1][0].strip()
															for i in range(len(SCREAMING_SNAKE_CASE_  )  ):
																				if seq[i] not in residue_constants.restypes:
																									__lowerCAmelCase            = "X"  # FIXME: strings are immutable
															__lowerCAmelCase            = np.array(
															    [residue_constants.restype_order.get(SCREAMING_SNAKE_CASE_   ,   residue_constants.restype_num  ) for res_symbol in seq]  )
										elif "[TERTIARY]" == g[0]:
															__lowerCAmelCase            = []
															for axis in range(3  ):
																				tertiary.append(list(map(SCREAMING_SNAKE_CASE_   ,   g[1][axis].split()  )  )  )
															__lowerCAmelCase            = np.array(SCREAMING_SNAKE_CASE_  )
															__lowerCAmelCase            = np.zeros((len(tertiary[0]  ) // 3, residue_constants.atom_type_num, 3)  ).astype(np.floataa  )
															for i, atom in enumerate(SCREAMING_SNAKE_CASE_  ):
																				__lowerCAmelCase            = np.transpose(tertiary_np[:, i::3]  )
															atom_positions *= PICO_TO_ANGSTROM
										elif "[MASK]" == g[0]:
															__lowerCAmelCase            = np.array(list(map({"-": 0, "+": 1}.get   ,   g[1][0].strip()  )  )  )
															__lowerCAmelCase            = np.zeros(
															    (
															        len(SCREAMING_SNAKE_CASE_  ),
															        residue_constants.atom_type_num,
															    )  ).astype(np.floataa  )
															for i, atom in enumerate(SCREAMING_SNAKE_CASE_  ):
																				__lowerCAmelCase            = 1
															atom_mask *= mask[..., None]

					assert aatype is not None

					return Protein(
					    atom_positions=SCREAMING_SNAKE_CASE_   ,   atom_mask=SCREAMING_SNAKE_CASE_   ,   aatype=SCREAMING_SNAKE_CASE_   ,   residue_index=np.arange(len(SCREAMING_SNAKE_CASE_  )  )   ,   b_factors=SCREAMING_SNAKE_CASE_   ,   )







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Protein   ,   SCREAMING_SNAKE_CASE_							:       int = 0  ):
					__lowerCAmelCase            = []

					__lowerCAmelCase            = prot.remark
					if remark is not None:
										pdb_headers.append(F"""REMARK {remark}"""  )

					__lowerCAmelCase            = prot.parents
					__lowerCAmelCase            = prot.parents_chain_index
					if parents is not None and parents_chain_index is not None:
										__lowerCAmelCase            = [p for i, p in zip(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  ) if i == chain_id]

					if parents is None or len(SCREAMING_SNAKE_CASE_  ) == 0:
										__lowerCAmelCase            = ["N/A"]

					pdb_headers.append(F"""PARENT {' '.join(SCREAMING_SNAKE_CASE_  )}"""  )

					return pdb_headers







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Protein   ,   SCREAMING_SNAKE_CASE_							:       str  ):
					__lowerCAmelCase            = []
					__lowerCAmelCase            = pdb_str.split("\n"  )

					__lowerCAmelCase            = prot.remark
					if remark is not None:
										out_pdb_lines.append(F"""REMARK {remark}"""  )

					__lowerCAmelCase            = 42
					if prot.parents is not None and len(prot.parents  ) > 0:
										__lowerCAmelCase            = []
										if prot.parents_chain_index is not None:
															__lowerCAmelCase            = {}
															for p, i in zip(prot.parents   ,   prot.parents_chain_index  ):
																				parent_dict.setdefault(str(SCREAMING_SNAKE_CASE_  )   ,   []  )
																				parent_dict[str(SCREAMING_SNAKE_CASE_  )].append(SCREAMING_SNAKE_CASE_  )

															__lowerCAmelCase            = max([int(SCREAMING_SNAKE_CASE_  ) for chain_idx in parent_dict]  )
															for i in range(max_idx + 1  ):
																				__lowerCAmelCase            = parent_dict.get(str(SCREAMING_SNAKE_CASE_  )   ,   ["N/A"]  )
																				parents_per_chain.append(SCREAMING_SNAKE_CASE_  )
										else:
															parents_per_chain.append(list(prot.parents  )  )
					else:
										__lowerCAmelCase            = [["N/A"]]

					def make_parent_line(SCREAMING_SNAKE_CASE_							:       Sequence[str]  ) -> str:
										return F"""PARENT {' '.join(SCREAMING_SNAKE_CASE_  )}"""

					out_pdb_lines.append(make_parent_line(parents_per_chain[0]  )  )

					__lowerCAmelCase            = 0
					for i, l in enumerate(SCREAMING_SNAKE_CASE_  ):
										if "PARENT" not in l and "REMARK" not in l:
															out_pdb_lines.append(SCREAMING_SNAKE_CASE_  )
										if "TER" in l and "END" not in lines[i + 1]:
															chain_counter += 1
															if not chain_counter >= len(SCREAMING_SNAKE_CASE_  ):
																				__lowerCAmelCase            = parents_per_chain[chain_counter]
															else:
																				__lowerCAmelCase            = ["N/A"]

															out_pdb_lines.append(make_parent_line(SCREAMING_SNAKE_CASE_  )  )

					return "\n".join(SCREAMING_SNAKE_CASE_  )







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Protein  ):
					__lowerCAmelCase            = residue_constants.restypes + ["X"]

					def res_atoa(SCREAMING_SNAKE_CASE_							:       int  ) -> str:
										return residue_constants.restype_atoa.get(restypes[r]   ,   "UNK"  )

					__lowerCAmelCase            = residue_constants.atom_types

					__lowerCAmelCase            = []

					__lowerCAmelCase            = prot.atom_mask
					__lowerCAmelCase            = prot.aatype
					__lowerCAmelCase            = prot.atom_positions
					__lowerCAmelCase            = prot.residue_index.astype(np.intaa  )
					__lowerCAmelCase            = prot.b_factors
					__lowerCAmelCase            = prot.chain_index

					if np.any(aatype > residue_constants.restype_num  ):
										raise ValueError("Invalid aatypes."  )

					__lowerCAmelCase            = get_pdb_headers(SCREAMING_SNAKE_CASE_  )
					if len(SCREAMING_SNAKE_CASE_  ) > 0:
										pdb_lines.extend(SCREAMING_SNAKE_CASE_  )

					__lowerCAmelCase            = aatype.shape[0]
					__lowerCAmelCase            = 1
					__lowerCAmelCase            = 0
					__lowerCAmelCase            = string.ascii_uppercase
					__lowerCAmelCase            = None
					# Add all atom sites.
					for i in range(SCREAMING_SNAKE_CASE_  ):
										__lowerCAmelCase            = res_atoa(aatype[i]  )
										for atom_name, pos, mask, b_factor in zip(SCREAMING_SNAKE_CASE_   ,   atom_positions[i]   ,   atom_mask[i]   ,   b_factors[i]  ):
															if mask < 0.5:
																				continue

															__lowerCAmelCase            = "ATOM"
															__lowerCAmelCase            = atom_name if len(SCREAMING_SNAKE_CASE_  ) == 4 else F""" {atom_name}"""
															__lowerCAmelCase            = ""
															__lowerCAmelCase            = ""
															__lowerCAmelCase            = 1.00
															__lowerCAmelCase            = atom_name[0]  # Protein supports only C, N, O, S, this works.
															__lowerCAmelCase            = ""

															__lowerCAmelCase            = "A"
															if chain_index is not None:
																				__lowerCAmelCase            = chain_tags[chain_index[i]]

															# PDB is a columnar format, every space matters here!
															__lowerCAmelCase            = (
															    F"""{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}"""
															    F"""{res_name_a:>3} {chain_tag:>1}"""
															    F"""{residue_index[i]:>4}{insertion_code:>1}   """
															    F"""{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}"""
															    F"""{occupancy:>6.2f}{b_factor:>6.2f}          """
															    F"""{element:>2}{charge:>2}"""
															)
															pdb_lines.append(SCREAMING_SNAKE_CASE_  )
															atom_index += 1

										__lowerCAmelCase            = i == n - 1
										if chain_index is not None:
															if i != n - 1 and chain_index[i + 1] != prev_chain_index:
																				__lowerCAmelCase            = True
																				__lowerCAmelCase            = chain_index[i + 1]

										if should_terminate:
															# Close the chain.
															__lowerCAmelCase            = "TER"
															__lowerCAmelCase            = (
															    F"""{chain_end:<6}{atom_index:>5}      {res_atoa(aatype[i]  ):>3} {chain_tag:>1}{residue_index[i]:>4}"""
															)
															pdb_lines.append(SCREAMING_SNAKE_CASE_  )
															atom_index += 1

															if i != n - 1:
																				# "prev" is a misnomer here. This happens at the beginning of
																				# each new chain.
																				pdb_lines.extend(get_pdb_headers(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  )  )

					pdb_lines.append("END"  )
					pdb_lines.append(""  )
					return "\n".join(SCREAMING_SNAKE_CASE_  )







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Protein  ):
					return residue_constants.STANDARD_ATOM_MASK[prot.aatype]







def 			_a   (		SCREAMING_SNAKE_CASE_							:       FeatureDict   ,   SCREAMING_SNAKE_CASE_							:       ModelOutput   ,   SCREAMING_SNAKE_CASE_							:       Optional[np.ndarray] = None   ,   SCREAMING_SNAKE_CASE_							:       Optional[np.ndarray] = None   ,   SCREAMING_SNAKE_CASE_							:       Optional[str] = None   ,   SCREAMING_SNAKE_CASE_							:       Optional[Sequence[str]] = None   ,   SCREAMING_SNAKE_CASE_							:       Optional[Sequence[int]] = None   ,   ):
					return Protein(
					    aatype=features["aatype"]   ,   atom_positions=result["final_atom_positions"]   ,   atom_mask=result["final_atom_mask"]   ,   residue_index=features["residue_index"] + 1   ,   b_factors=b_factors if b_factors is not None else np.zeros_like(result["final_atom_mask"]  )   ,   chain_index=SCREAMING_SNAKE_CASE_   ,   remark=SCREAMING_SNAKE_CASE_   ,   parents=SCREAMING_SNAKE_CASE_   ,   parents_chain_index=SCREAMING_SNAKE_CASE_   ,   )










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


UpperCamelCase__					    =    logging.get_logger(__name__)

UpperCamelCase__					    =    {
    """facebook/nllb-moe-54B""": """https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json""",
}
class     a__					(      snake_case__   ):
						_a							:					Tuple      =					"""nllb-moe"""
						_a							:					Dict      =					["""past_key_values"""]
						_a							:					Optional[int]      =					{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






						def __init__(     self	,    _A=1_2_8_1_1_2	,    _A=1_0_2_4	,    _A=1_2	,    _A=4_0_9_6	,    _A=1_6	,    _A=1_2	,    _A=4_0_9_6	,    _A=1_6	,    _A=0.05	,    _A=0.05	,    _A=True	,    _A=True	,    _A="relu"	,    _A=1_0_2_4	,    _A=0.1	,    _A=0.1	,    _A=0.0	,    _A=0.02	,    _A=2	,    _A=True	,    _A=False	,    _A="float32"	,    _A=False	,    _A=1_2_8	,    _A=6_4	,    _A=4	,    _A=4	,    _A=0.0_01	,    _A=0.0_01	,    _A="all"	,    _A=False	,    _A=False	,    _A=1.0	,    _A=0.2	,    _A=1	,    _A=0	,    _A=2	,    _A=False	,    **_A	,    ):




											"""simple docstring"""



											__lowerCAmelCase            = vocab_size
											__lowerCAmelCase            = max_position_embeddings
											__lowerCAmelCase            = d_model
											__lowerCAmelCase            = encoder_ffn_dim
											__lowerCAmelCase            = encoder_layers
											__lowerCAmelCase            = encoder_attention_heads
											__lowerCAmelCase            = decoder_ffn_dim
											__lowerCAmelCase            = decoder_layers
											__lowerCAmelCase            = decoder_attention_heads
											__lowerCAmelCase            = dropout
											__lowerCAmelCase            = attention_dropout
											__lowerCAmelCase            = activation_dropout
											__lowerCAmelCase            = activation_function
											__lowerCAmelCase            = init_std
											__lowerCAmelCase            = encoder_layerdrop
											__lowerCAmelCase            = decoder_layerdrop
											__lowerCAmelCase            = use_cache
											__lowerCAmelCase            = encoder_layers
											__lowerCAmelCase            = scale_embedding  # scale factor will be sqrt(d_model) if True
											__lowerCAmelCase            = router_z_loss_coef
											__lowerCAmelCase            = router_aux_loss_coef
											__lowerCAmelCase            = decoder_sparse_step
											__lowerCAmelCase            = encoder_sparse_step
											__lowerCAmelCase            = num_experts
											__lowerCAmelCase            = expert_capacity
											__lowerCAmelCase            = router_bias
											if router_dtype not in ["float32", "float16", "bfloat16"]:
																raise ValueError(f"""`router_dtype` must be one of 'float32', 'float16' or 'bfloat16', got {router_dtype}"""    )
											__lowerCAmelCase            = router_dtype

											__lowerCAmelCase            = router_ignore_padding_tokens
											__lowerCAmelCase            = batch_prioritized_routing
											__lowerCAmelCase            = second_expert_policy
											__lowerCAmelCase            = normalize_router_prob_before_dropping
											__lowerCAmelCase            = moe_eval_capacity_token_fraction
											__lowerCAmelCase            = moe_token_dropout
											__lowerCAmelCase            = output_router_logits
											super().__init__(
											    pad_token_id=_A	,    bos_token_id=_A	,    eos_token_id=_A	,    is_encoder_decoder=_A	,    decoder_start_token_id=_A	,    **_A	,    )




