From a589ec34d1bf0efdcad86ee7cf25ac83b4dec4e1 Mon Sep 17 00:00:00 2001
From: electroglyph <admin@terminoid.com>
Date: Tue, 16 Dec 2025 02:50:37 -0800
Subject: [PATCH] refactor pooling detection, add missing pooling types

---
 unsloth/models/sentence_transformer.py | 111 +++++++++++++------------
 1 file changed, 59 insertions(+), 52 deletions(-)

diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py
index 36bc14e203..3e8850cf39 100644
--- a/unsloth/models/sentence_transformer.py
+++ b/unsloth/models/sentence_transformer.py
@@ -21,6 +21,64 @@ from huggingface_hub import hf_hub_download
 
 
 class FastSentenceTransformer(FastModel):
+    @staticmethod
+    def read_pooling_mode(model_name, token):
+        try:
+            if os.path.exists(model_name) and os.path.exists(
+                os.path.join(model_name, "modules.json")
+            ):
+                modules_json_path = os.path.join(model_name, "modules.json")
+            else:
+                modules_json_path = hf_hub_download(
+                    model_name, "modules.json", token = token
+                )
+
+            with open(modules_json_path, "r") as f:
+                modules_config = json.load(f)
+
+            pooling_config_path = None
+            for module in modules_config:
+                if module.get("type", "") == "sentence_transformers.models.Pooling":
+                    pooling_path = module.get("path", "")
+                    if pooling_path:
+                        # try to find config.json for pooling module
+                        if os.path.exists(model_name) and os.path.exists(
+                            os.path.join(model_name, pooling_path, "config.json")
+                        ):
+                            pooling_config_path = os.path.join(
+                                model_name, pooling_path, "config.json"
+                            )
+                        else:
+                            pooling_config_path = hf_hub_download(
+                                model_name,
+                                os.path.join(pooling_path, "config.json"),
+                                token = token,
+                            )
+                        break
+
+            if pooling_config_path:
+                with open(pooling_config_path, "r") as f:
+                    pooling_config = json.load(f)
+                    pooling_map = {
+                        "pooling_mode_cls_token": "cls",
+                        "pooling_mode_mean_tokens": "mean",
+                        "pooling_mode_max_tokens": "max",
+                        "pooling_mode_mean_sqrt_len_tokens": "mean_sqrt_len",
+                        "pooling_mode_weightedmean_tokens": "weightedmean",
+                        "pooling_mode_lasttoken": "lasttoken",
+                    }
+                    for config_key, mode in pooling_map.items():
+                        if pooling_config.get(config_key):
+                            if mode != "mean":
+                                print(f"Pooling mode detected as {mode}, updating...")
+                            return mode
+
+        except Exception as e:
+            print(
+                f"Failed to detect pooling mode: {e}, defaulting to mean pooling."
+            )
+            return "mean"
+    
     @staticmethod
     def from_pretrained(
         model_name,
@@ -152,58 +210,7 @@ class FastSentenceTransformer(FastModel):
 
         # detect pooling mode if not specified/default
         if pooling_mode == "mean":
-            try:
-                if os.path.exists(model_name) and os.path.exists(
-                    os.path.join(model_name, "modules.json")
-                ):
-                    modules_json_path = os.path.join(model_name, "modules.json")
-                else:
-                    modules_json_path = hf_hub_download(
-                        model_name, "modules.json", token = token
-                    )
-
-                with open(modules_json_path, "r") as f:
-                    modules_config = json.load(f)
-
-                pooling_config_path = None
-                for module in modules_config:
-                    if module.get("type", "") == "sentence_transformers.models.Pooling":
-                        pooling_path = module.get("path", "")
-                        if pooling_path:
-                            # try to find config.json for pooling module
-                            if os.path.exists(model_name) and os.path.exists(
-                                os.path.join(model_name, pooling_path, "config.json")
-                            ):
-                                pooling_config_path = os.path.join(
-                                    model_name, pooling_path, "config.json"
-                                )
-                            else:
-                                pooling_config_path = hf_hub_download(
-                                    model_name,
-                                    os.path.join(pooling_path, "config.json"),
-                                    token = token,
-                                )
-                            break
-
-                if pooling_config_path:
-                    with open(pooling_config_path, "r") as f:
-                        pooling_config = json.load(f)
-                        pooling_map = {
-                            "pooling_mode_cls_token": "cls",
-                            "pooling_mode_mean_tokens": "mean",
-                            "pooling_mode_max_tokens": "max",
-                            "pooling_mode_mean_sqrt_len_tokens": "mean_sqrt_len",
-                        }
-                        for config_key, mode in pooling_map.items():
-                            if pooling_config.get(config_key):
-                                print(f"Pooling mode detected as {mode}, updating...")
-                                pooling_mode = mode
-                                break
-
-            except Exception as e:
-                print(
-                    f"Failed to detect pooling mode: {e}, defaulting to mean pooling."
-                )
+            pooling_mode = FastSentenceTransformer.read_pooling_mode(model_name, token)
 
         pooling_module = Pooling(
             word_embedding_dimension = hidden_size,