| | from dataclasses import dataclass |
| |
|
| | from .supported_models import currently_supported_models, standard_models, experimental_models |
| |
|
| |
|
| | @dataclass |
| | class BaseModelArguments: |
| | def __init__(self, model_names: list[str] = None, model_paths: list[str] = None, model_types: list[str] = None, model_dtype=None, **kwargs): |
| | if model_paths is not None: |
| | assert model_types is not None, "model_types is required when model_paths is provided." |
| | assert len(model_paths) == len(model_types), f"model_paths ({len(model_paths)}) and model_types ({len(model_types)}) must have the same length." |
| | self.model_names = [p.split('/')[-1] for p in model_paths] |
| | self._model_types = list(model_types) |
| | self._model_paths = list(model_paths) |
| | else: |
| | assert model_names is not None, "Either model_names or model_paths/model_types must be provided." |
| | if model_names[0] == 'standard': |
| | self.model_names = standard_models |
| | elif 'exp' in model_names[0].lower(): |
| | self.model_names = experimental_models |
| | else: |
| | self.model_names = model_names |
| | self._model_types = None |
| | self._model_paths = None |
| | self.model_dtype = model_dtype |
| |
|
| | def model_entries(self): |
| | """Yields (display_name, dispatch_type, model_path) tuples for each model. |
| | |
| | In preset mode: dispatch_type is the preset name, model_path is None. |
| | In path mode: dispatch_type is the model type keyword, model_path is the explicit path. |
| | """ |
| | if self._model_paths is not None: |
| | for name, mtype, mpath in zip(self.model_names, self._model_types, self._model_paths): |
| | yield name, mtype, mpath |
| | else: |
| | for name in self.model_names: |
| | yield name, name, None |
| |
|
| |
|
| | def get_base_model(model_name: str, masked_lm: bool = False, dtype=None, model_path: str = None): |
| | if 'random' in model_name.lower(): |
| | from .random import build_random_model |
| | return build_random_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'esm2' in model_name.lower() and model_name.lower().count('esm2') == 1: |
| | from .esm2 import build_esm2_model |
| | return build_esm2_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'dsm' in model_name.lower(): |
| | from .esm2 import build_esm2_model |
| | return build_esm2_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'esmc' in model_name.lower(): |
| | from .esmc import build_esmc_model |
| | return build_esmc_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'protbert' in model_name.lower(): |
| | from .protbert import build_protbert_model |
| | return build_protbert_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'prott5' in model_name.lower(): |
| | from .prott5 import build_prott5_model |
| | return build_prott5_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'ankh' in model_name.lower(): |
| | from .ankh import build_ankh_model |
| | return build_ankh_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'glm' in model_name.lower(): |
| | from .glm import build_glm2_model |
| | return build_glm2_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'dplm2' in model_name.lower(): |
| | from .dplm2 import build_dplm2_model |
| | return build_dplm2_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'dplm' in model_name.lower(): |
| | from .dplm import build_dplm_model |
| | return build_dplm_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'protclm' in model_name.lower(): |
| | from .protCLM import build_protCLM |
| | return build_protCLM(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'onehot' in model_name.lower(): |
| | from .one_hot import build_one_hot_model |
| | return build_one_hot_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'amplify' in model_name.lower(): |
| | from .amplify import build_amplify_model |
| | return build_amplify_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'e1' in model_name.lower(): |
| | from .e1 import build_e1_model |
| | return build_e1_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'vec2vec' in model_name.lower(): |
| | from .vec2vec import build_vec2vec_model |
| | return build_vec2vec_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'calm' in model_name.lower(): |
| | from .calm import build_calm_model |
| | return build_calm_model(model_name, masked_lm=masked_lm, dtype=dtype, model_path=model_path) |
| | elif 'custom' in model_name.lower(): |
| | from .custom_model import build_custom_model |
| | assert model_path is not None, "model_path is required for custom models. Use --model_paths and --model_types custom." |
| | return build_custom_model(model_path, masked_lm=masked_lm, dtype=dtype) |
| | else: |
| | raise ValueError(f"Model {model_name} not supported") |
| |
|
| |
|
| | def get_base_model_for_training(model_name: str, tokenwise: bool = False, num_labels: int = None, hybrid: bool = False, dtype=None, model_path: str = None): |
| | if 'esm2' in model_name.lower() or 'dsm' in model_name.lower(): |
| | from .esm2 import get_esm2_for_training |
| | return get_esm2_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | elif 'esmc' in model_name.lower(): |
| | from .esmc import get_esmc_for_training |
| | return get_esmc_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | elif 'protbert' in model_name.lower(): |
| | from .protbert import get_protbert_for_training |
| | return get_protbert_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | elif 'prott5' in model_name.lower(): |
| | from .prott5 import get_prott5_for_training |
| | return get_prott5_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | elif 'ankh' in model_name.lower(): |
| | from .ankh import get_ankh_for_training |
| | return get_ankh_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | elif 'glm' in model_name.lower(): |
| | from .glm import get_glm2_for_training |
| | return get_glm2_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | elif 'dplm2' in model_name.lower(): |
| | from .dplm2 import get_dplm2_for_training |
| | return get_dplm2_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | elif 'dplm' in model_name.lower(): |
| | from .dplm import get_dplm_for_training |
| | return get_dplm_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | elif 'e1' in model_name.lower(): |
| | from .e1 import get_e1_for_training |
| | return get_e1_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | elif 'protclm' in model_name.lower(): |
| | from .protCLM import get_protCLM_for_training |
| | return get_protCLM_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | elif 'amplify' in model_name.lower(): |
| | from .amplify import get_amplify_for_training |
| | return get_amplify_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | elif 'calm' in model_name.lower(): |
| | from .calm import get_calm_for_training |
| | return get_calm_for_training(model_name, tokenwise, num_labels, hybrid, dtype=dtype, model_path=model_path) |
| | else: |
| | raise ValueError(f"Model {model_name} not supported") |
| |
|
| |
|
| | def get_tokenizer(model_name: str, model_path: str = None): |
| | if 'custom' in model_name.lower(): |
| | from .custom_model import build_custom_tokenizer |
| | assert model_path is not None, "model_path is required for custom models. Use --model_paths and --model_types custom." |
| | return build_custom_tokenizer(model_path) |
| | if 'esm2' in model_name.lower() or 'random' in model_name.lower() or 'dsm' in model_name.lower(): |
| | from .esm2 import get_esm2_tokenizer |
| | return get_esm2_tokenizer(model_name, model_path=model_path) |
| | elif 'esmc' in model_name.lower(): |
| | from .esmc import get_esmc_tokenizer |
| | return get_esmc_tokenizer(model_name, model_path=model_path) |
| | elif 'protbert' in model_name.lower(): |
| | from .protbert import get_protbert_tokenizer |
| | return get_protbert_tokenizer(model_name, model_path=model_path) |
| | elif 'prott5' in model_name.lower(): |
| | from .prott5 import get_prott5_tokenizer |
| | return get_prott5_tokenizer(model_name, model_path=model_path) |
| | elif 'ankh' in model_name.lower(): |
| | from .ankh import get_ankh_tokenizer |
| | return get_ankh_tokenizer(model_name, model_path=model_path) |
| | elif 'glm' in model_name.lower(): |
| | from .glm import get_glm2_tokenizer |
| | return get_glm2_tokenizer(model_name, model_path=model_path) |
| | elif 'dplm2' in model_name.lower(): |
| | from .dplm2 import get_dplm2_tokenizer |
| | return get_dplm2_tokenizer(model_name, model_path=model_path) |
| | elif 'dplm' in model_name.lower(): |
| | from .dplm import get_dplm_tokenizer |
| | return get_dplm_tokenizer(model_name, model_path=model_path) |
| | elif 'e1' in model_name.lower(): |
| | from .e1 import get_e1_tokenizer |
| | return get_e1_tokenizer(model_name, model_path=model_path) |
| | elif 'protclm' in model_name.lower(): |
| | from .protCLM import get_protCLM_tokenizer |
| | return get_protCLM_tokenizer(model_name, model_path=model_path) |
| | elif 'onehot' in model_name.lower(): |
| | from .one_hot import get_one_hot_tokenizer |
| | return get_one_hot_tokenizer(model_name, model_path=model_path) |
| | elif 'amplify' in model_name.lower(): |
| | from .amplify import get_amplify_tokenizer |
| | return get_amplify_tokenizer(model_name, model_path=model_path) |
| | elif 'calm' in model_name.lower(): |
| | from .calm import get_calm_tokenizer |
| | return get_calm_tokenizer(model_name, model_path=model_path) |
| | else: |
| | raise ValueError(f"Model {model_name} not supported") |
| |
|
| |
|
| | if __name__ == '__main__': |
| | |
| | import sys |
| | import argparse |
| | |
| | parser = argparse.ArgumentParser(description='Download and list supported models') |
| | parser.add_argument('--download', action='store_true', help='Download all standard models') |
| | parser.add_argument('--list', action='store_true', help='List all supported models with descriptions') |
| | args = parser.parse_args() |
| | |
| | if len(sys.argv) == 1: |
| | parser.print_help() |
| | sys.exit(1) |
| | |
| | if args.list: |
| | try: |
| | from resource_info import model_descriptions |
| | print("\n=== Currently Supported Models ===\n") |
| | |
| | max_name_len = max(len(name) for name in currently_supported_models) |
| | max_type_len = max(len(model_descriptions.get(name, {}).get('type', 'Unknown')) for name in currently_supported_models if name in model_descriptions) |
| | max_size_len = max(len(model_descriptions.get(name, {}).get('size', 'Unknown')) for name in currently_supported_models if name in model_descriptions) |
| | |
| | |
| | print(f"{'Model':<{max_name_len+2}}{'Type':<{max_type_len+2}}{'Size':<{max_size_len+2}}Description") |
| | print("-" * (max_name_len + max_type_len + max_size_len + 50)) |
| | |
| | for model_name in currently_supported_models: |
| | if model_name in model_descriptions: |
| | model_info = model_descriptions[model_name] |
| | print(f"{model_name:<{max_name_len+2}}{model_info.get('type', 'Unknown'):<{max_type_len+2}}{model_info.get('size', 'Unknown'):<{max_size_len+2}}{model_info.get('description', 'No description available')}") |
| | else: |
| | print(f"{model_name:<{max_name_len+2}}{'Unknown':<{max_type_len+2}}{'Unknown':<{max_size_len+2}}No description available") |
| | |
| | print("\n=== Standard Models ===\n") |
| | for model_name in standard_models: |
| | print(f"- {model_name}") |
| | |
| | except ImportError: |
| | print("Model descriptions file not found. Only listing model names.") |
| | print("\n=== Currently Supported Models ===\n") |
| | for model_name in currently_supported_models: |
| | print(f"- {model_name}") |
| | |
| | print("\n=== Standard Models ===\n") |
| | for model_name in standard_models: |
| | print(f"- {model_name}") |
| | |
| | if args.download: |
| | |
| | from torchinfo import summary |
| | from ..utils import clear_screen |
| | download_args = BaseModelArguments(model_names=['standard']) |
| | for model_name in download_args.model_names: |
| | model, tokenizer = get_base_model(model_name) |
| | print(f'Downloaded {model_name}') |
| | tokenized = tokenizer('MEKVQYLTRSAIRRASTIEMPQQARQKLQNLFINFCLILICLLLICIIVMLL', return_tensors='pt').input_ids |
| | summary(model, input_data=tokenized) |
| | clear_screen() |
| |
|
