code: 
  
  
  
  
  
  
  """simple docstring"""
  
  import argparse
  import pathlib
  from pathlib import Path
  from tempfile import TemporaryDirectory
  
  import esm as esm_module
  import torch
  from esm.esmfold.va.misc import batch_encode_sequences as esmfold_encode_sequences
  from esm.esmfold.va.pretrained import esmfold_va
  
  from transformers.models.esm.configuration_esm import EsmConfig, EsmFoldConfig
  from transformers.models.esm.modeling_esm import (
      EsmForMaskedLM,
      EsmForSequenceClassification,
      EsmIntermediate,
      EsmLayer,
      EsmOutput,
      EsmSelfAttention,
      EsmSelfOutput,
  )
  from transformers.models.esm.modeling_esmfold import EsmForProteinFolding
  from transformers.models.esm.tokenization_esm import EsmTokenizer
  from transformers.utils import logging
  
  
  logging.set_verbosity_info()
  lowercase__      =				logging.get_logger(__name__)
  
  lowercase__      =				[
      (
          """protein1""",
          """MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA""",
      ),
      ("""protein2""", """MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA"""),
      ("""protein3""", """MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG"""),
      ("""protein4""", """MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA"""),
  ]
  
  lowercase__      =				{
      """esm1b_t33_650M_UR50S""": esm_module.pretrained.esm1b_t33_650M_UR50S,
      """esm1v_t33_650M_UR90S_1""": esm_module.pretrained.esm1v_t33_650M_UR90S_1,
      """esm1v_t33_650M_UR90S_2""": esm_module.pretrained.esm1v_t33_650M_UR90S_2,
      """esm1v_t33_650M_UR90S_3""": esm_module.pretrained.esm1v_t33_650M_UR90S_3,
      """esm1v_t33_650M_UR90S_4""": esm_module.pretrained.esm1v_t33_650M_UR90S_4,
      """esm1v_t33_650M_UR90S_5""": esm_module.pretrained.esm1v_t33_650M_UR90S_5,
      """esm2_t48_15B_UR50D""": esm_module.pretrained.esm2_t48_15B_UR50D,
      """esm2_t36_3B_UR50D""": esm_module.pretrained.esm2_t36_3B_UR50D,
      """esm2_t33_650M_UR50D""": esm_module.pretrained.esm2_t33_650M_UR50D,
      """esm2_t30_150M_UR50D""": esm_module.pretrained.esm2_t30_150M_UR50D,
      """esm2_t12_35M_UR50D""": esm_module.pretrained.esm2_t12_35M_UR50D,
      """esm2_t6_8M_UR50D""": esm_module.pretrained.esm2_t6_8M_UR50D,
      """esmfold_v1""": esmfold_va,
  }
  
  lowercase__      =				list("""ARNDCQEGHILKMFPSTWYV""")
  
  lowercase__      =				restypes + ["""X"""]
  lowercase__      =				restypes_with_x + ["""<pad>""", """<mask>""", """<cls>""", """<sep>""", """<eos>"""]
  
  def 	__lowerCamelCase			(						)		->					int:
  
  		"""simple docstring"""
  
  
  
  
  
  
  
  		with TemporaryDirectory() as tempdir:
  				lowerCAmelCase_     :							Any             =    "\n".join(__UpperCamelCase )
  				lowerCAmelCase_     :							List[Any]             =    Path(__UpperCamelCase ) / "vocab.txt"
  				vocab_file.write_text(__UpperCamelCase )
  				lowerCAmelCase_     :							Optional[Any]             =    EsmTokenizer(vocab_file=str(__UpperCamelCase ) )
  		lowerCAmelCase_     :							Any             =    0  # Overlaps with 'A' but that seems to be what they want
  		return hf_tokenizer
  
  def 	__lowerCamelCase			(						__UpperCamelCase   , __UpperCamelCase )		->					List[str]:
  
  		"""simple docstring"""
  
  
  
  
  
  
  
  		lowerCAmelCase_     :							Optional[Any]             =    our_module.load_state_dict(original_module.state_dict() )
  		if status.missing_keys:
  				raise ValueError(f'''Missing keys: {status.missing_keys}''' )
  		if status.unexpected_keys:
  				raise ValueError(f'''Unexpected keys: {status.unexpected_keys}''' )
  
  def 	__lowerCamelCase			(						__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase )		->					Optional[Any]:
  
  		"""simple docstring"""
  
  
  
  
  
  
  
  		if model.startswith("esmfold" ):
  				lowerCAmelCase_     :							Any             =    MODEL_MAPPING[model]()
  		else:
  				lowerCAmelCase_       ,			lowerCAmelCase_     :							Dict             =    MODEL_MAPPING[model]()
  		esm.eval()  # disable dropout
  
  		if model.startswith("esmfold" ):
  				lowerCAmelCase_     :							Optional[Any]             =    esm.esm.embed_dim
  				lowerCAmelCase_     :							Union[str, Any]             =    esm.esm.num_layers
  				lowerCAmelCase_     :							Optional[Any]             =    esm.esm.attention_heads
  				lowerCAmelCase_     :							str             =    4 * embed_dim
  				lowerCAmelCase_     :							Tuple             =    esm.esm.token_dropout
  				lowerCAmelCase_     :							str             =    False  # This code path does not exist in ESM-2
  				lowerCAmelCase_     :							List[Any]             =    "rotary"
  				lowerCAmelCase_     :							Dict             =    True
  				lowerCAmelCase_     :							Union[str, Any]             =    EsmFoldConfig()
  				for key, val in esm.cfg.items():
  						if hasattr(__UpperCamelCase   , __UpperCamelCase ) and key != "trunk":
  								setattr(__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase )
  				for key, val in esm.cfg.trunk.items():
  						if hasattr(esmfold_config.trunk   , __UpperCamelCase ) and key != "structure_module":
  								setattr(esmfold_config.trunk   , __UpperCamelCase   , __UpperCamelCase )
  				for key, val in esm.cfg.trunk.structure_module.items():
  						if hasattr(esmfold_config.trunk.structure_module   , __UpperCamelCase ):
  								setattr(esmfold_config.trunk.structure_module   , __UpperCamelCase   , __UpperCamelCase )
  		elif hasattr(__UpperCamelCase   , "args" ):
  				# Indicates an ESM-1b or ESM-1v model
  				lowerCAmelCase_     :							List[Any]             =    esm.args.embed_dim
  				lowerCAmelCase_     :							int             =    esm.args.layers
  				lowerCAmelCase_     :							List[str]             =    esm.args.attention_heads
  				lowerCAmelCase_     :							Optional[Any]             =    esm.args.ffn_embed_dim
  				lowerCAmelCase_     :							Optional[Any]             =    esm.args.token_dropout
  				lowerCAmelCase_     :							Tuple             =    True if esm.emb_layer_norm_before else False
  				lowerCAmelCase_     :							Optional[int]             =    "absolute"
  				lowerCAmelCase_     :							List[str]             =    False
  				lowerCAmelCase_     :							Any             =    None
  		else:
  				# Indicates an ESM-2 model
  				lowerCAmelCase_     :							Any             =    esm.embed_dim
  				lowerCAmelCase_     :							Any             =    esm.num_layers
  				lowerCAmelCase_     :							Optional[int]             =    esm.attention_heads
  				lowerCAmelCase_     :							Union[str, Any]             =    4 * embed_dim  # This is hardcoded in ESM-2
  				lowerCAmelCase_     :							str             =    esm.token_dropout
  				lowerCAmelCase_     :							Tuple             =    False  # This code path does not exist in ESM-2
  				lowerCAmelCase_     :							List[str]             =    "rotary"
  				lowerCAmelCase_     :							List[Any]             =    False
  				lowerCAmelCase_     :							int             =    None
  
  		if is_folding_model:
  				lowerCAmelCase_     :							int             =    esm.esm.alphabet
  		lowerCAmelCase_     :							Tuple             =    tuple(alphabet.all_toks )
  		lowerCAmelCase_     :							str             =    alphabet.mask_idx
  		lowerCAmelCase_     :							str             =    alphabet.padding_idx
  
  		if is_folding_model:
  				lowerCAmelCase_     :							Dict             =    esm.esm
  		else:
  				lowerCAmelCase_     :							int             =    esm
  
  		lowerCAmelCase_     :							Union[str, Any]             =    EsmConfig(
  		    vocab_size=original_esm_model.embed_tokens.num_embeddings   , mask_token_id=__UpperCamelCase   , hidden_size=__UpperCamelCase   , num_hidden_layers=__UpperCamelCase   , num_attention_heads=__UpperCamelCase   , intermediate_size=__UpperCamelCase   , max_position_embeddings=1026   , layer_norm_eps=1e-5   , attention_probs_dropout_prob=0.0   , hidden_dropout_prob=0.0   , pad_token_id=__UpperCamelCase   , emb_layer_norm_before=__UpperCamelCase   , token_dropout=__UpperCamelCase   , position_embedding_type=__UpperCamelCase   , is_folding_model=__UpperCamelCase   , esmfold_config=__UpperCamelCase   , vocab_list=__UpperCamelCase   , )
  		if classification_head:
  				lowerCAmelCase_     :							int             =    esm.classification_heads["mnli"].out_proj.weight.shape[0]
  		print("Our ESM config:"   , __UpperCamelCase )
  
  		if model.startswith("esmfold" ):
  				lowerCAmelCase_     :							Optional[int]             =    EsmForProteinFolding
  		elif classification_head:
  				lowerCAmelCase_     :							Union[str, Any]             =    EsmForSequenceClassification
  		else:
  				lowerCAmelCase_     :							Tuple             =    EsmForMaskedLM
  		lowerCAmelCase_     :							int             =    model_class(__UpperCamelCase )
  		model.eval()
  
  		# Now let's copy all the weights.
  		# Embeddings
  		lowerCAmelCase_     :							Optional[Any]             =    original_esm_model.embed_tokens.weight
  		if position_embedding_type == "absolute":
  				lowerCAmelCase_     :							str             =    original_esm_model.embed_positions.weight
  
  		if config.emb_layer_norm_before:
  				lowerCAmelCase_     :							List[Any]             =    original_esm_model.emb_layer_norm_before.weight
  				lowerCAmelCase_     :							Dict             =    original_esm_model.emb_layer_norm_before.bias
  
  		lowerCAmelCase_     :							Union[str, Any]             =    original_esm_model.emb_layer_norm_after.weight
  		lowerCAmelCase_     :							Union[str, Any]             =    original_esm_model.emb_layer_norm_after.bias
  
  		for i in range(config.num_hidden_layers ):
  				# Encoder: start of layer
  				lowerCAmelCase_     :							EsmLayer             =    model.esm.encoder.layer[i]
  				# esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
  				lowerCAmelCase_     :							Union[str, Any]             =    original_esm_model.layers[i]
  
  				# self attention
  				lowerCAmelCase_     :							EsmSelfAttention             =    layer.attention.self
  				assert (
  				    esm_layer.self_attn.k_proj.weight.data.shape
  				    == esm_layer.self_attn.q_proj.weight.data.shape
  				    == esm_layer.self_attn.v_proj.weight.data.shape
  				    == torch.Size((config.hidden_size, config.hidden_size) )
  				)
  
  				lowerCAmelCase_     :							Any             =    esm_layer.self_attn.q_proj.weight
  				lowerCAmelCase_     :							Optional[Any]             =    esm_layer.self_attn.q_proj.bias
  				lowerCAmelCase_     :							Any             =    esm_layer.self_attn.k_proj.weight
  				lowerCAmelCase_     :							Optional[int]             =    esm_layer.self_attn.k_proj.bias
  				lowerCAmelCase_     :							Dict             =    esm_layer.self_attn.v_proj.weight
  				lowerCAmelCase_     :							Optional[int]             =    esm_layer.self_attn.v_proj.bias
  
  				if getattr(esm_layer.self_attn   , "rot_emb"   , __UpperCamelCase ) is not None:
  						# Matt: Although inv_freq is not a trainable weight, it is computed at model init and cached.
  						# During the training of ESM-2 the model was converted to float16 precision, which also converts
  						# the inv_freq tensor, and the loss of precision remains even if the model is loaded later as float32.
  						# If we recompute inv_freq without this loss of precision then we will get subtly different rotary
  						# embeddings, which are enough to cause significant discrepancies in model outputs. To avoid this,
  						# we make sure the new model copies the data from the old inv_freq.
  						lowerCAmelCase_     :							Optional[Any]             =    esm_layer.self_attn.rot_emb.inv_freq
  
  				# LayerNorm changes for pre-activation
  				lowerCAmelCase_     :							List[Any]             =    esm_layer.self_attn_layer_norm.weight
  				lowerCAmelCase_     :							str             =    esm_layer.self_attn_layer_norm.bias
  				lowerCAmelCase_     :							int             =    esm_layer.final_layer_norm.weight
  				lowerCAmelCase_     :							Optional[int]             =    esm_layer.final_layer_norm.bias
  
  				# self-attention output
  				lowerCAmelCase_     :							EsmSelfOutput             =    layer.attention.output
  				assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
  				lowerCAmelCase_     :							Any             =    esm_layer.self_attn.out_proj.weight
  				lowerCAmelCase_     :							Dict             =    esm_layer.self_attn.out_proj.bias
  
  				# intermediate
  				lowerCAmelCase_     :							EsmIntermediate             =    layer.intermediate
  				assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
  				lowerCAmelCase_     :							List[Any]             =    esm_layer.fca.weight
  				lowerCAmelCase_     :							str             =    esm_layer.fca.bias
  
  				# output
  				lowerCAmelCase_     :							EsmOutput             =    layer.output
  				assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
  				lowerCAmelCase_     :							List[Any]             =    esm_layer.fca.weight
  				lowerCAmelCase_     :							Any             =    esm_layer.fca.bias
  				# end of layer
  
  		if is_folding_model:
  				lowerCAmelCase_     :							Union[str, Any]             =    esm.esm_s_combine.data
  				lowerCAmelCase_     :							str             =    esm.afa_to_esm.data
  				transfer_and_check_weights(esm.embedding   , model.embedding )
  				transfer_and_check_weights(esm.esm_s_mlp   , model.esm_s_mlp )
  				transfer_and_check_weights(esm.trunk   , model.trunk )
  				transfer_and_check_weights(esm.distogram_head   , model.distogram_head )
  				transfer_and_check_weights(esm.ptm_head   , model.ptm_head )
  				transfer_and_check_weights(esm.lm_head   , model.lm_head )
  				transfer_and_check_weights(esm.lddt_head   , model.lddt_head )
  
  		elif classification_head:
  				lowerCAmelCase_     :							str             =    esm.esm.classification_heads["mnli"].dense.weight
  				lowerCAmelCase_     :							List[str]             =    esm.classification_heads["mnli"].dense.bias
  				lowerCAmelCase_     :							Dict             =    esm.classification_heads["mnli"].out_proj.weight
  				lowerCAmelCase_     :							List[str]             =    esm.classification_heads["mnli"].out_proj.bias
  		else:
  				# LM Head
  				lowerCAmelCase_     :							List[str]             =    esm.lm_head.dense.weight
  				lowerCAmelCase_     :							Optional[int]             =    esm.lm_head.dense.bias
  				lowerCAmelCase_     :							int             =    esm.lm_head.layer_norm.weight
  				lowerCAmelCase_     :							Union[str, Any]             =    esm.lm_head.layer_norm.bias
  				lowerCAmelCase_     :							str             =    esm.lm_head.weight
  				lowerCAmelCase_     :							List[Any]             =    esm.lm_head.bias
  
  		# Contact prediction head
  		transfer_and_check_weights(esm.contact_head   , model.esm.contact_head )
  
  		# Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
  		if is_folding_model:
  				# Folding models aren't trained on masked inputs and don't like mask tokens.
  				lowerCAmelCase_     :							List[Any]             =    SAMPLE_DATA[:2]
  		else:
  				lowerCAmelCase_     :							Tuple             =    SAMPLE_DATA
  
  		if is_folding_model:
  				lowerCAmelCase_     :							List[Any]             =    get_esmfold_tokenizer()
  				lowerCAmelCase_     :							int             =    hf_tokenizer(
  				    [row[1] for row in sample_data]   , return_tensors="pt"   , padding=__UpperCamelCase   , add_special_tokens=__UpperCamelCase )
  				lowerCAmelCase_       ,			lowerCAmelCase_       ,			lowerCAmelCase_       ,			lowerCAmelCase_       ,			lowerCAmelCase_     :							List[Any]             =    esmfold_encode_sequences([row[1] for row in sample_data] )
  				lowerCAmelCase_     :							Dict             =    torch.all(hf_tokens["input_ids"] == esmfold_aas ) and torch.all(
  				    hf_tokens["attention_mask"] == esmfold_mask )
  		else:
  				# Let's check that we get the same results.
  				lowerCAmelCase_     :							Union[str, Any]             =    alphabet.get_batch_converter()
  				lowerCAmelCase_       ,			lowerCAmelCase_       ,			lowerCAmelCase_     :							Any             =    batch_converter(__UpperCamelCase )
  				# Prepare tokenizer and make sure it matches
  				with TemporaryDirectory() as tempdir:
  						lowerCAmelCase_     :							str             =    "\n".join(alphabet.all_toks )
  						lowerCAmelCase_     :							Optional[int]             =    Path(__UpperCamelCase ) / "vocab.txt"
  						vocab_file.write_text(__UpperCamelCase )
  						lowerCAmelCase_     :							int             =    EsmTokenizer(vocab_file=str(__UpperCamelCase ) )
  
  				lowerCAmelCase_     :							Union[str, Any]             =    hf_tokenizer([row[1] for row in sample_data]   , return_tensors="pt"   , padding=__UpperCamelCase )
  				lowerCAmelCase_     :							List[str]             =    torch.all(hf_tokens["input_ids"] == batch_tokens )
  
  		print("Do both models tokenizers output the same tokens?"   , "🔥" if success else "💩" )
  		if not success:
  				raise Exception("Tokenization does not match!" )
  
  		with torch.no_grad():
  				if is_folding_model:
  						# Let's test the model in parts
  						# ESMFold always converts the ESM stem to float16, which requires float16 ops
  						# that don't exist on CPU. Therefore, to test it we need to run it on GPU. However,
  						# ESMFold is what we in the community call a "big boy" and so we desperately avoid putting both the
  						# original and the converted model on the GPU at the same time.
  						lowerCAmelCase_     :							Optional[int]             =    esm.cuda().infer([row[1] for row in sample_data] )
  						lowerCAmelCase_     :							Union[str, Any]             =    model.cuda()(
  						    input_ids=hf_tokens["input_ids"].cuda()   , attention_mask=hf_tokens["attention_mask"].cuda() )
  				else:
  						lowerCAmelCase_     :							List[str]             =    model(**__UpperCamelCase   , output_hidden_states=__UpperCamelCase )
  						lowerCAmelCase_     :							List[Any]             =    our_output["logits"]
  						if classification_head:
  								lowerCAmelCase_     :							Union[str, Any]             =    esm.model.classification_heads["mnli"](esm.extract_features(__UpperCamelCase ) )
  						else:
  								lowerCAmelCase_     :							Dict             =    esm(hf_tokens["input_ids"]   , repr_layers=list(range(999 ) ) )
  								lowerCAmelCase_     :							Optional[int]             =    their_output["logits"]
  
  				if is_folding_model:
  						lowerCAmelCase_     :							Union[str, Any]             =    torch.max(torch.abs(our_output["positions"] - their_output["positions"] ) ).item()
  						lowerCAmelCase_     :							Union[str, Any]             =    torch.allclose(our_output["positions"]   , their_output["positions"]   , atol=1e-5 )
  				else:
  						lowerCAmelCase_     :							Union[str, Any]             =    torch.max(torch.abs(our_output - their_output ) ).item()
  						lowerCAmelCase_     :							Union[str, Any]             =    torch.allclose(__UpperCamelCase   , __UpperCamelCase   , atol=1e-5 )
  
  				print(f'''max_absolute_diff = {max_absolute_diff}''' )  # ~ 1e-5
  				print("Do both models output the same tensors?"   , "🔥" if success else "💩" )
  
  				if not success:
  						raise Exception("Something went wRoNg" )
  
  				if not is_folding_model:
  						# Let's check contact prediction too
  						lowerCAmelCase_     :							Optional[int]             =    model.predict_contacts(hf_tokens["input_ids"]   , hf_tokens["attention_mask"] )
  						lowerCAmelCase_     :							Optional[int]             =    esm.predict_contacts(hf_tokens["input_ids"] )
  						lowerCAmelCase_     :							str             =    torch.max(torch.abs(our_output - their_output ) ).item()
  						lowerCAmelCase_     :							List[Any]             =    torch.allclose(__UpperCamelCase   , __UpperCamelCase   , atol=1e-5 )
  
  						print("Contact prediction testing:" )
  						print(f'''max_absolute_diff = {max_absolute_diff}''' )  # ~ 1e-5
  						print("Do both models output the same tensors?"   , "🔥" if success else "💩" )
  
  						if not success:
  								raise Exception("Something went wRoNg" )
  
  				pathlib.Path(__UpperCamelCase ).mkdir(parents=__UpperCamelCase   , exist_ok=__UpperCamelCase )
  				print(f'''Saving model to {pytorch_dump_folder_path}''' )
  				model.save_pretrained(__UpperCamelCase )
  
  				del esm  # Free up some memory before continuing
  
  		print(f'''Saving tokenizer to {pytorch_dump_folder_path}''' )
  		hf_tokenizer.save_pretrained(__UpperCamelCase )
  
  		if push_to_repo:
  				model.push_to_hub(repo_id=__UpperCamelCase   , use_auth_token=__UpperCamelCase )
  				hf_tokenizer.push_to_hub(repo_id=__UpperCamelCase   , use_auth_token=__UpperCamelCase )
  
  
  if __name__ == "__main__":
  					lowercase__      =				argparse.ArgumentParser()
  					# Required parameters
  					parser.add_argument(
  					    """--pytorch_dump_folder_path""", type=str, required=True, help="""Path to the output PyTorch model."""
  					)
  					parser.add_argument(
  					    """--classification_head""", action="""store_true""", help="""Whether to convert a final classification head."""
  					)
  					parser.add_argument("""--model""", default=None, type=str, required=True, help="""Name of model to convert.""")
  					parser.add_argument("""--push_to_repo""", type=str, help="""Repo to upload to (including username!).""")
  					parser.add_argument("""--auth_token""", type=str, help="""HuggingFace auth token.""")
  					lowercase__      =				parser.parse_args()
  					convert_esm_checkpoint_to_pytorch(
  					    args.model, args.pytorch_dump_folder_path, args.classification_head, args.push_to_repo, args.auth_token
  					)
  
style_context: 
  
  
  
  
  
  
  """simple docstring"""
  
  import warnings
  
  import numpy as np
  import pandas as pd
  from matplotlib import pyplot as plt
  from sklearn.metrics import pairwise_distances
  
  warnings.filterwarnings("""ignore""")
  
  lowercase__      =				"""K-MEANS-CLUST/ """
  
  def 	__lowerCamelCase			(						__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase=None )		->					Optional[Any]:
  
  		"""simple docstring"""
  
  
  
  
  
  
  
  		if seed is not None:  # useful for obtaining consistent results
  				np.random.seed(__UpperCamelCase )
  		lowerCAmelCase_     :							List[Any]             =    data.shape[0]  # number of data points
  
  		# Pick K indices from range [0, N).
  		lowerCAmelCase_     :							Optional[Any]             =    np.random.randint(0   , __UpperCamelCase   , __UpperCamelCase )
  
  		# Keep centroids as dense format, as many entries will be nonzero due to averaging.
  		# As long as at least one document in a cluster contains a word,
  		# it will carry a nonzero weight in the TF-IDF vector of the centroid.
  		lowerCAmelCase_     :							List[str]             =    data[rand_indices, :]
  
  		return centroids
  
  def 	__lowerCamelCase			(						__UpperCamelCase   , __UpperCamelCase )		->					Dict:
  
  		"""simple docstring"""
  
  
  
  
  
  
  
  		return pairwise_distances(__UpperCamelCase   , __UpperCamelCase   , metric="euclidean" )
  
  def 	__lowerCamelCase			(						__UpperCamelCase   , __UpperCamelCase )		->					List[str]:
  
  		"""simple docstring"""
  
  
  
  
  
  
  
  		lowerCAmelCase_     :							Optional[Any]             =    centroid_pairwise_dist(__UpperCamelCase   , __UpperCamelCase )
  
  		# Compute cluster assignments for each data point:
  		# Fill in the blank (RHS only)
  		lowerCAmelCase_     :							str             =    np.argmin(__UpperCamelCase   , axis=1 )
  
  		return cluster_assignment
  
  def 	__lowerCamelCase			(						__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase )		->					Dict:
  
  		"""simple docstring"""
  
  
  
  
  
  
  
  		lowerCAmelCase_     :							Tuple             =    []
  		for i in range(__UpperCamelCase ):
  				# Select all data points that belong to cluster i. Fill in the blank (RHS only)
  				lowerCAmelCase_     :							List[Any]             =    data[cluster_assignment == i]
  				# Compute the mean of the data points. Fill in the blank (RHS only)
  				lowerCAmelCase_     :							List[str]             =    member_data_points.mean(axis=0 )
  				new_centroids.append(__UpperCamelCase )
  		lowerCAmelCase_     :							Tuple             =    np.array(__UpperCamelCase )
  
  		return new_centroids
  
  def 	__lowerCamelCase			(						__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase )		->					int:
  
  		"""simple docstring"""
  
  
  
  
  
  
  
  		lowerCAmelCase_     :							Union[str, Any]             =    0.0
  		for i in range(__UpperCamelCase ):
  				# Select all data points that belong to cluster i. Fill in the blank (RHS only)
  				lowerCAmelCase_     :							List[str]             =    data[cluster_assignment == i, :]
  
  				if member_data_points.shape[0] > 0:  # check if i-th cluster is non-empty
  						# Compute distances from centroid to data points (RHS only)
  						lowerCAmelCase_     :							int             =    pairwise_distances(
  						    __UpperCamelCase   , [centroids[i]]   , metric="euclidean" )
  						lowerCAmelCase_     :							Optional[int]             =    distances**2
  						heterogeneity += np.sum(__UpperCamelCase )
  
  		return heterogeneity
  
  def 	__lowerCamelCase			(						__UpperCamelCase   , __UpperCamelCase )		->					Dict:
  
  		"""simple docstring"""
  
  
  
  
  
  
  
  		plt.figure(figsize=(7, 4) )
  		plt.plot(__UpperCamelCase   , linewidth=4 )
  		plt.xlabel("# Iterations" )
  		plt.ylabel("Heterogeneity" )
  		plt.title(f'''Heterogeneity of clustering over time, K={k:d}''' )
  		plt.rcParams.update({"font.size": 16} )
  		plt.show()
  
  def 	__lowerCamelCase			(						__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase=500   , __UpperCamelCase=None   , __UpperCamelCase=False )		->					List[Any]:
  
  		"""simple docstring"""
  
  
  
  
  
  
  
  		lowerCAmelCase_     :							Dict             =    initial_centroids[:]
  		lowerCAmelCase_     :							Optional[int]             =    None
  
  		for itr in range(__UpperCamelCase ):
  				if verbose:
  						print(__UpperCamelCase   , end="" )
  
  				# 1. Make cluster assignments using nearest centroids
  				lowerCAmelCase_     :							Optional[Any]             =    assign_clusters(__UpperCamelCase   , __UpperCamelCase )
  
  				# 2. Compute a new centroid for each of the k clusters, averaging all data
  				#    points assigned to that cluster.
  				lowerCAmelCase_     :							Tuple             =    revise_centroids(__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase )
  
  				# Check for convergence: if none of the assignments changed, stop
  				if (
  				    prev_cluster_assignment is not None
  				    and (prev_cluster_assignment == cluster_assignment).all()
  				):
  						break
  
  				# Print number of new assignments
  				if prev_cluster_assignment is not None:
  						lowerCAmelCase_     :							Union[str, Any]             =    np.sum(prev_cluster_assignment != cluster_assignment )
  						if verbose:
  								print(
  								    f'''    {num_changed:5d} elements changed their cluster assignment.''' )
  
          # Record heterogeneity convergence metric
  				if record_heterogeneity is not None:
  						# YOUR CODE HERE
  						lowerCAmelCase_     :							Tuple             =    compute_heterogeneity(__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase )
  						record_heterogeneity.append(__UpperCamelCase )
  
  				lowerCAmelCase_     :							str             =    cluster_assignment[:]
  
  		return centroids, cluster_assignment
  
  
  # Mock test below
  if False:  # change to true to run this test case.
  					from sklearn import datasets as ds
  
  					lowercase__      =				ds.load_iris()
  					lowercase__      =				3
  					lowercase__      =				[]
  					lowercase__      =				get_initial_centroids(dataset["""data"""], k, seed=0)
  					lowercase__ ,     lowercase__      =				kmeans(
  					    dataset["""data"""],
  					    k,
  					    initial_centroids,
  					    maxiter=400,
  					    record_heterogeneity=heterogeneity,
  					    verbose=True,
  					)
  					plot_heterogeneity(heterogeneity, k)
  
  def 	__lowerCamelCase			(						__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase=None )		->					pd.DataFrame:
  
  		"""simple docstring"""
  
  
  
  
  
  
  
  		if fill_missing_report:
  				lowerCAmelCase_     :							str             =    df.fillna(value=__UpperCamelCase )
  		lowerCAmelCase_     :							Optional[int]             =    1
  		lowerCAmelCase_     :							Optional[int]             =    df.select_dtypes(np.number ).columns
  		lowerCAmelCase_     :							Dict             =    (
  		    df.groupby(["Cluster"] )[  # construct report dataframe
  		        numeric_cols
  		    ]  # group by cluster number
  		    .agg(
  		        [
  		            ("sum", np.sum),
  		            ("mean_with_zeros", lambda __UpperCamelCase : np.mean(np.nan_to_num(__UpperCamelCase ) )),
  		            ("mean_without_zeros", lambda __UpperCamelCase : x.replace(0   , np.NaN ).mean()),
  		            (
  		                "mean_25-75",
  		                lambda __UpperCamelCase : np.mean(
  		                    np.nan_to_num(
  		                        sorted(__UpperCamelCase )[
  		                            round(len(__UpperCamelCase ) * 25 / 100 ) : round(len(__UpperCamelCase ) * 75 / 100 )
  		                        ] ) ),
  		            ),
  		            ("mean_with_na", np.mean),
  		            ("min", lambda __UpperCamelCase : x.min()),
  		            ("5%", lambda __UpperCamelCase : x.quantile(0.05 )),
  		            ("25%", lambda __UpperCamelCase : x.quantile(0.25 )),
  		            ("50%", lambda __UpperCamelCase : x.quantile(0.50 )),
  		            ("75%", lambda __UpperCamelCase : x.quantile(0.75 )),
  		            ("95%", lambda __UpperCamelCase : x.quantile(0.95 )),
  		            ("max", lambda __UpperCamelCase : x.max()),
  		            ("count", lambda __UpperCamelCase : x.count()),
  		            ("stdev", lambda __UpperCamelCase : x.std()),
  		            ("mode", lambda __UpperCamelCase : x.mode()[0]),
  		            ("median", lambda __UpperCamelCase : x.median()),
  		            ("# > 0", lambda __UpperCamelCase : (x > 0).sum()),
  		        ] )
  		    .T.reset_index()
  		    .rename(index=__UpperCamelCase   , columns={"level_0": "Features", "level_1": "Type"} )
  		)  # rename columns
  		# calculate the size of cluster(count of clientID's)
  		lowerCAmelCase_     :							Optional[Any]             =    report[
  		    (report["Features"] == "dummy") & (report["Type"] == "count")
  		].copy()  # avoid SettingWithCopyWarning
  		lowerCAmelCase_     :							List[Any]             =    (
  		    "ClusterSize"  # rename created cluster df to match report column names
  		)
  		lowerCAmelCase_     :							Dict             =    "# of Customers"
  		lowerCAmelCase_     :							List[str]             =    pd.DataFrame(
  		    clustersize.iloc[:, 2:].values
  		    / clustersize.iloc[:, 2:].values.sum() )
  		lowerCAmelCase_     :							Optional[Any]             =    "% of Customers"  # rename created cluster df to match report column names
  		lowerCAmelCase_     :							Dict             =    "ClusterProportion"
  		lowerCAmelCase_     :							Optional[Any]             =    clusterproportion.columns.tolist()
  		lowerCAmelCase_     :							Optional[Any]             =    cols[-2:] + cols[:-2]
  		lowerCAmelCase_     :							Dict             =    clusterproportion[cols]  # rearrange columns to match report
  		lowerCAmelCase_     :							Union[str, Any]             =    report.columns
  		lowerCAmelCase_     :							Dict             =    pd.DataFrame(
  		    abs(
  		        report[report["Type"] == "count"].iloc[:, 2:].values
  		        - clustersize.iloc[:, 2:].values ) )  # generating df with count of nan values
  		lowerCAmelCase_     :							Union[str, Any]             =    0
  		lowerCAmelCase_     :							Any             =    "# of nan"
  		lowerCAmelCase_     :							Dict             =    report[
  		    report["Type"] == "count"
  		].Features.tolist()  # filling values in order to match report
  		lowerCAmelCase_     :							Optional[Any]             =    a.columns.tolist()
  		lowerCAmelCase_     :							Optional[Any]             =    cols[-2:] + cols[:-2]
  		lowerCAmelCase_     :							str             =    a[cols]  # rearrange columns to match report
  		lowerCAmelCase_     :							Tuple             =    report.columns  # rename columns to match report
  		lowerCAmelCase_     :							List[str]             =    report.drop(
  		    report[report.Type == "count"].index )  # drop count values except cluster size
  		lowerCAmelCase_     :							Tuple             =    pd.concat(
  		    [report, a, clustersize, clusterproportion]   , axis=0 )  # concat report with clustert size and nan values
  		lowerCAmelCase_     :							List[Any]             =    report["Features"].isin(__UpperCamelCase )
  		lowerCAmelCase_     :							Tuple             =    report.columns.tolist()
  		lowerCAmelCase_     :							List[str]             =    cols[0:2] + cols[-1:] + cols[2:-1]
  		lowerCAmelCase_     :							List[str]             =    report[cols]
  		lowerCAmelCase_     :							str             =    {
  		    "ClusterSize": 9,
  		    "ClusterProportion": 8,
  		    "mean_with_zeros": 7,
  		    "mean_with_na": 6,
  		    "max": 5,
  		    "50%": 4,
  		    "min": 3,
  		    "25%": 2,
  		    "75%": 1,
  		    "# of nan": 0,
  		    "# > 0": -1,
  		    "sum_with_na": -2,
  		}
  		lowerCAmelCase_     :							List[str]             =    (
  		    report.assign(
  		        Sortera=lambda __UpperCamelCase : x.Type.map(__UpperCamelCase )   , Sortera=lambda __UpperCamelCase : list(reversed(range(len(__UpperCamelCase ) ) ) )   , )
  		    .sort_values(["Sorter1", "Mark", "Sorter2"]   , ascending=__UpperCamelCase )
  		    .drop(["Sorter1", "Sorter2"]   , axis=1 )
  		)
  		lowerCAmelCase_     :							str             =    ""
  		lowerCAmelCase_     :							Any             =    report.reset_index()
  		lowerCAmelCase_     :							Any             =    report.drop(columns=["index"] )
  		return report
  
  
  if __name__ == "__main__":
  					import doctest
  
  					doctest.testmod()
  