diff --git a/README.md b/README.md
index 64b6a07..35a6126 100644
--- a/README.md
+++ b/README.md
@@ -35,7 +35,7 @@
-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! ✨
@@ -53,8 +53,8 @@ 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
@@ -62,7 +62,7 @@ GITHUB_TOKEN=
REPO_NAME=
```
-*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:*
```
diff --git a/app/main.py b/app/main.py
index c5bbdaf..2261bac 100644
--- a/app/main.py
+++ b/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)