guofangandClaude Sonnet 5.5
822b55f26f
Replace readline with a prompt_toolkit session: the completion menu opens as you type, history suggestions show in gray (accept with Right), and history lives in .cache/history. Enter on a highlighted candidate accepts it instead of submitting the line. Ctrl+C now exits cleanly. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
305 lines
11 KiB
Python
Executable File
305 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
print("Loading torch/transformers, please wait...", flush=True)
|
|
|
|
import os
|
|
|
|
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)
|
|
# 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"
|
|
_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 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
|
|
|
|
|
|
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:
|
|
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
|
|
save_last_model(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 "<key> <value>"; 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 <pooling|max_length|batch_size> <value>")
|
|
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 <text>")
|
|
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 <text1> | <text2>")
|
|
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 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 "):
|
|
return [m for m in MODELS if m.startswith(text)]
|
|
if 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") 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 []
|
|
|
|
|
|
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),
|
|
)
|
|
|
|
|
|
HELP = """Commands:
|
|
/lsmod list models (* marks the current one)
|
|
/select <model> switch model (Tab completes)
|
|
/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
|
|
/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()
|
|
|
|
session = make_session()
|
|
select_model(load_last_model())
|
|
print(HELP)
|
|
while True:
|
|
try:
|
|
text = session.prompt("> ").strip()
|
|
except (EOFError, KeyboardInterrupt):
|
|
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 <model>")
|
|
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()
|