Tom Aarsen commited on
Commit
e82960d
1 Parent(s): 6c6aac5

Fix embedding dimensions if Dense module exists

Browse files

e.g. for https://huggingface.co/aspire/acge_text_embedding

Files changed (1) hide show
  1. app.py +9 -6
app.py CHANGED
@@ -1168,12 +1168,15 @@ with open("EXTERNAL_MODEL_RESULTS.json", "w") as f:
1168
  def get_dim_seq_size(model):
1169
  filenames = [sib.rfilename for sib in model.siblings]
1170
  dim, seq, size = "", "", ""
1171
- if "1_Pooling/config.json" in filenames:
1172
- st_config_path = hf_hub_download(model.modelId, filename="1_Pooling/config.json")
1173
- dim = json.load(open(st_config_path)).get("word_embedding_dimension", "")
1174
- elif "2_Pooling/config.json" in filenames:
1175
- st_config_path = hf_hub_download(model.modelId, filename="2_Pooling/config.json")
1176
- dim = json.load(open(st_config_path)).get("word_embedding_dimension", "")
 
 
 
1177
  if "config.json" in filenames:
1178
  config_path = hf_hub_download(model.modelId, filename="config.json")
1179
  config = json.load(open(config_path))
 
1168
  def get_dim_seq_size(model):
1169
  filenames = [sib.rfilename for sib in model.siblings]
1170
  dim, seq, size = "", "", ""
1171
+ for filename in filenames:
1172
+ if re.match("\d+_Pooling/config.json", filename):
1173
+ st_config_path = hf_hub_download(model.modelId, filename=filename)
1174
+ dim = json.load(open(st_config_path)).get("word_embedding_dimension", "")
1175
+ break
1176
+ for filename in filenames:
1177
+ if re.match("\d+_Dense/config.json", filename):
1178
+ st_config_path = hf_hub_download(model.modelId, filename=filename)
1179
+ dim = json.load(open(st_config_path)).get("out_features", dim)
1180
  if "config.json" in filenames:
1181
  config_path = hf_hub_download(model.modelId, filename="config.json")
1182
  config = json.load(open(config_path))