update app.py

app.py

@@ -0,0 +1,103 @@
+import gradio as gr
+from datasets import load_dataset
+import re
+
+
+def load_dataset_demo(name):
+    try:
+        dataset = load_dataset(name)["train"].filter(lambda x: x["flags"])
+    except Exception as _:
+        dataset = load_dataset(name)["train"]
+    return dataset
+
+
+NAME_DATASETS = [
+    "Self-GRIT/selfrag_dataset-embed_query_instruct-Meta-Llama-3-70B-Instruct_temp-0.01",
+    "Self-GRIT/selfrag_dataset_mini-embed_query_instruct-Meta-Llama-3-70B-Instruct_temp-0.01",
+    "Self-GRIT/selfrag_dataset_mini-embed_query_instruct-Meta-Llama-3-8B-Instruct",
+]
+
+DATASETS = {name: load_dataset_demo(name) for name in NAME_DATASETS}
+INSTRUCTION_COL = "instruction"
+OUTPUT_COL = "output"
+OUTPUT_ORIGIN = "output_origin"
+
+
+def extract_pairs(text):
+    # Regex pattern to match <embed>...</embed><passage>...</passage> pairs
+    pattern = r"<embed>(.*?)</embed><passage>(.*?)</passage>"
+    # Find all matches in the text
+    matches = re.findall(pattern, text, re.DOTALL)
+    return matches
+
+
+def preprocess_qa_pairs(text):
+    qa_pairs = extract_pairs(text)
+    response = ""
+    if len(qa_pairs) == 0:
+        response = "No query-passage pairs found."
+    else:
+        for i, (query, passage) in enumerate(qa_pairs):
+            response += f"========================== QP-Pair {i+1} =============================\n"
+            response += f"Query:\n{query.strip()}\n"
+            response += f"Passage:\n{passage.strip()}\n\n"
+    return response
+
+
+def output_fn(dropdown, slider):
+    dataset = DATASETS[dropdown]
+    example = dataset[int(slider)]
+    return (
+        example[INSTRUCTION_COL],
+        example[OUTPUT_ORIGIN],
+        example[OUTPUT_COL],
+        preprocess_qa_pairs(example[OUTPUT_COL]),
+    )
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Explore Self-RAG Datasets")
+    with gr.Group():
+        with gr.Row():
+            with gr.Column():
+                dropdown = gr.Dropdown(
+                    NAME_DATASETS,
+                    value=[NAME_DATASETS[0]],
+                    multiselect=False,
+                    label="Dataset",
+                    info="Select the dataset name",
+                )
+            with gr.Column():
+                slider = gr.Slider(
+                    minimum=0,
+                    maximum=max([len(dataset) for _, dataset in DATASETS.items()]),
+                    step=1,
+                    label="#example",
+                    value=0,
+                )
+                button = gr.Button(value="Submit", variant="primary")
+    with gr.Group():
+        with gr.Row():
+            output_instruction = gr.Textbox(
+                label="Instruction", placeholder="Instruction", type="text"
+            )
+        with gr.Row():
+            with gr.Row():
+                output_self_rag = gr.Textbox(
+                    label="SELG-RAG output", placeholder="SELG-RAG output", type="text"
+                )
+                output_self_grit = gr.Textbox(
+                    label="SELF-GRIT output",
+                    placeholder="SELF-GRIT output",
+                    type="text",
+                )
+    with gr.Group():
+        output_qps = gr.Textbox(
+            label="Query-Passage Pairs", placeholder="Query-Passage Pairs", type="text"
+        )
+    button.click(
+        fn=output_fn,
+        inputs=[dropdown, slider],
+        outputs=[output_instruction, output_self_rag, output_self_grit, output_qps],
+    )
+    demo.launch(share=True, debug=True)