MODELS = {
"pretrained": {
"<=4B": [
"microsoft/phi-1",
"microsoft/phi-1_5",
"microsoft/phi-2",
"Qwen/Qwen1.5-0.5B",
"Qwen/Qwen1.5-1.8B",
"Qwen/Qwen1.5-4B",
"google/gemma-2b",
"allenai/OLMo-1B",
],
"<=7B": [
"google/gemma-7b",
"mistralai/Mistral-7B-v0.1",
"Qwen/Qwen1.5-7B",
"01-ai/Yi-6B",
"meta-llama/Llama-2-7b-hf",
"codellama/CodeLlama-7b-hf",
"EleutherAI/llemma_7b",
"allenai/OLMo-7B",
"mistral-community/Mistral-7B-v0.2",
],
"<=14B": [
"Qwen/Qwen1.5-14B",
"meta-llama/Llama-2-13b-hf",
"codellama/CodeLlama-13b-hf",
"upstage/SOLAR-10.7B-v1.0",
"meta-llama/Meta-Llama-3-8B",
],
"<=50B": [
"01-ai/Yi-34B",
"EleutherAI/llemma_34b",
"codellama/CodeLlama-34b-hf",
"mistralai/Mixtral-8x7B-v0.1",
"Qwen/Qwen1.5-32B",
],
"<=75B": [
"meta-llama/Llama-2-70b-hf",
"codellama/CodeLlama-70b-hf",
"meta-llama/Meta-Llama-3-70B",
"Qwen/Qwen1.5-72B",
],
"<=175B": [
"mistral-community/Mixtral-8x22B-v0.1-AWQ",
],
},
"instruction_tuned": {
"<=4B": [
"Qwen/Qwen1.5-0.5B-Chat",
"Qwen/Qwen1.5-1.8B-Chat",
"Qwen/Qwen1.5-4B-Chat",
"google/gemma-2b-it",
"google/gemma-1.1-2b-it",
"microsoft/Phi-3-mini-4k-instruct",
"microsoft/Phi-3-mini-128k-instruct",
],
"<=7B": [
"google/gemma-7b-it",
"mistralai/Mistral-7B-Instruct-v0.2",
"Qwen/Qwen1.5-7B-Chat",
"01-ai/Yi-6B-Chat",
"meta-llama/Llama-2-7b-chat-hf",
"codellama/CodeLlama-7b-Instruct-hf",
"allenai/OLMo-7B-SFT",
"allenai/OLMo-7B-Instruct",
"allenai/tulu-2-7b",
"allenai/tulu-2-dpo-7b",
"allenai/codetulu-2-7b",
"microsoft/Orca-2-7b",
"openchat/openchat-3.5-0106",
"teknium/OpenHermes-2-Mistral-7B",
"teknium/OpenHermes-2.5-Mistral-7B",
"NousResearch/Nous-Hermes-2-Mistral-7B-DPO",
"HuggingFaceH4/zephyr-7b-beta",
"berkeley-nest/Starling-LM-7B-alpha",
"Nexusflow/Starling-LM-7B-beta",
"kaist-ai/mistral-orpo-alpha",
"kaist-ai/mistral-orpo-beta",
"google/gemma-1.1-7b-it",
],
"<=14B": [
"Qwen/Qwen1.5-14B-Chat",
"meta-llama/Llama-2-13b-chat-hf",
"codellama/CodeLlama-13b-Instruct-hf",
"allenai/tulu-2-13b",
"allenai/tulu-2-dpo-13b",
"allenai/codetulu-2-13b",
"microsoft/Orca-2-13b",
"upstage/SOLAR-10.7B-Instruct-v1.0",
"meta-llama/Meta-Llama-3-8B-Instruct",
"CohereForAI/aya-101",
],
"<=50B": [
"01-ai/Yi-34B-Chat",
"codellama/CodeLlama-34b-Instruct-hf",
"allenai/codetulu-2-34b",
"mistralai/Mixtral-8x7B-Instruct-v0.1",
"NousResearch/Nous-Hermes-2-Mixtral-8x7B-SFT",
"NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
"NousResearch/Nous-Hermes-2-Yi-34B",
"CohereForAI/c4ai-command-r-v01",
"Qwen/Qwen1.5-32B-Chat",
],
"<=75B": [
"meta-llama/Llama-2-70b-chat-hf",
"codellama/CodeLlama-70b-Instruct-hf",
"Qwen/Qwen1.5-72B-Chat",
"allenai/tulu-2-dpo-70b",
"meta-llama/Meta-Llama-3-70B-Instruct",
],
"<=175B": [
"alpindale/c4ai-command-r-plus-GPTQ",
"MaziyarPanahi/zephyr-orpo-141b-A35b-v0.1-AWQ",
"MaziyarPanahi/Mixtral-8x22B-Instruct-v0.1-AWQ",
],
},
}
API_MODELS = [
"gpt-3.5-turbo-0125",
"gpt-3.5-turbo-1106",
"gpt-4-0125-preview",
"gpt-4-1106-preview",
"gpt-4-turbo-2024-04-09",
"gpt-4o-2024-05-13",
"claude-3-haiku-20240307",
"claude-3-opus-20240229",
"claude-3-sonnet-20240229",
"mistral-large",
"mistral-medium",
"gemini-1.0-pro",
"gemini-pro-1.5",
"google/gemini-flash-1.5",
"qwen/qwen-110b-chat",
]
ORDERED_MODELS = [
"microsoft/phi-1",
"microsoft/phi-1_5",
"microsoft/phi-2",
"Qwen/Qwen1.5-0.5B",
"Qwen/Qwen1.5-1.8B",
"Qwen/Qwen1.5-4B",
"google/gemma-2b",
"allenai/OLMo-1B",
"Qwen/Qwen1.5-0.5B-Chat",
"Qwen/Qwen1.5-1.8B-Chat",
"Qwen/Qwen1.5-4B-Chat",
"microsoft/Phi-3-mini-4k-instruct",
"microsoft/Phi-3-mini-128k-instruct",
"google/gemma-2b-it",
"google/gemma-1.1-2b-it",
"google/gemma-7b",
"mistralai/Mistral-7B-v0.1",
"mistral-community/Mistral-7B-v0.2",
"Qwen/Qwen1.5-7B",
"01-ai/Yi-6B",
"meta-llama/Llama-2-7b-hf",
"codellama/CodeLlama-7b-hf",
"meta-llama/Meta-Llama-3-8B",
"EleutherAI/llemma_7b",
"allenai/OLMo-7B",
"google/gemma-7b-it",
"google/gemma-1.1-7b-it",
"mistralai/Mistral-7B-Instruct-v0.2",
"Qwen/Qwen1.5-7B-Chat",
"01-ai/Yi-6B-Chat",
"meta-llama/Llama-2-7b-chat-hf",
"codellama/CodeLlama-7b-Instruct-hf",
"meta-llama/Meta-Llama-3-8B-Instruct",
"allenai/OLMo-7B-SFT",
"allenai/OLMo-7B-Instruct",
"allenai/tulu-2-7b",
"allenai/tulu-2-dpo-7b",
"allenai/codetulu-2-7b",
"microsoft/Orca-2-7b",
"openchat/openchat-3.5-0106",
"teknium/OpenHermes-2-Mistral-7B",
"teknium/OpenHermes-2.5-Mistral-7B",
"NousResearch/Nous-Hermes-2-Mistral-7B-DPO",
"Starling-LM-7B-alpha",
"Starling-LM-7B-beta",
"kaist-ai/mistral-orpo-alpha",
"kaist-ai/mistral-orpo-beta",
"HuggingFaceH4/zephyr-7b-beta",
"Qwen/Qwen1.5-14B",
"meta-llama/Llama-2-13b-hf",
"codellama/CodeLlama-13b-hf",
"upstage/SOLAR-10.7B-v1.0",
"Qwen/Qwen1.5-14B-Chat",
"upstage/SOLAR-10.7B-Instruct-v1.0",
"CohereForAI/aya-101",
"meta-llama/Llama-2-13b-chat-hf",
"codellama/CodeLlama-13b-Instruct-hf",
"allenai/tulu-2-13b",
"allenai/tulu-2-dpo-13b",
"allenai/codetulu-2-13b",
"microsoft/Orca-2-13b",
"01-ai/Yi-34B",
"EleutherAI/llemma_34b",
"Qwen/Qwen1.5-32B",
"codellama/CodeLlama-34b-hf",
"mistralai/Mixtral-8x7B-v0.1",
"01-ai/Yi-34B-Chat",
"NousResearch/Nous-Hermes-2-Yi-34B",
"codellama/CodeLlama-34b-Instruct-hf",
"allenai/codetulu-2-34b",
"Qwen/Qwen1.5-32B-Chat",
"mistralai/Mixtral-8x7B-Instruct-v0.1",
"NousResearch/Nous-Hermes-2-Mixtral-8x7B-SFT",
"NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
"CohereForAI/c4ai-command-r-v01",
"meta-llama/Llama-2-70b-hf",
"codellama/CodeLlama-70b-hf",
"mistral-community/Mixtral-8x22B-v0.1-AWQ",
"meta-llama/Meta-Llama-3-70B",
"Qwen/Qwen1.5-72B",
"meta-llama/Llama-2-70b-chat-hf",
"codellama/CodeLlama-70b-Instruct-hf",
"allenai/tulu-2-dpo-70b",
"alpindale/c4ai-command-r-plus-GPTQ",
"meta-llama/Meta-Llama-3-70B-Instruct",
"MaziyarPanahi/Mixtral-8x22B-Instruct-v0.1-AWQ",
"MaziyarPanahi/zephyr-orpo-141b-A35b-v0.1-AWQ",
"Qwen/Qwen1.5-72B-Chat",
"qwen/qwen-110b-chat",
"gpt-3.5-turbo-1106",
"gpt-3.5-turbo-0125",
"gpt-4-1106-preview",
"gpt-4-0125-preview",
"gpt-4-turbo-2024-04-09",
"gpt-4o-2024-05-13",
"mistral-medium",
"mistral-large",
"gemini-1.0-pro",
"gemini-pro-1.5",
"google/gemini-flash-1.5",
"claude-3-haiku-20240307",
"claude-3-sonnet-20240229",
"claude-3-opus-20240229",
]
bgb_trained_models = [
"microsoft/phi-1",
"microsoft/phi-1_5",
"microsoft/phi-2",
"Qwen/Qwen1.5-0.5B",
"Qwen/Qwen1.5-1.8B",
"Qwen/Qwen1.5-4B",
"google/gemma-2b",
"allenai/OLMo-1B",
"google/gemma-7b",
"mistralai/Mistral-7B-v0.1",
"Qwen/Qwen1.5-7B",
"01-ai/Yi-6B",
"meta-llama/Llama-2-7b-hf",
"codellama/CodeLlama-7b-hf",
"EleutherAI/llemma_7b",
"allenai/OLMo-7B",
"Qwen/Qwen1.5-14B",
"meta-llama/Llama-2-13b-hf",
"codellama/CodeLlama-13b-hf",
"upstage/SOLAR-10.7B-v1.0",
"01-ai/Yi-34B",
"EleutherAI/llemma_34b",
"codellama/CodeLlama-34b-hf",
"mistralai/Mixtral-8x7B-v0.1",
"meta-llama/Llama-2-70b-hf",
"codellama/CodeLlama-70b-hf",
"Qwen/Qwen1.5-72B",
"Qwen/Qwen1.5-0.5B-Chat",
"Qwen/Qwen1.5-1.8B-Chat",
"Qwen/Qwen1.5-4B-Chat",
"google/gemma-2b-it",
"google/gemma-7b-it",
"mistralai/Mistral-7B-Instruct-v0.2",
"Qwen/Qwen1.5-7B-Chat",
"01-ai/Yi-6B-Chat",
"meta-llama/Llama-2-7b-chat-hf",
"codellama/CodeLlama-7b-Instruct-hf",
"allenai/OLMo-7B-SFT",
"allenai/OLMo-7B-Instruct",
"allenai/tulu-2-7b",
"allenai/tulu-2-dpo-7b",
"allenai/codetulu-2-7b",
"microsoft/Orca-2-7b",
"openchat/openchat-3.5-0106",
"teknium/OpenHermes-2-Mistral-7B",
"teknium/OpenHermes-2.5-Mistral-7B",
"NousResearch/Nous-Hermes-2-Mistral-7B-DPO",
"HuggingFaceH4/zephyr-7b-beta",
"Qwen/Qwen1.5-14B-Chat",
"meta-llama/Llama-2-13b-chat-hf",
"codellama/CodeLlama-13b-Instruct-hf",
"allenai/tulu-2-13b",
"allenai/tulu-2-dpo-13b",
"allenai/codetulu-2-13b",
"microsoft/Orca-2-13b",
"01-ai/Yi-34B-Chat",
"codellama/CodeLlama-34b-Instruct-hf",
"allenai/codetulu-2-34b",
"mistralai/Mixtral-8x7B-Instruct-v0.1",
"NousResearch/Nous-Hermes-2-Mistral-8x7B-SFT",
"NousResearch/Nous-Hermes-2-Mistral-8x7B-DPO",
"NousResearch/Nous-Hermes-2-Yi-34B",
"meta-llama/Llama-2-70b-chat-hf",
"codellama/CodeLlama-70b-Instruct-hf",
"Qwen/Qwen1.5-72B-Chat",
"allenai/tulu-2-dpo-72b",
]
MODEL_MAPPING = {
"microsoft/phi-1": [1.3, "Base"],
"microsoft/phi-1_5": [1.3, "Base"],
"microsoft/phi-2": [2.7, "Base"],
"Qwen/Qwen1.5-0.5B": [0.5, "Base"],
"Qwen/Qwen1.5-1.8B": [1.8, "Base"],
"Qwen/Qwen1.5-4B": [4.0, "Base"],
"google/gemma-2b": [2.0, "Base"],
"allenai/OLMo-1B": [1.0, "Base"],
"Qwen/Qwen1.5-0.5B-Chat": [0.5, "Chat", "Qwen/Qwen1.5-0.5B"],
"Qwen/Qwen1.5-1.8B-Chat": [1.8, "Chat", "Qwen/Qwen1.5-1.8B"],
"Qwen/Qwen1.5-4B-Chat": [4.0, "Chat", "Qwen/Qwen1.5-4B"],
"microsoft/Phi-3-mini-4k-instruct": [3.8, "Chat"],
"microsoft/Phi-3-mini-128k-instruct": [3.8, "Chat"],
"google/gemma-2b-it": [2.0, "Chat", "google/gemma-2b"],
"google/gemma-1.1-2b-it": [2.0, "Chat"],
"google/gemma-7b": [7.0, "Base"],
"mistralai/Mistral-7B-v0.1": [7.0, "Base"],
"mistral-community/Mistral-7B-v0.2": [7.0, "Base"],
"Qwen/Qwen1.5-7B": [7.0, "Base"],
"01-ai/Yi-6B": [6.0, "Base"],
"meta-llama/Llama-2-7b-hf": [7.0, "Base"],
"codellama/CodeLlama-7b-hf": [7.0, "Base"],
"meta-llama/Meta-Llama-3-8B": [8.0, "Base"],
"EleutherAI/llemma_7b": [7.0, "Base"],
"allenai/OLMo-7B": [7.0, "Base"],
"google/gemma-7b-it": [7.0, "Chat", "google/gemma-7b"],
"google/gemma-1.1-7b-it": [7.0, "Chat"],
"mistralai/Mistral-7B-Instruct-v0.2": [7.0, "Chat", "mistral-community/Mistral-7B-v0.2"],
"Qwen/Qwen1.5-7B-Chat": [7.0, "Chat", "Qwen/Qwen1.5-7B"],
"01-ai/Yi-6B-Chat": [6.0, "Chat", "01-ai/Yi-6B"],
"meta-llama/Llama-2-7b-chat-hf": [7.0, "Chat", "meta-llama/Llama-2-7b-hf"],
"codellama/CodeLlama-7b-Instruct-hf": [7.0, "Chat", "codellama/CodeLlama-7b-hf"],
"meta-llama/Meta-Llama-3-8B-Instruct": [8.0, "Chat", "meta-llama/Meta-Llama-3-8B"],
"allenai/OLMo-7B-SFT": [7.0, "Chat", "allenai/OLMo-7B"],
"allenai/OLMo-7B-Instruct": [7.0, "Chat", "allenai/OLMo-7B"],
"allenai/tulu-2-7b": [7.0, "Chat", "meta-llama/Llama-2-7b-hf"],
"allenai/tulu-2-dpo-7b": [7.0, "Chat", "meta-llama/Llama-2-7b-hf"],
"allenai/codetulu-2-7b": [7.0, "Chat", "codellama/CodeLlama-7b-hf"],
"microsoft/Orca-2-7b": [7.0, "Chat", "meta-llama/Llama-2-7b-hf"],
"openchat/openchat-3.5-0106": [7.0, "Chat", "mistralai/Mistral-7B-v0.1"],
"teknium/OpenHermes-2-Mistral-7B": [7.0, "Chat", "mistralai/Mistral-7B-v0.1"],
"teknium/OpenHermes-2.5-Mistral-7B": [7.0, "Chat", "mistralai/Mistral-7B-v0.1"],
"NousResearch/Nous-Hermes-2-Mistral-7B-DPO": [7.0, "Chat", "mistralai/Mistral-7B-v0.1"],
"Starling-LM-7B-alpha": [7.0, "Chat"],
"Starling-LM-7B-beta": [7.0, "Chat"],
"kaist-ai/mistral-orpo-alpha": [7.0, "Chat", "mistralai/Mistral-7B-v0.1"],
"kaist-ai/mistral-orpo-beta": [7.0, "Chat", "mistralai/Mistral-7B-v0.1"],
"HuggingFaceH4/zephyr-7b-beta": [7.0, "Chat", "mistralai/Mistral-7B-v0.1"],
"Qwen/Qwen1.5-14B": [14.0, "Base"],
"meta-llama/Llama-2-13b-hf": [13.0, "Base"],
"codellama/CodeLlama-13b-hf": [13.0, "Base"],
"upstage/SOLAR-10.7B-v1.0": [10.7, "Base"],
"Qwen/Qwen1.5-14B-Chat": [14.0, "Chat", "Qwen/Qwen1.5-14B"],
"upstage/SOLAR-10.7B-Instruct-v1.0": [10.7, "Chat", "upstage/SOLAR-10.7B-v1.0"],
"CohereForAI/aya-101": [13.0, "Chat"],
"meta-llama/Llama-2-13b-chat-hf": [13.0, "Chat", "meta-llama/Llama-2-13b-hf"],
"codellama/CodeLlama-13b-Instruct-hf": [13.0, "Chat", "codellama/CodeLlama-13b-hf"],
"allenai/tulu-2-13b": [13.0, "Chat", "meta-llama/Llama-2-13b-hf"],
"allenai/tulu-2-dpo-13b": [13.0, "Chat", "meta-llama/Llama-2-13b-hf"],
"allenai/codetulu-2-13b": [13.0, "Chat", "codellama/CodeLlama-13b-hf"],
"microsoft/Orca-2-13b": [13.0, "Chat", "meta-llama/Llama-2-13b-hf"],
"01-ai/Yi-34B": [34.0, "Base"],
"EleutherAI/llemma_34b": [34.0, "Base"],
"Qwen/Qwen1.5-32B": [32.0, "Base"],
"codellama/CodeLlama-34b-hf": [34.0, "Base"],
"mistralai/Mixtral-8x7B-v0.1": [46.7, "Base"],
"01-ai/Yi-34B-Chat": [34.0, "Chat", "01-ai/Yi-34B"],
"NousResearch/Nous-Hermes-2-Yi-34B": [34.0, "Chat", "01-ai/Yi-34B"],
"codellama/CodeLlama-34b-Instruct-hf": [34.0, "Chat", "codellama/CodeLlama-34b-hf"],
"allenai/codetulu-2-34b": [34.0, "Chat", "codellama/CodeLlama-34b-hf"],
"Qwen/Qwen1.5-32B-Chat": [32.0, "Chat", "Qwen/Qwen1.5-32B"],
"mistralai/Mixtral-8x7B-Instruct-v0.1": [46.7, "Chat", "mistralai/Mixtral-8x7B-v0.1"],
"NousResearch/Nous-Hermes-2-Mixtral-8x7B-SFT": [46.7, "Chat", "mistralai/Mixtral-8x7B-v0.1"],
"NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO": [46.7, "Chat", "mistralai/Mixtral-8x7B-v0.1"],
"CohereForAI/c4ai-command-r-v01": [35.0, "Chat"],
"meta-llama/Llama-2-70b-hf": [70.0, "Base"],
"codellama/CodeLlama-70b-hf": [70.0, "Base"],
"mistral-community/Mixtral-8x22B-v0.1-AWQ": [141, "Base"],
"meta-llama/Meta-Llama-3-70B": [70.0, "Base"],
"Qwen/Qwen1.5-72B": [72.0, "Base"],
"meta-llama/Llama-2-70b-chat-hf": [70.0, "Chat", "meta-llama/Llama-2-70b-hf"],
"codellama/CodeLlama-70b-Instruct-hf": [70.0, "Chat", "codellama/CodeLlama-70b-hf"],
"allenai/tulu-2-dpo-70b": [70.0, "Chat", "meta-llama/Llama-2-70b-hf"],
"alpindale/c4ai-command-r-plus-GPTQ": [104, "Chat"],
"meta-llama/Meta-Llama-3-70B-Instruct": [70.0, "Chat", "meta-llama/Meta-Llama-3-70B"],
"MaziyarPanahi/Mixtral-8x22B-Instruct-v0.1-AWQ": ["141", "Chat", "mistral-community/Mixtral-8x22B-v0.1-AWQ"],
"MaziyarPanahi/zephyr-orpo-141b-A35b-v0.1-AWQ": ["141", "Chat", "mistral-community/Mixtral-8x22B-v0.1-AWQ"],
"Qwen/Qwen1.5-72B-Chat": [72.0, "Chat", "Qwen/Qwen1.5-72B"],
"qwen/qwen-110b-chat": [110.0, "Chat", None],
"gpt-3.5-turbo-1106": [None, "Proprietary"],
"gpt-3.5-turbo-0125": [None, "Proprietary"],
"gpt-4-1106-preview": [None, "Proprietary"],
"gpt-4-0125-preview": [None, "Proprietary"],
"gpt-4-turbo-2024-04-09": [None, "Proprietary"],
"gpt-4o-2024-05-13": [None, "Proprietary"],
"mistral-medium": [None, "Proprietary"],
"mistral-large": [None, "Proprietary"],
"gemini-1.0-pro": [None, "Proprietary"],
"gemini-pro-1.5": [None, "Proprietary"],
"google/gemini-flash-1.5": [None, "Proprietary"],
"claude-3-haiku-20240307": [None, "Proprietary"],
"claude-3-sonnet-20240229": [None, "Proprietary"],
"claude-3-opus-20240229": [None, "Proprietary"],
}
MODEL_SHORT_TO_LONG = {model.split("/")[-1]: model for model in ORDERED_MODELS}
def get_model_type(model_name: str) -> str:
for _, model_list in MODELS["pretrained"].items():
if model_name in model_list:
return "base"
for _, model_list in MODELS["instruction_tuned"].items():
if model_name in model_list:
return "instruct"
if model_name in API_MODELS:
return "api"
raise ValueError(f"Model {model_name} not found in model_list.py")
return None
def get_open_model_list() -> list:
all_models = []
for _, model_list in MODELS["pretrained"].items():
all_models.extend(model_list)
for _, model_list in MODELS["instruction_tuned"].items():
all_models.extend(model_list)
return all_models
def get_all_model_list() -> list:
all_models = []
for _, model_list in MODELS["pretrained"].items():
all_models.extend(model_list)
for _, model_list in MODELS["instruction_tuned"].items():
all_models.extend(model_list)
all_models.extend(API_MODELS)
return all_models
def get_pretrained_models() -> list:
all_models = []
for _, model_list in MODELS["pretrained"].items():
all_models.extend(model_list)
return all_models
def get_instruct_models() -> list:
all_models = []
for _, model_list in MODELS["instruction_tuned"].items():
all_models.extend(model_list)
return all_models
def get_model_params(model_name: str) -> int:
for size_range, model_list in MODELS["pretrained"].items():
if model_name in model_list:
return int(size_range.split("B")[0].replace("<=", ""))
for size_range, model_list in MODELS["instruction_tuned"].items():
if model_name in model_list:
return int(size_range.split("B")[0].replace("<=", ""))
raise ValueError(f"Model {model_name} not found in model_list.py")
def get_model_num_gpus(model_name: str) -> int:
model_params = get_model_params(model_name)
num_gpus = {
4: 1,
7: 1,
14: 2,
50: 4,
75: 8,
175: 4,
}[model_params]
return num_gpus
def get_not_trained_models() -> list:
all_models = get_all_model_list()
trained_models = bgb_trained_models
not_trained_models = [model for model in all_models if model not in trained_models]
return not_trained_models
def is_trained_model(model_name: str) -> bool:
return model_name in bgb_trained_models
if __name__ == "__main__":
assert get_model_type("microsoft/phi-1"), "base"
assert get_model_params("microsoft/phi-2"), 4
models = get_all_model_list()
model_list_str = ""
for model in models:
model_list_str += f'"{model}"\n'
print(model_list_str)
print(f"{len(models)} models found in src/model_list.py")
print(get_not_trained_models())