Initial commit

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

.gitignore

@@ -1,3 +1 @@
 .venv/
-__pycache__/
-.cache/

embedlab.py

@@ -2,49 +2,29 @@
 
 print("Loading torch/transformers, please wait...", flush=True)
 
-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
 
-# 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)
+# Model -> default parameters. The bge family uses CLS pooling.
+#   pooling:    "cls" (first token) or "mean" (average of all tokens)
 #   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"]
 
-if torch.cuda.is_available():
-    device = "cuda"
-elif torch.backends.mps.is_available():  # Apple Silicon GPU
-    device = "mps"
-else:
-    device = "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 _loaded = {}      # loaded models: name -> (tokenizer, model)
-_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)
+_params = {name: dict(d) for name, d in MODELS.items()}  # live, per-model parameters
 current = None    # currently selected model
 
 
@@ -54,94 +34,27 @@ def is_cached(name):
 
     return any(
         isinstance(try_to_load_from_cache(name, f), str)
-        for f in ("model.safetensors", "model.safetensors.index.json", "pytorch_model.bin")
+        for f in ("model.safetensors", "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:
-        with open(LAST_MODEL_FILE) as f:
-            name = f.read().strip()
-    except OSError:
-        return DEFAULT_MODEL
-    return name if name in MODELS else DEFAULT_MODEL  # MODELS was refreshed at startup
-
-
-def save_last_model(name):
-    try:
-        os.makedirs(os.path.dirname(LAST_MODEL_FILE), exist_ok=True)
-        with open(LAST_MODEL_FILE, "w") as f:
-            f.write(name)
-    except OSError:
-        pass
-
-
 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, dtype="auto").to(device).eval()  # dtype from the model's config
+        mod = AutoModel.from_pretrained(name).to(device).eval()
         _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}")
 
 
 def params():
     for k in PARAM_KEYS:
-        v, d = _params[current][k], _defaults[current][k]
+        v, d = _params[current][k], MODELS[current][k]
         print(f"  {k:<11} = {v}" + ("" if v == d else f"   (default: {d})"))
 
 
@@ -153,8 +66,8 @@ def set_param(arg):
         return
     key, val = parts
     if key == "pooling":
-        if val not in ("cls", "mean", "last"):
-            print("pooling must be 'cls', 'mean' or 'last'")
+        if val not in ("cls", "mean"):
+            print("pooling must be 'cls' or 'mean'")
             return
     else:
         try:
@@ -173,7 +86,7 @@ def set_param(arg):
 
 
 def reset_params():
-    _params[current] = dict(_defaults[current])
+    _params[current] = dict(MODELS[current])
     print(f"[{current}] parameters reset to defaults")
 
 
@@ -186,10 +99,6 @@ 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}")
@@ -212,7 +121,6 @@ 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"
@@ -237,69 +145,41 @@ 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.float(), p=2, dim=1).cpu())
+        outs.append(F.normalize(emb, p=2, dim=1).cpu())
     return torch.cat(outs)
 
 
-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()
+def completer(text, state):
+    """Tab completion: commands at line start, model names after /select."""
+    line = readline.get_line_buffer().lstrip()
     if line.startswith("/select "):
-        return [m for m in MODELS if m.startswith(text)]
-    if line.startswith("/set "):
+        options = [m for m in MODELS if m.startswith(text)]
+    elif line.startswith("/set "):
         words = line.split(" ")
         if len(words) <= 2:  # completing the key
-            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 []
+            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
 
 
-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 make_session():
-    try:
-        os.makedirs(os.path.dirname(HISTORY_FILE), exist_ok=True)
-    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),
-    )
+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")
 
 
 HELP = """Commands:
@@ -308,7 +188,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|last), max_length, batch_size
+  /set <key> <value>  change a parameter: pooling (cls|mean), 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)
@@ -323,14 +203,13 @@ def main():
     parser.add_argument("--full", action="store_true", help="print the full vector")
     args = parser.parse_args()
 
-    session = make_session()
-    refresh_models()
-    select_model(load_last_model())
+    setup_readline()
+    select_model(DEFAULT_MODEL)
     print(HELP)
     while True:
         try:
-            text = session.prompt("> ").strip()
-        except (EOFError, KeyboardInterrupt):
+            text = input("> ").strip()
+        except EOFError:
             break
         if not text or text in ("/quit", "/exit"):
             break

requirements.txt

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