From e838f591aa76688f033621295bc51b3067b48c2a Mon Sep 17 00:00:00 2001 From: Cyberes Date: Fri, 20 Oct 2023 14:00:24 -0600 Subject: [PATCH] fix keyerror? --- llm_server/routes/helpers/model.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/llm_server/routes/helpers/model.py b/llm_server/routes/helpers/model.py index ca35867..bf18b66 100644 --- a/llm_server/routes/helpers/model.py +++ b/llm_server/routes/helpers/model.py @@ -5,9 +5,11 @@ def estimate_model_size(config: dict): :param config: :return: """ - vocab_size = config['vocab_size'] - hidden_size = config['hidden_size'] - num_hidden_layers = config['num_hidden_layers'] - intermediate_size = config['intermediate_size'] - total_params = (vocab_size * hidden_size) + (num_hidden_layers * ((hidden_size * intermediate_size * 4) + (hidden_size * hidden_size * 3))) - return int(total_params / 1e9) + vocab_size = config.get('vocab_size') + hidden_size = config.get('hidden_size') + num_hidden_layers = config.get('num_hidden_layers') + intermediate_size = config.get('intermediate_size') + if vocab_size and hidden_size and num_hidden_layers and intermediate_size: + total_params = (vocab_size * hidden_size) + (num_hidden_layers * ((hidden_size * intermediate_size * 4) + (hidden_size * hidden_size * 3))) + return int(total_params / 1e9) + return 0