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,33 +2,49 @@
 
 print("Loading torch/transformers, please wait...", flush=True)
 
-import atexit
+import json
 import os
-import readline
 
 import torch
 import torch.nn.functional as F
+from prompt_toolkit import PromptSession
+from prompt_toolkit.application import get_app
+from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
+from prompt_toolkit.completion import Completer, Completion
+from prompt_toolkit.filters import Condition
+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.
-#   pooling:    "cls" (first token) or "mean" (average of all tokens)
+# 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},
-}
 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")
 HISTORY_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".cache", "history")
 COMMANDS = ["/select", "/sim", "/tokens", "/params", "/set", "/reset", "/lsmod", "/help", "/quit", "/exit"]
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():  # Apple Silicon GPU
+    device = "mps"
+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
 
 
@@ -38,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:
@@ -49,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):
@@ -63,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)
-        mod = AutoModel.from_pretrained(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}")
@@ -78,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})"))
 
 
@@ -90,8 +153,8 @@ def set_param(arg):
         return
     key, val = parts
     if key == "pooling":
-        if val not in ("cls", "mean"):
-            print("pooling must be 'cls' or 'mean'")
+        if val not in ("cls", "mean", "last"):
+            print("pooling must be 'cls', 'mean' or 'last'")
             return
     else:
         try:
@@ -110,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")
 
 
@@ -123,6 +186,10 @@ def show_tokens(text, preview=50):
     tokenizer = _loaded[current][0]
     ids = tokenizer(text, truncation=False)["input_ids"]  # includes special tokens
     toks = tokenizer.convert_ids_to_tokens(ids)
+    if type(getattr(tokenizer.backend_tokenizer, "decoder", None)).__name__ == "ByteLevel":
+        # byte-level BPE (Qwen) stores raw bytes as odd characters; decode each token to readable text.
+        # A token holding only part of a multi-byte character can't be decoded, so keep its raw form.
+        toks = [d if d and "\ufffd" not in d else t for t, d in zip(toks, map(tokenizer.decode, ([i] for i in ids)))]
     shown = ", ".join(f"{t}({i})" for t, i in zip(toks[:preview], ids[:preview]))
     more = f", ... (+{len(ids) - preview} more)" if len(ids) > preview else ""
     print(f"[{current}] {len(ids)} tokens: {shown}{more}")
@@ -145,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"
@@ -169,55 +237,69 @@ def encode(texts):
         hidden = model(**inputs).last_hidden_state
         if pooling == "cls":
             emb = hidden[:, 0]
+        elif pooling == "last":
+            if tokenizer.padding_side == "left":
+                emb = hidden[:, -1]
+            else:  # right padding: pick the last non-pad token of each row
+                last = inputs["attention_mask"].sum(1) - 1
+                emb = hidden[torch.arange(hidden.size(0), device=hidden.device), last]
         else:  # mean
             mask = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)
             emb = (hidden * mask).sum(1) / mask.sum(1)
-        outs.append(F.normalize(emb, p=2, dim=1).cpu())
+        outs.append(F.normalize(emb.float(), p=2, dim=1).cpu())
     return torch.cat(outs)
 
 
-def completer(text, state):
-    """Tab completion: commands at line start, model names after /select."""
-    line = readline.get_line_buffer().lstrip()
+def candidates(line, text):
+    """Completions for the word `text` at the end of `line`: commands at line start, model names after /select."""
+    line = line.lstrip()
     if line.startswith("/select "):
-        options = [m for m in MODELS if m.startswith(text)]
-    elif line.startswith("/set "):
+        return [m for m in MODELS if m.startswith(text)]
+    if line.startswith("/set "):
         words = line.split(" ")
         if len(words) <= 2:  # completing the key
-            options = [k + " " for k in PARAM_KEYS if k.startswith(text)]
-        elif len(words) == 3 and words[1] == "pooling":
-            options = [v for v in ("cls", "mean") if v.startswith(text)]
-        else:
-            options = []
-    elif " " not in line:
-        options = [c + " " if c in ("/select", "/sim", "/tokens", "/set") else c for c in COMMANDS if c.startswith(text)]
-    else:
-        options = []
-    return options[state] if state < len(options) else None
+            return [k + " " for k in PARAM_KEYS if k.startswith(text)]
+        if len(words) == 3 and words[1] == "pooling":
+            return [v for v in ("cls", "mean", "last") if v.startswith(text)]
+        return []
+    if " " not in line:
+        return [c + " " if c in ("/select", "/sim", "/tokens", "/set") else c for c in COMMANDS if c.startswith(text)]
+    return []
 
 
-def setup_readline():
-    readline.set_completer_delims(" \t\n")  # do not treat / and - as delimiters
-    readline.set_completer(completer)
-    doc = readline.__doc__ or ""
-    if "libedit" in doc:  # macOS
-        readline.parse_and_bind("bind ^I rl_complete")
-    else:
-        readline.parse_and_bind("tab: complete")
-    try:
-        readline.read_history_file(HISTORY_FILE)
-    except OSError:
-        pass
-    readline.set_history_length(1000)
-    atexit.register(save_history)
+class InputCompleter(Completer):
+    """Completes while typing."""
+
+    def get_completions(self, document, complete_event):
+        text = document.get_word_before_cursor(WORD=True)
+        for option in candidates(document.text_before_cursor, text):
+            yield Completion(option, start_position=-len(text))
 
 
-def save_history():
+def make_session():
     try:
         os.makedirs(os.path.dirname(HISTORY_FILE), exist_ok=True)
-        readline.write_history_file(HISTORY_FILE)
     except OSError:
         pass
+    bindings = KeyBindings()
+
+    @Condition
+    def completion_selected():
+        state = get_app().current_buffer.complete_state
+        return state is not None and state.current_completion is not None
+
+    @bindings.add("enter", filter=completion_selected)
+    def _(event):
+        # Enter on a highlighted candidate accepts it (already inserted) instead of submitting
+        event.current_buffer.complete_state = None
+
+    return PromptSession(
+        key_bindings=bindings,
+        completer=InputCompleter(),
+        complete_while_typing=True,
+        auto_suggest=AutoSuggestFromHistory(),
+        history=FileHistory(HISTORY_FILE),
+    )
 
 
 HELP = """Commands:
@@ -226,7 +308,7 @@ HELP = """Commands:
   /sim <a> | <b>      cosine similarity of two texts (separated by |)
   /tokens <text>      show how the text is split into tokens
   /params             show the current model's parameters
-  /set <key> <value>  change a parameter: pooling (cls|mean), max_length, batch_size
+  /set <key> <value>  change a parameter: pooling (cls|mean|last), max_length, batch_size
   /reset              restore the current model's default parameters
   /help               show this help
   /quit, /exit        quit (empty line or Ctrl+D also works)
@@ -241,13 +323,14 @@ def main():
     parser.add_argument("--full", action="store_true", help="print the full vector")
     args = parser.parse_args()
 
-    setup_readline()
+    session = make_session()
+    refresh_models()
     select_model(load_last_model())
     print(HELP)
     while True:
         try:
-            text = input("> ").strip()
-        except EOFError:
+            text = session.prompt("> ").strip()
+        except (EOFError, KeyboardInterrupt):
             break
         if not text or text in ("/quit", "/exit"):
             break

requirements.txt

@@ -1,3 +1,4 @@
 torch
 transformers
 huggingface_hub
+prompt_toolkit