mrdbourke commited on
Commit
146fa95
Β·
verified Β·
1 Parent(s): 1a5e77e

uploading multimodal RAG with nemotron

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ kitchen_bench.png filter=lfs diff=lfs merge=lfs -text
README.md CHANGED
@@ -1,12 +1,32 @@
1
  ---
2
- title: Multimodal Rag With Nemotron
3
- emoji: 🌍
4
  colorFrom: green
5
- colorTo: purple
6
  sdk: gradio
7
- sdk_version: 6.3.0
8
  app_file: app.py
9
  pinned: false
 
10
  ---
11
 
12
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ title: Multimodal RAG with Nemotron Embed and Rerank VL
3
+ emoji: πŸ‘οΈπŸ“‘
4
  colorFrom: green
5
+ colorTo: green
6
  sdk: gradio
 
7
  app_file: app.py
8
  pinned: false
9
+ license: apache-2.0
10
  ---
11
 
12
+ # Multimodal RAG with Nemotron Embed and Rerank VL Models
13
+
14
+ Search for recipes using input queries of text or images.
15
+
16
+ Dataset is 10,000 recipes from [Food.com](https://www.kaggle.com/datasets/shuyangli94/food-com-recipes-and-user-interactions) turned into images and markdown and embedded into the same feature space.
17
+
18
+ Options:
19
+
20
+ * Default is to retrieve most similar recipes via visual or text search
21
+ * Optional: Rerank initial results with reranker model (note: this increases latency)
22
+ * Optional: Generate a summary of the returned recipes (note: this increases latency)
23
+
24
+ Details:
25
+
26
+ * **Dataset used:** https://huggingface.co/datasets/mrdbourke/recipe-synthetic-images-10k
27
+ * **Embedding model used:** https://huggingface.co/nvidia/llama-nemotron-embed-vl-1b-v2
28
+ * **Note:** By default we use the image + text embeddings as we have access to image and text pairs in our dataset, and according to the [launch blog post](https://huggingface.co/blog/nvidia/llama-nemotron-vl-1b), these work the best.
29
+ * **Rerank model used:** https://huggingface.co/nvidia/llama-nemotron-rerank-vl-1b-v2
30
+ * **Generation model used:** https://huggingface.co/Qwen/Qwen3-VL-2B-Instruct (note: you could use a larger model such as [Nemotron v3](https://huggingface.co/collections/nvidia/nvidia-nemotron-v3), however, this will require more compute resources)
31
+
32
+ **Note:** This README.md was authored in a [live tutorial recorded for YouTube](https://youtube.com/live/XWmqcFiBvpU?feature=share).
app.py ADDED
@@ -0,0 +1,649 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Multimodal RAG Demo with Nemotron Embed VL and Rerank VL
3
+
4
+ A Gradio demo for multimodal retrieval augmented generation using:
5
+ - Dataset: mrdbourke/recipe-synthetic-images-10k
6
+ - Embedding model: nvidia/llama-nemotron-embed-vl-1b-v2
7
+ - Rerank model: nvidia/llama-nemotron-rerank-vl-1b-v2
8
+ - Generation model: Qwen/Qwen3-VL-2B-Instruct
9
+ """
10
+
11
+ import time
12
+ import torch
13
+ import gradio as gr
14
+ from PIL import Image
15
+ from datasets import load_dataset
16
+ from safetensors.torch import load_file
17
+ from transformers import (
18
+ AutoModel,
19
+ AutoModelForSequenceClassification,
20
+ AutoProcessor,
21
+ Qwen3VLForConditionalGeneration,
22
+ )
23
+
24
+ # ============================================================================
25
+ # Configuration
26
+ # ============================================================================
27
+
28
+ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
29
+
30
+ # Model paths and commit hashes (required for sdpa attention on HF Spaces)
31
+ EMBED_MODEL_PATH = "nvidia/llama-nemotron-embed-vl-1b-v2"
32
+ EMBED_COMMIT_HASH = "5b5ca69c35bf6ec1484d2d5ff238626e67a745e2"
33
+
34
+ RERANK_MODEL_PATH = "nvidia/llama-nemotron-rerank-vl-1b-v2"
35
+ RERANK_COMMIT_HASH = "47e5a355d1a050c3e5f69d53f14964b1d34bcd9d"
36
+
37
+ GENERATION_MODEL_ID = "Qwen/Qwen3-VL-2B-Instruct"
38
+
39
+ # ============================================================================
40
+ # Load Dataset and Embeddings
41
+ # ============================================================================
42
+
43
+ print("[INFO] Loading dataset...")
44
+ dataset = load_dataset(path="mrdbourke/recipe-synthetic-images-10k")
45
+ print(f"[INFO] Dataset loaded with {len(dataset['train'])} samples")
46
+
47
+ print("[INFO] Loading embeddings...")
48
+ image_text_embeddings = load_file("image_text_embeddings_10k.safetensors")
49
+ image_text_embeddings = image_text_embeddings["image_text_embeddings"].to(DEVICE)
50
+ print(f"[INFO] Embeddings loaded: {image_text_embeddings.shape}")
51
+
52
+ # ============================================================================
53
+ # Load Models
54
+ # ============================================================================
55
+ modality_to_tokens = {
56
+ "image": 2048,
57
+ "image_text": 10240,
58
+ "text": 8192
59
+ }
60
+
61
+ print(f"[INFO] Loading embedding model from: {EMBED_MODEL_PATH} with commit: {EMBED_COMMIT_HASH}")
62
+ embed_model = AutoModel.from_pretrained(
63
+ EMBED_MODEL_PATH,
64
+ revision=EMBED_COMMIT_HASH,
65
+ dtype=torch.bfloat16,
66
+ trust_remote_code=True,
67
+ attn_implementation="sdpa",
68
+ device_map="auto",
69
+ ).eval()
70
+
71
+ # Set embed processor kwargs
72
+ # Note: These are the suggest settings from the embed model card
73
+ embed_modality = "image_text"
74
+ embed_processor_kwargs = {
75
+ "trust_remote_code": True,
76
+ "max_input_tiles": 6,
77
+ "use_thumbnail": True,
78
+ "p_max_length": modality_to_tokens[embed_modality]
79
+ }
80
+
81
+ embed_processor = AutoProcessor.from_pretrained(
82
+ EMBED_MODEL_PATH,
83
+ revision=EMBED_COMMIT_HASH,
84
+ trust_remote_code=True,
85
+ **embed_processor_kwargs
86
+ )
87
+ print(f"[INFO] Embedding model loaded!")
88
+ print(f"[INFO] Embed processor using p_max_length: {embed_processor.p_max_length}")
89
+
90
+ print(f"[INFO] Loading rerank model from: {RERANK_MODEL_PATH} with commit: {RERANK_COMMIT_HASH}")
91
+ rerank_model = AutoModelForSequenceClassification.from_pretrained(
92
+ RERANK_MODEL_PATH,
93
+ revision=RERANK_COMMIT_HASH,
94
+ dtype=torch.bfloat16,
95
+ trust_remote_code=True,
96
+ attn_implementation="sdpa",
97
+ device_map="auto",
98
+ ).eval()
99
+
100
+ # Set rerank processor kwargs
101
+ # Note: These are the suggest settings from the rerank model card
102
+ rerank_modality = "image_text"
103
+ rereank_processor_kwargs = {
104
+ "trust_remote_code": True,
105
+ "max_input_tiles": 6,
106
+ "use_thumbnail": True,
107
+ "rerank_max_length": modality_to_tokens[rerank_modality]
108
+ }
109
+
110
+ rerank_processor = AutoProcessor.from_pretrained(
111
+ RERANK_MODEL_PATH,
112
+ revision=RERANK_COMMIT_HASH,
113
+ trust_remote_code=True,
114
+ **rereank_processor_kwargs
115
+ )
116
+ print(f"[INFO] Rerank processor using rerank_max_length: {rerank_processor.rerank_max_length}")
117
+
118
+ print(f"[INFO] Rerank model loaded!")
119
+
120
+ print("[INFO] Loading generation model...")
121
+ qwen_model = Qwen3VLForConditionalGeneration.from_pretrained(
122
+ GENERATION_MODEL_ID,
123
+ dtype="auto",
124
+ device_map="auto"
125
+ )
126
+ qwen_processor = AutoProcessor.from_pretrained(GENERATION_MODEL_ID)
127
+ print(f"[INFO] Generation model loaded")
128
+
129
+ # ============================================================================
130
+ # Helper Functions
131
+ # ============================================================================
132
+
133
+ def _l2_normalize(x: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
134
+ """L2 normalize a tensor along the last dimension."""
135
+ return x / (x.norm(p=2, dim=-1, keepdim=True) + eps)
136
+
137
+
138
+ def match_query_to_embeddings(
139
+ query: str | Image.Image,
140
+ target_embeddings_to_match: torch.Tensor,
141
+ top_k: int = 100
142
+ ) -> tuple[torch.Tensor, torch.Tensor]:
143
+ """
144
+ Match a query (text or image) to target embeddings.
145
+
146
+ Returns:
147
+ Tuple of (sorted_scores, sorted_indices)
148
+ """
149
+ with torch.inference_mode():
150
+ if isinstance(query, Image.Image):
151
+ query_embeddings = embed_model.encode_documents(images=[query])
152
+ else:
153
+ query_embeddings = embed_model.encode_queries([query])
154
+
155
+ cos_sim = _l2_normalize(query_embeddings) @ _l2_normalize(target_embeddings_to_match).T
156
+ cos_sim_flat = cos_sim.flatten()
157
+ sorted_indices = torch.argsort(cos_sim_flat, descending=True)[:top_k]
158
+ sorted_scores = cos_sim_flat[sorted_indices][:top_k]
159
+
160
+ return sorted_scores, sorted_indices
161
+
162
+
163
+ def rerank_samples(
164
+ dataset,
165
+ query_text: str,
166
+ sorted_indices: list | torch.Tensor,
167
+ num_samples_to_rerank: int,
168
+ rerank_model,
169
+ rerank_processor,
170
+ device: str = DEVICE,
171
+ text_column: str = "recipe_markdown",
172
+ image_column: str = "image",
173
+ dataset_split: str = "train",
174
+ ) -> tuple:
175
+ """
176
+ Rerank top samples using the vision-language reranker model.
177
+
178
+ Returns:
179
+ Tuple of (dataset_samples_to_rerank, rerank_sorted_indices)
180
+ """
181
+ top_indices = sorted_indices[:num_samples_to_rerank]
182
+ dataset_samples_to_rerank = dataset[dataset_split].select(top_indices)
183
+
184
+ texts_to_rerank = dataset_samples_to_rerank[text_column]
185
+ images_to_rerank = dataset_samples_to_rerank[image_column]
186
+
187
+ samples_to_rerank = [
188
+ {"question": query_text, "doc_text": text, "doc_image": image}
189
+ for text, image in zip(texts_to_rerank, images_to_rerank)
190
+ ]
191
+
192
+ batch_dict_rerank = rerank_processor.process_queries_documents_crossencoder(samples_to_rerank)
193
+ batch_dict_rerank = {
194
+ k: v.to(device) if isinstance(v, torch.Tensor) else v
195
+ for k, v in batch_dict_rerank.items()
196
+ }
197
+
198
+ with torch.inference_mode():
199
+ rerank_outputs = rerank_model(**batch_dict_rerank, return_dict=True)
200
+
201
+ rerank_logits = rerank_outputs.logits.squeeze(-1)
202
+ rerank_sorted_indices = torch.argsort(rerank_logits, descending=True)
203
+
204
+ return dataset_samples_to_rerank, rerank_sorted_indices
205
+
206
+
207
+ def generate_recipe_summary(
208
+ recipe_texts: list[str],
209
+ model = None,
210
+ processor = None,
211
+ max_new_tokens: int = 512
212
+ ) -> str:
213
+ """Generate a markdown summary of multiple recipes."""
214
+ if model is None:
215
+ model = qwen_model
216
+ if processor is None:
217
+ processor = qwen_processor
218
+
219
+ recipes_combined = ""
220
+ for i, recipe in enumerate(recipe_texts[:3], 1):
221
+ recipes_combined += f"\n\n--- RECIPE {i} ---\n{recipe}"
222
+
223
+ prompt = f"""You are a helpful culinary assistant. Below are {len(recipe_texts[:3])} recipes.
224
+ Please provide a brief markdown summary with:
225
+ - A short 1-2 sentence overview of each recipe
226
+ - Key ingredients highlighted
227
+ - Estimated difficulty (Easy/Medium/Hard)
228
+ - Which recipe might be best for a quick weeknight dinner
229
+
230
+ For example use the following format:
231
+
232
+ ```markdown
233
+ # Recipe summary
234
+
235
+ ## <recipe_name>
236
+
237
+ [details]
238
+
239
+ ## <recipe_name>
240
+
241
+ [details]
242
+
243
+ ## <recipe_name>
244
+
245
+ [details]
246
+ ```
247
+
248
+ Keep the summary concise and well-formatted in markdown. Return in ```markdown``` tags so it can be easily parsed.
249
+
250
+ <recipes>
251
+ {recipes_combined}
252
+ </recipes>
253
+
254
+ ## Summary:"""
255
+
256
+ messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
257
+
258
+ inputs = processor.apply_chat_template(
259
+ messages,
260
+ tokenize=True,
261
+ add_generation_prompt=True,
262
+ return_dict=True,
263
+ return_tensors="pt"
264
+ )
265
+ inputs = inputs.to(model.device)
266
+
267
+ with torch.no_grad():
268
+ generated_ids = model.generate(
269
+ **inputs,
270
+ max_new_tokens=max_new_tokens,
271
+ do_sample=True,
272
+ temperature=0.7,
273
+ top_p=0.9
274
+ )
275
+
276
+ generated_ids_trimmed = [
277
+ out_ids[len(in_ids):]
278
+ for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
279
+ ]
280
+
281
+ output_text = processor.batch_decode(
282
+ generated_ids_trimmed,
283
+ skip_special_tokens=True,
284
+ clean_up_tokenization_spaces=False
285
+ )[0]
286
+
287
+ return output_text.strip()
288
+
289
+
290
+ def _markdown_to_simple_html(markdown_text: str, max_reviews: int = 1) -> str:
291
+ """Convert recipe markdown to a simple HTML card."""
292
+ lines = markdown_text.strip().split('\n')
293
+
294
+ title = ""
295
+ description = ""
296
+ recipe_id = ""
297
+ cook_time = ""
298
+ num_ratings = ""
299
+ ingredients = []
300
+ steps = []
301
+ tags = []
302
+ reviews = []
303
+
304
+ current_section = None
305
+ in_ingredients = False
306
+ in_steps = False
307
+ in_reviews = False
308
+ in_tags = False
309
+ review_count = 0
310
+
311
+ for line in lines:
312
+ line = line.strip()
313
+
314
+ if line.startswith('# ') and not title:
315
+ title = line[2:].strip()
316
+ continue
317
+
318
+ if line.startswith('**ID:**'):
319
+ recipe_id = line.replace('**ID:**', '').strip()
320
+ continue
321
+ if line.startswith('**Time:**'):
322
+ cook_time = line.replace('**Time:**', '').strip()
323
+ continue
324
+ if line.startswith('**Number of Ratings:**'):
325
+ num_ratings = line.replace('**Number of Ratings:**', '').strip()
326
+ continue
327
+
328
+ if line.startswith('## '):
329
+ section_name = line[3:].strip().lower()
330
+ in_ingredients = section_name == 'ingredients'
331
+ in_steps = section_name.startswith('steps')
332
+ in_reviews = section_name == 'reviews'
333
+ in_tags = section_name == 'tags'
334
+ current_section = section_name
335
+ continue
336
+
337
+ if current_section == 'description' and line and not line.startswith('#'):
338
+ description = line
339
+ continue
340
+
341
+ if in_ingredients and line.startswith('- '):
342
+ ingredients.append(line[2:].strip())
343
+ continue
344
+
345
+ if in_steps and line and line[0].isdigit():
346
+ step_text = line.split('. ', 1)[-1] if '. ' in line else line
347
+ steps.append(step_text.strip())
348
+ continue
349
+
350
+ if in_tags and line.startswith('`'):
351
+ tag_list = [t.strip().strip('`') for t in line.split(',')]
352
+ tags.extend(tag_list)
353
+ continue
354
+
355
+ if in_reviews and line.startswith('> ') and review_count < max_reviews:
356
+ reviews.append(line[2:].strip())
357
+ review_count += 1
358
+ continue
359
+
360
+ html = f'''
361
+ <div style="border: 1px solid #ddd; border-radius: 8px; padding: 16px; margin: 4px; background: #fff; font-family: system-ui, -apple-system, sans-serif; font-size: 12px; height: 400px; overflow-y: auto; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
362
+ <div style="font-weight: bold; font-size: 14px; color: #333; margin-bottom: 8px;">{title}</div>
363
+
364
+ <div style="display: flex; gap: 12px; font-size: 11px; color: #666; margin-bottom: 10px; flex-wrap: wrap;">
365
+ {f'<span>⏱️ {cook_time}</span>' if cook_time else ''}
366
+ {f'<span>⭐ {num_ratings} ratings</span>' if num_ratings else ''}
367
+ {f'<span style="color: #999;">ID: {recipe_id}</span>' if recipe_id else ''}
368
+ </div>
369
+
370
+ <div style="color: #555; margin-bottom: 12px; font-style: italic; line-height: 1.4;">{description[:150]}{"..." if len(description) > 150 else ""}</div>
371
+
372
+ <div style="margin-bottom: 12px;">
373
+ <div style="font-weight: bold; font-size: 11px; color: #333; margin-bottom: 4px;">πŸ“ Ingredients</div>
374
+ <div style="color: #444; line-height: 1.5;">{", ".join(ingredients[:8])}{"..." if len(ingredients) > 8 else ""}</div>
375
+ </div>
376
+
377
+ <div style="margin-bottom: 12px;">
378
+ <div style="font-weight: bold; font-size: 11px; color: #333; margin-bottom: 4px;">πŸ‘¨β€πŸ³ Steps ({len(steps)} total)</div>
379
+ <ol style="margin: 0; padding-left: 20px; color: #444; line-height: 1.5;">
380
+ {"".join(f'<li style="margin-bottom: 4px;">{step[:80]}{"..." if len(step) > 80 else ""}</li>' for step in steps[:4])}
381
+ {f'<li style="color: #999;">...and {len(steps) - 4} more steps</li>' if len(steps) > 4 else ''}
382
+ </ol>
383
+ </div>
384
+ '''
385
+
386
+ if tags:
387
+ display_tags = tags[:5]
388
+ html += f'''
389
+ <div style="margin-bottom: 12px;">
390
+ <div style="font-weight: bold; font-size: 11px; color: #333; margin-bottom: 4px;">🏷️ Tags</div>
391
+ <div style="display: flex; flex-wrap: wrap; gap: 4px;">
392
+ {"".join(f'<span style="background: #f0f0f0; padding: 2px 6px; border-radius: 4px; font-size: 10px;">{tag}</span>' for tag in display_tags)}
393
+ {f'<span style="color: #999; font-size: 10px;">+{len(tags) - 5} more</span>' if len(tags) > 5 else ''}
394
+ </div>
395
+ </div>
396
+ '''
397
+
398
+ if reviews:
399
+ html += f'''
400
+ <div style="border-top: 1px solid #eee; padding-top: 10px; margin-top: 10px;">
401
+ <div style="font-weight: bold; font-size: 11px; color: #333; margin-bottom: 4px;">πŸ’¬ Review</div>
402
+ <div style="color: #555; font-size: 11px; line-height: 1.4; background: #f9f9f9; padding: 8px; border-radius: 4px; font-style: italic;">"{reviews[0][:200]}{"..." if len(reviews[0]) > 200 else ""}"</div>
403
+ </div>
404
+ '''
405
+
406
+ html += '</div>'
407
+ return html
408
+
409
+
410
+ def create_recipe_cards_html(
411
+ scores_and_samples: list[dict],
412
+ num_results: int = 3,
413
+ text_key: str = "text",
414
+ max_reviews: int = 1
415
+ ) -> str:
416
+ """Generate combined HTML cards from scored recipe samples."""
417
+ recipe_cards_html = []
418
+
419
+ for item in scores_and_samples[:num_results]:
420
+ sample = item["sample"]
421
+ markdown_text = sample.get(text_key, "") or sample.get("markdown", "")
422
+ card_html = _markdown_to_simple_html(markdown_text, max_reviews=max_reviews)
423
+ recipe_cards_html.append(f'<div style="flex: 1; min-width: 0;">{card_html}</div>')
424
+
425
+ combined_html = f'''
426
+ <div style="margin-top: 16px;">
427
+ <h3 style="font-family: system-ui, -apple-system, sans-serif; font-size: 16px; font-weight: 600; color: #333; margin-bottom: 12px;">Retrieved Texts</h3>
428
+ <div style="display: flex; flex-direction: row; gap: 12px; width: 100%;">
429
+ {"".join(recipe_cards_html)}
430
+ </div>
431
+ </div>
432
+ '''
433
+
434
+ return combined_html
435
+
436
+
437
+ # ============================================================================
438
+ # Main Retrieve Function
439
+ # ============================================================================
440
+
441
+ def retrieve(
442
+ query_text: str | None,
443
+ query_image: Image.Image | None,
444
+ rerank_option: str,
445
+ generate_summary_option: str
446
+ ):
447
+ """
448
+ Main retrieval function for the Gradio interface.
449
+
450
+ Args:
451
+ query_text: Text query input
452
+ query_image: Image query input (PIL Image)
453
+ rerank_option: "True" or "False" to enable reranking
454
+ generate_summary_option: "True" or "False" to enable summary generation
455
+
456
+ Returns:
457
+ Tuple of (image_gallery, recipe_cards_html, summary, timing_dict)
458
+ """
459
+ # Determine input query (prefer text over image)
460
+ if query_text and query_text.strip():
461
+ input_query = query_text
462
+ elif query_image is not None:
463
+ input_query = query_image
464
+ else:
465
+ raise gr.Error("Please provide either a text query or an image query.")
466
+
467
+ # === Retrieval ===
468
+ start_time_query_to_embed_match = time.time()
469
+ result_sorted_scores, result_sorted_indices = match_query_to_embeddings(
470
+ query=input_query,
471
+ target_embeddings_to_match=image_text_embeddings,
472
+ top_k=20
473
+ )
474
+ end_time_query_to_embed_match = time.time()
475
+
476
+ top_dataset_results_to_show = dataset["train"].select(result_sorted_indices)
477
+
478
+ scores_and_samples = [
479
+ {"score": round(score.item(), 4), "sample": sample}
480
+ for score, sample in zip(result_sorted_scores, top_dataset_results_to_show)
481
+ ]
482
+
483
+ output_image_gallery = [
484
+ (item["sample"]["image"], f'Score: {item["score"]}')
485
+ for item in scores_and_samples[:3]
486
+ ]
487
+
488
+ output_recipe_cards_html = create_recipe_cards_html(
489
+ scores_and_samples=scores_and_samples,
490
+ num_results=3,
491
+ text_key="recipe_markdown",
492
+ max_reviews=1
493
+ )
494
+
495
+ # === Reranking (optional) ===
496
+ if rerank_option == "True":
497
+ start_time_reranking = time.time()
498
+ dataset_samples_to_rerank, rerank_sorted_indicies = rerank_samples(
499
+ sorted_indices=result_sorted_indices,
500
+ dataset=dataset,
501
+ dataset_split="train",
502
+ query_text=input_query,
503
+ num_samples_to_rerank=20,
504
+ rerank_model=rerank_model,
505
+ rerank_processor=rerank_processor
506
+ )
507
+ end_time_reranking = time.time()
508
+ rerank_time = round(end_time_reranking - start_time_reranking, 4)
509
+
510
+ top_dataset_results_to_show = dataset_samples_to_rerank.select(rerank_sorted_indicies)
511
+ samples_and_rerank_changes = []
512
+ for new_rank, (sample, original_rank) in enumerate(zip(top_dataset_results_to_show, rerank_sorted_indicies)):
513
+ movement = new_rank - original_rank
514
+ if movement == 0:
515
+ movement_string = f"{movement}"
516
+ else:
517
+ movement_string = f"+{abs(movement)}" if movement < 0 else f"-{movement}"
518
+ rerank_string = f"Original rank: {original_rank} | New rank: {new_rank} | Movement: {movement_string}"
519
+ samples_and_rerank_changes.append({"sample": sample, "rerank_string": rerank_string})
520
+
521
+ output_image_gallery = [
522
+ (item["sample"]["image"], item["rerank_string"])
523
+ for item in samples_and_rerank_changes[:3]
524
+ ]
525
+ output_recipe_cards_html = create_recipe_cards_html(
526
+ scores_and_samples=samples_and_rerank_changes,
527
+ num_results=3,
528
+ text_key="recipe_markdown",
529
+ max_reviews=1
530
+ )
531
+ else:
532
+ rerank_time = "Reranking turned off"
533
+
534
+ # === Generation (optional) ===
535
+ if generate_summary_option == "True":
536
+ start_time_generation_output = time.time()
537
+ if rerank_option == "True":
538
+ recipe_texts = [item["sample"]["recipe_markdown"] for item in samples_and_rerank_changes[:3]]
539
+ else:
540
+ recipe_texts = [item["sample"]["recipe_markdown"] for item in scores_and_samples[:3]]
541
+
542
+ summary = generate_recipe_summary(recipe_texts)
543
+ summary = summary.replace("```markdown", "").replace("```", "")
544
+ end_time_generation_output = time.time()
545
+ generation_time = round(end_time_generation_output - start_time_generation_output, 4)
546
+ else:
547
+ generation_time = "Generation turned off"
548
+ summary = "Generation turned off, no summary created"
549
+
550
+ timing_dict = {
551
+ "query_embed_and_match_time": round(end_time_query_to_embed_match - start_time_query_to_embed_match, 4),
552
+ "rerank_time": rerank_time,
553
+ "generation_time": generation_time
554
+ }
555
+
556
+ return output_image_gallery, output_recipe_cards_html, summary, timing_dict
557
+
558
+
559
+ # ============================================================================
560
+ # Gradio Interface
561
+ # ============================================================================
562
+
563
+ with gr.Blocks(title="Multimodal RAG Demo") as demo:
564
+ gr.Markdown("""# πŸ‘οΈπŸ“‘ Multimodal RAG Demo with Nemotron Embed VL and Rerank VL
565
+
566
+ Input an image or text about food and get recipe images/text back.
567
+
568
+ This is a scalable workflow that can lend itself to many use cases such as business document retrieval, technical manual look ups and more.
569
+
570
+ By default it returns the top 3 results from a database of 10,000+ recipes. We've limited it to 3 for the demo but in practice you could return as many as you like.
571
+
572
+ * **Dataset used:** https://huggingface.co/datasets/mrdbourke/recipe-synthetic-images-10k
573
+ * **Embedding model used:** https://huggingface.co/nvidia/llama-nemotron-embed-vl-1b-v2
574
+ * **Note:** By default we use the image + text embeddings as we have access to image and text pairs in our dataset, and according to the [launch blog post](https://huggingface.co/blog/nvidia/llama-nemotron-vl-1b), these work the best.
575
+ * **Rerank model used:** https://huggingface.co/nvidia/llama-nemotron-rerank-vl-1b-v2
576
+ * **Generation model used:** https://huggingface.co/Qwen/Qwen3-VL-2B-Instruct (note: you could use a larger model such as [Nemotron v3](https://huggingface.co/collections/nvidia/nvidia-nemotron-v3), however, this will require more compute resources)
577
+ """)
578
+
579
+ with gr.Row():
580
+ with gr.Column(scale=1):
581
+ gr.Markdown("### Query Input")
582
+
583
+ query_text = gr.Textbox(
584
+ label="Text Query",
585
+ placeholder="e.g. 'dinner recipes with tomatoes'",
586
+ lines=2
587
+ )
588
+
589
+ query_image = gr.Image(
590
+ label="Image Query (optional)",
591
+ type="pil",
592
+ height=200
593
+ )
594
+
595
+ generate_summary_option = gr.Radio(
596
+ choices=["True", "False"],
597
+ value="False",
598
+ label="Generate recipe summary"
599
+ )
600
+
601
+ rerank_option = gr.Radio(
602
+ choices=["True", "False"],
603
+ value="False",
604
+ label="Rerank initial results? (note: reranking is for text queries only)"
605
+ )
606
+
607
+ search_btn = gr.Button("Search", variant="primary", size="lg")
608
+
609
+ with gr.Column(scale=2):
610
+ gr.Markdown("### Retrieved Results")
611
+
612
+ gallery_output = gr.Gallery(
613
+ label="Retrieved Recipe Images",
614
+ columns=3,
615
+ height="auto",
616
+ object_fit="cover",
617
+ show_label=True
618
+ )
619
+
620
+ recipes_html = gr.HTML(label="Retrieved Recipe Texts")
621
+
622
+ summary_generation = gr.Markdown(
623
+ label="Retrieved Recipe Summary (generated from top results)"
624
+ )
625
+
626
+ timing_output = gr.JSON(label="Timings")
627
+
628
+ gr.Markdown("### Example Queries")
629
+
630
+ gr.Examples(
631
+ examples=[
632
+ ["best omelette recipes", None, "False", "False"],
633
+ ["best omelette recipes", None, "False", "True"],
634
+ ["best omelette recipes", None, "True", "True"],
635
+ ["eggplant dip", None, "True", "True"],
636
+ [None, "kitchen_bench.png", "False", "True"]
637
+ ],
638
+ inputs=[query_text, query_image, rerank_option, generate_summary_option],
639
+ label="Example Queries"
640
+ )
641
+
642
+ search_btn.click(
643
+ fn=retrieve,
644
+ inputs=[query_text, query_image, rerank_option, generate_summary_option],
645
+ outputs=[gallery_output, recipes_html, summary_generation, timing_output]
646
+ )
647
+
648
+ if __name__ == "__main__":
649
+ demo.launch()
image_text_embeddings_10k.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5a2046580ef9d85201f138a14c0ebc60c7e1fc7a5ca4c933d0077d545afb064f
3
+ size 41353320
kitchen_bench.png ADDED

Git LFS Details

  • SHA256: 27f56ad6d53b85e7fb30ff0ebb94268865d3283da046d51e62b7fec8568b5246
  • Pointer size: 132 Bytes
  • Size of remote file: 8.78 MB
requirements.txt ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ transformers
3
+ gradio
4
+ datasets
5
+ safetensors
6
+ Pillow
7
+ accelerate
8
+ qwen-vl-utils