fix keyerror?

This commit is contained in:
Cyberes 2023-10-20 14:00:24 -06:00
parent 0abd4b94fb
commit e838f591aa
1 changed files with 8 additions and 6 deletions

View File

@ -5,9 +5,11 @@ def estimate_model_size(config: dict):
:param config: :param config:
:return: :return:
""" """
vocab_size = config['vocab_size'] vocab_size = config.get('vocab_size')
hidden_size = config['hidden_size'] hidden_size = config.get('hidden_size')
num_hidden_layers = config['num_hidden_layers'] num_hidden_layers = config.get('num_hidden_layers')
intermediate_size = config['intermediate_size'] 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))) 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 int(total_params / 1e9)
return 0