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>
This commit is contained in:
1 parent
30d7be1a80
commit
0cb17d6638
1 file changed
+20
-9
+20
-9
@@ -15,8 +15,8 @@ from prompt_toolkit.history import FileHistory
|
|||||||
from prompt_toolkit.key_binding import KeyBindings
|
from prompt_toolkit.key_binding import KeyBindings
|
||||||
from transformers import AutoModel, AutoTokenizer
|
from transformers import AutoModel, AutoTokenizer
|
||||||
|
|
||||||
# Model -> default parameters. The bge family uses CLS pooling.
|
# Model -> default parameters. The bge family uses CLS pooling, Qwen3-Embedding uses last-token pooling.
|
||||||
# pooling: "cls" (first token) or "mean" (average of all tokens)
|
# pooling: "cls" (first token), "mean" (average of all tokens) or "last" (last token)
|
||||||
# max_length: max tokens per text; longer input is truncated
|
# max_length: max tokens per text; longer input is truncated
|
||||||
# batch_size: texts encoded per forward pass (speed/memory only, no effect on vectors)
|
# batch_size: texts encoded per forward pass (speed/memory only, no effect on vectors)
|
||||||
MODELS = {
|
MODELS = {
|
||||||
@@ -24,7 +24,12 @@ MODELS = {
|
|||||||
"BAAI/bge-base-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-large-zh-v1.5": {"pooling": "cls", "max_length": 512, "batch_size": 32},
|
||||||
"BAAI/bge-m3": {"pooling": "cls", "max_length": 8192, "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"]
|
PARAM_KEYS = ["pooling", "max_length", "batch_size"]
|
||||||
DEFAULT_MODEL = "BAAI/bge-small-zh-v1.5"
|
DEFAULT_MODEL = "BAAI/bge-small-zh-v1.5"
|
||||||
LAST_MODEL_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".cache", "last_model")
|
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
|
return
|
||||||
if name not in _loaded:
|
if name not in _loaded:
|
||||||
print(f"Loading {name} ...", flush=True)
|
print(f"Loading {name} ...", flush=True)
|
||||||
tok = AutoTokenizer.from_pretrained(name)
|
tok = AutoTokenizer.from_pretrained(name, **TOKENIZER_OPTS.get(name, {}))
|
||||||
mod = AutoModel.from_pretrained(name).to(device).eval()
|
mod = AutoModel.from_pretrained(name, **MODEL_OPTS.get(name, {})).to(device).eval()
|
||||||
_loaded[name] = (tok, mod)
|
_loaded[name] = (tok, mod)
|
||||||
current = name
|
current = name
|
||||||
save_last_model(name)
|
save_last_model(name)
|
||||||
@@ -100,8 +105,8 @@ def set_param(arg):
|
|||||||
return
|
return
|
||||||
key, val = parts
|
key, val = parts
|
||||||
if key == "pooling":
|
if key == "pooling":
|
||||||
if val not in ("cls", "mean"):
|
if val not in ("cls", "mean", "last"):
|
||||||
print("pooling must be 'cls' or 'mean'")
|
print("pooling must be 'cls', 'mean' or 'last'")
|
||||||
return
|
return
|
||||||
else:
|
else:
|
||||||
try:
|
try:
|
||||||
@@ -179,10 +184,16 @@ def encode(texts):
|
|||||||
hidden = model(**inputs).last_hidden_state
|
hidden = model(**inputs).last_hidden_state
|
||||||
if pooling == "cls":
|
if pooling == "cls":
|
||||||
emb = hidden[:, 0]
|
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
|
else: # mean
|
||||||
mask = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)
|
mask = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)
|
||||||
emb = (hidden * mask).sum(1) / mask.sum(1)
|
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)
|
return torch.cat(outs)
|
||||||
|
|
||||||
|
|
||||||
@@ -196,7 +207,7 @@ def candidates(line, text):
|
|||||||
if len(words) <= 2: # completing the key
|
if len(words) <= 2: # completing the key
|
||||||
return [k + " " for k in PARAM_KEYS if k.startswith(text)]
|
return [k + " " for k in PARAM_KEYS if k.startswith(text)]
|
||||||
if len(words) == 3 and words[1] == "pooling":
|
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 []
|
return []
|
||||||
if " " not in line:
|
if " " not in line:
|
||||||
return [c + " " if c in ("/select", "/sim", "/tokens", "/set") else c for c in COMMANDS if c.startswith(text)]
|
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 |)
|
/sim <a> | <b> cosine similarity of two texts (separated by |)
|
||||||
/tokens <text> show how the text is split into tokens
|
/tokens <text> show how the text is split into tokens
|
||||||
/params show the current model's parameters
|
/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
|
/reset restore the current model's default parameters
|
||||||
/help show this help
|
/help show this help
|
||||||
/quit, /exit quit (empty line or Ctrl+D also works)
|
/quit, /exit quit (empty line or Ctrl+D also works)
|
||||||
|
|||||||
Reference in new issue
Block a user