Read model parameters from the model instead of hard-coding

MODELS is now a list: a short recommended set plus sentence-transformers
models found in the Hugging Face cache, refreshed on startup, /lsmod and
/select. Defaults are read on first load: pooling from 1_Pooling/config.json,
max_length from the model and tokenizer limits (capped at 8192), batch_size
from the parameter count, and dtype from the model config. Also recognize
sharded weights in is_cached.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

embedlab.py

@@ -2,6 +2,7 @@
 
 print("Loading torch/transformers, please wait...", flush=True)
 
+import json
 import os
 
 import torch
@@ -15,21 +16,19 @@ from prompt_toolkit.history import FileHistory
 from prompt_toolkit.key_binding import KeyBindings
 from transformers import AutoModel, AutoTokenizer
 
-# Model -> default parameters. The bge family uses CLS pooling, Qwen3-Embedding uses last-token pooling.
+# Models offered even before they are downloaded. Any sentence-transformers style model found in
+# the local Hugging Face cache is listed too. Parameters are read from each model's own files.
+RECOMMENDED = [
+    "BAAI/bge-small-zh-v1.5",
+    "BAAI/bge-base-zh-v1.5",
+    "BAAI/bge-large-zh-v1.5",
+    "BAAI/bge-m3",
+    "Qwen/Qwen3-Embedding-8B",
+]
+# Parameters (defaults come from the model, see read_defaults):
 #   pooling:    "cls" (first token), "mean" (average of all tokens) or "last" (last token)
 #   max_length: max tokens per text; longer input is truncated
 #   batch_size: texts encoded per forward pass (speed/memory only, no effect on vectors)
-MODELS = {
-    "BAAI/bge-small-zh-v1.5": {"pooling": "cls", "max_length": 512, "batch_size": 32},
-    "BAAI/bge-base-zh-v1.5": {"pooling": "cls", "max_length": 512, "batch_size": 32},
-    "BAAI/bge-large-zh-v1.5": {"pooling": "cls", "max_length": 512, "batch_size": 32},
-    "BAAI/bge-m3": {"pooling": "cls", "max_length": 8192, "batch_size": 32},
-    "Qwen/Qwen3-Embedding-8B": {"pooling": "last", "max_length": 8192, "batch_size": 4},
-}
-# Extra loading options. Qwen3-Embedding is a decoder: left padding puts the last real token at
-# position -1, and bfloat16 halves the memory of the 8B weights.
-TOKENIZER_OPTS = {"Qwen/Qwen3-Embedding-8B": {"padding_side": "left"}}
-MODEL_OPTS = {"Qwen/Qwen3-Embedding-8B": {"torch_dtype": torch.bfloat16}}
 PARAM_KEYS = ["pooling", "max_length", "batch_size"]
 DEFAULT_MODEL = "BAAI/bge-small-zh-v1.5"
 LAST_MODEL_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".cache", "last_model")
@@ -43,7 +42,9 @@ elif torch.backends.mps.is_available():  # Apple Silicon GPU
 else:
     device = "cpu"
 _loaded = {}      # loaded models: name -> (tokenizer, model)
-_params = {name: dict(d) for name, d in MODELS.items()}  # live, per-model parameters
+_defaults = {}    # per-model default parameters, read when the model is first loaded
+_params = {}      # live, per-model parameters
+MODELS = list(RECOMMENDED)  # selectable models: recommended + cached (see refresh_models)
 current = None    # currently selected model
 
 
@@ -53,10 +54,52 @@ def is_cached(name):
 
     return any(
         isinstance(try_to_load_from_cache(name, f), str)
-        for f in ("model.safetensors", "pytorch_model.bin")
+        for f in ("model.safetensors", "model.safetensors.index.json", "pytorch_model.bin")
     )
 
 
+def refresh_models():
+    """Rebuild MODELS: the recommended list plus cached sentence-transformers models."""
+    from huggingface_hub import try_to_load_from_cache
+    from huggingface_hub.constants import HF_HUB_CACHE
+
+    found = []
+    try:
+        entries = sorted(os.listdir(HF_HUB_CACHE))
+    except OSError:
+        entries = []
+    for entry in entries:
+        kind, _, rest = entry.partition("--")
+        org, _, repo = rest.partition("--")
+        name = f"{org}/{repo}"
+        if kind == "models" and repo and isinstance(try_to_load_from_cache(name, "modules.json"), str):
+            found.append(name)
+    MODELS[:] = RECOMMENDED + [n for n in found if n not in RECOMMENDED]
+
+
+def read_defaults(name, tokenizer, model):
+    """Default parameters from the model's own files."""
+    from huggingface_hub import hf_hub_download
+
+    pooling = "mean"  # fallback when the repo has no sentence-transformers pooling config
+    try:
+        with open(hf_hub_download(name, "1_Pooling/config.json")) as f:
+            cfg = json.load(f)
+        for mode, key in (("cls", "pooling_mode_cls_token"), ("last", "pooling_mode_lasttoken"),
+                          ("mean", "pooling_mode_mean_tokens")):
+            if cfg.get(key):
+                pooling = mode
+                break
+    except Exception:
+        pass
+    limits = [getattr(model.config, "max_position_embeddings", 8192), 8192]
+    if tokenizer.model_max_length < 100_000:  # larger means the tokenizer reports no real limit
+        limits.append(tokenizer.model_max_length)
+    billions = sum(p.numel() for p in model.parameters()) / 1e9
+    return {"pooling": pooling, "max_length": min(limits),
+            "batch_size": 32 if billions < 1 else 8 if billions < 3 else 4}
+
+
 def load_last_model():
     """The model selected last time, or DEFAULT_MODEL if unknown."""
     try:
@@ -64,7 +107,7 @@ def load_last_model():
             name = f.read().strip()
     except OSError:
         return DEFAULT_MODEL
-    return name if name in MODELS else DEFAULT_MODEL
+    return name if name in MODELS else DEFAULT_MODEL  # MODELS was refreshed at startup
 
 
 def save_last_model(name):
@@ -78,14 +121,19 @@ def save_last_model(name):
 
 def select_model(name):
     global current
+    if name not in MODELS:
+        refresh_models()  # it may have been downloaded since startup
     if name not in MODELS:
         print(f"Unknown model: {name} (use /lsmod to list models)")
         return
     if name not in _loaded:
         print(f"Loading {name} ...", flush=True)
-        tok = AutoTokenizer.from_pretrained(name, **TOKENIZER_OPTS.get(name, {}))
+        tok = AutoTokenizer.from_pretrained(name)
-        mod = AutoModel.from_pretrained(name, **MODEL_OPTS.get(name, {})).to(device).eval()
+        mod = AutoModel.from_pretrained(name, dtype="auto").to(device).eval()  # dtype from the model's config
         _loaded[name] = (tok, mod)
+        _defaults[name] = read_defaults(name, tok, mod)
+        _params[name] = dict(_defaults[name])
+        refresh_models()
     current = name
     save_last_model(name)
     print(f"Current model: {name}")
@@ -93,7 +141,7 @@ def select_model(name):
 
 def params():
     for k in PARAM_KEYS:
-        v, d = _params[current][k], MODELS[current][k]
+        v, d = _params[current][k], _defaults[current][k]
         print(f"  {k:<11} = {v}" + ("" if v == d else f"   (default: {d})"))
 
 
@@ -125,7 +173,7 @@ def set_param(arg):
 
 
 def reset_params():
-    _params[current] = dict(MODELS[current])
+    _params[current] = dict(_defaults[current])
     print(f"[{current}] parameters reset to defaults")
 
 
@@ -164,6 +212,7 @@ def sim(arg):
 
 
 def lsmod():
+    refresh_models()
     for name in MODELS:
         mark = "*" if name == current else " "
         state = "cached" if is_cached(name) else "not downloaded"
@@ -275,6 +324,7 @@ def main():
     args = parser.parse_args()
 
     session = make_session()
+    refresh_models()
     select_model(load_last_model())
     print(HELP)
     while True: