Commit codes

app/components.py

@@ -0,0 +1,217 @@
+import collections
+import os
+from typing import Dict
+
+import streamlit as st
+from datasets import get_dataset_config_names
+from jinja2 import Environment, FileSystemLoader
+
+import utils
+from configuration import OPTIMIZERS_ACCELERATE, OPTIMIZERS_TRAINER, TASKS, TASKS_TO_PIPELINE_TAG
+from utils import (get_dataset_infos_dict, get_datasets, get_model_to_model_id,
+                   render_features)
+
+
+def show_API_component(inputs: Dict[str, str]) -> Dict[str, str]:
+    template_dict = collections.defaultdict()
+    template_dirs = [
+        f for f in os.scandir("templates") if f.is_dir() and f.name != "example"
+    ]
+    template_dirs = sorted(template_dirs, key=lambda e: e.name)
+    for template_dir in template_dirs:
+        template_dict[template_dir.name] = template_dir.path
+    st.write("## API")
+    inputs['api'] = st.selectbox(
+
+        "Which Hugging Face API do you want to use?", list(template_dict.keys())
+    )
+    inputs['template_dir'] = template_dict.get(inputs['api'])
+    return inputs
+
+
+def show_model_component(inputs: Dict[str, str]) -> Dict[str, str]:
+
+    model_info = get_model_to_model_id()
+    models = model_info['model_to_model_id']
+    models_pipeline = model_info["model_to_pipeline_tag"]
+    st.write("## Model")
+    models_for_task = []
+    for model in models:
+        if (models_pipeline[model] == inputs["nlp_task"]):
+            models_for_task.append(model)
+    model = st.selectbox("Which model?", list(models_for_task))
+    inputs["model_checkpoint"] = models.get(model)
+    inputs["pretrained"] = st.checkbox("Use pre-trained model")
+    return inputs
+
+
+def show_task_component(inputs: Dict[str, str]) -> Dict[str, str]:
+    st.write("## Task")
+    task = st.selectbox("Which task?", TASKS)
+    inputs["task"] = task
+    inputs["nlp_task"] = st.selectbox(
+        "Which NLP task?", TASKS_TO_PIPELINE_TAG[task])
+    return inputs
+
+
+def show_input_data_component(inputs: Dict[str, str]) -> Dict[str, str]:
+    st.write("## Input data")
+    english_datasets = get_datasets()
+    english_datasets_for_task = []
+
+    for dataset in english_datasets:
+        for task_category in english_datasets[dataset]:
+            if task_category == inputs["nlp_task"]:
+                english_datasets_for_task.append(dataset)
+                continue
+
+    inputs["dataset"] = st.selectbox(
+        "Which one?", tuple(english_datasets_for_task)
+    )
+
+    configs = get_dataset_config_names(inputs["dataset"])
+    inputs["subset"] = st.selectbox("Which subset?", list(configs))
+
+    data_info_dict = get_dataset_infos_dict(
+        inputs["dataset"], inputs["subset"])
+
+    assert data_info_dict.splits is not None
+    if 'train' in list(data_info_dict.splits.keys()):
+        train_index = list(data_info_dict.splits.keys()).index('train')
+    else:
+        train_index = 0
+
+    inputs["train"] = st.selectbox("Which split for training?", list(
+        data_info_dict.splits.keys()), index=train_index)
+
+    if 'validation' in list(data_info_dict.splits.keys()):
+        validation_index = list(
+            data_info_dict.splits.keys()).index('validation')
+    else:
+        validation_index = len(list(data_info_dict.splits.keys()))-1
+
+    inputs["validation"] = st.selectbox("Which split for validation?", list(
+        data_info_dict.splits.keys()), index=validation_index)
+
+    assert data_info_dict.features is not None
+    feature_index = 0
+    if inputs["nlp_task"] == 'translation':
+        if 'translation' in list(data_info_dict.features.keys()):
+            feature_index = list(
+                data_info_dict.features.keys()).index('translation')
+
+    inputs["feature"] = st.selectbox(
+        "Which data feature?", list(data_info_dict.features.keys()), feature_index)
+
+    if inputs["feature"] == 'translation':
+        inputs["source_language"] = st.selectbox(
+            "Which language for source?", list(data_info_dict.features['translation'].languages))
+        inputs["target_language"] = st.selectbox(
+            "Which language for target?", list(data_info_dict.features['translation'].languages))
+
+    return inputs
+
+
+def show_preprocessing_component(inputs: Dict[str, str]) -> Dict[str, str]:
+    st.write("## Preprocessing")
+    inputs["block_size"] = st.number_input(
+        "The length of each block (i.e. context size)", 1, None, 128)
+
+    if inputs["task"] == "MaskedLM":
+        inputs["mlm_probability"] = st.number_input(
+            "The probability with which to (randomly) mask tokens in the input", 0.0, 1.00, 0.15)
+        inputs["whole_word_masking"] = st.checkbox(
+            "Use whole word masking")
+    return inputs
+
+
+def show_training_comoponent(inputs: Dict[str, str]) -> Dict[str, str]:
+    st.write("## Training")
+
+    # inputs['with_tracker'] = st.selectbox(
+    #     "Loggers to monitor the training ", ["none", "all", "tensorboard", "wandb", "comet_ml"])
+    inputs["seed"] = st.number_input(
+        "Seed", 1, None, 4)
+    
+    if inputs['api'] == 'Accelerate':
+        optimizer_dict_to_use = OPTIMIZERS_ACCELERATE
+    else:
+        optimizer_dict_to_use = OPTIMIZERS_TRAINER
+        
+    inputs["optimizer"] = st.selectbox(
+            "Optimizer", list(optimizer_dict_to_use.keys()))
+    default_lr = optimizer_dict_to_use[inputs["optimizer"]]
+    inputs["lr"] = st.number_input(
+        "Learning rate", 0.000, None, default_lr, format="%f"
+    )
+    inputs["use_weight_decay"] = st.checkbox("Use weight decay")
+    if inputs["use_weight_decay"]:
+        inputs["weight_decay"] = st.number_input(
+            "Weight decay", 0.000, None, 0.01, format="%f"
+        )
+
+    inputs["gradient_accumulation_steps"] = st.number_input(
+        "Gradient Accumulation Steps", 1, None, 8)
+
+    inputs['lr_scheduler_type'] = st.selectbox(
+        "The scheduler type to use", ["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"])
+    inputs['num_warmup_steps'] = st.number_input(
+        "Num warmup steps", 0, None, 0)
+    inputs["batch_size"] = st.number_input("Batch size", 1, None, 32)
+    inputs["num_epochs"] = st.number_input("Epochs", 1, None, 3)
+    return inputs
+
+
+def show_datset_view_component(inputs: Dict[str, str]) -> Dict[str, str]:
+    data_info_dict = get_dataset_infos_dict(
+        inputs["dataset"], inputs["subset"])
+    st.write(f'## Dataset view: {inputs["dataset"]}/{inputs["subset"]}')
+    st.markdown(
+        "*Homepage*: "
+        + data_info_dict.homepage
+        + "\n\n*Dataset*: https://github.com/huggingface/datasets/blob/master/datasets/%s/%s.py"
+        % (inputs["dataset"], inputs["dataset"])
+    )
+    s = []
+    s .append('dataset' + "=" + inputs["dataset"])
+    s.append('config' + "=" + inputs["subset"])
+    st.markdown(
+        "*Permalink*: https://huggingface.co/datasets/viewer/?"
+        + "&".join(s)
+    )
+    # https://github.com/huggingface/datasets-viewer/blob/master/run.py#L282
+    st.write(f'{data_info_dict.description}')
+    st.write(render_features(data_info_dict.features))
+    # TODO make a conditional if the size of the data is too big, switch to streaming mode
+    # TODO cashe this part of the code
+    # selected_dataset = load_dataset(
+    #     inputs["dataset"], inputs["subset"], split=inputs["train"], streaming=True)
+    # print(selected_dataset)
+    # print(next(iter(selected_dataset)))
+    return inputs
+
+def show_code_component(inputs: Dict[str, str]) -> Dict[str, str]:
+    # Generate code and notebook based on template.py.jinja file in the template dir.
+    env = Environment(
+        loader=FileSystemLoader(inputs['template_dir']), trim_blocks=True, lstrip_blocks=True,
+    )
+
+    template = env.get_template(f'task_templates/{inputs["nlp_task"]}.py.jinja')
+    code = template.render(header=utils.code_header, notebook=False, **inputs)
+    notebook_code = template.render(
+        header=utils.notebook_header, notebook=True, **inputs)
+
+    notebook = utils.to_notebook(notebook_code)
+
+    st.write(f'## Code view: {inputs["api"]}')
+    st.write("")  # add vertical space
+    col1, col2 = st.beta_columns(2)
+    with col1:
+        utils.download_button(code, "generated-code.py", "🐍 Download (.py)")
+    with col2:
+        utils.download_button(
+            notebook, "generated-notebook.ipynb", "📓 Download (.ipynb)")
+    colab_error = st.empty()
+    # Display code.
+    st.code(code)
+    return inputs

app/configuration.py

@@ -0,0 +1,17 @@
+INCLUDED_USERS = ['google', 'EleutherAI',
+                  "Helsinki-NLP", "bigscience", "facebook", "openai", "microsoft"]
+
+#TODO create a tempalte for text2text-generation
+# TASKS_TO_PIPELINE_TAG = {
+#     "CausalLM": ['text-generation'], "MaskedLM": ["fill-mask"], "Seq2SeqLM": ['text2text-generation', 'translation']}
+TASKS_TO_PIPELINE_TAG = {
+    "CausalLM": ['text-generation'], "MaskedLM": ["fill-mask"], "Seq2SeqLM": ['translation']}
+
+
+TASKS = list(TASKS_TO_PIPELINE_TAG.keys())
+
+OPTIMIZERS_ACCELERATE = {
+    "AdamW": 0.0001, "Adadelta": 1.0, "Adagrad": 0.01, "Adam": 0.001, "SparseAdam": 0.001, "Adamax": 0.002, "ASGD": 0.01, "LBFGS": 1.0, "NAdam": 0.002, "RAdam": 0.001, "RMSprop": 0.01, "Rprop": 0.01, "SGD": 0.01
+}
+
+OPTIMIZERS_TRAINER = {'adamw_hf': 0.0001, 'adamw_torch': 0.0001, 'adamw_apex_fused': 0.0001, 'adafactor': 0.0001}

app/main.py

@@ -0,0 +1,36 @@
+import streamlit as st
+
+from components import (show_API_component, show_code_component,
+                        show_datset_view_component, show_input_data_component,
+                        show_model_component, show_preprocessing_component,
+                        show_task_component, show_training_comoponent)
+
+st.set_page_config(
+    page_title="Training Code Generator for Hugging Face Models ",  layout="wide"
+)
+
+st.markdown("<br>", unsafe_allow_html=True)
+
+"""
+# Training Code Generator for Hugging Face Models 🤗 
+"""
+st.markdown("<br>", unsafe_allow_html=True)
+"""
+---
+"""
+
+inputs = {}
+
+with st.sidebar:
+    st.info(
+        "**Select the configuration**"
+    )
+    inputs = show_API_component(inputs)
+    inputs = show_task_component(inputs)
+    inputs = show_model_component(inputs)
+    inputs = show_input_data_component(inputs)
+    inputs = show_preprocessing_component(inputs)
+    inputs = show_training_comoponent(inputs)
+
+inputs = show_datset_view_component(inputs)
+inputs = show_code_component(inputs)

app/utils.py

@@ -0,0 +1,216 @@
+import base64
+import importlib.util
+import math
+import re
+import uuid
+from types import ModuleType
+from typing import Dict
+
+import datasets
+import jupytext
+import requests
+import streamlit as st
+from datasets import DatasetInfo, get_dataset_infos
+from datasets.info import DatasetInfosDict
+
+from configuration import INCLUDED_USERS, TASKS_TO_PIPELINE_TAG
+
+
+def import_from_file(module_name: str, filepath: str) -> ModuleType:
+    """
+    Imports a module from file.
+    Args:
+        module_name (str): Assigned to the module's __name__ parameter (does not 
+            influence how the module is named outside of this function)
+        filepath (str): Path to the .py file
+    Returns:
+        The module
+    """
+    spec = importlib.util.spec_from_file_location(module_name, filepath)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def notebook_header(text: str):
+    """
+    Insert section header into a jinja file, formatted as notebook cell.
+
+    Leave 2 blank lines before the header.
+    """
+    return f"""# # {text}
+"""
+
+
+def code_header(text: str):
+    """
+    Insert section header into a jinja file, formatted as Python comment.
+
+    Leave 2 blank lines before the header.
+    """
+    seperator_len = (75 - len(text)) / 2
+    seperator_len_left = math.floor(seperator_len)
+    seperator_len_right = math.ceil(seperator_len)
+    return f"# {'-' * seperator_len_left} {text} {'-' * seperator_len_right}"
+
+
+def to_notebook(code: str) -> str:
+    """Converts Python code to Jupyter notebook format."""
+    notebook = jupytext.reads(code, fmt="py")
+    # print(jupytext.writes(notebook, fmt="ipynb"))
+    return jupytext.writes(notebook, fmt="ipynb")
+
+
+def download_button(
+    object_to_download: str, download_filename: str, button_text: str  # , pickle_it=False
+):
+    """
+    Generates a link to download the given object_to_download.
+
+    From: https://discuss.streamlit.io/t/a-download-button-with-custom-css/4220
+    Params:
+    ------
+    object_to_download:  The object to be downloaded.
+    download_filename (str): filename and extension of file. e.g. mydata.csv,
+    some_txt_output.txt download_link_text (str): Text to display for download
+    link.
+    button_text (str): Text to display on download button (e.g. 'click here to download file')
+    pickle_it (bool): If True, pickle file.
+    Returns:
+    -------
+    (str): the anchor tag to download object_to_download
+    Examples:
+    --------
+    download_link(your_df, 'YOUR_DF.csv', 'Click to download data!')
+    download_link(your_str, 'YOUR_STRING.txt', 'Click to download text!')
+    """
+
+    # try:
+    #     # some strings <-> bytes conversions necessary here
+    b64 = base64.b64encode(object_to_download.encode()).decode()
+    # except AttributeError:
+    #     b64 = base64.b64encode(object_to_download).decode()
+
+    button_uuid = str(uuid.uuid4()).replace("-", "")
+    button_id = re.sub("\d+", "", button_uuid)
+
+    custom_css = f""" 
+        <style>
+            #{button_id} {{
+                display: inline-flex;
+                align-items: center;
+                justify-content: center;
+                background-color: rgb(255, 255, 255);
+                color: rgb(38, 39, 48);
+                padding: .25rem .75rem;
+                position: relative;
+                text-decoration: none;
+                border-radius: 4px;
+                border-width: 1px;
+                border-style: solid;
+                border-color: rgb(230, 234, 241);
+                border-image: initial;
+            }} 
+            #{button_id}:hover {{
+                border-color: rgb(246, 51, 102);
+                color: rgb(246, 51, 102);
+            }}
+            #{button_id}:active {{
+                box-shadow: none;
+                background-color: rgb(246, 51, 102);
+                color: white;
+                }}
+        </style> """
+
+    dl_link = (
+        custom_css
+        + f'<a download="{download_filename}" id="{button_id}" href="data:file/txt;base64,{b64}">{button_text}</a><br><br>'
+    )
+
+    st.markdown(dl_link, unsafe_allow_html=True)
+
+
+@st.cache
+def get_model_to_model_id() -> Dict[str, Dict[str, str]]:
+    requests.get("https://huggingface.co")
+    response = requests.get("https://huggingface.co/api/models")
+    tags = response.json()
+    model_to_model_id = {}
+    model_to_pipeline_tag = {}
+
+    for model in tags:
+        model_name = model['modelId']
+        is_community_model = "/" in model_name
+        if is_community_model:
+            user = model_name.split("/")[0]
+            if user not in INCLUDED_USERS:
+                continue
+
+        # TODO Right now if pipiline is not defined, skip
+        if "pipeline_tag" in model:
+            model_to_model_id[model['id']] = model['modelId']
+            model_to_pipeline_tag[model['id']] = model["pipeline_tag"]
+    return {"model_to_model_id": model_to_model_id, "model_to_pipeline_tag": model_to_pipeline_tag}
+
+
+@st.cache
+def get_datasets() -> Dict[str, str]:
+    english_datasets = {}
+    response = requests.get(
+        "https://huggingface.co/api/datasets?full=true&languages=en")
+    tags = response.json()
+    for dataset in tags:
+        dataset_name = dataset["id"]
+
+        is_community_dataset = "/" in dataset_name
+        if is_community_dataset:
+            # user = dataset_name.split("/")[0]
+            # if user in INCLUDED_USERS:
+            #     english_datasets.append(dataset_name)
+            continue
+
+        if "cardData" not in dataset:
+            continue
+        metadata = dataset["cardData"]
+
+        if "languages" not in metadata:
+            continue
+
+        if "task_categories" not in metadata:
+            continue
+
+        task_is_valid = False
+        for task_category in metadata["task_categories"]:
+
+            if any(task_category in task for task in list(TASKS_TO_PIPELINE_TAG.values())):
+                task_is_valid = True
+        if not task_is_valid:
+            continue
+
+        languages = metadata["languages"]
+
+        if "en" in languages or "en-US" in languages:
+            english_datasets[dataset_name] = metadata["task_categories"]
+    return english_datasets
+
+
+@st.cache
+def get_dataset_infos_dict(dataset: str, subset: str) -> DatasetInfo:
+    return DatasetInfosDict(get_dataset_infos(dataset))[subset]
+
+# https://github.com/huggingface/datasets-viewer/blob/master/run.py#L49
+
+
+def render_features(features):
+    # TODO redner translation object with the languages tags
+    if isinstance(features, dict):
+        return {k: render_features(v) for k, v in features.items()}
+    if isinstance(features, datasets.features.ClassLabel):
+        return features.names
+
+    if isinstance(features, datasets.features.Value):
+        return features.dtype
+
+    if isinstance(features, datasets.features.Sequence):
+        return {"[]": render_features(features.feature)}
+    return features

requirements.txt

@@ -0,0 +1,4 @@
+streamlit
+datasets
+jupytext
+Jinja2

templates/Accelerate/task_templates/fill-mask.py.jinja

@@ -0,0 +1,271 @@
+# Before running, install required packages:
+{% if notebook %}
+
+!
+{%- else %}
+#
+{%- endif %}
+pip install datasets transformers[sentencepiece] accelerate
+
+import collections
+import logging
+import math
+
+import datasets
+import numpy as np
+import torch
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import set_seed
+from codecarbon import EmissionsTracker
+from datasets import load_dataset
+from torch.optim import {{ optimizer }}
+from torch.utils.data import DataLoader
+from torch.utils.data.dataloader import DataLoader
+from tqdm.auto import tqdm
+from transformers import (AutoConfig, AutoModelForMaskedLM, AutoTokenizer,
+                          DataCollatorForLanguageModeling, Trainer,
+                          TrainingArguments, default_data_collator,
+                          get_scheduler)
+from transformers.utils.versions import require_version
+
+{{ header("Setup") }}
+
+tracker = EmissionsTracker(log_level='error')
+tracker.start()
+
+logger = get_logger(__name__)
+require_version("datasets>=1.8.0")
+
+accelerator = Accelerator()
+set_seed({{ seed }})
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.ERROR,
+)
+logger.info(accelerator.state, main_process_only=False)
+if accelerator.is_local_main_process:
+    datasets.utils.logging.set_verbosity_warning()
+    transformers.utils.logging.set_verbosity_info()
+else:
+    datasets.utils.logging.set_verbosity_error()
+    transformers.utils.logging.set_verbosity_error()
+
+{{ header("Load model and dataset") }}
+
+{% if subset == 'default' %}
+datasets = load_dataset('{{dataset}}')
+{% else %}
+datasets = load_dataset('{{dataset}}', '{{ subset }}')
+{% endif %}
+model_checkpoint = "{{model_checkpoint}}"    
+tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
+{% if pretrained %}
+model = AutoModelFor{{task}}.from_pretrained(model_checkpoint)
+{% else %}
+config = AutoConfig.from_pretrained(model_checkpoint)
+model = AutoModelFor{{task}}.from_config(config)
+{% endif %}
+model.resize_token_embeddings(len(tokenizer))
+model_name = model_checkpoint.split("/")[-1]
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+{{ header("Preprocessing") }}
+
+def tokenize_function(examples):
+    result = tokenizer(examples["{{ feature }}"])
+    {% if task=="MaskedLM" %}
+    {% if whole_word_masking %}
+    if tokenizer.is_fast:
+        result["word_ids"] = [result.word_ids(i) for i in range(len(result["input_ids"]))]
+    {% endif %}
+    {% endif %}
+    return result
+
+with accelerator.main_process_first():
+    tokenized_datasets = datasets.map(tokenize_function, batched=True, num_proc=4, remove_columns=list(set(sum(list(datasets.column_names.values()),[]))), desc="Running tokenizer on dataset"
+    )
+
+block_size = {{ block_size }}
+
+def group_texts(examples):
+    # Concatenate all texts.
+    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+    total_length = len(concatenated_examples[list(examples.keys())[0]])
+    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+    total_length = (total_length // block_size) * block_size
+    # Split by chunks of max_len.
+    result = {
+        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+        for k, t in concatenated_examples.items()
+    }
+    result["labels"] = result["input_ids"].copy()
+    return result
+
+with accelerator.main_process_first():
+    lm_datasets = tokenized_datasets.map(
+        group_texts,
+        batched=True,
+        batch_size=1000,
+        num_proc=4,
+        desc=f"Grouping texts in chunks of {block_size}",
+    )
+
+{% if whole_word_masking %}
+def whole_word_masking_data_collator(features):
+    for feature in features:
+        word_ids = feature.pop("word_ids")
+
+        # Create a map between words and corresponding token indices
+        mapping = collections.defaultdict(list)
+        current_word_index = -1
+        current_word = None
+        for idx, word_id in enumerate(word_ids):
+            if word_id is not None:
+                if word_id != current_word:
+                    current_word = word_id
+                    current_word_index += 1
+                mapping[current_word_index].append(idx)
+
+        # Randomly mask words
+        wwm_probability = {{ mlm_probability }}
+        mask = np.random.binomial(1, wwm_probability, (len(mapping),))
+        input_ids = feature["input_ids"]
+        labels = feature["labels"]
+        new_labels = [-100] * len(labels)
+        for word_id in np.where(mask)[0]:
+            word_id = word_id.item()
+            for idx in mapping[word_id]:
+                new_labels[idx] = labels[idx]
+                input_ids[idx] = tokenizer.mask_token_id
+
+    return default_data_collator(features)
+
+data_collator = whole_word_masking_data_collator
+{% else %}
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability={{ mlm_probability }})
+{% endif %}
+
+def insert_random_mask(batch):
+    features = [dict(zip(batch, t)) for t in zip(*batch.values())]
+    masked_inputs = data_collator(features)
+    # Create a new "masked" column for each column in the dataset
+    return {"masked_" + k: v.numpy() for k, v in masked_inputs.items()}
+
+{% if whole_word_masking %}
+lm_datasetst = lm_datasets.remove_columns(["word_ids"])
+{% endif %}
+with accelerator.main_process_first():
+    eval_dataset = lm_datasets["{{ validation }}"].map(
+        insert_random_mask,
+        batched=True,
+        remove_columns=lm_datasets["{{ validation }}"].column_names,
+        desc="Inserting a random mask on eval dataset"
+    )
+
+eval_dataset = eval_dataset.rename_columns(
+    {
+        name: name.split('masked_')[1] for name in eval_dataset.features.keys()
+    }
+)
+
+
+batch_size = {{ batch_size }}
+train_dataloader = DataLoader(
+    lm_datasets["{{ train }}"],
+    shuffle=True,
+    batch_size=batch_size,
+    collate_fn=data_collator,
+)
+eval_dataloader = DataLoader(
+    eval_dataset, batch_size=batch_size, collate_fn=default_data_collator
+)
+
+{{ header("Training") }}
+
+{% if use_weight_decay %}
+weight_decay = {{ weight_decay }}
+def get_grouped_params(model, no_decay=["bias", "LayerNorm.weight"]):
+    params_with_wd, params_without_wd = [], []
+    for n, p in model.named_parameters():
+        if any(nd in n for nd in no_decay):
+            params_without_wd.append(p)
+        else:
+            params_with_wd.append(p)
+    return [
+        {"params": params_with_wd, "weight_decay": weight_decay},
+        {"params": params_without_wd, "weight_decay": 0.0},
+    ]
+
+optimizer = {{ optimizer }}(get_grouped_params(model), lr={{ lr }})
+{% else %}
+optimizer = {{ optimizer }}(model.parameters(), lr={{ lr }})
+{% endif %}
+
+accelerator = Accelerator()
+model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
+    model, optimizer, train_dataloader, eval_dataloader
+)
+
+num_train_epochs = {{ num_epochs }}
+gradient_accumulation_steps = {{ gradient_accumulation_steps }}
+num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
+max_train_steps = num_train_epochs * num_update_steps_per_epoch
+output_dir=f"{model_name}-finetuned"
+
+lr_scheduler = get_scheduler(
+    '{{ lr_scheduler_type }}',
+    optimizer=optimizer,
+    num_warmup_steps={{ num_warmup_steps }},
+    num_training_steps=max_train_steps,
+)
+
+progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
+for epoch in range(num_train_epochs):
+    # Training
+    model.train()
+    for step, batch in enumerate(train_dataloader):
+        outputs = model(**batch)
+        loss = outputs.loss / gradient_accumulation_steps
+        accelerator.backward(loss)
+        
+        if step % gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
+            #TODO Let the user decide on clip grad norm 
+            accelerator.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+            progress_bar.update(1)
+
+    # Evaluation
+    model.eval()
+    losses = []
+    for step, batch in enumerate(eval_dataloader):
+        with torch.no_grad():
+            outputs = model(**batch)
+
+        loss = outputs.loss
+        losses.append(accelerator.gather(loss.repeat(batch_size)))
+
+    losses = torch.cat(losses)
+    losses = losses[: len(eval_dataset)]
+    try:
+        eval_loss = torch.mean(losses)
+        perplexity = math.exp(eval_loss)
+    except OverflowError:
+        perplexity = float("inf")
+    accelerator.print({"loss/eval": eval_loss, "perplexity": perplexity})
+    model.train()
+    accelerator.wait_for_everyone()
+    unwrapped_model = accelerator.unwrap_model(model)
+    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
+    if accelerator.is_main_process:
+        tokenizer.save_pretrained(output_dir)
+
+emissions = tracker.stop()
+accelerator.print(f'Emissions: {emissions} kg')

templates/Accelerate/task_templates/text-generation.py.jinja

@@ -0,0 +1,207 @@
+# Before running, install required packages:
+{% if notebook %}
+
+!
+{%- else %}
+#
+{%- endif %}
+pip install datasets transformers[sentencepiece] accelerate sacrebleu==1.4.14 codecarbon sacremoses
+
+import collections
+import logging
+import math
+import random
+
+import datasets
+import numpy as np
+import torch
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import set_seed
+from codecarbon import EmissionsTracker
+from datasets import load_dataset
+from torch.optim import {{ optimizer }}
+from torch.utils.data import DataLoader
+from torch.utils.data.dataloader import DataLoader
+from tqdm.auto import tqdm
+from transformers import (AutoConfig, AutoModelForCausalLM, AutoModelForMaskedLM, AutoTokenizer,
+                          DataCollatorForLanguageModeling, Trainer,
+                          TrainingArguments, default_data_collator,
+                          get_scheduler)
+from transformers.utils.versions import require_version
+
+{{ header("Setup") }}
+
+tracker = EmissionsTracker(log_level='error')
+tracker.start()
+
+logger = get_logger(__name__)
+require_version("datasets>=1.8.0")
+
+accelerator = Accelerator()
+set_seed({{ seed }})
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.ERROR,
+)
+logger.info(accelerator.state, main_process_only=False)
+if accelerator.is_local_main_process:
+    datasets.utils.logging.set_verbosity_warning()
+    transformers.utils.logging.set_verbosity_info()
+else:
+    datasets.utils.logging.set_verbosity_error()
+    transformers.utils.logging.set_verbosity_error()
+
+{{ header("Load model and dataset") }}
+
+{% if subset == 'default' %}
+datasets = load_dataset('{{dataset}}')
+{% else %}
+datasets = load_dataset('{{dataset}}', '{{ subset }}')
+{% endif %}
+model_checkpoint = "{{model_checkpoint}}"    
+tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
+{% if pretrained %}
+model = AutoModelFor{{task}}.from_pretrained(model_checkpoint)
+{% else %}
+config = AutoConfig.from_pretrained(model_checkpoint)
+model = AutoModelFor{{task}}.from_config(config)
+{% endif %}
+model.resize_token_embeddings(len(tokenizer))
+model_name = model_checkpoint.split("/")[-1]
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+{{ header("Preprocessing") }}
+
+def tokenize_function(examples):
+    result = tokenizer(examples["{{ feature }}"])
+    {% if task=="MaskedLM" %}
+    {% if whole_word_masking %}
+    if tokenizer.is_fast:
+        result["word_ids"] = [result.word_ids(i) for i in range(len(result["input_ids"]))]
+    {% endif %}
+    {% endif %}
+    return result
+
+with accelerator.main_process_first():
+    tokenized_datasets = datasets.map(tokenize_function, batched=True, num_proc=4, remove_columns=list(set(sum(list(datasets.column_names.values()),[]))), desc="Running tokenizer on dataset"
+    )
+
+block_size = {{ block_size }}
+
+def group_texts(examples):
+    # Concatenate all texts.
+    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+    total_length = len(concatenated_examples[list(examples.keys())[0]])
+    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+    total_length = (total_length // block_size) * block_size
+    # Split by chunks of max_len.
+    result = {
+        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+        for k, t in concatenated_examples.items()
+    }
+    result["labels"] = result["input_ids"].copy()
+    return result
+
+with accelerator.main_process_first():
+    lm_datasets = tokenized_datasets.map(
+        group_texts,
+        batched=True,
+        batch_size=1000,
+        num_proc=4,
+        desc=f"Grouping texts in chunks of {block_size}",
+    )
+
+data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
+batch_size = {{ batch_size }}
+train_dataloader = DataLoader(lm_datasets["{{ train }}"], batch_size=batch_size, shuffle=True, collate_fn=data_collator)
+eval_dataloader = DataLoader(lm_datasets["{{ validation }}"], batch_size=batch_size, collate_fn=data_collator)
+
+{{ header("Training") }}
+
+{% if use_weight_decay %}
+weight_decay = {{ weight_decay }}
+def get_grouped_params(model, no_decay=["bias", "LayerNorm.weight"]):
+    params_with_wd, params_without_wd = [], []
+    for n, p in model.named_parameters():
+        if any(nd in n for nd in no_decay):
+            params_without_wd.append(p)
+        else:
+            params_with_wd.append(p)
+    return [
+        {"params": params_with_wd, "weight_decay": weight_decay},
+        {"params": params_without_wd, "weight_decay": 0.0},
+    ]
+
+optimizer = {{ optimizer }}(get_grouped_params(model), lr={{ lr }})
+{% else %}
+optimizer = {{ optimizer }}(model.parameters(), lr={{ lr }})
+{% endif %}
+
+model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
+    model, optimizer, train_dataloader, eval_dataloader
+)
+
+num_train_epochs = {{ num_epochs }}
+gradient_accumulation_steps = {{ gradient_accumulation_steps }}
+num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
+max_train_steps = num_train_epochs * num_update_steps_per_epoch
+output_dir=f"{model_name}-finetuned"
+
+lr_scheduler = get_scheduler(
+    '{{ lr_scheduler_type }}',
+    optimizer=optimizer,
+    num_warmup_steps={{ num_warmup_steps }},
+    num_training_steps=max_train_steps,
+)
+
+progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
+for epoch in range(num_train_epochs):
+    # Training
+    model.train()
+    for step, batch in enumerate(train_dataloader):
+        outputs = model(**batch)
+        loss = outputs.loss / gradient_accumulation_steps
+        accelerator.backward(loss)
+        if step % gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
+            #TODO Let the user decide on clip grad norm 
+            accelerator.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+            progress_bar.update(1)
+ 
+    # Evaluation
+    model.eval()
+    losses = []
+    for step, batch in enumerate(eval_dataloader):
+        with torch.no_grad():
+            outputs = model(**batch)
+
+        loss = outputs.loss
+        losses.append(accelerator.gather(loss.repeat(batch_size)))
+
+    losses = torch.cat(losses)
+    losses = losses[: len(eval_dataloader.dataset)]
+    try:
+        eval_loss = torch.mean(losses)
+        perplexity = math.exp(eval_loss)
+    except OverflowError:
+        perplexity = float("inf")
+
+    accelerator.print(f"epoch {epoch}: perplexity: {perplexity} eval_loss: {eval_loss}")
+    model.train()
+    accelerator.wait_for_everyone()
+    unwrapped_model = accelerator.unwrap_model(model)
+    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
+    if accelerator.is_main_process:
+        tokenizer.save_pretrained(output_dir)
+
+emissions = tracker.stop()
+accelerator.print(f'Emissions: {emissions} kg')

templates/Accelerate/task_templates/translation.py.jinja

@@ -0,0 +1,287 @@
+# Before running, install required packages:
+{% if notebook %}
+
+!
+{%- else %}
+#
+{%- endif %}
+pip install datasets transformers[sentencepiece] accelerate sacrebleu==1.4.14 codecarbon sacremoses
+
+import collections
+import logging
+import math
+import random
+
+import babel
+import datasets
+import numpy as np
+import torch
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import set_seed
+from codecarbon import EmissionsTracker
+from datasets import load_dataset, load_metric
+from torch.optim import {{ optimizer }}
+from torch.utils.data import DataLoader
+from torch.utils.data.dataloader import DataLoader
+from tqdm.auto import tqdm
+from transformers import (AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer,
+                          DataCollatorForLanguageModeling,
+                          DataCollatorForSeq2Seq, MBartTokenizer,
+                          MBartTokenizerFast, Trainer, TrainingArguments,
+                          default_data_collator, get_scheduler)
+from transformers.utils.versions import require_version
+
+{{ header("Setup") }}
+
+tracker = EmissionsTracker(log_level='error')
+tracker.start()
+
+logger = get_logger(__name__)
+require_version("datasets>=1.8.0")
+
+accelerator = Accelerator()
+set_seed({{ seed }})
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.ERROR,
+)
+logger.info(accelerator.state, main_process_only=False)
+if accelerator.is_local_main_process:
+    datasets.utils.logging.set_verbosity_warning()
+    transformers.utils.logging.set_verbosity_info()
+else:
+    datasets.utils.logging.set_verbosity_error()
+    transformers.utils.logging.set_verbosity_error()
+
+{{ header("Load model and dataset") }}
+
+{% if subset == 'default' %}
+datasets = load_dataset('{{dataset}}')
+{% else %}
+datasets = load_dataset('{{dataset}}', '{{ subset }}')
+{% endif %}
+metric = load_metric("sacrebleu")
+model_checkpoint = "{{model_checkpoint}}"    
+tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
+{% if pretrained %}
+model = AutoModelFor{{task}}.from_pretrained(model_checkpoint)
+{% else %}
+config = AutoConfig.from_pretrained(model_checkpoint)
+model = AutoModelFor{{task}}.from_config(config)
+{% endif %}
+model.resize_token_embeddings(len(tokenizer))
+model_name = model_checkpoint.split("/")[-1]
+
+{{ header("Preprocessing") }}
+
+source_lang = '{{ source_language }}'
+target_lang = '{{ target_language }}'
+{% if 'mbart' in model_checkpoint %}
+
+# Set decoder_start_token_id
+if model.config.decoder_start_token_id is None and isinstance(tokenizer, (MBartTokenizer, MBartTokenizerFast)):
+    assert (
+        target_lang is not None and source_lang is not None
+    ), "mBart requires --target_lang and --source_lang"
+    if isinstance(tokenizer, MBartTokenizer):
+        model.config.decoder_start_token_id = tokenizer.lang_code_to_id[target_lang]
+    else:
+        model.config.decoder_start_token_id = tokenizer.convert_tokens_to_ids(target_lang)
+
+{% endif %}
+{% if 't5' in model_checkpoint %}
+if model_checkpoint in ["t5-small", "t5-base", "t5-larg", "t5-3b", "t5-11b"]:
+    for language in (source_lang, target_lang):
+        if language != language[:2]:
+            logging.warning(
+                'Extended language code %s not supported. Falling back on %s.',
+                language, language[:2]
+            )
+    lang_id_to_string = {
+        source_lang: babel.Locale(source_lang[:2]).english_name,
+        target_lang: babel.Locale(target_lang[:2]).english_name,
+    }
+    src_str = 'translate {}'.format(lang_id_to_string[source_lang])
+    tgt_str = ' to {}: '.format(lang_id_to_string[target_lang])
+    prefix = src_str + tgt_str
+else:
+    prefix = ""
+{% else %}
+prefix = ""
+{% endif %}
+{% if 'mbart' in model_checkpoint %}
+
+# For translation we set the codes of our source and target languages (only useful for mBART, the others will
+# ignore those attributes).
+if isinstance(tokenizer, (MBartTokenizer, MBartTokenizerFast)):
+    label = ['ar_AR', 'cs_CZ', 'de_DE', 'en_XX', 'es_XX', 'et_EE', 'fi_FI', 'fr_XX', 'gu_IN', 'hi_IN', 'it_IT', 'ja_XX', 'kk_KZ', 'ko_KR', 'lt_LT', 'lv_LV', 'my_MM', 'ne_NP', 'nl_XX', 'ro_RO', 'ru_RU', 'si_LK', 'tr_TR', 'vi_VN', 'zh_CN']
+    source_code = [item for item in label if item.startswith(source_lang)][0]
+    target_code = [item for item in label if item.startswith(target_lang)][0]
+    if source_lang is not None:
+        tokenizer.src_lang = source_code
+    if target_lang is not None:
+        tokenizer.tgt_lang = target_code
+{% endif %}
+max_input_length = {{ block_size }}
+max_target_length = {{ block_size }}
+
+def preprocess_function(examples):
+    inputs = [prefix + ex[source_lang] for ex in examples["translation"]]
+    targets = [ex[target_lang] for ex in examples["translation"]]
+    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)
+
+    # Setup the tokenizer for targets
+    with tokenizer.as_target_tokenizer():
+        labels = tokenizer(targets, max_length=max_target_length, truncation=True)
+
+    model_inputs["labels"] = labels["input_ids"]
+    return model_inputs
+
+with accelerator.main_process_first():
+    tokenized_datasets = datasets.map(preprocess_function, batched=True, num_proc=4, remove_columns=list(
+        set(sum(list(datasets.column_names.values()), []))), desc="Running tokenizer on dataset")
+
+data_collator = DataCollatorForSeq2Seq(tokenizer, model=model, pad_to_multiple_of=8 if accelerator.use_fp16 else None)
+batch_size = {{ batch_size }}
+train_dataloader = DataLoader(tokenized_datasets["{{ train }}"], batch_size=batch_size, shuffle=True, collate_fn=data_collator)
+eval_dataloader = DataLoader(tokenized_datasets["{{ validation }}"], batch_size=batch_size, collate_fn=data_collator)
+
+{{ header("Training") }}
+
+def compute_metrics(eval_preds):
+    preds, labels = eval_preds
+    # In case the model returns more than the prediction logits
+    if isinstance(preds, tuple):
+        preds = preds[0]
+
+    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
+
+    # Replace -100s in the labels as we can't decode them
+    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+    # Some simple post-processing
+    decoded_preds = [pred.strip() for pred in decoded_preds]
+    decoded_labels = [[label.strip()] for label in decoded_labels]
+
+    result = metric.compute(predictions=decoded_preds,
+                            references=decoded_labels)
+    return {"bleu": result["score"]}
+
+
+def postprocess(predictions, labels):
+    predictions = predictions.cpu().numpy()
+    labels = labels.cpu().numpy()
+
+    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+
+    # Replace -100 in the labels as we can't decode them.
+    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+    # Some simple post-processing
+    decoded_preds = [pred.strip() for pred in decoded_preds]
+    decoded_labels = [[label.strip()] for label in decoded_labels]
+    return decoded_preds, decoded_labels
+
+{% if use_weight_decay %}
+weight_decay = {{ weight_decay }}
+def get_grouped_params(model, no_decay=["bias", "LayerNorm.weight"]):
+    params_with_wd, params_without_wd = [], []
+    for n, p in model.named_parameters():
+        if any(nd in n for nd in no_decay):
+            params_without_wd.append(p)
+        else:
+            params_with_wd.append(p)
+    return [
+        {"params": params_with_wd, "weight_decay": weight_decay},
+        {"params": params_without_wd, "weight_decay": 0.0},
+    ]
+
+optimizer = {{ optimizer }}(get_grouped_params(model), lr={{ lr }})
+{% else %}
+optimizer = {{ optimizer }}(model.parameters(), lr={{ lr }})
+{% endif %}
+model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
+    model, optimizer, train_dataloader, eval_dataloader
+)
+
+num_train_epochs = {{ num_epochs }}
+gradient_accumulation_steps = {{ gradient_accumulation_steps }}
+num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
+max_train_steps = num_train_epochs * num_update_steps_per_epoch
+output_dir=f"{model_name}-finetuned"
+
+lr_scheduler = get_scheduler(
+    '{{ lr_scheduler_type }}',
+    optimizer=optimizer,
+    num_warmup_steps={{ num_warmup_steps }},
+    num_training_steps=max_train_steps,
+)
+
+progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
+for epoch in range(num_train_epochs):
+    # Training
+    model.train()
+    for step, batch in enumerate(train_dataloader):
+        outputs = model(**batch)
+        loss = outputs.loss
+        loss = loss / gradient_accumulation_steps
+        accelerator.backward(loss)
+
+        if step % gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+            progress_bar.update(1)
+            
+
+    # Evaluation
+    model.eval()
+    samples_seen = 0
+    for step, batch in enumerate(eval_dataloader):
+        with torch.no_grad():
+            generated_tokens = accelerator.unwrap_model(model).generate(
+                batch["input_ids"],
+                attention_mask=batch["attention_mask"],
+                max_length=128,
+            )
+        labels = batch["labels"]
+
+        # Necessary to pad predictions and labels for being gathered
+        generated_tokens = accelerator.pad_across_processes(
+            generated_tokens, dim=1, pad_index=tokenizer.pad_token_id
+        )
+        labels = accelerator.pad_across_processes(labels, dim=1, pad_index=-100)
+
+        predictions_gathered = accelerator.gather(generated_tokens)
+        labels_gathered = accelerator.gather(labels)
+
+        decoded_preds, decoded_labels = postprocess(predictions_gathered, labels_gathered)
+
+        if accelerator.num_processes > 1:
+            if step == len(eval_dataloader) - 1:
+                decoded_preds = decoded_preds[: len(
+                    eval_dataloader.dataset) - samples_seen]
+                decoded_labels = decoded_labels[: len(
+                    eval_dataloader.dataset) - samples_seen]
+            else:
+                samples_seen += len(decoded_labels)
+        metric.add_batch(predictions=decoded_preds, references=decoded_labels)
+
+    results = metric.compute()
+    print(f"epoch {epoch}, BLEU score: {results['score']:.2f}")
+
+    # Save and upload
+    accelerator.wait_for_everyone()
+    unwrapped_model = accelerator.unwrap_model(model)
+    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
+    if accelerator.is_main_process:
+        tokenizer.save_pretrained(output_dir)
+
+emissions = tracker.stop()
+print(f'Emissions: {emissions} kg')

templates/Trainer/task_templates/fill-mask.py.jinja

@@ -0,0 +1,185 @@
+# Before running, install required packages:
+{% if notebook %}
+
+!
+{%- else %}
+#
+{%- endif %}
+pip install datasets transformers
+
+import collections
+import math
+import logging
+
+import numpy as np
+import transformers
+import datasets
+from datasets import load_dataset
+from transformers import (AutoConfig, AutoModelForMaskedLM, AutoTokenizer,
+                          DataCollatorForLanguageModeling, Trainer,
+                          TrainingArguments, default_data_collator, set_seed)
+from transformers.testing_utils import CaptureLogger
+from transformers.utils.versions import require_version
+
+{{ header("Setup") }}
+
+
+logger = logging.getLogger(__name__)
+require_version("datasets>=1.8.0")
+set_seed({{ seed }})
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.ERROR,
+)
+datasets.utils.logging.set_verbosity_warning()
+transformers.utils.logging.set_verbosity_info()
+
+
+{{ header("Load model and dataset") }}
+
+{% if subset == 'default' %}
+datasets = load_dataset('{{dataset}}')
+{% else %}
+datasets = load_dataset('{{dataset}}', '{{ subset }}')
+{% endif %}
+model_checkpoint = "{{model_checkpoint}}"
+tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
+{% if pretrained %}
+model = AutoModelFor{{task}}.from_pretrained(model_checkpoint)
+{% else %}
+config = AutoConfig.from_pretrained(model_checkpoint)
+model = AutoModelFor{{task}}.from_config(config)
+{% endif %}
+model.resize_token_embeddings(len(tokenizer))
+model_name = model_checkpoint.split("/")[-1]
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+{{ header("Preprocessing") }}
+
+# since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function
+tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
+def tokenize_function(examples):
+    with CaptureLogger(tok_logger) as cl:
+        result = tokenizer(examples["{{ feature }}"])
+    if "Token indices sequence length is longer than the" in cl.out:
+        tok_logger.warning(
+                "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits"
+                " before being passed to the model."
+            )
+    if tokenizer.is_fast:
+        result["word_ids"] = [result.word_ids(i) for i in range(len(result["input_ids"]))]
+    return result
+
+tokenized_datasets = datasets.map(tokenize_function, batched=True, num_proc=4, remove_columns=list(set(sum(list(datasets.column_names.values()),[]))), desc="Running tokenizer on dataset"
+    )
+block_size = {{ block_size }}
+
+def group_texts(examples):
+    # Concatenate all texts.
+    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+    total_length = len(concatenated_examples[list(examples.keys())[0]])
+    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+    total_length = (total_length // block_size) * block_size
+    # Split by chunks of max_len.
+    result = {
+        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+        for k, t in concatenated_examples.items()
+    }
+    result["labels"] = result["input_ids"].copy()
+    return result
+
+lm_datasets = tokenized_datasets.map(
+    group_texts,
+    batched=True,
+    batch_size=1000,
+    num_proc=4,
+    desc=f"Grouping texts in chunks of {block_size}",
+)
+
+{{ header("Training") }}
+
+training_args = TrainingArguments(
+    output_dir=f"{model_name}-finetuned",
+    per_device_train_batch_size={{ batch_size }},
+    per_device_eval_batch_size={{ batch_size }},
+    evaluation_strategy='epoch',
+    logging_strategy='epoch',
+    save_strategy='epoch',
+    optim='{{ optimizer }}',
+    learning_rate={{ lr }},
+    num_train_epochs={{ num_epochs }},
+    gradient_accumulation_steps={{ gradient_accumulation_steps }},
+    lr_scheduler_type='{{ lr_scheduler_type }}',
+    warmup_steps={{ num_warmup_steps }},
+    {% if use_weight_decay%}
+    weight_decay={{ weight_decay }},
+    {% endif %}
+    push_to_hub=False,
+    dataloader_num_workers=0,
+    {% if task=="MaskedLM" %}
+    {% if whole_word_masking %}
+    remove_unused_columns=False,
+    {% endif %}
+    {% endif %}
+    load_best_model_at_end=True,
+    log_level='error' 
+)
+
+
+{% if whole_word_masking %}
+def whole_word_masking_data_collator(features):
+    for feature in features:
+        word_ids = feature.pop("word_ids")
+
+        # Create a map between words and corresponding token indices
+        mapping = collections.defaultdict(list)
+        current_word_index = -1
+        current_word = None
+        for idx, word_id in enumerate(word_ids):
+            if word_id is not None:
+                if word_id != current_word:
+                    current_word = word_id
+                    current_word_index += 1
+                mapping[current_word_index].append(idx)
+
+        # Randomly mask words
+        wwm_probability = {{ mlm_probability }}
+        mask = np.random.binomial(1, wwm_probability, (len(mapping),))
+        input_ids = feature["input_ids"]
+        labels = feature["labels"]
+        new_labels = [-100] * len(labels)
+        for word_id in np.where(mask)[0]:
+            word_id = word_id.item()
+            for idx in mapping[word_id]:
+                new_labels[idx] = labels[idx]
+                input_ids[idx] = tokenizer.mask_token_id
+
+    return default_data_collator(features)
+data_collator = whole_word_masking_data_collator
+    {% else %}
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability={{ mlm_probability }})
+{% endif %}
+
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=lm_datasets["{{ train }}"],
+    eval_dataset=lm_datasets["{{ validation }}"],
+    data_collator=data_collator,
+)
+
+train_result = trainer.train()
+trainer.save_model()
+trainer.log_metrics("train", train_result.metrics)
+trainer.save_metrics("train", train_result.metrics)
+trainer.save_state()
+eval_results = trainer.evaluate()
+eval_results["perplexity"] = math.exp(eval_results['eval_loss'])
+print(f"Perplexity: {math.exp(eval_results['eval_loss']):.2f}")
+trainer.log_metrics("eval", eval_results)
+trainer.save_metrics("eval", eval_results)

templates/Trainer/task_templates/text-generation.py.jinja

@@ -0,0 +1,152 @@
+# Before running, install required packages:
+{% if notebook %}
+
+!
+{%- else %}
+#
+{%- endif %}
+pip install datasets transformers
+
+import collections
+import math
+import logging
+
+import numpy as np
+import transformers
+import datasets
+from datasets import load_dataset
+from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer,
+                          DataCollatorForLanguageModeling, Trainer,
+                          TrainingArguments, default_data_collator, set_seed)
+from transformers.testing_utils import CaptureLogger
+from transformers.utils.versions import require_version
+
+{{ header("Setup") }}
+
+
+logger = logging.getLogger(__name__)
+require_version("datasets>=1.8.0")
+set_seed({{ seed }})
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.ERROR,
+)
+datasets.utils.logging.set_verbosity_warning()
+transformers.utils.logging.set_verbosity_info()
+
+
+{{ header("Load model and dataset") }}
+
+{% if subset == 'default' %}
+datasets = load_dataset('{{dataset}}')
+{% else %}
+datasets = load_dataset('{{dataset}}', '{{ subset }}')
+{% endif %}
+model_checkpoint = "{{model_checkpoint}}"
+tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
+{% if pretrained %}
+model = AutoModelFor{{task}}.from_pretrained(model_checkpoint)
+{% else %}
+config = AutoConfig.from_pretrained(model_checkpoint)
+model = AutoModelFor{{task}}.from_config(config)
+{% endif %}
+model.resize_token_embeddings(len(tokenizer))
+model_name = model_checkpoint.split("/")[-1]
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+{{ header("Preprocessing") }}
+
+# since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function
+tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
+def tokenize_function(examples):
+    with CaptureLogger(tok_logger) as cl:
+        result = tokenizer(examples["{{ feature }}"])
+    if "Token indices sequence length is longer than the" in cl.out:
+        tok_logger.warning(
+                "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits"
+                " before being passed to the model."
+            )
+    if tokenizer.is_fast:
+        result["word_ids"] = [result.word_ids(i) for i in range(len(result["input_ids"]))]
+    return result
+
+tokenized_datasets = datasets.map(tokenize_function, batched=True, num_proc=4, remove_columns=list(set(sum(list(datasets.column_names.values()),[]))), desc="Running tokenizer on dataset"
+    )
+block_size = {{ block_size }}
+
+def group_texts(examples):
+    # Concatenate all texts.
+    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+    total_length = len(concatenated_examples[list(examples.keys())[0]])
+    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+    total_length = (total_length // block_size) * block_size
+    # Split by chunks of max_len.
+    result = {
+        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+        for k, t in concatenated_examples.items()
+    }
+    result["labels"] = result["input_ids"].copy()
+    return result
+
+lm_datasets = tokenized_datasets.map(
+    group_texts,
+    batched=True,
+    batch_size=1000,
+    num_proc=4,
+    desc=f"Grouping texts in chunks of {block_size}",
+)
+
+{{ header("Training") }}
+
+training_args = TrainingArguments(
+    output_dir=f"{model_name}-finetuned",
+    per_device_train_batch_size={{ batch_size }},
+    per_device_eval_batch_size={{ batch_size }},
+    evaluation_strategy='epoch',
+    logging_strategy='epoch',
+    save_strategy='epoch',
+    optim='{{ optimizer }}',
+    learning_rate={{ lr }},
+    num_train_epochs={{ num_epochs }},
+    gradient_accumulation_steps={{ gradient_accumulation_steps }},
+    lr_scheduler_type='{{ lr_scheduler_type }}',
+    warmup_steps={{ num_warmup_steps }},
+    {% if use_weight_decay%}
+    weight_decay={{ weight_decay }},
+    {% endif %}
+    push_to_hub=False,
+    dataloader_num_workers=0,
+    {% if task=="MaskedLM" %}
+    {% if whole_word_masking %}
+    remove_unused_columns=False,
+    {% endif %}
+    {% endif %}
+    load_best_model_at_end=True,
+    log_level='error' 
+)
+
+data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
+
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=lm_datasets["{{ train }}"],
+    eval_dataset=lm_datasets["{{ validation }}"],
+    data_collator=data_collator,
+)
+
+train_result = trainer.train()
+trainer.save_model()
+trainer.log_metrics("train", train_result.metrics)
+trainer.save_metrics("train", train_result.metrics)
+trainer.save_state()
+eval_results = trainer.evaluate()
+eval_results["perplexity"] = math.exp(eval_results['eval_loss'])
+print(f"Perplexity: {math.exp(eval_results['eval_loss']):.2f}")
+trainer.log_metrics("eval", eval_results)
+trainer.save_metrics("eval", eval_results)

templates/Trainer/task_templates/translation.py.jinja

@@ -0,0 +1,218 @@
+# Before running, install required packages:
+{% if notebook %}
+
+!
+{%- else %}
+#
+{%- endif %}
+pip install datasets transformers[sentencepiece] accelerate sacrebleu==1.4.14 sacremoses
+
+import collections
+import logging
+import math
+import random
+
+import babel
+import datasets
+import numpy as np
+import torch
+import transformers
+from datasets import load_dataset, load_metric
+from torch.utils.data import DataLoader
+from torch.utils.data.dataloader import DataLoader
+from tqdm.auto import tqdm
+from transformers import (AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer,
+                          DataCollatorForLanguageModeling,
+                          DataCollatorForSeq2Seq, MBartTokenizer,
+                          MBartTokenizerFast, Seq2SeqTrainer, Seq2SeqTrainingArguments,
+                          default_data_collator, get_scheduler)
+from transformers.utils.versions import require_version
+
+{{ header("Setup") }}
+
+
+logger = logging.getLogger(__name__)
+require_version("datasets>=1.8.0")
+set_seed({{ seed }})
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.ERROR,
+)
+datasets.utils.logging.set_verbosity_warning()
+transformers.utils.logging.set_verbosity_info()
+
+{{ header("Load model and dataset") }}
+
+{% if subset == 'default' %}
+datasets = load_dataset('{{dataset}}')
+{% else %}
+datasets = load_dataset('{{dataset}}', '{{ subset }}')
+{% endif %}
+metric = load_metric("sacrebleu")
+model_checkpoint = "{{model_checkpoint}}"    
+tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
+{% if pretrained %}
+model = AutoModelFor{{task}}.from_pretrained(model_checkpoint)
+{% else %}
+config = AutoConfig.from_pretrained(model_checkpoint)
+model = AutoModelFor{{task}}.from_config(config)
+{% endif %}
+model.resize_token_embeddings(len(tokenizer))
+model_name = model_checkpoint.split("/")[-1]
+
+{{ header("Preprocessing") }}
+
+source_lang = '{{ source_language }}'
+target_lang = '{{ target_language }}'
+{% if 'mbart' in model_checkpoint %}
+
+# Set decoder_start_token_id
+if model.config.decoder_start_token_id is None and isinstance(tokenizer, (MBartTokenizer, MBartTokenizerFast)):
+    assert (
+        target_lang is not None and source_lang is not None
+    ), "mBart requires --target_lang and --source_lang"
+    if isinstance(tokenizer, MBartTokenizer):
+        model.config.decoder_start_token_id = tokenizer.lang_code_to_id[target_lang]
+    else:
+        model.config.decoder_start_token_id = tokenizer.convert_tokens_to_ids(target_lang)
+
+{% endif %}
+{% if 't5' in model_checkpoint %}
+if model_checkpoint in ["t5-small", "t5-base", "t5-larg", "t5-3b", "t5-11b"]:
+    for language in (source_lang, target_lang):
+        if language != language[:2]:
+            logging.warning(
+                'Extended language code %s not supported. Falling back on %s.',
+                language, language[:2]
+            )
+    lang_id_to_string = {
+        source_lang: babel.Locale(source_lang[:2]).english_name,
+        target_lang: babel.Locale(target_lang[:2]).english_name,
+    }
+    src_str = 'translate {}'.format(lang_id_to_string[source_lang])
+    tgt_str = ' to {}: '.format(lang_id_to_string[target_lang])
+    prefix = src_str + tgt_str
+else:
+    prefix = ""
+{% else %}
+prefix = ""
+{% endif %}
+{% if 'mbart' in model_checkpoint %}
+
+# For translation we set the codes of our source and target languages (only useful for mBART, the others will
+# ignore those attributes).
+if isinstance(tokenizer, (MBartTokenizer, MBartTokenizerFast)):
+    label = ['ar_AR', 'cs_CZ', 'de_DE', 'en_XX', 'es_XX', 'et_EE', 'fi_FI', 'fr_XX', 'gu_IN', 'hi_IN', 'it_IT', 'ja_XX', 'kk_KZ', 'ko_KR', 'lt_LT', 'lv_LV', 'my_MM', 'ne_NP', 'nl_XX', 'ro_RO', 'ru_RU', 'si_LK', 'tr_TR', 'vi_VN', 'zh_CN']
+    source_code = [item for item in label if item.startswith(source_lang)][0]
+    target_code = [item for item in label if item.startswith(target_lang)][0]
+    if source_lang is not None:
+        tokenizer.src_lang = source_code
+    if target_lang is not None:
+        tokenizer.tgt_lang = target_code
+{% endif %}
+max_input_length = {{ block_size }}
+max_target_length = {{ block_size }}
+
+def preprocess_function(examples):
+    inputs = [prefix + ex[source_lang] for ex in examples["translation"]]
+    targets = [ex[target_lang] for ex in examples["translation"]]
+    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)
+
+    # Setup the tokenizer for targets
+    with tokenizer.as_target_tokenizer():
+        labels = tokenizer(targets, max_length=max_target_length, truncation=True)
+
+    model_inputs["labels"] = labels["input_ids"]
+    return model_inputs
+
+
+tokenized_datasets = datasets.map(preprocess_function, batched=True, num_proc=4, remove_columns=list(
+    set(sum(list(datasets.column_names.values()), []))), desc="Running tokenizer on dataset")
+
+data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+batch_size = {{ batch_size }}
+
+{{ header("Training") }}
+
+def compute_metrics(eval_preds):
+    preds, labels = eval_preds
+    # In case the model returns more than the prediction logits
+    if isinstance(preds, tuple):
+        preds = preds[0]
+
+    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
+
+    # Replace -100s in the labels as we can't decode them
+    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+    # Some simple post-processing
+    decoded_preds = [pred.strip() for pred in decoded_preds]
+    decoded_labels = [[label.strip()] for label in decoded_labels]
+
+    result = metric.compute(predictions=decoded_preds,
+                            references=decoded_labels)
+    return {"bleu": result["score"]}
+
+
+def postprocess(predictions, labels):
+    predictions = predictions.cpu().numpy()
+    labels = labels.cpu().numpy()
+
+    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+
+    # Replace -100 in the labels as we can't decode them.
+    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+    # Some simple post-processing
+    decoded_preds = [pred.strip() for pred in decoded_preds]
+    decoded_labels = [[label.strip()] for label in decoded_labels]
+    return decoded_preds, decoded_labels
+
+
+
+training_args = Seq2SeqTrainingArguments(
+    output_dir=f"{model_name}-finetuned",
+    per_device_train_batch_size={{ batch_size }},
+    per_device_eval_batch_size={{ batch_size }},
+    evaluation_strategy='epoch',
+    logging_strategy='epoch',
+    save_strategy='epoch',
+    optim='{{ optimizer }}',
+    learning_rate={{ lr }},
+    num_train_epochs={{ num_epochs }},
+    gradient_accumulation_steps={{ gradient_accumulation_steps }},
+    lr_scheduler_type='{{ lr_scheduler_type }}',
+    warmup_steps={{ num_warmup_steps }},
+    {% if use_weight_decay%}
+    weight_decay={{ weight_decay }},
+    {% endif %}
+    push_to_hub=False,
+    dataloader_num_workers=0,
+    {% if task=="MaskedLM" %}
+    {% if whole_word_masking %}
+    remove_unused_columns=False,
+    {% endif %}
+    {% endif %}
+    load_best_model_at_end=True,
+    log_level='error' 
+)
+
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=lm_datasets["{{ train }}"],
+    eval_dataset=lm_datasets["{{ validation }}"],
+    data_collator=data_collator,
+)
+
+train_result = trainer.train()
+trainer.save_model()
+trainer.log_metrics("train", train_result.metrics)
+trainer.save_metrics("train", train_result.metrics)
+trainer.save_state()
+eval_results = trainer.evaluate()
+trainer.log_metrics("eval", eval_results)
+trainer.save_metrics("eval", eval_results)