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)