From 26282b545a0d25b038ba98b6109a8c6231c6e94a Mon Sep 17 00:00:00 2001 From: Guofang Date: Mon, 9 Nov 2026 21:40:57 +0800 Subject: [PATCH] 1 --- .gitignore | 1 + bge_demo.py | 171 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 172 insertions(+) create mode 100644 .gitignore create mode 100755 bge_demo.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bee8a64 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__ diff --git a/bge_demo.py b/bge_demo.py new file mode 100755 index 0000000..1ddc956 --- /dev/null +++ b/bge_demo.py @@ -0,0 +1,171 @@ +#!/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 -> (pooling, max tokens). The bge family uses CLS pooling. +MODELS = { + "BAAI/bge-small-zh-v1.5": ("cls", 512), + "BAAI/bge-base-zh-v1.5": ("cls", 512), + "BAAI/bge-large-zh-v1.5": ("cls", 512), + "BAAI/bge-m3": ("cls", 8192), +} +DEFAULT_MODEL = "BAAI/bge-small-zh-v1.5" +COMMANDS = ["/select", "/sim", "/lsmod", "/help", "/quit", "/exit"] + +device = "cuda" if torch.cuda.is_available() else "cpu" +_loaded = {} # loaded models: name -> (tokenizer, model) +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 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, batch_size=32): + """Encode with the current model; returns L2-normalized vectors, shape (N, dim).""" + tokenizer, model = _loaded[current] + pooling, max_length = MODELS[current] + 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 " " not in line: + options = [c + " " if c in ("/select", "/sim") 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 |) + /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 == "/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()