| | import copy |
| | import os |
| | import torch |
| | from huggingface_hub import HfApi, login |
| | from transformers import AutoModelForMaskedLM, AutoTokenizer, AutoConfig |
| |
|
| | from esm_plusplus.load_official import load_official_model |
| | from esm_plusplus.modeling_esm_plusplus import ESMplusplusForMaskedLM |
| | from weight_parity_utils import assert_state_dict_equal, assert_model_parameters_fp32 |
| |
|
| |
|
| | MODEL_DICT = { |
| | "Synthyra/ESMplusplus_small": "esmc-300", |
| | "Synthyra/ESMplusplus_large": "esmc-600", |
| | } |
| |
|
| |
|
| | def _resolve_repo_items(repo_ids: list[str] | None) -> list[tuple[str, str]]: |
| | if repo_ids is None or len(repo_ids) == 0: |
| | return list(MODEL_DICT.items()) |
| |
|
| | selected_items: list[tuple[str, str]] = [] |
| | for repo_id in repo_ids: |
| | assert repo_id in MODEL_DICT, ( |
| | f"Unknown repo_id {repo_id}. " |
| | f"Valid options: {sorted(MODEL_DICT.keys())}" |
| | ) |
| | selected_items.append((repo_id, MODEL_DICT[repo_id])) |
| | return selected_items |
| |
|
| |
|
| | if __name__ == "__main__": |
| | |
| | import argparse |
| |
|
| | parser = argparse.ArgumentParser() |
| | parser.add_argument("--hf_token", type=str, default=None) |
| | parser.add_argument("--repo_ids", nargs="*", type=str, default=None) |
| | parser.add_argument("--dry_run", action="store_true") |
| | parser.add_argument("--skip-weights", action="store_true") |
| | args = parser.parse_args() |
| | api = HfApi() |
| |
|
| | if args.hf_token is not None: |
| | assert len(args.hf_token) > 0, "--hf_token cannot be empty." |
| | login(token=args.hf_token) |
| |
|
| | script_root = os.path.dirname(os.path.abspath(__file__)) |
| |
|
| | for repo_id, esmc_model_key in _resolve_repo_items(args.repo_ids): |
| | official_model, tokenizer = load_official_model(esmc_model_key, device=torch.device("cpu"), dtype=torch.float32) |
| | |
| | config = AutoConfig.from_pretrained(repo_id, trust_remote_code=True) |
| | config.auto_map = { |
| | "AutoConfig": "modeling_esm_plusplus.ESMplusplusConfig", |
| | "AutoModel": "modeling_esm_plusplus.ESMplusplusModel", |
| | "AutoModelForMaskedLM": "modeling_esm_plusplus.ESMplusplusForMaskedLM", |
| | "AutoModelForSequenceClassification": "modeling_esm_plusplus.ESMplusplusForSequenceClassification", |
| | "AutoModelForTokenClassification": "modeling_esm_plusplus.ESMplusplusForTokenClassification", |
| | } |
| | config.tie_word_embeddings = False |
| | if args.skip_weights: |
| | if args.dry_run: |
| | print(f"[skip-weights][dry-run] validated config+tokenizer parity for {repo_id}") |
| | continue |
| | config.push_to_hub(repo_id) |
| | tokenizer.push_to_hub(repo_id) |
| | print(f"[skip-weights] uploaded config+tokenizer for {repo_id}") |
| | continue |
| | model = ESMplusplusForMaskedLM(config=config).eval().cpu().to(torch.float32) |
| | load_result = model.load_state_dict(official_model.model.state_dict(), strict=True) |
| |
|
| | |
| | model.sequence_head[0].weight = copy.deepcopy(official_model.model.sequence_head[0].weight) |
| | model.sequence_head[0].bias = copy.deepcopy(official_model.model.sequence_head[0].bias) |
| | model.sequence_head[2].weight = copy.deepcopy(official_model.model.sequence_head[2].weight) |
| | model.sequence_head[2].bias = copy.deepcopy(official_model.model.sequence_head[2].bias) |
| | model.sequence_head[3].weight = copy.deepcopy(official_model.model.sequence_head[3].weight) |
| | model.sequence_head[3].bias = copy.deepcopy(official_model.model.sequence_head[3].bias) |
| | |
| | assert_model_parameters_fp32( |
| | model=model, |
| | model_name=f"mapped ESM++ model ({esmc_model_key})", |
| | ) |
| | assert_model_parameters_fp32( |
| | model=official_model.model, |
| | model_name=f"official ESM++ model ({esmc_model_key})", |
| | ) |
| | assert_state_dict_equal( |
| | reference_state_dict=official_model.model.state_dict(), |
| | candidate_state_dict=model.state_dict(), |
| | context=f"ESMC/ESM++ weight parity ({esmc_model_key})", |
| | ) |
| |
|
| | if args.dry_run: |
| | print(f"[dry_run] validated ESM++ parity for {repo_id} <- {esmc_model_key}") |
| | continue |
| |
|
| | tokenizer.push_to_hub(repo_id) |
| | model.push_to_hub(repo_id) |
| | api.upload_file( |
| | path_or_fileobj="esm_plusplus/modeling_esm_plusplus.py", |
| | path_in_repo="modeling_esm_plusplus.py", |
| | repo_id=repo_id, |
| | repo_type="model", |
| | ) |
| | downloaded_model = AutoModelForMaskedLM.from_pretrained( |
| | repo_id, |
| | dtype=torch.float32, |
| | device_map="cpu", |
| | force_download=True, |
| | trust_remote_code=True, |
| | ) |
| | assert_state_dict_equal( |
| | reference_state_dict=official_model.model.state_dict(), |
| | candidate_state_dict=downloaded_model.state_dict(), |
| | context=f"ESMC/ESM++ weight parity post-download ({repo_id})", |
| | ) |
| |
|
