@@ -1850,20 +1850,28 @@ def _try_set_pooling_type(self) -> None:
18501850 with open(module_path, encoding="utf-8") as f:
18511851 modules = json.load(f)
18521852 for mod in modules:
1853- if mod["type"] == "sentence_transformers.models. Pooling":
1853+ if mod["type"].endswith(" Pooling") :
18541854 pooling_path = mod["path"]
18551855 break
18561856
1857+ mode_mapping = {
1858+ "mean": gguf.PoolingType.MEAN,
1859+ "cls": gguf.PoolingType.CLS,
1860+ "lasttoken": gguf.PoolingType.LAST,
1861+ }
1862+
18571863 # get pooling type
18581864 if pooling_path is not None:
18591865 with open(self.dir_model / pooling_path / "config.json", encoding="utf-8") as f:
18601866 pooling = json.load(f)
1861- if pooling[ "pooling_mode_mean_tokens"] :
1867+ if pooling.get( "pooling_mode_mean_tokens") :
18621868 pooling_type = gguf.PoolingType.MEAN
1863- elif pooling[ "pooling_mode_cls_token"] :
1869+ elif pooling.get( "pooling_mode_cls_token") :
18641870 pooling_type = gguf.PoolingType.CLS
1865- elif pooling[ "pooling_mode_lasttoken"] :
1871+ elif pooling.get( "pooling_mode_lasttoken") :
18661872 pooling_type = gguf.PoolingType.LAST
1873+ elif (pooling_mode := pooling.get("pooling_mode")) in mode_mapping:
1874+ pooling_type = mode_mapping[pooling_mode]
18671875 else:
18681876 raise NotImplementedError("Only MEAN, CLS, and LAST pooling types supported")
18691877 self.gguf_writer.add_pooling_type(pooling_type)
@@ -7180,7 +7188,7 @@ def __init__(self, *args, **kwargs):
71807188 with open(modules_file, encoding="utf-8") as modules_json_file:
71817189 mods = json.load(modules_json_file)
71827190 for mod in mods:
7183- if mod["type"] == "sentence_transformers.models. Dense":
7191+ if mod["type"].endswith(" Dense") :
71847192 mod_path = mod["path"]
71857193 # check if model.safetensors file for Dense layer exists
71867194 model_tensors_file = self.dir_model / mod_path / "model.safetensors"
0 commit comments