From 0cb17d663839284dd6bdcab9ec5f2933b3a1e471 Mon Sep 17 00:00:00 2001 From: Guofang Date: Tue, 10 Nov 2026 22:16:50 +0800 Subject: [PATCH] Add Qwen3-Embedding-8B and last-token pooling Qwen3-Embedding is a decoder model: it pools the last token, so add a "last" pooling mode, pad on the left, and load the weights in bfloat16. Co-Authored-By: Claude Sonnet 5.5 --- embedlab.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/embedlab.py b/embedlab.py index 9db35ba..1f43bb0 100755 --- a/embedlab.py +++ b/embedlab.py @@ -15,8 +15,8 @@ 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) +# Model -> default parameters. The bge family uses CLS pooling, Qwen3-Embedding uses last-token pooling. +# 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 = { @@ -24,7 +24,12 @@ MODELS = { "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") @@ -78,8 +83,8 @@ def select_model(name): 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() + tok = AutoTokenizer.from_pretrained(name, **TOKENIZER_OPTS.get(name, {})) + mod = AutoModel.from_pretrained(name, **MODEL_OPTS.get(name, {})).to(device).eval() _loaded[name] = (tok, mod) current = name save_last_model(name) @@ -100,8 +105,8 @@ def set_param(arg): return key, val = parts if key == "pooling": - if val not in ("cls", "mean"): - print("pooling must be 'cls' or 'mean'") + if val not in ("cls", "mean", "last"): + print("pooling must be 'cls', 'mean' or 'last'") return else: try: @@ -179,10 +184,16 @@ 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, p=2, dim=1).cpu()) + outs.append(F.normalize(emb.float(), p=2, dim=1).cpu()) return torch.cat(outs) @@ -196,7 +207,7 @@ def candidates(line, text): 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 [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)] @@ -244,7 +255,7 @@ HELP = """Commands: /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 + /set change a parameter: pooling (cls|mean|last), 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)