neuml/txtai
1
0
Fork 0
mirror of https://github.com/neuml/txtai.git synced 2026-10-12 02:19:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support pooling_mode with models trained with recent Sentence Transformers, closes #1292

Browse Source
This commit is contained in:
davidmezzetti
2026-09-29 10:23:16 -04:00
parent e313248ff2
commit f87574de4c
1 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/python/txtai/models/pooling/factory.py
+7 -3
View File
@@ -85,16 +85,20 @@ class PoolingFactory:
# Load 1_Pooling/config.json file
config = PoolingFactory.load(path, "1_Pooling/config.json")
# Models trained with Sentence Transformers v5.4+ use `pooling_mode` parameter
# Support both the new and legacy parameters
mode = config.get("pooling_mode")
# Set to CLS pooling if it's enabled and mean pooling is disabled
if config and config.get("pooling_mode_cls_token") and not config.get("pooling_mode_mean_tokens"):
if config and config.get("pooling_mode_cls_token", mode == "cls") and not config.get("pooling_mode_mean_tokens", mode == "mean"):
method = "clspooling"
# Set to last token pooling if it's enabled and mean pooling is disabled
if config and config.get("pooling_mode_lasttoken") and not config.get("pooling_mode_mean_tokens"):
if config and config.get("pooling_mode_lasttoken", mode == "lasttoken") and not config.get("pooling_mode_mean_tokens", mode == "mean"):
method = "lastpooling"
# Set to max pooling if it's enabled and mean pooling is disabled
if config and config.get("pooling_mode_max_tokens") and not config.get("pooling_mode_mean_tokens"):
if config and config.get("pooling_mode_max_tokens", mode == "max") and not config.get("pooling_mode_mean_tokens", mode == "mean"):
method = "maxpooling"
# Check for late interaction pooling