Initial Gradient Notebook support

README.md

@@ -35,7 +35,7 @@
 
 <br>
 
-Generate custom template code for PyTorch & sklearn, using a simple web UI built with [streamlit](https://www.streamlit.io/). traingenerator offers multiple options for preprocessing, model setup, training, and visualization (using Tensorboard or comet.ml). It exports to .py, Jupyter Notebook, or  [Google Colab](https://colab.research.google.com/). The perfect tool to jumpstart your next machine learning project! ✨
+Generate custom template code for PyTorch & sklearn, using a simple web UI built with [streamlit](https://www.streamlit.io/). traingenerator offers multiple options for preprocessing, model setup, training, and visualization (using Tensorboard or comet.ml). It exports to .py, Jupyter Notebook, [Google Colab](https://colab.research.google.com/), or [Paperspace Gradient](https://gradient.paperspace.com/). The perfect tool to jumpstart your next machine learning project! ✨
 
 <br>
 
@@ -53,16 +53,16 @@ cd traingenerator
 pip install -r requirements.txt
 ```
 
-*Optional: For the "Open in Colab" button to work you need to set up a Github repo 
-where the notebook files can be stored (Colab can only open public files if 
+*Optional: For the "Open in Colab" and "Open in Gradient" buttons to work you need to set up a Github repo 
+where the notebook files can be stored (Colab and Gradient can only open public files if 
 they are on Github). After setting up the repo, create a file `.env` with content:*
 
 ```bash
 GITHUB_TOKEN=<your-github-access-token>
 REPO_NAME=<user/notebooks-repo>
 ```
 
-*If you don't set this up, the app will still work but the "Open in Colab" button 
+*If you don't set this up, the app will still work but the "Open in Colab" and "Open in Gradient" buttons 
 will only show an error message.*
 
 
@@ -92,7 +92,7 @@ To update the deployed app, commit your changes and run:
 git push heroku main
 ```
 
-*Optional: If you set up a Github repo to enable the "Open in Colab" button (see above),
+*Optional: If you set up a Github repo to enable the "Open in Colab" or "Open in Gradient" buttons (see above),
 you also need to run:*
 
 ```

app/main.py

@@ -24,12 +24,13 @@
 )
 
 
-# Set up github access for "Open in Colab" button.
+# Set up github access for "Open in Colab" and "Open in Gradient" buttons.
 load_dotenv()  # load environment variables from .env file
 if os.getenv("GITHUB_TOKEN") and os.getenv("REPO_NAME"):
     g = Github(os.getenv("GITHUB_TOKEN"))
     repo = g.get_repo(os.getenv("REPO_NAME"))
     colab_enabled = True
+    gradient_enabled = True
 
     def add_to_colab(notebook):
         """Adds notebook to Colab by pushing it to Github repo and returning Colab link."""
@@ -42,9 +43,21 @@ def add_to_colab(notebook):
         colab_link = f"http://colab.research.google.com/github/{os.getenv('REPO_NAME')}/blob/main/notebooks/{notebook_id}/generated-notebook.ipynb"
         return colab_link
 
+    def add_to_gradient(notebook):
+        """Adds notebook to Colab by pushing it to Github repo and returning Colab link."""
+        notebook_id = str(uuid.uuid4())
+        repo.create_file(
+            f"notebooks/{notebook_id}/generated-notebook.ipynb",
+            f"Added notebook {notebook_id}",
+            notebook,
+        )
+        gradient_link = f"http://console.paperspace.com/github/{os.getenv('REPO_NAME')}/blob/main/notebooks/{notebook_id}/generated-notebook.ipynb"
+        return gradient_link
+
 
 else:
     colab_enabled = False
+    gradient_enabled = False
 
 
 # Display header.
@@ -88,6 +101,7 @@ def add_to_colab(notebook):
     st.write("")  # add vertical space
     col1, col2, col3 = st.beta_columns(3)
     open_colab = col1.button("🚀 Open in Colab")  # logic handled further down
+    open_gradient = col1.button("🚀 Open in Gradient")  # logic handled further down
     with col2:
         utils.download_button(code, "generated-code.py", "🐍 Download (.py)")
     with col3:
@@ -102,6 +116,18 @@ def add_to_colab(notebook):
 
     # Handle "Open Colab" button. Down here because to open the new web page, it
     # needs to create a temporary element, which we don't want to show above.
+
+    if open_gradient:
+        if gradient_enabled:
+            gradient_link = add_to_gradient(notebook)
+            utils.open_link(gradient_link)
+        else:
+            colab_error.error(
+                """
+                **Gradient support is disabled.** (If you are hosting this: Create a Github 
+                repo to store notebooks and register it via a .env file)
+                """
+            )
     if open_colab:
         if colab_enabled:
             colab_link = add_to_colab(notebook)