Show readable tokens for byte-level BPE tokenizers

Qwen tokens were printed as raw byte characters (e.g. 你好 as ä½łå¥½).
Decode each token for byte-level tokenizers, keeping the raw form for
tokens that hold only part of a multi-byte character.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

embedlab.py

@@ -2,16 +2,21 @@
 
 print("Loading torch/transformers, please wait...", flush=True)
 
-import atexit
 import os
-import readline
 
 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)
+# 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 = {
@@ -19,14 +24,24 @@ 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")
 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"
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():  # Apple Silicon GPU
+    device = "mps"
+else:
+    device = "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
@@ -68,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)
@@ -90,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:
@@ -123,6 +138,10 @@ def show_tokens(text, preview=50):
     tokenizer = _loaded[current][0]
     ids = tokenizer(text, truncation=False)["input_ids"]  # includes special tokens
     toks = tokenizer.convert_ids_to_tokens(ids)
+    if type(getattr(tokenizer.backend_tokenizer, "decoder", None)).__name__ == "ByteLevel":
+        # byte-level BPE (Qwen) stores raw bytes as odd characters; decode each token to readable text.
+        # A token holding only part of a multi-byte character can't be decoded, so keep its raw form.
+        toks = [d if d and "\ufffd" not in d else t for t, d in zip(toks, map(tokenizer.decode, ([i] for i in 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}")
@@ -169,55 +188,69 @@ 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)
 
 
-def completer(text, state):
-    """Tab completion: commands at line start, model names after /select."""
-    line = readline.get_line_buffer().lstrip()
+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 "):
-        options = [m for m in MODELS if m.startswith(text)]
-    elif line.startswith("/set "):
+        return [m for m in MODELS if m.startswith(text)]
+    if 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
+            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", "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)]
+    return []
 
 
-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")
-    try:
-        readline.read_history_file(HISTORY_FILE)
-    except OSError:
-        pass
-    readline.set_history_length(1000)
-    atexit.register(save_history)
+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 save_history():
+def make_session():
     try:
         os.makedirs(os.path.dirname(HISTORY_FILE), exist_ok=True)
-        readline.write_history_file(HISTORY_FILE)
     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:
@@ -226,7 +259,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)
@@ -241,13 +274,13 @@ def main():
     parser.add_argument("--full", action="store_true", help="print the full vector")
     args = parser.parse_args()
 
-    setup_readline()
+    session = make_session()
     select_model(load_last_model())
     print(HELP)
     while True:
         try:
-            text = input("> ").strip()
-        except EOFError:
+            text = session.prompt("> ").strip()
+        except (EOFError, KeyboardInterrupt):
             break
         if not text or text in ("/quit", "/exit"):
             break

requirements.txt

@@ -1,3 +1,4 @@
 torch
 transformers
 huggingface_hub
+prompt_toolkit