commit 85b230681430d9a564df981634fe29907908b914 Author: Guofang Date: Mon Nov 9 21:40:57 2026 +0800 Initial commit Co-Authored-By: Claude Sonnet 5.5 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..21d0b89 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.venv/ diff --git a/embedlab.py b/embedlab.py new file mode 100755 index 0000000..5dd0ca0 --- /dev/null +++ b/embedlab.py @@ -0,0 +1,253 @@ +#!/usr/bin/env python3 + +print("Loading torch/transformers, please wait...", flush=True) + +import readline + +import torch +import torch.nn.functional as F +from transformers import AutoModel, AutoTokenizer + +# 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" +COMMANDS = ["/select", "/sim", "/tokens", "/params", "/set", "/reset", "/lsmod", "/help", "/quit", "/exit"] + +device = "cuda" if torch.cuda.is_available() else "cpu" +_loaded = {} # loaded models: name -> (tokenizer, model) +_params = {name: dict(d) for name, d in MODELS.items()} # live, per-model parameters +current = None # currently selected model + + +def is_cached(name): + """Whether the model weights exist in the local Hugging Face cache.""" + from huggingface_hub import try_to_load_from_cache + + return any( + isinstance(try_to_load_from_cache(name, f), str) + for f in ("model.safetensors", "pytorch_model.bin") + ) + + +def select_model(name): + global current + 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() + _loaded[name] = (tok, mod) + current = name + print(f"Current model: {name}") + + +def params(): + for k in PARAM_KEYS: + v, d = _params[current][k], MODELS[current][k] + print(f" {k:<11} = {v}" + ("" if v == d else f" (default: {d})")) + + +def set_param(arg): + """arg is " "; validates and applies it to the current model.""" + parts = arg.split() + if len(parts) != 2 or parts[0] not in PARAM_KEYS: + print("Usage: /set ") + return + key, val = parts + if key == "pooling": + if val not in ("cls", "mean"): + print("pooling must be 'cls' or 'mean'") + return + else: + try: + val = int(val) + except ValueError: + print(f"{key} must be an integer") + return + limit = _loaded[current][0].model_max_length if key == "max_length" else 4096 + if key == "max_length" and limit > 100_000: # tokenizer reports no real limit + limit = 8192 + if not 1 <= val <= limit: + print(f"{key} must be between 1 and {limit}") + return + _params[current][key] = val + print(f"[{current}] {key} = {val}") + + +def reset_params(): + _params[current] = dict(MODELS[current]) + print(f"[{current}] parameters reset to defaults") + + +def show_tokens(text, preview=50): + """Show how the current model's tokenizer splits text.""" + text = text.strip() + if not text: + print("Usage: /tokens ") + return + tokenizer = _loaded[current][0] + ids = tokenizer(text, truncation=False)["input_ids"] # includes special tokens + toks = tokenizer.convert_ids_to_tokens(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}") + max_length = _params[current]["max_length"] + if len(ids) > max_length: + print(f" note: exceeds max_length={max_length}; " + f"the last {len(ids) - max_length} tokens will be truncated") + + +def sim(arg): + """arg is "text1 | text2"; prints their cosine similarity.""" + a, sep, b = arg.partition("|") + a, b = a.strip(), b.strip() + if not sep or not a or not b: + print("Usage: /sim | ") + return + va, vb = encode([a, b]) + # vectors are L2-normalized, so the dot product equals cosine similarity + print(f"[{current}] similarity: {(va @ vb).item():.4f}") + + +def lsmod(): + for name in MODELS: + mark = "*" if name == current else " " + state = "cached" if is_cached(name) else "not downloaded" + print(f"{mark} {name} [{state}]") + + +@torch.no_grad() +def encode(texts): + """Encode with the current model; returns L2-normalized vectors, shape (N, dim).""" + tokenizer, model = _loaded[current] + p = _params[current] + pooling, max_length, batch_size = p["pooling"], p["max_length"], p["batch_size"] + outs = [] + for i in range(0, len(texts), batch_size): + inputs = tokenizer( + texts[i : i + batch_size], + padding=True, + truncation=True, + max_length=max_length, + return_tensors="pt", + ).to(device) + hidden = model(**inputs).last_hidden_state + if pooling == "cls": + emb = hidden[:, 0] + 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()) + return torch.cat(outs) + + +def completer(text, state): + """Tab completion: commands at line start, model names after /select.""" + line = readline.get_line_buffer().lstrip() + if line.startswith("/select "): + options = [m for m in MODELS if m.startswith(text)] + elif 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 + + +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: + /lsmod list models (* marks the current one) + /select switch model (Tab completes) + /sim | cosine similarity of two texts (separated by |) + /tokens show how the text is split into tokens + /params show the current model's parameters + /set 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) + Any other text (including "lsmod" without /) prints its vector. + To embed text that starts with /, escape it with //, e.g. //lsmod""" + + +def main(): + import argparse + + parser = argparse.ArgumentParser(description="Enter text, get its embedding vector") + parser.add_argument("--full", action="store_true", help="print the full vector") + args = parser.parse_args() + + setup_readline() + select_model(DEFAULT_MODEL) + print(HELP) + while True: + try: + text = input("> ").strip() + except EOFError: + break + if not text or text in ("/quit", "/exit"): + break + if text.startswith("//"): + text = text[1:] # escape: treat text starting with / as plain text + cmd = None + else: + cmd = text if text.startswith("/") else None + if cmd == "/lsmod": + lsmod() + elif cmd == "/help": + print(HELP) + elif cmd == "/tokens" or (cmd and cmd.startswith("/tokens ")): + show_tokens(cmd[len("/tokens"):]) + elif cmd == "/params": + params() + elif cmd == "/reset": + reset_params() + elif cmd == "/set" or (cmd and cmd.startswith("/set ")): + set_param(cmd[len("/set"):]) + elif cmd == "/sim" or (cmd and cmd.startswith("/sim ")): + sim(cmd[len("/sim"):]) + elif cmd == "/select" or (cmd and cmd.startswith("/select ")): + name = cmd[len("/select"):].strip() + if name: + select_model(name) + else: + print("Usage: /select ") + elif cmd: + print(f"Unknown command: {cmd} (type /help)") + else: + vec = encode([text])[0] + print(f"[{current}] dim: {vec.numel()} norm: {vec.norm().item():.4f}") + if args.full: + print(vec.tolist()) + else: + print("first 8 values:", [round(x, 4) for x in vec[:8].tolist()]) + + +if __name__ == "__main__": + main() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..60be314 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +torch +transformers +huggingface_hub diff --git a/setup.sh b/setup.sh new file mode 100755 index 0000000..948e728 --- /dev/null +++ b/setup.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +# 创建 venv 并安装依赖 +set -euo pipefail +cd "$(dirname "$0")" +python3 -m venv .venv +.venv/bin/pip install -r requirements.txt +echo "完成。运行: .venv/bin/python embedlab.py"