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 from huggingface_hub import InferenceClient import json import re import pandas as pd from gradio.data_classes import FileData from utils.prompts import ( generate_mapping_prompt, generate_eda_prompt, generate_embedding_prompt, ) """ TODOs: - Need feedback on the output commands to validate if operations are appropiate to data types - Refactor - Make the notebook generation more dynamic, add loading components to do not freeze the UI - Fix errors: - When generating output - When parsing output - When pushing notebook - Add target tasks to choose for the notebook: - Exploratory data analysis - Auto training - RAG - etc. - Enable 'generate notebook' button only if dataset is available and supports library - First get compatible-libraries and let user choose the library """ # Configuration BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co" HEADERS = {"Accept": "application/json", "Content-Type": "application/json"} client = Client(headers=HEADERS) inference_client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct") logging.basicConfig(level=logging.INFO) def get_compatible_libraries(dataset: str): resp = client.get( f"{BASE_DATASETS_SERVER_URL}/compatible-libraries?dataset={dataset}" ) resp.raise_for_status() return resp.json() def create_notebook_file(cell_commands, notebook_name): nb = nbf.v4.new_notebook() nb["cells"] = [ nbf.v4.new_code_cell(command["source"]) if command["cell_type"] == "code" else nbf.v4.new_markdown_cell(command["source"]) for command in cell_commands ] with open(notebook_name, "w") as f: nbf.write(nb, f) logging.info(f"Notebook {notebook_name} created successfully") def push_notebook(file_path, dataset_id, token): notebook_name = "dataset_analysis.ipynb" api = HfApi(token=token) try: api.upload_file( path_or_fileobj=file_path, path_in_repo=notebook_name, repo_id=dataset_id, repo_type="dataset", ) link = f"https://huggingface.co/datasets/{dataset_id}/blob/main/{notebook_name}" return gr.HTML( value=f'See notebook', visible=True, ) except Exception as err: logging.error(f"Failed to push notebook: {err}") return gr.HTML(value="Failed to push notebook", visible=True) def get_first_rows_as_df(dataset: str, config: str, split: str, limit: int): 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) features = content["features"] features_dict = {feature["name"]: feature["type"] for feature in features} return features_dict, first_rows_df def get_txt_from_output(output): extracted_text = content_from_output(output) content = json.loads(extracted_text) logging.info(content) return content def content_from_output(output): pattern = r"`json(.*?)`" match = re.search(pattern, output, re.DOTALL) if not match: pattern = r"```(.*?)```" match = re.search(pattern, output, re.DOTALL) if not match: try: index = output.index("```json") logging.info(f"Index: {index}") return output[index + 7 :] except: pass raise Exception("Unable to generate jupyter notebook.") return match.group(1) def generate_eda_cells(dataset_id): for messages in generate_cells(dataset_id, generate_eda_prompt): yield messages, gr.update(visible=False), None # Keep button hidden yield messages, gr.update(visible=True), f"{dataset_id.replace('/', '-')}.ipynb" def generate_embedding_cells(dataset_id): for messages in generate_cells(dataset_id, generate_embedding_prompt): yield messages, gr.update(visible=False), None # Keep button hidden yield messages, gr.update(visible=True), f"{dataset_id.replace('/', '-')}.ipynb" def push_to_hub( history, dataset_id, notebook_file, profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None, ): logging.info(f"Pushing notebook to hub: {dataset_id} on file {notebook_file}") if not profile or not oauth_token: yield history + [ gr.ChatMessage(role="assistant", content="⏳ _Login to push to hub..._") ] logging.info(f"Profile: {profile}, token: {oauth_token.token}") notebook_name = "dataset_analysis.ipynb" api = HfApi(token=oauth_token.token) try: logging.info(f"About to push {notebook_file} - {notebook_name} - {dataset_id}") api.upload_file( path_or_fileobj=notebook_file, path_in_repo=notebook_name, repo_id=dataset_id, repo_type="dataset", ) link = f"https://huggingface.co/datasets/{dataset_id}/blob/main/{notebook_name}" logging.info(f"Notebook pushed to hub: {link}") yield history + [ gr.ChatMessage( role="assistant", content=f"[Here is the generated notebook]({link})" ) ] except Exception as err: logging.info("Failed to push notebook", err) yield history + [gr.ChatMessage(role="assistant", content=err)] def generate_cells(dataset_id, prompt_fn): 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 [] if not libraries: gr.Error("Dataset not compatible with pandas library.") logging.error(f"Dataset not compatible with pandas library") return gr.File(visible=False), gr.Row.update(visible=False) pandas_library = next( (lib for lib in libraries.get("libraries", []) if lib["library"] == "pandas"), None, ) if not pandas_library: gr.Error("Dataset not compatible with pandas library.") return [] 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] logging.info(f"First config: {first_config} - first split: {first_split}") first_file = f"hf://datasets/{dataset_id}/{first_config_loading_code['arguments']['splits'][first_split]}" logging.info(f"First split file: {first_file}") features, df = get_first_rows_as_df(dataset_id, first_config, first_split, 3) sample_data = df.head(5).to_dict(orient="records") prompt = prompt_fn(features, sample_data, first_code) messages = [gr.ChatMessage(role="user", content=prompt)] yield messages + [gr.ChatMessage(role="assistant", content="⏳ _Starting task..._")] prompt_messages = [{"role": "user", "content": prompt}] output = inference_client.chat_completion( messages=prompt_messages, stream=True, max_tokens=2500 ) generated_text = "" current_line = "" for chunk in output: current_line += chunk.choices[0].delta.content if current_line.endswith("\n"): generated_text += current_line messages.append(gr.ChatMessage(role="assistant", content=current_line)) current_line = "" yield messages yield messages logging.info("---> Formated prompt") formatted_prompt = generate_mapping_prompt(generated_text) logging.info(formatted_prompt) prompt_messages = [{"role": "user", "content": formatted_prompt}] yield messages + [ gr.ChatMessage(role="assistant", content="⏳ _Generating notebook..._") ] output = inference_client.chat_completion( messages=prompt_messages, stream=False, max_tokens=2500 ) cells_txt = output.choices[0].message.content logging.info("---> Model output") logging.info(cells_txt) commands = get_txt_from_output(cells_txt) html_code = f"" # Adding dataset viewer on the first part commands.insert( 0, { "cell_type": "code", "source": f'from IPython.display import HTML\n\ndisplay(HTML("{html_code}"))', }, ) commands.insert(0, {"cell_type": "markdown", "source": "# Dataset Viewer"}) notebook_name = f"{dataset_id.replace('/', '-')}.ipynb" create_notebook_file(commands, notebook_name=notebook_name) messages.append( gr.ChatMessage(role="user", content="Here is the generated notebook") ) yield messages messages.append( gr.ChatMessage( role="user", content=FileData(path=notebook_name, mime_type="application/x-ipynb+json"), ) ) yield messages def comming_soon_message(): gr.Info("Comming soon") with gr.Blocks(fill_height=True) as demo: gr.Markdown("# 🤖 Dataset notebook creator 🕵️") with gr.Row(): with gr.Column(scale=1): dataset_name = HuggingfaceHubSearch( label="Hub Dataset ID", placeholder="Search for dataset id on Huggingface", search_type="dataset", value="", ) @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) with gr.Row(): generate_eda_btn = gr.Button("Generate EDA notebook") generate_embedding_btn = gr.Button("Generate Embeddings notebook") generate_training_btn = gr.Button("Generate Training notebook") with gr.Column(): chatbot = gr.Chatbot( label="Results", type="messages", avatar_images=( None, None, ), ) with gr.Row(): login_btn = gr.LoginButton() push_btn = gr.Button("Push to hub", visible=False) notebook_file = gr.File(visible=False) generate_eda_btn.click( generate_eda_cells, inputs=[dataset_name], outputs=[chatbot, push_btn, notebook_file], ) generate_embedding_btn.click( generate_embedding_cells, inputs=[dataset_name], outputs=[chatbot, push_btn, notebook_file], ) generate_training_btn.click(comming_soon_message, inputs=[], outputs=[]) push_btn.click( push_to_hub, inputs=[ chatbot, dataset_name, notebook_file, ], outputs=[chatbot], ) demo.launch()