Skip to content

Commit 88af1f0

Browse files
Blaizzyclaude
Blaizzy
and
claude
committed
qwen3_vl: use a text-to-image retrieval demo as the README example
Replaces the toy "embed 4 mixed inputs and print a 4x4 similarity" snippet with a real retrieval workflow: embed an image gallery once, score multiple text queries, and plot the similarity heatmap + top-K image grid. - README embedding example swapped for the retrieval script excerpt, with a rendered PNG of the plot. - examples/qwen3_vl_retrieval.py: runnable demo (saves one PNG). - examples/compare_mlx_vs_transformers.py: side-by-side comparison against the HF Qwen3VLEmbedder reference (runs via uv ephemeral env so the torch-free main venv stays clean). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 15208cb commit 88af1f0

4 files changed

Lines changed: 333 additions & 23 deletions

File tree

‎README.md‎

Lines changed: 55 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -39,39 +39,71 @@ pip install mlx-embeddings
3939

4040
Qwen3-VL uses a model-specific processor and a high-level `model.process(...)` API for multimodal embedding and reranking.
4141

42-
#### Multimodal Embedding
42+
#### Multimodal Retrieval
43+
44+
Text-to-image retrieval over a small gallery — embed images once, then score any number of text queries against them. Full script: [`examples/qwen3_vl_retrieval.py`](examples/qwen3_vl_retrieval.py).
4345

4446
```python
47+
from io import BytesIO
48+
49+
import matplotlib.pyplot as plt
4550
import mlx.core as mx
46-
from mlx_embeddings import load
51+
import numpy as np
52+
import requests
53+
from PIL import Image
4754

48-
model, processor = load("Qwen/Qwen3-VL-Embedding-2B")
55+
from mlx_embeddings import load
4956

50-
inputs = [
51-
{
52-
"text": "A woman playing with her dog on a beach at sunset.",
53-
"instruction": "Retrieve images or text relevant to the user's query.",
54-
},
55-
{
56-
"text": "A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset."
57-
},
58-
{
59-
"image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
60-
},
61-
{
62-
"text": "A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset.",
63-
"image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
64-
},
57+
GALLERY = [
58+
("woman with dog on beach",
59+
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"),
60+
("two cats on a couch",
61+
"http://images.cocodataset.org/val2017/000000039769.jpg"),
62+
("tennis player on a court",
63+
"http://images.cocodataset.org/val2017/000000000872.jpg"),
64+
("bear in the wild",
65+
"http://images.cocodataset.org/val2017/000000000285.jpg"),
66+
("dark train tunnel",
67+
"http://images.cocodataset.org/val2017/000000001268.jpg"),
68+
("group of people standing together",
69+
"http://images.cocodataset.org/val2017/000000001000.jpg"),
70+
]
71+
QUERIES = [
72+
"a person spending time with their pet outdoors at sunset",
73+
"sleepy cats relaxing indoors",
74+
"someone playing a racquet sport",
75+
"wildlife in a natural habitat",
76+
"the inside of a transit tunnel",
77+
"a crowd of people gathered outside",
6578
]
79+
INSTRUCTION = "Retrieve images that match the user's query."
80+
81+
def fetch(src):
82+
if src.startswith(("http://", "https://")):
83+
return Image.open(BytesIO(requests.get(src, timeout=30).content)).convert("RGB")
84+
return Image.open(src).convert("RGB")
85+
86+
labels, urls = zip(*GALLERY)
87+
images = [fetch(u) for u in urls]
6688

67-
embeddings = model.process(inputs, processor=processor)
68-
similarity = embeddings @ embeddings.T
89+
model, processor = load("Qwen/Qwen3-VL-Embedding-2B")
90+
img_embeds = model.process([{"image": i} for i in images], processor=processor)
91+
txt_embeds = model.process(
92+
[{"text": q, "instruction": INSTRUCTION} for q in QUERIES], processor=processor,
93+
)
94+
sim = np.array((txt_embeds @ img_embeds.T).astype(mx.float32))
6995

70-
mx.eval(embeddings, similarity)
71-
print(embeddings.shape) # (4, 2048)
72-
print(similarity)
96+
for qi, q in enumerate(QUERIES):
97+
top = np.argsort(-sim[qi])[:3]
98+
print(f"q{qi}: {q}")
99+
for k, idx in enumerate(top):
100+
print(f" #{k + 1} {sim[qi, idx]:.3f} {labels[idx]}")
73101
```
74102

103+
Output (`examples/qwen3_vl_retrieval.py` additionally plots the full query × image similarity heatmap and the top-K image grid):
104+
105+
![Qwen3-VL retrieval](examples/qwen3_vl_retrieval.png)
106+
75107
#### Multimodal Reranking
76108

77109
```python

‎examples/compare_mlx_vs_transformers.py‎

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
"""
2+
Compare Qwen3-VL-Embedding-2B outputs between mlx-embeddings (this repo)
3+
and the official HF transformers reference implementation.
4+
5+
Loads both, runs the same GALLERY + QUERIES through both, and prints the
6+
two similarity matrices side-by-side with max/mean abs diff.
7+
8+
Run (with torch + transformers available):
9+
uv run --with torch --with transformers --with qwen-vl-utils \\
10+
--with pillow --with requests --with matplotlib --with mlx \\
11+
--with mlx-vlm --with huggingface-hub --with accelerate \\
12+
examples/compare_mlx_vs_transformers.py
13+
14+
Or inside a venv that has those deps installed.
15+
"""
16+
from __future__ import annotations
17+
18+
import importlib.util
19+
import sys
20+
from io import BytesIO
21+
from pathlib import Path
22+
23+
import numpy as np
24+
import requests
25+
from PIL import Image
26+
27+
GALLERY = [
28+
("woman with dog on beach",
29+
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"),
30+
("two cats on a couch",
31+
"http://images.cocodataset.org/val2017/000000039769.jpg"),
32+
("tennis player on a court",
33+
"http://images.cocodataset.org/val2017/000000000872.jpg"),
34+
("bear in the wild",
35+
"http://images.cocodataset.org/val2017/000000000285.jpg"),
36+
("dark train tunnel",
37+
"http://images.cocodataset.org/val2017/000000001268.jpg"),
38+
("group of people standing together",
39+
"http://images.cocodataset.org/val2017/000000001000.jpg"),
40+
]
41+
42+
QUERIES = [
43+
"a person spending time with their pet outdoors at sunset",
44+
"sleepy cats relaxing indoors",
45+
"someone playing a racquet sport",
46+
"wildlife in a natural habitat",
47+
"the inside of a transit tunnel",
48+
"a crowd of people gathered outside",
49+
]
50+
51+
INSTRUCTION = "Retrieve images that match the user's query."
52+
MODEL_ID = "Qwen/Qwen3-VL-Embedding-2B"
53+
54+
55+
def fetch_image(src: str) -> Image.Image:
56+
if src.startswith(("http://", "https://")):
57+
return Image.open(BytesIO(requests.get(src, timeout=30).content)).convert("RGB")
58+
return Image.open(src).convert("RGB")
59+
60+
61+
def run_mlx(images, queries):
62+
import mlx.core as mx
63+
from mlx_embeddings import load
64+
65+
model, processor = load(MODEL_ID)
66+
image_embeds = model.process(
67+
[{"image": img} for img in images], processor=processor
68+
)
69+
mx.eval(image_embeds)
70+
text_embeds = model.process(
71+
[{"text": q, "instruction": INSTRUCTION} for q in queries],
72+
processor=processor,
73+
)
74+
mx.eval(text_embeds)
75+
sim = (text_embeds @ image_embeds.T).astype(mx.float32)
76+
mx.eval(sim)
77+
return np.array(sim)
78+
79+
80+
def _load_hf_reference():
81+
"""Import the model's bundled scripts/qwen3_vl_embedding.py (the reference
82+
implementation shipped with the checkpoint) without needing it on sys.path.
83+
"""
84+
from huggingface_hub import snapshot_download
85+
import transformers.utils.generic as _tgeneric
86+
87+
# The reference script was written against transformers 4.x and imports
88+
# `check_model_inputs`, which was removed in 5.x. The one call site using
89+
# the decorator is already commented out, so a no-op is safe.
90+
if not hasattr(_tgeneric, "check_model_inputs"):
91+
def check_model_inputs(fn):
92+
return fn
93+
_tgeneric.check_model_inputs = check_model_inputs
94+
95+
local = Path(snapshot_download(MODEL_ID))
96+
script_path = local / "scripts" / "qwen3_vl_embedding.py"
97+
spec = importlib.util.spec_from_file_location("qwen3_vl_embedding", script_path)
98+
mod = importlib.util.module_from_spec(spec)
99+
sys.modules["qwen3_vl_embedding"] = mod
100+
spec.loader.exec_module(mod)
101+
return mod
102+
103+
104+
def run_transformers(images, queries):
105+
import torch
106+
107+
ref = _load_hf_reference()
108+
embedder = ref.Qwen3VLEmbedder(MODEL_ID, torch_dtype=torch.bfloat16)
109+
110+
image_inputs = [{"image": img} for img in images]
111+
text_inputs = [{"text": q, "instruction": INSTRUCTION} for q in queries]
112+
113+
image_embeds = embedder.process(image_inputs, normalize=True).to(torch.float32).cpu()
114+
text_embeds = embedder.process(text_inputs, normalize=True).to(torch.float32).cpu()
115+
sim = (text_embeds @ image_embeds.T).numpy()
116+
return sim
117+
118+
119+
def format_matrix(name: str, sim: np.ndarray, row_labels, col_labels) -> str:
120+
header = " " + " ".join(f"{c:>6}" for c in col_labels)
121+
rows = [f" q{i}: " + " ".join(f"{v:+.3f}" for v in sim[i]) + f" <- {r}"
122+
for i, r in enumerate(row_labels)]
123+
return f"=== {name} ===\n{header}\n" + "\n".join(rows)
124+
125+
126+
def top_k_table(sim: np.ndarray, labels, queries, k: int = 3) -> str:
127+
lines = []
128+
for i, q in enumerate(queries):
129+
order = np.argsort(-sim[i])[:k]
130+
picks = ", ".join(f"{labels[j]}({sim[i, j]:+.3f})" for j in order)
131+
lines.append(f" q{i} {q!r}\n {picks}")
132+
return "\n".join(lines)
133+
134+
135+
def main() -> None:
136+
labels, urls = zip(*GALLERY)
137+
short_labels = [f"img{i}" for i in range(len(GALLERY))]
138+
print(f"Fetching {len(GALLERY)} gallery images…")
139+
images = [fetch_image(u) for u in urls]
140+
141+
print("\n[MLX] running mlx-embeddings path…")
142+
sim_mlx = run_mlx(images, QUERIES)
143+
144+
print("[HF ] running transformers reference path…")
145+
sim_hf = run_transformers(images, QUERIES)
146+
147+
print()
148+
print(format_matrix("mlx-embeddings", sim_mlx, QUERIES, short_labels))
149+
print()
150+
print(format_matrix("transformers reference", sim_hf, QUERIES, short_labels))
151+
152+
diff = sim_mlx - sim_hf
153+
print("\n=== diff (mlx - hf) ===")
154+
print(f" max |diff|: {np.abs(diff).max():.4f}")
155+
print(f" mean |diff|: {np.abs(diff).mean():.4f}")
156+
print(f" per-row max |diff|: "
157+
f"{np.array2string(np.abs(diff).max(axis=1), precision=4)}")
158+
159+
def rank(s):
160+
return np.argsort(-s, axis=1)
161+
162+
ranks_mlx = rank(sim_mlx)
163+
ranks_hf = rank(sim_hf)
164+
top1_agree = (ranks_mlx[:, 0] == ranks_hf[:, 0]).mean()
165+
top3_agree = np.mean([
166+
set(ranks_mlx[i, :3]) == set(ranks_hf[i, :3])
167+
for i in range(len(QUERIES))
168+
])
169+
print(f"\n top-1 agreement across queries: {top1_agree * 100:.0f}%")
170+
print(f" top-3 set agreement: {top3_agree * 100:.0f}%")
171+
172+
print("\n=== mlx top-3 per query ===")
173+
print(top_k_table(sim_mlx, labels, QUERIES))
174+
print("\n=== hf top-3 per query ===")
175+
print(top_k_table(sim_hf, labels, QUERIES))
176+
177+
178+
if __name__ == "__main__":
179+
main()

‎examples/qwen3_vl_retrieval.png‎

965 KB
Loading

‎examples/qwen3_vl_retrieval.py‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
"""Qwen3-VL text-to-image retrieval demo.
2+
3+
Embeds a small image gallery, embeds a batch of text queries, then plots
4+
the query × image similarity heatmap and the top-K images per query.
5+
6+
Run:
7+
.venv/bin/python examples/qwen3_vl_retrieval.py
8+
9+
Saves one PNG next to this file.
10+
"""
11+
from io import BytesIO
12+
from pathlib import Path
13+
14+
import matplotlib.pyplot as plt
15+
import mlx.core as mx
16+
import numpy as np
17+
import requests
18+
from PIL import Image
19+
20+
from mlx_embeddings import load
21+
22+
GALLERY = [
23+
("woman with dog on beach",
24+
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"),
25+
("two cats on a couch",
26+
"http://images.cocodataset.org/val2017/000000039769.jpg"),
27+
("tennis player on a court",
28+
"http://images.cocodataset.org/val2017/000000000872.jpg"),
29+
("bear in the wild",
30+
"http://images.cocodataset.org/val2017/000000000285.jpg"),
31+
("dark train tunnel",
32+
"http://images.cocodataset.org/val2017/000000001268.jpg"),
33+
("group of people standing together",
34+
"http://images.cocodataset.org/val2017/000000001000.jpg"),
35+
]
36+
QUERIES = [
37+
"a person spending time with their pet outdoors at sunset",
38+
"sleepy cats relaxing indoors",
39+
"someone playing a racquet sport",
40+
"wildlife in a natural habitat",
41+
"the inside of a transit tunnel",
42+
"a crowd of people gathered outside",
43+
]
44+
INSTRUCTION = "Retrieve images that match the user's query."
45+
TOP_K = 3
46+
OUT = Path(__file__).with_name("qwen3_vl_retrieval.png")
47+
48+
49+
def fetch(src):
50+
if src.startswith(("http://", "https://")):
51+
return Image.open(BytesIO(requests.get(src, timeout=30).content)).convert("RGB")
52+
return Image.open(src).convert("RGB")
53+
54+
55+
labels, urls = zip(*GALLERY)
56+
images = [fetch(u) for u in urls]
57+
58+
model, processor = load("Qwen/Qwen3-VL-Embedding-2B")
59+
img_embeds = model.process([{"image": i} for i in images], processor=processor)
60+
txt_embeds = model.process(
61+
[{"text": q, "instruction": INSTRUCTION} for q in QUERIES], processor=processor,
62+
)
63+
sim = np.array((txt_embeds @ img_embeds.T).astype(mx.float32))
64+
65+
fig, (ax_hm, ax_tk) = plt.subplots(
66+
2, 1, figsize=(2 + TOP_K * 3, len(QUERIES) * 2.2),
67+
gridspec_kw={"height_ratios": [1, 2]},
68+
)
69+
70+
# Heatmap with text labels
71+
im = ax_hm.imshow(sim, cmap="viridis", aspect="auto")
72+
ax_hm.set_xticks(range(len(GALLERY)), labels, rotation=30, ha="right", fontsize=8)
73+
ax_hm.set_yticks(range(len(QUERIES)), [f"q{i}" for i in range(len(QUERIES))], fontsize=8)
74+
ax_hm.set_title("text ↔ image similarity")
75+
thresh = sim.mean()
76+
for i in range(sim.shape[0]):
77+
for j in range(sim.shape[1]):
78+
ax_hm.text(j, i, f"{sim[i, j]:.2f}", ha="center", va="center", fontsize=7,
79+
color="white" if sim[i, j] < thresh else "black")
80+
fig.colorbar(im, ax=ax_hm, shrink=0.8)
81+
82+
# Top-K images per query
83+
ax_tk.axis("off")
84+
gs = ax_tk.get_subplotspec().subgridspec(len(QUERIES), TOP_K + 1, wspace=0.1, hspace=0.4)
85+
for qi, q in enumerate(QUERIES):
86+
lbl_ax = fig.add_subplot(gs[qi, 0])
87+
lbl_ax.axis("off")
88+
lbl_ax.text(1.0, 0.5, f"q{qi}: {q}", ha="right", va="center", fontsize=9, wrap=True)
89+
for k, idx in enumerate(np.argsort(-sim[qi])[:TOP_K]):
90+
ax = fig.add_subplot(gs[qi, k + 1])
91+
ax.imshow(images[idx])
92+
ax.set_xticks([]); ax.set_yticks([])
93+
ax.set_title(f"#{k + 1} {sim[qi, idx]:.2f}\n{labels[idx]}", fontsize=8,
94+
color="tab:green" if k == 0 else "black")
95+
96+
fig.suptitle("Qwen3-VL retrieval", fontsize=13)
97+
fig.tight_layout(rect=[0, 0, 1, 0.98])
98+
fig.savefig(OUT, dpi=140, bbox_inches="tight")
99+
print(f"saved -> {OUT}")

0 commit comments

Comments
 (0)