Show retrieved prompt for each NN1/NN2/NN3

app.py

@@ -14,9 +14,9 @@ TOP1_ID = SAMPLE_IDS[0] if SAMPLE_IDS else None
 
 
 def build_noise_demo(sample_id: str | None):
-    """Returns (baseline, nn1, nn2, nn3) each: (spec, bg_wav, fg_wav, m_wav)."""
+    """Returns (baseline, nn1, nn2, nn3) each: (spec, bg_wav, fg_wav, m_wav); nn1/nn2/nn3 also have prompt."""
     if not sample_id:
-        return (None,) * 16
+        return (None,) * 19
     data = get_noise_demo_paths(sample_id)
     out = []
     for key in ("baseline", "nn1", "nn2", "nn3"):
@@ -27,6 +27,9 @@ def build_noise_demo(sample_id: str | None):
             block.get("fg_wav"),
             block.get("m_wav"),
         ])
+        if key.startswith("nn"):
+            prompt = block.get("prompt", "") or ""
+            out.append(f"**Prompt:** {prompt}" if prompt else "")
     return tuple(out)
 
 
@@ -51,33 +54,37 @@ with gr.Blocks(title="NearestNeighbor Audio Demo", css=".gradio-container { max-
         label="Noise (ID)",
     )
 
-    def block_ui(title: str):
+    def block_ui(title: str, with_prompt: bool = False):
         with gr.Group():
             gr.Markdown(f"### {title}")
+            if with_prompt:
+                prompt_md = gr.Markdown(value="", elem_id=f"{title}_prompt")
             img = gr.Image(label=f"{title}", show_label=True)
             with gr.Row():
                 abg = gr.Audio(label="BG", show_label=True)
                 afg = gr.Audio(label="FG", show_label=True)
                 am = gr.Audio(label="Mix", show_label=True)
+        if with_prompt:
+            return prompt_md, img, abg, afg, am
         return img, abg, afg, am
 
     # Baseline
     bl_img, bl_bg, bl_fg, bl_m = block_ui("Baseline")
     gr.Markdown("---")
     # NN1
-    nn1_img, nn1_bg, nn1_fg, nn1_m = block_ui("NN1")
+    nn1_prompt, nn1_img, nn1_bg, nn1_fg, nn1_m = block_ui("NN1", with_prompt=True)
     gr.Markdown("---")
     # NN2
-    nn2_img, nn2_bg, nn2_fg, nn2_m = block_ui("NN2")
+    nn2_prompt, nn2_img, nn2_bg, nn2_fg, nn2_m = block_ui("NN2", with_prompt=True)
     gr.Markdown("---")
     # NN3
-    nn3_img, nn3_bg, nn3_fg, nn3_m = block_ui("NN3")
+    nn3_prompt, nn3_img, nn3_bg, nn3_fg, nn3_m = block_ui("NN3", with_prompt=True)
 
     all_outputs = [
         bl_img, bl_bg, bl_fg, bl_m,
-        nn1_img, nn1_bg, nn1_fg, nn1_m,
-        nn2_img, nn2_bg, nn2_fg, nn2_m,
-        nn3_img, nn3_bg, nn3_fg, nn3_m,
+        nn1_img, nn1_bg, nn1_fg, nn1_m, nn1_prompt,
+        nn2_img, nn2_bg, nn2_fg, nn2_m, nn2_prompt,
+        nn3_img, nn3_bg, nn3_fg, nn3_m, nn3_prompt,
     ]
 
     def on_noise_select(sid):

dataset_loader.py

@@ -105,13 +105,20 @@ def get_noise_demo_paths(bid: str) -> dict:
     """
     One block per method: baseline, nn1, nn2, nn3.
     Each block: one combined image (spec) + 3 audios (bg_wav, fg_wav, m_wav).
-    Returns { "baseline": {spec, bg_wav, fg_wav, m_wav}, "nn1": {...}, "nn2": {...}, "nn3": {...} }.
+    nn1/nn2/nn3 also include "prompt" (retrieved text). 
+    Returns { "baseline": {spec, bg_wav, fg_wav, m_wav}, "nn1": {..., prompt}, ... }.
     """
     inner = f"{ROOT_PREFIX}{bid}/{bid}"
     files = _find_files(inner)
     baseline_inner = f"{inner}/baseline"
     baseline_files = _find_files(baseline_inner) if any(f.startswith(baseline_inner) for f in files) else []
 
+    prompts = _load_json_from_repo(f"{inner}/temp_retrieval.json")
+    if not prompts:
+        prompts = _load_json_from_repo(f"{inner}/natural_prompts.json")
+    if not prompts:
+        prompts = []
+
     def collect_block(file_list: list, folder_prefix: str) -> dict:
         """From files under folder_prefix, get spec + bg_wav, fg_wav, m_wav."""
         spec = bg = fg = m = None
@@ -144,10 +151,12 @@ def get_noise_demo_paths(bid: str) -> dict:
                 break
     baseline_block = collect_block(baseline_files, baseline_prefix) if baseline_prefix else {"spec": None, "bg_wav": None, "fg_wav": None, "m_wav": None}
 
-    # NN1, NN2, NN3: generated_01_, generated_02_, generated_03_
+    # NN1, NN2, NN3: generated_01_, generated_02_, generated_03_ + prompt from retrieval JSON
     result = {"baseline": baseline_block}
     for i in range(1, 4):
         rel_prefix = f"generated_{i:02d}_"
         nn_files = [f for f in files if f.replace(inner + "/", "").startswith(rel_prefix)]
-        result[f"nn{i}"] = collect_block(nn_files, rel_prefix)
+        block = collect_block(nn_files, rel_prefix)
+        block["prompt"] = prompts[i - 1].get("prompt", "") if i <= len(prompts) else ""
+        result[f"nn{i}"] = block
     return result