alozowski commited on
Commit
24c603a
1 Parent(s): 37e5956

Set unknown model size to -1 and improve logging

Browse files
Files changed (1) hide show
  1. src/submission/check_validity.py +12 -6
src/submission/check_validity.py CHANGED
@@ -1,6 +1,7 @@
1
  import json
2
  import os
3
  import re
 
4
  from collections import defaultdict
5
  from datetime import datetime, timedelta, timezone
6
 
@@ -75,28 +76,33 @@ def is_model_on_hub(
75
  return False, f"was not found or misconfigured on the hub! Error raised was {e.args[0]}", None
76
 
77
 
78
- def get_model_size(model_info: ModelInfo, precision: str):
79
  size_pattern = re.compile(r"(\d+\.)?\d+(b|m)")
80
  safetensors = None
 
81
  try:
82
  safetensors = get_safetensors_metadata(model_info.id)
83
  except Exception as e:
84
- print(e)
85
 
86
  if safetensors is not None:
87
  model_size = round(sum(safetensors.parameter_count.values()) / 1e9, 3)
88
  else:
89
  try:
90
  size_match = re.search(size_pattern, model_info.id.lower())
91
- model_size = size_match.group(0)
92
- model_size = round(float(model_size[:-1]) if model_size[-1] == "b" else float(model_size[:-1]) / 1e3, 3)
 
 
 
93
  except AttributeError:
94
- return 0 # Unknown model sizes are indicated as 0, see NUMERIC_INTERVALS in app.py
 
95
 
96
  size_factor = 8 if (precision == "GPTQ" or "gptq" in model_info.id.lower()) else 1
97
  model_size = size_factor * model_size
98
- return model_size
99
 
 
100
 
101
  def get_model_arch(model_info: ModelInfo):
102
  return model_info.config.get("architectures", "Unknown")
 
1
  import json
2
  import os
3
  import re
4
+ import logging
5
  from collections import defaultdict
6
  from datetime import datetime, timedelta, timezone
7
 
 
76
  return False, f"was not found or misconfigured on the hub! Error raised was {e.args[0]}", None
77
 
78
 
79
+ def get_model_size(model_info: ModelInfo, precision: str) -> float:
80
  size_pattern = re.compile(r"(\d+\.)?\d+(b|m)")
81
  safetensors = None
82
+
83
  try:
84
  safetensors = get_safetensors_metadata(model_info.id)
85
  except Exception as e:
86
+ logging.error(f"Failed to get safetensors metadata for model {model_info.id}: {str(e)}")
87
 
88
  if safetensors is not None:
89
  model_size = round(sum(safetensors.parameter_count.values()) / 1e9, 3)
90
  else:
91
  try:
92
  size_match = re.search(size_pattern, model_info.id.lower())
93
+ if size_match:
94
+ model_size = size_match.group(0)
95
+ model_size = round(float(model_size[:-1]) if model_size[-1] == "b" else float(model_size[:-1]) / 1e3, 3)
96
+ else:
97
+ return -1 # Unknown model size
98
  except AttributeError:
99
+ logging.warning(f"Unable to parse model size from ID: {model_info.id}")
100
+ return -1 # Unknown model size
101
 
102
  size_factor = 8 if (precision == "GPTQ" or "gptq" in model_info.id.lower()) else 1
103
  model_size = size_factor * model_size
 
104
 
105
+ return model_size
106
 
107
  def get_model_arch(model_info: ModelInfo):
108
  return model_info.config.get("architectures", "Unknown")