ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor pooling detection, add missing pooling types

Browse source
This commit is contained in:
electroglyph 2025-12-16 02:50:37 -08:00
commit a589ec34d1
1 changed files with 59 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

111
unsloth/models/sentence_transformer.py
View file

@ -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,