Spaces:
Sleeping
Sleeping
| 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_user_prompt, | |
| generate_rag_system_prompt, | |
| generate_eda_system_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): | |
| 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(cell_commands, 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 cell_commands | |
| ] | |
| with open(notebook_name, "w") as f: | |
| nbf.write(nb, f) | |
| logging.info(f"Notebook {notebook_name} created successfully") | |
| 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) | |
| features = content["features"] | |
| features_dict = {feature["name"]: feature["type"] for feature in features} | |
| return features_dict, first_rows_df | |
| except Exception as e: | |
| logging.error(f"Error fetching first rows: {e}") | |
| raise | |
| def get_txt_from_output(output): | |
| try: | |
| extracted_text = extract_content_from_output(output) | |
| content = json.loads(extracted_text) | |
| logging.info(content) | |
| return content | |
| except Exception as e: | |
| gr.Error("Error when parsing notebook, try again.") | |
| logging.error(f"Failed to fetch compatible libraries: {e}") | |
| raise | |
| def extract_content_from_output(output): | |
| patterns = [r"`json(.*?)`", r"```(.*?)```"] | |
| for pattern in patterns: | |
| match = re.search(pattern, output, re.DOTALL) | |
| if match: | |
| return match.group(1) | |
| try: | |
| index = output.index("```json") | |
| logging.info(f"Index: {index}") | |
| return output[index + 7 :] | |
| except ValueError: | |
| logging.error("Unable to generate Jupyter notebook.") | |
| raise | |
| 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, profile: gr.OAuthProfile | None): | |
| for messages in generate_cells(dataset_id, generate_eda_system_prompt, "eda"): | |
| yield messages, gr.update(visible=False), None # Keep button hidden | |
| yield ( | |
| messages, | |
| gr.update(visible=profile and dataset_id.split("/")[0] == profile.username), | |
| f"{dataset_id.replace('/', '-')}-eda.ipynb", | |
| ) | |
| def generate_rag_cells(dataset_id, profile: gr.OAuthProfile | None): | |
| for messages in generate_cells(dataset_id, generate_rag_system_prompt, "rag"): | |
| yield messages, gr.update(visible=False), None # Keep button hidden | |
| yield ( | |
| messages, | |
| gr.update(visible=profile and dataset_id.split("/")[0] == profile.username), | |
| f"{dataset_id.replace('/', '-')}-rag.ipynb", | |
| ) | |
| def generate_embedding_cells(dataset_id, profile: gr.OAuthProfile | None): | |
| for messages in generate_cells(dataset_id, generate_embedding_prompt, "embedding"): | |
| yield messages, gr.update(visible=False), None # Keep button hidden | |
| yield ( | |
| messages, | |
| gr.update(visible=profile and dataset_id.split("/")[0] == profile.username), | |
| f"{dataset_id.replace('/', '-')}-embedding.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..._") | |
| ] | |
| return | |
| 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="user", | |
| content=f"[See the notebook on the Hub]({link})", | |
| ) | |
| ] | |
| except Exception as e: | |
| logging.info("Failed to push notebook", e) | |
| yield history + [gr.ChatMessage(role="assistant", content=e)] | |
| def generate_cells(dataset_id, prompt_fn, notebook_type="eda"): | |
| 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] | |
| features, df = get_first_rows_as_df(dataset_id, first_config, first_split, 3) | |
| prompt = generate_user_prompt( | |
| features, df.head(5).to_dict(orient="records"), first_code | |
| ) | |
| messages = [gr.ChatMessage(role="user", content=prompt)] | |
| yield messages + [gr.ChatMessage(role="assistant", content="⏳ _Starting task..._")] | |
| prompt_messages = [ | |
| {"role": "system", "content": prompt_fn()}, | |
| {"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"<iframe src='https://huggingface.co/datasets/{dataset_id}/embed/viewer' width='80%' height='560px'></iframe>" | |
| 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('/', '-')}-{notebook_type}.ipynb" | |
| create_notebook_file(commands, notebook_name=notebook_name) | |
| messages.append( | |
| gr.ChatMessage(role="user", content="Here is the generated notebook file") | |
| ) | |
| yield messages | |
| messages.append( | |
| gr.ChatMessage( | |
| role="user", | |
| content=FileData(path=notebook_name, mime_type="application/x-ipynb+json"), | |
| ) | |
| ) | |
| yield messages | |
| def coming_soon_message(): | |
| return gr.Info("Coming 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="", | |
| ) | |
| def embed(name): | |
| if not name: | |
| return gr.Markdown("### No dataset provided") | |
| html_code = f""" | |
| <iframe | |
| src="https://huggingface.co/datasets/{name}/embed/viewer/default/train" | |
| frameborder="0" | |
| width="100%" | |
| height="350px" | |
| ></iframe> | |
| """ | |
| 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_rag_btn = gr.Button("Generate RAG 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_rag_btn.click( | |
| generate_rag_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(coming_soon_message, inputs=[], outputs=[]) | |
| push_btn.click( | |
| push_to_hub, | |
| inputs=[ | |
| chatbot, | |
| dataset_name, | |
| notebook_file, | |
| ], | |
| outputs=[chatbot], | |
| ) | |
| demo.launch() | |