ninjals commited on
Commit
a3cb9f2
Β·
verified Β·
1 Parent(s): 4442a18

Uploading Gradio multimodal RAG demo

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