import gradio as gr from gradio_huggingfacehub_search import HuggingfaceHubSearch import nbformat as nbf from huggingface_hub import HfApi from httpx import Client import logging import pandas as pd from utils.notebook_utils import ( replace_wildcards, load_json_files_from_folder, ) from dotenv import load_dotenv import os from nbconvert import HTMLExporter import uuid load_dotenv() HF_TOKEN = os.getenv("HF_TOKEN") assert HF_TOKEN is not None, "You need to set HF_TOKEN in your environment variables" NOTEBOOKS_REPOSITORY = os.getenv("NOTEBOOKS_REPOSITORY") assert ( NOTEBOOKS_REPOSITORY is not None ), "You need to set NOTEBOOKS_REPOSITORY in your environment variables" URL = "https://huggingface.co/spaces/asoria/auto-notebook-creator" BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co" HEADERS = {"Accept": "application/json", "Content-Type": "application/json"} client = Client(headers=HEADERS) logging.basicConfig(level=logging.INFO) # TODO: Validate notebook templates format folder_path = "notebooks" notebook_templates = load_json_files_from_folder(folder_path) logging.info(f"Available notebooks {notebook_templates.keys()}") def get_compatible_libraries(dataset: str): try: response = client.get( f"{BASE_DATASETS_SERVER_URL}/compatible-libraries?dataset={dataset}" ) response.raise_for_status() return response.json() except Exception as e: logging.error(f"Error fetching compatible libraries: {e}") raise def create_notebook_file(cells, notebook_name): nb = nbf.v4.new_notebook() nb["cells"] = [ nbf.v4.new_code_cell( cmd["source"] if isinstance(cmd["source"], str) else "\n".join(cmd["source"]) ) if cmd["cell_type"] == "code" else nbf.v4.new_markdown_cell(cmd["source"]) for cmd in cells ] with open(notebook_name, "w") as f: nbf.write(nb, f) logging.info(f"Notebook {notebook_name} created successfully") html_exporter = HTMLExporter() html_data, _ = html_exporter.from_notebook_node(nb) return html_data def get_first_rows_as_df(dataset: str, config: str, split: str, limit: int): try: resp = client.get( f"{BASE_DATASETS_SERVER_URL}/first-rows?dataset={dataset}&config={config}&split={split}" ) resp.raise_for_status() content = resp.json() rows = content["rows"] rows = [row["row"] for row in rows] first_rows_df = pd.DataFrame.from_dict(rows).sample(frac=1).head(limit) return first_rows_df except Exception as e: logging.error(f"Error fetching first rows: {e}") raise def longest_string_column(df): longest_col = None max_length = 0 for col in df.select_dtypes(include=["object", "string"]): max_col_length = df[col].str.len().max() if max_col_length > max_length: max_length = max_col_length longest_col = col return longest_col def _push_to_hub( dataset_id, notebook_file, ): logging.info(f"Pushing notebook to hub: {dataset_id} on file {notebook_file}") notebook_name = notebook_file.split("/")[-1] api = HfApi(token=HF_TOKEN) try: logging.info(f"About to push {notebook_file} - {dataset_id}") api.upload_file( path_or_fileobj=notebook_file, path_in_repo=notebook_name, repo_id=NOTEBOOKS_REPOSITORY, repo_type="dataset", ) except Exception as e: logging.info("Failed to push notebook", e) raise def generate_cells(dataset_id, notebook_title): logging.info(f"Generating {notebook_title} notebook for dataset {dataset_id}") cells = notebook_templates[notebook_title]["notebook_template"] notebook_type = notebook_templates[notebook_title]["notebook_type"] dataset_types = notebook_templates[notebook_title]["dataset_types"] try: libraries = get_compatible_libraries(dataset_id) except Exception as err: gr.Error("Unable to retrieve dataset info from HF Hub.") logging.error(f"Failed to fetch compatible libraries: {err}") return "", "## ❌ This dataset is not accessible from the Hub ❌" if not libraries: logging.error(f"Dataset not compatible with pandas library - not libraries") return "", "## ❌ This dataset is not compatible with pandas library ❌" pandas_library = next( (lib for lib in libraries.get("libraries", []) if lib["library"] == "pandas"), None, ) if not pandas_library: logging.error("Dataset not compatible with pandas library - not pandas library") return "", "## ❌ This dataset is not compatible with pandas library ❌" first_config_loading_code = pandas_library["loading_codes"][0] first_code = first_config_loading_code["code"] first_config = first_config_loading_code["config_name"] first_split = list(first_config_loading_code["arguments"]["splits"].keys())[0] df = get_first_rows_as_df(dataset_id, first_config, first_split, 3) longest_col = longest_string_column(df) html_code = f"" wildcards = ["{dataset_name}", "{first_code}", "{html_code}", "{longest_col}"] replacements = [dataset_id, first_code, html_code, longest_col] has_numeric_columns = len(df.select_dtypes(include=["number"]).columns) > 0 has_categoric_columns = len(df.select_dtypes(include=["object"]).columns) > 0 valid_dataset = False if "text" in dataset_types and has_categoric_columns: valid_dataset = True if "numeric" in dataset_types and has_numeric_columns: valid_dataset = True if not valid_dataset: logging.error( f"Dataset does not have the column types needed for this notebook which expects to have {dataset_types} data types." ) return ( "", f"## ❌ This dataset does not have {dataset_types} columns, which are required for this notebook type ❌", ) cells = replace_wildcards( cells, wildcards, replacements, has_numeric_columns, has_categoric_columns ) notebook_name = ( f"{dataset_id.replace('/', '-')}-{notebook_type}-{uuid.uuid4()}.ipynb" ) html_content = create_notebook_file(cells, notebook_name=notebook_name) _push_to_hub(dataset_id, notebook_name) notebook_link = f"https://colab.research.google.com/#fileId=https%3A//huggingface.co/datasets/{NOTEBOOKS_REPOSITORY}/blob/main/{notebook_name}" return ( html_content, f"[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)]({notebook_link})", ) css = """ .prose :where(pre):not(:where([class~=not-prose],[class~=not-prose] *)) { background-color: var(--table-even-background-fill); /* Fix dark mode */ } """ with gr.Blocks(css=css) as demo: gr.Markdown("# 🤖 Dataset notebook creator 🕵️") gr.Markdown(f"[![Notebooks: {len(notebook_templates)}](https://img.shields.io/badge/Notebooks-{len(notebook_templates)}-blue.svg)]({URL}/tree/main/notebooks)") gr.Markdown(f"[![Contribute a Notebook](https://img.shields.io/badge/Contribute%20a%20Notebook-8A2BE2)]({URL}/blob/main/CONTRIBUTING.md)") text_input = gr.Textbox(label="Suggested notebook type", visible=False) gr.Markdown("## 1. Select and preview a dataset from Huggingface Hub") dataset_name = HuggingfaceHubSearch( label="Hub Dataset ID", placeholder="Search for dataset id on Huggingface", search_type="dataset", value="", ) dataset_samples = gr.Examples( examples=[ [ "scikit-learn/iris", "Try this dataset for Exploratory Data Analysis", ], [ "infinite-dataset-hub/GlobaleCuisineRecipes", "Try this dataset for Embeddings generation", ], [ "infinite-dataset-hub/GlobalBestSellersSummaries", "Try this dataset for RAG generation", ], ], inputs=[dataset_name, text_input], cache_examples=False, ) @gr.render(inputs=dataset_name) def embed(name): if not name: return gr.Markdown("### No dataset provided") html_code = f""" """ return gr.HTML(value=html_code, elem_classes="viewer") gr.Markdown("## 2. Select the type of notebook you want to generate") notebook_type = gr.Dropdown( choices=notebook_templates.keys(), label="Notebook type", value="Text Embeddings", ) generate_button = gr.Button("Generate Notebook", variant="primary") gr.Markdown("## 3. Notebook result + Open in Colab") go_to_notebook = gr.Markdown() code_component = gr.HTML() generate_button.click( generate_cells, inputs=[dataset_name, notebook_type], outputs=[code_component, go_to_notebook], ) gr.Markdown("🚧 Note: Some code may not be compatible with datasets that contain binary data or complex structures. 🚧") demo.launch()