with open(module_path, encoding="utf-8") as f:
modules = json.load(f)
for mod in modules:
- if mod["type"] == "sentence_transformers.models.Pooling":
+ if mod["type"].endswith("Pooling"):
pooling_path = mod["path"]
break
+ mode_mapping = {
+ "mean": gguf.PoolingType.MEAN,
+ "cls": gguf.PoolingType.CLS,
+ "lasttoken": gguf.PoolingType.LAST,
+ }
+
# get pooling type
if pooling_path is not None:
with open(self.dir_model / pooling_path / "config.json", encoding="utf-8") as f:
pooling = json.load(f)
- if pooling["pooling_mode_mean_tokens"]:
+ if pooling.get("pooling_mode_mean_tokens"):
pooling_type = gguf.PoolingType.MEAN
- elif pooling["pooling_mode_cls_token"]:
+ elif pooling.get("pooling_mode_cls_token"):
pooling_type = gguf.PoolingType.CLS
- elif pooling["pooling_mode_lasttoken"]:
+ elif pooling.get("pooling_mode_lasttoken"):
pooling_type = gguf.PoolingType.LAST
+ elif (pooling_mode := pooling.get("pooling_mode")) in mode_mapping:
+ pooling_type = mode_mapping[pooling_mode]
else:
raise NotImplementedError("Only MEAN, CLS, and LAST pooling types supported")
self.gguf_writer.add_pooling_type(pooling_type)
with open(modules_file, encoding="utf-8") as modules_json_file:
mods = json.load(modules_json_file)
for mod in mods:
- if mod["type"] == "sentence_transformers.models.Dense":
+ if mod["type"].endswith("Dense"):
mod_path = mod["path"]
# check if model.safetensors file for Dense layer exists
model_tensors_file = self.dir_model / mod_path / "model.safetensors"