visualization/components/example_browser.py

"""
Gradio widgets for browsing datasets and examples in the interpretability app.

What this module provides:
- create_example_browser: a Dataset-like widget to preview/select examples (dev/public modes).
- create_dataset_selector: a dropdown to pick which dataset to browse.

Expected backends:
- loader.data.list_datasets()
- loader.data.get_examples(dataset, n=...)
"""

import gradio as gr
from typing import List, Dict


def create_example_browser(mode: str = "public") -> gr.Dataset:
    """
    Create an example browser (grid/list) for selecting samples.

    Args:
        mode: "dev" or "public" — in "public" you typically show curated/precomputed
              examples; in "dev" you may show more free-form or larger lists.

    Expected usage:
        datasets = list_datasets()             # from loader/data.py
        examples  = get_examples(dataset, n=10)

    Returns:
        A gr.Dataset component configured for example selection.
    """
    mode = (mode or "public").lower()
    label = "Examples" if mode == "public" else "Example Browser (dev)"
    samples_per_page = 6 if mode == "public" else 12
    return gr.Dataset(
        components=[
            gr.Textbox(label="Context", lines=4, interactive=False),
            gr.Textbox(label="Prompt", lines=2, interactive=False),
            gr.Textbox(label="Answer", lines=2, interactive=False),
        ],
        headers=["context", "prompt", "answer"],
        label=label,
        samples=None,
        samples_per_page=samples_per_page,
        type="index",
        elem_id=f"example-browser-{mode}",
    )


def create_dataset_selector() -> gr.Dropdown:
    """
    Create a dataset selector dropdown.

    Expected usage:
        datasets = list_datasets()             # from loader/data.py
        dropdown.choices = datasets

    Returns:
        A gr.Dropdown component for choosing a dataset.
    """
    return gr.Dropdown(
        choices=[],
        label="Dataset",
        interactive=True,
        allow_custom_value=False,
        elem_id="dataset-selector",
    )