From db8e536254953390b7ea55a2738a45bbcf9c88c0 Mon Sep 17 00:00:00 2001 From: Guofang Date: Tue, 10 Nov 2026 22:33:02 +0800 Subject: [PATCH] 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 --- embedlab.py | 88 +++++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 69 insertions(+), 19 deletions(-) diff --git a/embedlab.py b/embedlab.py index 29b6f36..61d05af 100755 --- a/embedlab.py +++ b/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, {})) - mod = AutoModel.from_pretrained(name, **MODEL_OPTS.get(name, {})).to(device).eval() + tok = AutoTokenizer.from_pretrained(name) + 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: