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 <noreply@anthropic.com>

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 <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
+  /set <key> <value>  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)