librarian-bots/function-calling-datasets

1import os2import re3 4import gradio as gr5from huggingface_hub import get_collection6 7 8def extract_collection_id(input_text):9    if url_match := re.match(r"https://huggingface\.co/collections/(.+)$", input_text):10        return url_match[1]11 12    # Check if input is already in the correct format13    return input_text if re.match(r"^[\w-]+/[\w-]+", input_text) else None14 15 16def load_collection():17    collection_input = os.getenv("COLLECTION_SLUG_OR_URL")18    if not collection_input:19        raise ValueError("COLLECTION_SLUG_OR_URL environment variable is not set.")20 21    collection_id = extract_collection_id(collection_input)22    if not collection_id:23        raise ValueError(24            "Invalid collection ID or URL in COLLECTION_SLUG_OR_URL environment variable."25        )26 27    collection = get_collection(collection_id)28    if dataset_ids := [29        item.item_id for item in collection.items if item.item_type == "dataset"30    ]:31        return dataset_ids, collection_id32    else:33        raise ValueError("No datasets found in this collection.")34 35 36def display_dataset(dataset_ids, index):37    dataset_id = dataset_ids[index]38    return gr.HTML(f"""<iframe39    src="https://huggingface.co/datasets/{dataset_id}/embed/viewer"40    frameborder="0"41    width="100%"42    height="560px"43></iframe>""")44 45 46def navigate_dataset(dataset_ids, index, direction):47    new_index = (index + direction) % len(dataset_ids)48    return (49        new_index,50        f"Dataset {new_index + 1} of {len(dataset_ids)}: {dataset_ids[new_index]}",51    )52 53 54def get_display_name(collection_id):55    # Pattern to match username/repo-name with an optional ID of 16 or more hexadecimal characters56    pattern = r"^(.+?)-([a-f0-9]{16,})$"57    if match := re.match(pattern, collection_id):58        return match[1]59    else:60        # If no match, return the original61        return collection_id62 63 64try:65    dataset_ids, collection_id = load_collection()66    display_name = get_display_name(collection_id)67 68    with gr.Blocks() as demo:69        gr.Markdown(f"<h1>Dataset Viewer for Collection: {display_name}</h1>")70        gr.Markdown(71            f"[View full collection on Hugging Face](https://huggingface.co/collections/{collection_id})"72        )73 74        gr.Markdown("""75        This app allows you to browse and view datasets from a specific Hugging Face collection. 76        Use the 'Previous' and 'Next' buttons to navigate through the datasets in the collection. 77        See below for how to set up this app for a different collection.""")78 79        index_state = gr.State(value=0)80 81        with gr.Row():82            left_btn = gr.Button("Previous")83            right_btn = gr.Button("Next")84 85        dataset_info = gr.Markdown(f"Dataset 1 of {len(dataset_ids)}: {dataset_ids[0]}")86        iframe_output = gr.HTML()87        gr.Markdown("""**Note**: This space is currently set up to display datasets from a specific collection. 88        If you'd like to use it for a different collection:89        1. Duplicate this space90        2. In your duplicated space, set the `COLLECTION_SLUG_OR_URL` environment variable to your desired collection ID or URL91        3. Your new space will then display datasets from your chosen collection!92        Checkout the [docs](https://huggingface.co/docs/hub/datasets-viewer-embed) for other ways to use the iframe viewer.93        """)94        left_btn.click(95            navigate_dataset,96            inputs=[gr.State(dataset_ids), index_state, gr.Number(-1, visible=False)],97            outputs=[index_state, dataset_info],98        )99        right_btn.click(100            navigate_dataset,101            inputs=[gr.State(dataset_ids), index_state, gr.Number(1, visible=False)],102            outputs=[index_state, dataset_info],103        )104 105        index_state.change(106            display_dataset,107            inputs=[gr.State(dataset_ids), index_state],108            outputs=[iframe_output],109        )110 111        # Initialize the display with the first dataset112        demo.load(113            fn=lambda: display_dataset(dataset_ids, 0),114            inputs=None,115            outputs=[iframe_output],116        )117 118    if __name__ == "__main__":119        demo.launch()120 121except Exception as e:122    print(f"Error: {str(e)}")123    print(124        "Please set the COLLECTION_SLUG_OR_URL environment variable with a valid collection ID or URL."125    )126