🧙  A web app to generate template code for machine learning ✨


Code style: black

🎉 Traingenerator is now live! 🎉

Try it out:

Generate custom template code for PyTorch & sklearn, using a simple web UI built with streamlit. Traingenerator offers multiple options for preprocessing, model setup, training, and visualization (using Tensorboard or It exports to .py, Jupyter Notebook, or Google Colab. The perfect tool to jumpstart your next machine learning project!

For updates, follow me on Twitter, and if you like this project, please consider sponsoring

Adding new templates

You can add your own template in 4 easy steps (see below), without changing any code
in the app itself. Your new template will be automatically discovered by Traingenerator
and shown in the sidebar. That’s it! 🎈

Want to share your magic? 🧙 PRs are welcome! Please have a look
at and write on Gitter.

Some ideas for new templates: Keras/Tensorflow, Pytorch Lightning, object detection,
segmentation, text classification, …

  1. Create a folder under ./templates.
    The folder name should be the task that your template solves (e.g.
    Image classification). Optionally, you can add a framework name (e.g.
    Image classification_PyTorch). Both names are automatically shown in the first two
    dropdowns in the sidebar (see image).
    Tip: Copy the example template to get started more quickly.
  2. Add a file to the folder (see example).
    It needs to contain a method show(), which displays all template-specific streamlit
    components in the sidebar (i.e. everything below Task) and returns a dictionary of
    user inputs.
  3. Add a file to the folder (see example).
    This Jinja2 template is used
    to generate the code. You can write normal Python code in it and modify it
    (through Jinja) based on the user inputs in the sidebar (e.g. insert a parameter
    value from the sidebar or show different code parts based on the user’s selection).
  4. Optional: Add a file test-inputs.yml to the folder (see example).
    This simple YAML file should define a few possible user inputs that can be used for
    testing. If you run pytest (see below), it will automatically pick up this file, render
    the code template with its values, and check that the generated code runs without
    errors. This file is optional – but it’s required if you want to contribute your
    template to this repo.


Note: You only need to install Traingenerator if you want to contribute or run it
locally. If you just want to use it, go here.

git clone
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
they are on Github). After setting up the repo, create a file .env with content:


If you don’t set this up, the app will still work but the “Open in Colab” button
will only show an error message.

Running locally

streamlit run app/

Make sure to run always from the traingenerator dir (not from the app dir),
otherwise the app will not be able to find the templates.

Deploying to Heroku

First, install heroku and login.
To create a new deployment, run inside traingenerator:

heroku create
git push heroku main
heroku open

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),
you also need to run:

heroku config:set GITHUB_TOKEN=<your-github-access-token>
heroku config:set REPO_NAME=<user/notebooks-repo>


First, install pytest and required plugins via:

pip install -r requirements-dev.txt

To run all tests:

pytest ./tests

Note that this only tests the code templates (i.e. it renders them with different
input values and makes sure that the code executes without error). The streamlit app
itself is not tested at the moment.

You can also test an individual template by passing the name of the template dir to
--template, e.g.:

pytest ./tests --template "Image classification_scikit-learn"

The mage image used in Traingenerator is from
Twitter’s Twemoji library and
released under Creative Commons Attribution 4.0 International Public License.


View Github