download

app.py

@@ -3,34 +3,61 @@ from datasets import load_dataset
 import nbformat
 from nbconvert import HTMLExporter
 from traitlets.config import Config
+import os
+import shutil
 
-# we want to export the outputs too:
+# Configuration for HTMLExporter
 config = Config()
 config.HTMLExporter.preprocessors = ["nbconvert.preprocessors.ExtractOutputPreprocessor"]
-
-#  Instantiate the exporter. We use the `classic` template for now; we'll get into more details
-# later about how to customize the exporter further.
 html_exporter = HTMLExporter(config=config, template_name="classic")
 
 ds = load_dataset("data-agents/kaggle-notebooks", split="train", streaming=True)
 ds_iter = iter(ds)
 
+TMP_DIR = './tmp/'
+
+def reset_tmp_folder():
+    if os.path.exists(TMP_DIR):
+        shutil.rmtree(TMP_DIR)
+    os.makedirs(TMP_DIR)
+
 def parse_notebook():
-    notebook_string = next(ds_iter)["text"]
+    reset_tmp_folder()
+    
+    notebook_data = next(ds_iter)
+    notebook_string = notebook_data["text"]
+    notebook_id = notebook_data["id"].split("/")[-1]
+    
+    # Save the notebook string to a file
+    with open(os.path.join(TMP_DIR, notebook_id), 'w') as f:
+        f.write(notebook_string)
+    
     notebook_parsed = nbformat.reads(notebook_string, as_version=4)
     (notebook_body, resources) = html_exporter.from_notebook_node(notebook_parsed)
+    
     print("Resources:", resources["outputs"])
     return notebook_body
 
+def download_notebook():
+    for filename in os.listdir(TMP_DIR):
+        if filename.endswith('.ipynb'):
+            path = os.path.join(TMP_DIR, filename)
+            print(path)
+            return path
 
 with gr.Blocks() as demo:
     gr.Markdown("# Kaggle Notebooks")
-    button = gr.Button("Show next!")
-
+    
+    with gr.Row():
+        button = gr.Button("Show next!")
+        button_download = gr.Button("Download notebook")
+    
     with gr.Row():
         with gr.Column():
             html = gr.HTML("")
-
+    
     button.click(fn=parse_notebook, inputs=[], outputs=[html])
+    button_download.click(fn=download_notebook, inputs=[], outputs=[])
     demo.load(fn=parse_notebook, inputs=[], outputs=[html])
+
 demo.launch()