Integration with compatible-libraries and other commands

app.py

@@ -3,30 +3,37 @@ from gradio_huggingfacehub_search import HuggingfaceHubSearch
 import nbformat as nbf
 from huggingface_hub import HfApi
 from httpx import Client
+import logging
+
 
 """
 TODOs:
-- Handle erros
 - Add more commands to the notebook
 - Parametrize the commands (Move to another file)
-- How to handle configs and splits? -> Got from /compatible-libraries
-- Let user choose the framework
-- Use an LLM to suggest commands
+- Let user choose the framework and get if from /compatible-libraries
+- Use an LLM to suggest commands by column types
 - Add commands for auto training
-- Improve logs
 - Enable 'generate notebook' button only if dataset is available and supports library
 """
 
+# Configuration
 BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co"
-headers = {"Accept": "application/json", "Content-Type": "application/json"}
-client = Client(headers=headers)
+HEADERS = {"Accept": "application/json", "Content-Type": "application/json"}
+client = Client(headers=HEADERS)
+
+logging.basicConfig(level=logging.INFO)
 
 
 def get_compatible_libraries(dataset: str):
-    resp = client.get(
-        f"{BASE_DATASETS_SERVER_URL}/compatible-libraries?dataset={dataset}"
-    )
-    return resp.json()
+    try:
+        resp = client.get(
+            f"{BASE_DATASETS_SERVER_URL}/compatible-libraries?dataset={dataset}"
+        )
+        resp.raise_for_status()
+        return resp.json()
+    except Exception as err:
+        logging.error(f"Failed to fetch compatible libraries: {err}")
+        return None
 
 
 def create_notebook_file(cell_commands, notebook_name):
@@ -35,41 +42,56 @@ def create_notebook_file(cell_commands, notebook_name):
 
     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)
-    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}"
-    html = f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline; text-decoration-style: dotted;">See notebook</a>'
-
-    return gr.HTML(value=html, visible=True)
+    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'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline; text-decoration-style: dotted;">See notebook</a>',
+            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 generate_notebook(dataset_id):
     first_code = f"import pandas as pd\n\ndf = pd.read_parquet('hf://datasets/{dataset_id}/data/train-00000-of-00001.parquet')"
-    try:
-        libraries = get_compatible_libraries(dataset_id)["libraries"]
-    except Exception as err:
-        print(f"Error: {err}")
+    libraries = get_compatible_libraries(dataset_id)
+
+    if not libraries:
         return gr.File(visible=False), gr.Row.update(visible=False)
-    if pandas_library := next(
-        (element for element in libraries if element["library"] == "pandas"), None
-    ):
+
+    pandas_library = next(
+        (lib for lib in libraries.get("libraries", []) if lib["library"] == "pandas"),
+        None,
+    )
+    if pandas_library:
         first_code = pandas_library["loading_codes"][0]["code"]
     else:
         return gr.File(visible=False), gr.Row.update(visible=False)
 
+    html_code = f"<iframe src='https://huggingface.co/datasets/{dataset_id}/embed/viewer' width='80%' height='560px'></iframe>"
     commands = [
         "!pip install pandas",
         first_code,
         "df.head()",
+        f'from IPython.display import HTML\n\ndisplay(HTML("{html_code}"))',
+        "print(df.shape)",
+        "df.columns",
+        "df.describe()",
+        "df.info()",
+        # TODO: Generate more commands according to column types for EDA and then for auto training?
     ]
     notebook_name = f"{dataset_id.replace('/', '-')}.ipynb"
     create_notebook_file(commands, notebook_name=notebook_name)
@@ -103,7 +125,7 @@ with gr.Blocks() as demo:
     download_link = gr.File(label="Download notebook", visible=False)
     with gr.Row(visible=False) as auth_page:
         with gr.Column():
-            auth_title = gr.Markdown(
+            gr.Markdown(
                 "Want to push to hub? Enter your token ([settings](https://huggingface.co/settings/tokens)):"
             )
             token_box = gr.Textbox(