| """ |
| Data-driven HuggingFace upload script for all FastPLMs models. |
| |
| Runs weight conversion scripts for each family, then uploads |
| modeling code, embedding_mixin, entrypoint_setup, readmes, and licenses |
| to each HF repo. |
| |
| Usage: |
| py -m update_HF |
| py -m update_HF --hf_token YOUR_TOKEN |
| py -m update_HF --families esm2 dplm |
| py -m update_HF --skip-weights |
| py -m update_HF --files-only |
| """ |
|
|
| import argparse |
| import platform |
| import subprocess |
|
|
| from huggingface_hub import HfApi, login |
|
|
|
|
| MODEL_REGISTRY = [ |
| { |
| "family": "e1", |
| "repo_ids": [ |
| "Synthyra/Profluent-E1-150M", |
| "Synthyra/Profluent-E1-300M", |
| "Synthyra/Profluent-E1-600M", |
| ], |
| "files": { |
| "e1_fastplms/modeling_e1.py": "modeling_e1.py", |
| "e1_fastplms/tokenizer.json": "tokenizer.json", |
| }, |
| "readme_map": { |
| "Synthyra/Profluent-E1-150M": "readmes/e1_readme.md", |
| "Synthyra/Profluent-E1-300M": "readmes/e1_readme.md", |
| "Synthyra/Profluent-E1-600M": "readmes/e1_readme.md", |
| }, |
| "license": "LICENSE", |
| "weight_module": "e1_fastplms.get_e1_weights", |
| }, |
| { |
| "family": "esmplusplus", |
| "repo_ids": [ |
| "Synthyra/ESMplusplus_small", |
| "Synthyra/ESMplusplus_large", |
| ], |
| "files": { |
| "esm_plusplus/modeling_esm_plusplus.py": "modeling_esm_plusplus.py", |
| }, |
| "readme_map": { |
| "Synthyra/ESMplusplus_small": "readmes/esm_plusplus_small_readme.md", |
| "Synthyra/ESMplusplus_large": "readmes/esm_plusplus_large_readme.md", |
| }, |
| "license": "LICENSE", |
| "weight_module": "esm_plusplus.get_esmc_weights", |
| }, |
| { |
| "family": "esm2", |
| "repo_ids": [ |
| "Synthyra/ESM2-8M", |
| "Synthyra/ESM2-35M", |
| "Synthyra/ESM2-150M", |
| "Synthyra/ESM2-650M", |
| "Synthyra/ESM2-3B", |
| "Synthyra/FastESM2_650", |
| ], |
| "files": { |
| "esm2/modeling_fastesm.py": "modeling_fastesm.py", |
| }, |
| "readme_map": { |
| "Synthyra/ESM2-8M": "readmes/fastesm2_readme.md", |
| "Synthyra/ESM2-35M": "readmes/fastesm2_readme.md", |
| "Synthyra/ESM2-150M": "readmes/fastesm2_readme.md", |
| "Synthyra/ESM2-650M": "readmes/fastesm2_readme.md", |
| "Synthyra/ESM2-3B": "readmes/fastesm2_readme.md", |
| "Synthyra/FastESM2_650": "readmes/fastesm_650_readme.md", |
| }, |
| "license": "LICENSE", |
| "weight_module": "esm2.get_esm2_weights", |
| }, |
| { |
| "family": "dplm", |
| "repo_ids": [ |
| "Synthyra/DPLM-150M", |
| "Synthyra/DPLM-650M", |
| "Synthyra/DPLM-3B", |
| ], |
| "files": { |
| "dplm_fastplms/modeling_dplm.py": "modeling_dplm.py", |
| }, |
| "readme_map": { |
| "Synthyra/DPLM-150M": "readmes/dplm_readme.md", |
| "Synthyra/DPLM-650M": "readmes/dplm_readme.md", |
| "Synthyra/DPLM-3B": "readmes/dplm_readme.md", |
| }, |
| "license": "LICENSE", |
| "weight_module": "dplm_fastplms.get_dplm_weights", |
| }, |
| { |
| "family": "dplm2", |
| "repo_ids": [ |
| "Synthyra/DPLM2-150M", |
| "Synthyra/DPLM2-650M", |
| "Synthyra/DPLM2-3B", |
| ], |
| "files": { |
| "dplm2_fastplms/modeling_dplm2.py": "modeling_dplm2.py", |
| }, |
| "readme_map": { |
| "Synthyra/DPLM2-150M": "readmes/dplm2_readme.md", |
| "Synthyra/DPLM2-650M": "readmes/dplm2_readme.md", |
| "Synthyra/DPLM2-3B": "readmes/dplm2_readme.md", |
| }, |
| "license": "LICENSE", |
| "weight_module": "dplm2_fastplms.get_dplm2_weights", |
| }, |
| { |
| "family": "ankh", |
| "repo_ids": [ |
| "Synthyra/ANKH_base", |
| "Synthyra/ANKH_large", |
| "Synthyra/ANKH2_large", |
| ], |
| "files": {}, |
| "readme_map": {}, |
| "license": "LICENSE", |
| "weight_module": None, |
| }, |
| { |
| "family": "boltz", |
| "repo_ids": [ |
| "Synthyra/Boltz2", |
| ], |
| "files": { |
| "boltz_fastplms/modeling_boltz2.py": "modeling_boltz2.py", |
| "boltz_fastplms/__init__.py": "__init__.py", |
| "boltz_fastplms/minimal_featurizer.py": "minimal_featurizer.py", |
| "boltz_fastplms/minimal_structures.py": "minimal_structures.py", |
| "boltz_fastplms/cif_writer.py": "cif_writer.py", |
| "boltz_fastplms/vb_const.py": "vb_const.py", |
| "boltz_fastplms/vb_layers_attention.py": "vb_layers_attention.py", |
| "boltz_fastplms/vb_layers_attentionv2.py": "vb_layers_attentionv2.py", |
| "boltz_fastplms/vb_layers_confidence_utils.py": "vb_layers_confidence_utils.py", |
| "boltz_fastplms/vb_layers_dropout.py": "vb_layers_dropout.py", |
| "boltz_fastplms/vb_layers_initialize.py": "vb_layers_initialize.py", |
| "boltz_fastplms/vb_layers_outer_product_mean.py": "vb_layers_outer_product_mean.py", |
| "boltz_fastplms/vb_layers_pair_averaging.py": "vb_layers_pair_averaging.py", |
| "boltz_fastplms/vb_layers_pairformer.py": "vb_layers_pairformer.py", |
| "boltz_fastplms/vb_layers_transition.py": "vb_layers_transition.py", |
| "boltz_fastplms/vb_layers_triangular_mult.py": "vb_layers_triangular_mult.py", |
| "boltz_fastplms/vb_loss_diffusionv2.py": "vb_loss_diffusionv2.py", |
| "boltz_fastplms/vb_modules_confidencev2.py": "vb_modules_confidencev2.py", |
| "boltz_fastplms/vb_modules_diffusion_conditioning.py": "vb_modules_diffusion_conditioning.py", |
| "boltz_fastplms/vb_modules_diffusionv2.py": "vb_modules_diffusionv2.py", |
| "boltz_fastplms/vb_modules_encodersv2.py": "vb_modules_encodersv2.py", |
| "boltz_fastplms/vb_modules_transformersv2.py": "vb_modules_transformersv2.py", |
| "boltz_fastplms/vb_modules_trunkv2.py": "vb_modules_trunkv2.py", |
| "boltz_fastplms/vb_modules_utils.py": "vb_modules_utils.py", |
| "boltz_fastplms/vb_potentials_potentials.py": "vb_potentials_potentials.py", |
| "boltz_fastplms/vb_potentials_schedules.py": "vb_potentials_schedules.py", |
| "boltz_fastplms/vb_tri_attn_attention.py": "vb_tri_attn_attention.py", |
| "boltz_fastplms/vb_tri_attn_primitives.py": "vb_tri_attn_primitives.py", |
| "boltz_fastplms/vb_tri_attn_utils.py": "vb_tri_attn_utils.py", |
| }, |
| "readme_map": { |
| "Synthyra/Boltz2": "readmes/boltz2_readme.md", |
| }, |
| "license": "LICENSE", |
| "weight_module": "boltz_fastplms.get_boltz2_weights", |
| }, |
| ] |
|
|
| SHARED_FILES = { |
| "embedding_mixin.py": "embedding_mixin.py", |
| "entrypoint_setup.py": "entrypoint_setup.py", |
| } |
|
|
|
|
| def _run_weight_scripts( |
| families: list[str] | None, hf_token: str | None, skip_weights: bool |
| ) -> None: |
| python_cmd = "python" if platform.system().lower() == "linux" else "py" |
| for entry in MODEL_REGISTRY: |
| if families is not None and entry["family"] not in families: |
| continue |
| module = entry["weight_module"] |
| if module is None: |
| continue |
| command = [python_cmd, "-m", module] |
| if hf_token is not None: |
| command.extend(["--hf_token", hf_token]) |
| if skip_weights: |
| command.append("--skip-weights") |
| print(f"Running: {' '.join(command)}") |
| subprocess.run(command, check=True) |
|
|
|
|
| def _upload_files(api: HfApi, families: list[str] | None) -> None: |
| for entry in MODEL_REGISTRY: |
| if families is not None and entry["family"] not in families: |
| continue |
|
|
| for repo_id in entry["repo_ids"]: |
| print(f"\nUploading to {repo_id}") |
|
|
| for local_path, repo_path in entry["files"].items(): |
| api.upload_file( |
| path_or_fileobj=local_path, |
| path_in_repo=repo_path, |
| repo_id=repo_id, |
| repo_type="model", |
| ) |
|
|
| for local_path, repo_path in SHARED_FILES.items(): |
| api.upload_file( |
| path_or_fileobj=local_path, |
| path_in_repo=repo_path, |
| repo_id=repo_id, |
| repo_type="model", |
| ) |
|
|
| if entry["license"]: |
| api.upload_file( |
| path_or_fileobj=entry["license"], |
| path_in_repo="LICENSE", |
| repo_id=repo_id, |
| repo_type="model", |
| ) |
|
|
| readme_path = entry["readme_map"].get(repo_id) |
| if readme_path: |
| api.upload_file( |
| path_or_fileobj=readme_path, |
| path_in_repo="README.md", |
| repo_id=repo_id, |
| repo_type="model", |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="Upload FastPLMs models to HuggingFace") |
| parser.add_argument("--hf_token", type=str, default=None) |
| parser.add_argument("--families", nargs="+", default=None) |
| parser.add_argument( |
| "--skip-weights", |
| action="store_true", |
| help="Run weight scripts without downloading/pushing model weights", |
| ) |
| parser.add_argument("--files-only", action="store_true", help="Only upload files, skip weight conversion") |
| args = parser.parse_args() |
|
|
| if args.hf_token: |
| login(token=args.hf_token) |
|
|
| if not args.files_only: |
| _run_weight_scripts(args.families, args.hf_token, args.skip_weights) |
|
|
| api = HfApi() |
| _upload_files(api, args.families) |
| print("\nDone.") |
|
|
