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']
+    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