Introduce code from Rust monorepo. This is a mess

tts/NeMo/examples/asr/asr_webapp/Dockerfile

@@ -0,0 +1,41 @@
+# syntax=docker/dockerfile:experimental
+
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+ARG BASE_IMAGE=nvcr.io/nvidia/nemo:1.0.0rc1
+
+# build an image that includes only the nemo dependencies, ensures that dependencies
+# are included first for optimal caching, and useful for building a development
+# image (by specifying build target as `nemo-deps`)
+FROM ${BASE_IMAGE} as nemo-deps
+
+# copy asr service source into a temp directory
+WORKDIR /tmp/nemo-service
+COPY . .
+
+# override nemo installation with dependency from requirements.txt
+RUN /bin/echo "export BASE_IMAGE=${BASE_IMAGE}" >> /root/.bashrc
+RUN cd /tmp/nemo-service && pip install -r "requirements.txt"
+
+# copy webapp into container for end user
+WORKDIR /workspace/nemo-service
+COPY . /workspace/nemo-service/
+
+# pre-import some asr dependencies which would take time on each worker
+RUN python -c "import nemo.collections.asr as nemo_asr"
+
+# expose port 8000 and launch gunicorn
+EXPOSE 8000
+CMD ["gunicorn", "wsgi:app"]

tts/NeMo/examples/asr/asr_webapp/README.rst

@@ -0,0 +1,69 @@
+**ASR EVALUATION SERVICE **
+===========================
+
+Usage
+-----
+
+A simple service that will load a pre-trained model from NeMo's ASR collection and after uploading some files, will transcribe them. Any pre-trained model (or model which has been uploaded to the service) can be selected and if the local / docker environment supports GPUs then it will be possible to transcribe the audio segments with GPU.
+
+All uploaded files are immediately deleted after transcriptions are obtained.
+
+ .. note::
+
+    When using `Gunicorn <https://gunicorn.org/>`_, you might notice that each pretrained checkpoint takes a long time for the first transcription. This is because that worker is instantiating the model into memory and then moving the model onto the GPU. For large models, this step might take a significant amount of time. However, this step is cached, and subsequent transcription requests should be much faster (especially if a GPU is utilized).
+
+From inside the ``asr_webapp`` directory, there are three options to run the ASR Service to transcribe audio -
+
+Local Setup
+-----------
+
+1) Install all the dependencies (``pip install -r requirements.txt``). Note : It will utilize the latest branch set in the requirements and will override current nemo installation.
+
+2) Simply run ``python asr_service.py``. This will launch a single service worker that can be run locally.
+
+3) The app will run at ``127.0.0.1:5000`` by default.
+
+Gunicorn Setup
+--------------
+
+1) Follow Step (1) of the ``Local Setup`` section.
+
+2) Edit the configuration of ``gunicorn`` inside ``gunicorn.conf.py`` if you wish to change the port or number of workers (though this can be achieved via command line overrides as well).
+
+3) Simply run ``gunicorn wsgi:app``. This will launch two workers (default) in a gunicorn environment.
+
+4) The app will run at ``0.0.0.0:8000`` by default.
+
+Docker Setup
+------------
+
+The cleanest approach of the three, and requires simply building and running a docker container.
+
+1) Build the docker by executing ``bash docker_container_build.sh``. This will build a container using the latest branch of nemo.
+
+2) Run the container by executing ``bash docker_container_run.sh``. This will run a detached container that can be used by visiting ``0.0.0.0:8000`` on a modern browser.
+
+Note About Uploading Models
+---------------------------
+
+Uploading models is a useful method to evaluate models quickly without having to write scripts.
+By design, models uploaded via the browser are visible to **all** users who have access to the web app - so that there only needs to be one user who uploads the model and everyone else who can access this service will be able to evaluate with that model.
+
+This also means that models are uploaded indefinitely - and cannot be removed easily without explicitly deleting the ``models`` directory or by shutting down the container (if a container is used).
+
+There are several reasons this approach was taken -
+
+* It allows for easy evaluation amongst peers without duplicating model uploads (which may be very big).
+* When using ``gunicorn``, every worker has access to a central cache of models which can be loaded into memory.
+* Model upload mode is meant to be used in a local, protected environment.
+
+If you wish to disable model uploading, please open ``templates/main.html`` and comment out the following section:
+
+.. code-block:: html
+
+    <!-- Load .nemo checkpoint -->
+    <!-- Comment this section out to remove ability to upload files -->
+    <br>
+    NeMo File :
+    <input type = "file" name = "nemo_model" accept=".nemo" id="nemo_model_file" />
+    <!-- Comment out upto here to remove ability to upload models -->

tts/NeMo/examples/asr/asr_webapp/asr_service.py

@@ -0,0 +1,374 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import atexit
+import glob
+import os
+import shutil
+import time
+from html import unescape
+from uuid import uuid4
+
+import model_api
+import torch
+import werkzeug
+from flask import Flask, make_response, render_template, request, url_for
+from flask_cors import CORS
+from werkzeug.utils import secure_filename
+
+from nemo.utils import logging
+
+app = Flask(__name__)
+CORS(app)
+
+# Upload folder for audio files; models are stored in permanent cache
+# which gets deleted once the container shuts down
+app.config[f'UPLOAD_FOLDER'] = f"tmp/"
+
+
+@app.route('/initialize_model', methods=['POST'])
+def initialize_model():
+    """
+    API Endpoint to instantiate a model
+
+    Loads ASR model by its pretrained checkpoint name or upload ASR model that is provided by the user,
+    then load that checkpoint into the cache.
+
+    Loading of the model into cache is done once per worker. Number of workers should be limited
+    so as not to exhaust the GPU memory available on device (if GPU is being used).
+    """
+    logging.info("Starting ASR service")
+    if torch.cuda.is_available():
+        logging.info("CUDA is available. Running on GPU")
+    else:
+        logging.info("CUDA is not available. Defaulting to CPUs")
+
+    # get form fields
+    model_name = request.form['model_names_select']
+    use_gpu_if_available = request.form.get('use_gpu_ckbx', "off")
+
+    # get nemo model from user (if not none)
+    nemo_model_file = request.files.get('nemo_model', '')
+
+    # if nemo model is not None, upload it to model cache
+    if nemo_model_file != '':
+        model_name = _store_model(nemo_model_file)
+
+        # Alert user that model has been uploaded into the model cache,
+        # and they should refresh the page to access the model
+        result = render_template(
+            'toast_msg.html', toast_message=f"Model {model_name} has been uploaded. " f"Refresh page !", timeout=5000
+        )
+
+    else:
+        # Alert user that model has been loaded onto a workers memory
+        result = render_template(
+            'toast_msg.html', toast_message=f"Model {model_name} has been initialized !", timeout=2000
+        )
+
+    # Load model into memory cache
+    model_api.initialize_model(model_name=model_name)
+
+    # reset file banner
+    reset_nemo_model_file_script = """
+        <script>
+            document.getElementById('nemo_model_file').value = ""
+        </script>
+    """
+
+    result = result + reset_nemo_model_file_script
+    result = make_response(result)
+
+    # set cookies
+    result.set_cookie("model_name", model_name)
+    result.set_cookie("use_gpu", use_gpu_if_available)
+    return result
+
+
+def _store_model(nemo_model_file):
+    """
+    Preserve the model supplied by user into permanent cache
+    This cache needs to be manually deleted (if run locally), or gets deleted automatically
+    (when the container gets shutdown / killed).
+
+    Args:
+        nemo_model_file: User path to .nemo checkpoint.
+
+    Returns:
+        A file name (with a .nemo) at the end - to signify this is an uploaded checkpoint.
+    """
+    filename = secure_filename(nemo_model_file.filename)
+    file_basename = os.path.basename(filename)
+    model_dir = os.path.splitext(file_basename)[0]
+
+    model_store = os.path.join('models', model_dir)
+    if not os.path.exists(model_store):
+        os.makedirs(model_store)
+
+    # upload model
+    model_path = os.path.join(model_store, filename)
+    nemo_model_file.save(model_path)
+    return file_basename
+
+
+@app.route('/upload_audio_files', methods=['POST'])
+def upload_audio_files():
+    """
+    API Endpoint to upload audio files for inference.
+
+    The uploaded files must be wav files, 16 KHz sample rate, mono-channel audio samples.
+    """
+    # Try to get one or more files from form
+    try:
+        f = request.files.getlist('file')
+    except werkzeug.exceptions.BadRequestKeyError:
+        f = None
+
+    # If user did not select any file to upload, notify them.
+    if f is None or len(f) == 0:
+        toast = render_template('toast_msg.html', toast_message="No file has been selected to upload !", timeout=2000)
+        result = render_template('updates/upload_files_failed.html', pre_exec=toast, url=url_for('upload_audio_files'))
+        result = unescape(result)
+        return result
+
+    # temporary id to store data
+    uuid = str(uuid4())
+    data_store = os.path.join(app.config[f'UPLOAD_FOLDER'], uuid)
+
+    # If the user attempt to upload another set of files without first transcribing them,
+    # delete the old cache of files and create a new cache entirely
+    _remove_older_files_if_exists()
+
+    # Save each file into this unique cache
+    for fn in f:
+        filename = secure_filename(fn.filename)
+        if not os.path.exists(data_store):
+            os.makedirs(data_store)
+
+        fn.save(os.path.join(data_store, filename))
+        logging.info(f"Saving file : {fn.filename}")
+
+    # Update user that N files were uploaded.
+    msg = f"{len(f)} file(s) uploaded. Click to upload more !"
+    toast = render_template('toast_msg.html', toast_message=f"{len(f)} file(s) uploaded !", timeout=2000)
+    result = render_template(
+        'updates/upload_files_successful.html', pre_exec=toast, msg=msg, url=url_for('upload_audio_files')
+    )
+    result = unescape(result)
+
+    result = make_response(result)
+    result.set_cookie("uuid", uuid)
+    return result
+
+
+def _remove_older_files_if_exists():
+    """
+    Helper method to prevent cache leakage when user attempts to upload another set of files
+    without first transcribing the files already uploaded.
+    """
+    # remove old data store (if exists)
+    old_uuid = request.cookies.get('uuid', '')
+    if old_uuid is not None and old_uuid != '':
+        # delete old data store
+        old_data_store = os.path.join(app.config[f'UPLOAD_FOLDER'], old_uuid)
+
+        logging.info("Tried uploading more data without using old uploaded data. Purging data cache.")
+        shutil.rmtree(old_data_store, ignore_errors=True)
+
+
+@app.route('/remove_audio_files', methods=['POST'])
+def remove_audio_files():
+    """
+    API Endpoint for removing audio files
+
+    # Note: Sometimes data may persist due to set of circumstances:
+
+        - User uploads audio then closes app without transcribing anything
+
+    In such a case, the files will be deleted when gunicorn shutsdown, or container is stopped.
+    However the data may not be automatically deleted if the flast server is used as is.
+    """
+    # Get the unique cache id from cookie
+    uuid = request.cookies.get("uuid", "")
+    data_store = os.path.join(app.config[f'UPLOAD_FOLDER'], uuid)
+
+    # If the data does not exist (cache is empty), notify user
+    if not os.path.exists(data_store) or uuid == "":
+        files_dont_exist = render_template(
+            'toast_msg.html', toast_message="No files have been uploaded !", timeout=2000
+        )
+        result = render_template(
+            'updates/remove_files.html', pre_exec=files_dont_exist, url=url_for('remove_audio_files')
+        )
+        result = unescape(result)
+        return result
+
+    else:
+        # delete data that exists in cache
+        shutil.rmtree(data_store, ignore_errors=True)
+
+        logging.info("Removed all data")
+
+        # Notify user that cache was deleted.
+        toast = render_template('toast_msg.html', toast_message="All files removed !", timeout=2000)
+        result = render_template('updates/remove_files.html', pre_exec=toast, url=url_for('remove_audio_files'))
+        result = unescape(result)
+
+        result = make_response(result)
+        result.set_cookie("uuid", '', expires=0)
+        return result
+
+
+@app.route('/transcribe', methods=['POST'])
+def transcribe():
+    """
+    API Endpoint to transcribe a set of audio files.
+
+    The files are sorted according to their name, so order may not be same as upload order.
+
+    Utilizing the cached info inside the cookies, a model with selected name will be loaded into memory,
+    and maybe onto a GPU (if it is supported on the device).
+
+    Then the transcription api will be called from the model_api. If all is successful, a template is updated
+    with results. If some issue occurs (memory ran out, file is invalid format), notify the user.
+    """
+    # load model name from cookie
+    model_name = request.cookies.get('model_name')
+    logging.info(f"Model name : {model_name}")
+
+    # If model name is not selected via Load Model, notify user.
+    if model_name is None or model_name == '':
+        result = render_template('toast_msg.html', toast_message="Model has not been initialized !", timeout=2000)
+        return result
+
+    # load whether gpu should be used
+    use_gpu_if_available = request.cookies.get('use_gpu') == 'on'
+    gpu_used = torch.cuda.is_available() and use_gpu_if_available
+
+    # Load audio from paths
+    uuid = request.cookies.get("uuid", "")
+    data_store = os.path.join(app.config[f'UPLOAD_FOLDER'], uuid)
+
+    files = list(glob.glob(os.path.join(data_store, "*.wav")))
+
+    # If no files found in cache, notify user
+    if len(files) == 0:
+        result = render_template('toast_msg.html', toast_message="No audio files were found !", timeout=2000)
+        return result
+
+    # transcribe file via model api
+    t1 = time.time()
+    transcriptions = model_api.transcribe_all(files, model_name, use_gpu_if_available=use_gpu_if_available)
+    t2 = time.time()
+
+    # delete all transcribed files immediately
+    for fp in files:
+        try:
+            os.remove(fp)
+        except FileNotFoundError:
+            logging.info(f"Failed to delete transcribed file : {os.path.basename(fp)}")
+
+    # delete temporary transcription directory
+    shutil.rmtree(data_store, ignore_errors=True)
+
+    # If something happened during transcription, and it failed, notify user.
+    if type(transcriptions) == str and transcriptions == model_api.TAG_ERROR_DURING_TRANSCRIPTION:
+        toast = render_template(
+            'toast_msg.html',
+            toast_message=f"Failed to transcribe files due to unknown reason. "
+            f"Please provide 16 KHz Monochannel wav files onle.",
+            timeout=5000,
+        )
+        transcriptions = ["" for _ in range(len(files))]
+
+    else:
+        # Transcriptions obtained successfully, notify user.
+        toast = render_template(
+            'toast_msg.html',
+            toast_message=f"Transcribed {len(files)} files using {model_name} (gpu={gpu_used}), "
+            f"in {(t2 - t1): 0.2f} s",
+            timeout=5000,
+        )
+
+    # Write results to data table
+    results = []
+    for filename, transcript in zip(files, transcriptions):
+        results.append(dict(filename=os.path.basename(filename), transcription=transcript))
+
+    result = render_template('transcripts.html', transcripts=results)
+    result = toast + result
+    result = unescape(result)
+
+    result = make_response(result)
+    result.set_cookie("uuid", "", expires=0)
+    return result
+
+
+def remove_tmp_dir_at_exit():
+    """
+    Helper method to attempt a deletion of audio file cache on flask api exit.
+    Gunicorn and Docker container (based on gunicorn) will delete any remaining files on
+    shutdown of the gunicorn server or the docker container.
+
+    This is a patch that might not always work for Flask server, but in general should ensure
+    that local audio file cache is deleted.
+
+    This does *not* impact the model cache. Flask and Gunicorn servers will *never* delete uploaded models.
+    Docker container will delete models *only* when the container is killed (since models are uploaded to
+    local storage path inside container).
+    """
+    try:
+        uuid = request.cookies.get("uuid", "")
+
+        if uuid is not None or uuid != "":
+            cache_dir = os.path.join(os.path.join(app.config[f'UPLOAD_FOLDER'], uuid))
+            logging.info(f"Removing cache file for worker : {os.getpid()}")
+
+            if os.path.exists(cache_dir):
+                shutil.rmtree(cache_dir, ignore_errors=True)
+                logging.info(f"Deleted tmp folder : {cache_dir}")
+
+    except RuntimeError:
+        # Working outside of request context (probably shutdown)
+        # simply delete entire tmp folder
+        shutil.rmtree(app.config[f'UPLOAD_FOLDER'], ignore_errors=True)
+
+
+@app.route('/')
+def main():
+    """
+    API Endpoint for ASR Service.
+    """
+    nemo_model_names, local_model_names = model_api.get_model_names()
+    model_names = []
+    model_names.extend(local_model_names)  # prioritize local models
+    model_names.extend(nemo_model_names)  # attach all other pretrained models
+
+    # page initializations
+    result = render_template('main.html', model_names=model_names)
+    result = make_response(result)
+
+    # Reset cookies
+    result.set_cookie("model_name", '', expires=0)  # model name from pretrained model list
+    result.set_cookie("use_gpu", '', expires=0)  # flag to use gpu (if available)
+    result.set_cookie("uuid", '', expires=0)  # session id
+    return result
+
+
+# Register hook to delete file cache (for flask server only)
+atexit.register(remove_tmp_dir_at_exit)
+
+
+if __name__ == '__main__':
+    app.run(debug=True)

tts/NeMo/examples/asr/asr_webapp/docker_container_build.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# build container
+DOCKER_BUILDKIT=1 docker build -f Dockerfile --progress=plain -t nemo-asr-service:latest .
+

tts/NeMo/examples/asr/asr_webapp/docker_container_run.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# run container
+docker run -d --gpus=all -p 8000:8000 --shm-size=8g --ulimit memlock=-1 --ulimit \
+  stack=67108864 nemo-asr-service:latest
+

tts/NeMo/examples/asr/asr_webapp/gunicorn.conf.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# https://docs.gunicorn.org/en/stable/settings.html
+
+# NOTE: Do not import nemo / torch code here
+# Gunicorn creates forked processes - and CUDA cannot be used in forked multiprocessing environment.
+
+import shutil
+
+# General config
+bind = "0.0.0.0:8000"
+workers = 2
+
+# Worker specific config
+worker_connections = 1000
+timeout = 180  # 3 minutes of timeout
+
+
+def on_exit(server):
+    # delete tmp dir
+    print("Shutting down server ...")
+    print("Deleteing tmp directory ...")
+    shutil.rmtree('tmp/', ignore_errors=True)
+    print("Tmp directory deleted !")

tts/NeMo/examples/asr/asr_webapp/model_api.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import glob
+import os
+
+import torch
+
+import nemo.collections.asr as nemo_asr
+from nemo.utils import logging, model_utils
+
+# setup AMP (optional)
+if torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'):
+    logging.info("AMP enabled!\n")
+    autocast = torch.cuda.amp.autocast
+else:
+
+    @contextlib.contextmanager
+    def autocast():
+        yield
+
+
+MODEL_CACHE = {}
+
+# Special tags for fallbacks / user notifications
+TAG_ERROR_DURING_TRANSCRIPTION = "<ERROR_DURING_TRANSCRIPTION>"
+
+
+def get_model_names():
+    # Populate local copy of models
+    local_model_paths = glob.glob(os.path.join('models', "**", "*.nemo"), recursive=True)
+    local_model_names = list(sorted([os.path.basename(path) for path in local_model_paths]))
+
+    # Populate with pretrained nemo checkpoint list
+    nemo_model_names = set()
+    for model_info in nemo_asr.models.ASRModel.list_available_models():
+        for superclass in model_info.class_.mro():
+            if 'CTC' in superclass.__name__ or 'RNNT' in superclass.__name__:
+                nemo_model_names.add(model_info.pretrained_model_name)
+                break
+    nemo_model_names = list(sorted(nemo_model_names))
+    return nemo_model_names, local_model_names
+
+
+def initialize_model(model_name):
+    # load model
+    if model_name not in MODEL_CACHE:
+        if '.nemo' in model_name:
+            # use local model
+            model_name_no_ext = os.path.splitext(model_name)[0]
+            model_path = os.path.join('models', model_name_no_ext, model_name)
+
+            # Extract config
+            model_cfg = nemo_asr.models.ASRModel.restore_from(restore_path=model_path, return_config=True)
+            classpath = model_cfg.target  # original class path
+            imported_class = model_utils.import_class_by_path(classpath)  # type: ASRModel
+            logging.info(f"Restoring local model : {imported_class.__name__}")
+
+            # load model from checkpoint
+            model = imported_class.restore_from(restore_path=model_path, map_location='cpu')  # type: ASRModel
+
+        else:
+            # use pretrained model
+            model = nemo_asr.models.ASRModel.from_pretrained(model_name, map_location='cpu')
+
+        model.freeze()
+
+        # cache model
+        MODEL_CACHE[model_name] = model
+
+    model = MODEL_CACHE[model_name]
+    return model
+
+
+def transcribe_all(filepaths, model_name, use_gpu_if_available=True):
+    # instantiate model
+    if model_name in MODEL_CACHE:
+        model = MODEL_CACHE[model_name]
+    else:
+        model = initialize_model(model_name)
+
+    if torch.cuda.is_available() and use_gpu_if_available:
+        model = model.cuda()
+
+    # transcribe audio
+    logging.info("Begin transcribing audio...")
+    try:
+        with autocast():
+            with torch.no_grad():
+                transcriptions = model.transcribe(filepaths, batch_size=32)
+
+    except RuntimeError:
+        # Purge the cache to clear some memory
+        MODEL_CACHE.clear()
+
+        logging.info("Ran out of memory on device - performing inference on CPU for now")
+
+        try:
+            model = model.cpu()
+
+            with torch.no_grad():
+                transcriptions = model.transcribe(filepaths, batch_size=32)
+
+        except Exception as e:
+            logging.info(f"Exception {e} occured while attemting to transcribe audio. Returning error message")
+            return TAG_ERROR_DURING_TRANSCRIPTION
+
+    logging.info(f"Finished transcribing {len(filepaths)} files !")
+
+    # If RNNT models transcribe, they return a tuple (greedy, beam_scores)
+    if type(transcriptions[0]) == list and len(transcriptions) == 2:
+        # get greedy transcriptions only
+        transcriptions = transcriptions[0]
+
+    # Force onto CPU
+    model = model.cpu()
+    MODEL_CACHE[model_name] = model
+
+    return transcriptions

tts/NeMo/examples/asr/asr_webapp/requirements.txt

@@ -0,0 +1,4 @@
+-e git+git://github.com/NVIDIA/NeMo.git@v1.0.2#egg=nemo_toolkit[all]
+flask
+flask-cors
+gunicorn

tts/NeMo/examples/asr/asr_webapp/templates/main.html

@@ -0,0 +1,150 @@
+<!--
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+-->
+<!doctype html>
+<html lang="en">
+<head>
+    <link rel="stylesheet" href="http://fonts.googleapis.com/css?family=Roboto:300,400,500,700" type="text/css">
+    <meta charset="utf-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    <title>Automatic Speech Recognition</title>
+    <link rel="stylesheet" href="//code.jquery.com/ui/1.12.1/themes/base/jquery-ui.css">
+
+    <link rel="stylesheet" href="https://fonts.googleapis.com/icon?family=Material+Icons">
+    <link rel="stylesheet" href="https://code.getmdl.io/1.3.0/material.light_green-lime.min.css">
+    <script defer src="https://code.getmdl.io/1.3.0/material.min.js"></script>
+    <style>
+    .ui-resizable-se {
+      bottom: 17px;
+    }
+
+    </style>
+    <script src="https://code.jquery.com/jquery-1.12.4.js"></script>
+    <script src="https://code.jquery.com/ui/1.12.1/jquery-ui.js"></script>
+
+    <!-- Load from https://htmx.org/ -->
+    <script src="https://unpkg.com/htmx.org@1.3.1"></script>
+</head>
+<body>
+
+<!-- Always shows a header, even in smaller screens. -->
+<div class="mdl-layout mdl-js-layout mdl-layout--fixed-header">
+    <header class="mdl-layout__header">
+        <div class="mdl-layout__header-row">
+            <!-- Title -->
+            <span class="mdl-layout-title">NeMo Automatic Speech Recognition Demo</span>
+            <!-- Add spacer, to align navigation to the right -->
+            <div class="mdl-layout-spacer"></div>
+            <!-- Navigation. We hide it in small screens. -->
+            <nav class="mdl-navigation mdl-layout--large-screen-only">
+                <a class="mdl-navigation__link" href="#top">Top</a>
+            </nav>
+        </div>
+    </header>
+    <div class="mdl-layout__drawer">
+        <span class="mdl-layout-title">ASR Transcription</span>
+        <nav class="mdl-navigation">
+            <a class="mdl-navigation__link" href="https://github.com/NVIDIA/NeMo">Github</a>
+        </nav>
+    </div>
+    <main class="mdl-layout__content">
+        <div class="page-content">
+
+            <!--Placeholder for top of app-->
+            <div id="top"></div>
+
+            <form id="model_names_form" action="" enctype="multipart/form-data" style="display: inline;">
+
+                <h5> Model Name :
+
+                    <!-- Format selectable list of models -->
+                    <div class="mdl-textfield mdl-js-textfield mdl-textfield--floating-label" style="width: 25%">
+                        <select name="model_names_select" class="mdl-textfield__input" id="model_names_select_id">
+                            {% for x in model_names %}
+                            <option value="{{ x }}" {% if loop.first %} SELECTED{% endif %}>{{ x }}</option>
+                            {% endfor %}
+                        </select>
+                    </div>
+
+                    <!-- Load Model -->
+                    <button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
+                            hx-post="{{ url_for('initialize_model') }}" hx-target="this" hx-swap="afterend"
+                            hx-encoding="multipart/form-data">
+                        Load Model
+                    </button>
+
+                    <!-- Load .nemo checkpoint -->
+                    <!-- Comment this section out to remove ability to upload files -->
+
+                    &nbsp; <br>
+                    NeMo File :
+                    <input type="file" name="nemo_model" accept=".nemo" id="nemo_model_file"/>
+
+                    <!-- Comment out upto here to remove ability to upload models -->
+
+                </h5>
+
+                <!-- GPU checkbox -->
+                <label class="mdl-checkbox mdl-js-checkbox mdl-js-ripple-effect" for="use_gpu_ckbx_id">
+                    <input type="checkbox" id="use_gpu_ckbx_id" name="use_gpu_ckbx" class="mdl-checkbox__input" checked>
+                    <span class="mdl-checkbox__label">Use GPU (if available)</span>
+                </label>
+
+            </form>
+
+            <hr>
+            <br>
+
+            <!-- Form input for files -->
+            <form action="" enctype="multipart/form-data">
+                <input type="file" name="file" accept=".wav" multiple=""/>
+
+                <button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
+                        hx-post="{{ url_for('upload_audio_files') }}"
+                        hx-target="this" hx-swap="outerHTML" hx-encoding="multipart/form-data">
+                    Upload audio file(s)
+                </button>
+
+                <button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
+                        hx-post="{{ url_for('remove_audio_files') }}"
+                        hx-target="this" hx-swap="outerHTML">
+                    Remove all files
+                </button>
+            </form>
+            <hr>
+            <br>
+
+            <!-- Transcribe Button -->
+            <button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
+                    hx-post="{{ url_for('transcribe') }}" hx-target="#transcribeResults">
+                Transcribe all
+            </button>
+
+            <div id="transcribeResults"></div>
+
+
+            <!--Material Toast -->
+
+            <div id="toast" class="mdl-js-snackbar mdl-snackbar">
+                <div class="mdl-snackbar__text"></div>
+                <button class="mdl-snackbar__action" type="button"></button>
+            </div>
+
+        </div>
+    </main>
+</div>
+
+</body>
+</html>

tts/NeMo/examples/asr/asr_webapp/templates/toast_msg.html

@@ -0,0 +1,28 @@
+<!--
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+-->
+
+<script>
+(function() {
+  'use strict';
+  var snackbarContainer = document.querySelector('#toast');
+  var data = {
+      message: '{{ toast_message }}',
+      timeout: {{ timeout }},
+  };
+  snackbarContainer.MaterialSnackbar.showSnackbar(data);
+}());
+
+</script>

tts/NeMo/examples/asr/asr_webapp/templates/transcripts.html

@@ -0,0 +1,43 @@
+<!--
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+-->
+
+<table class="transcript-view-table mdl-data-table mdl-js-data-table mdl-shadow--4dp">
+    <thead>
+    <tr>
+        <th class="mdl-data-table__cell--non-numeric">Filename</th>
+        <th class="mdl-data-table__cell--non-numeric">Transcript</th>
+    </tr>
+    </thead>
+    <tbody>
+    {% for transcript in transcripts %}
+    <tr>
+        <td class="property-value-col mdl-data-table__cell--non-numeric" style="width: 10%">
+            <div class="col-value" style="overflow-x:auto;">
+                {{ transcript.filename }}
+            </div>
+        </td>
+        <td class="property-value-col mdl-data-table__cell--non-numeric">
+            <div class="col-value mdl-textfield mdl-js-textfield"
+                 style="width: 100%; display: block;">
+          <textarea class="mdl-textfield__input" type="text" rows="2" readonly>
+            {{ transcript.transcription }}
+          </textarea>
+            </div>
+        </td>
+    </tr>
+    {% endfor %}
+    </tbody>
+</table>

tts/NeMo/examples/asr/asr_webapp/templates/updates/remove_files.html

@@ -0,0 +1,22 @@
+<!--
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+-->
+
+{{ pre_exec }}
+<button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
+        hx-post="{{ url }}"
+        hx-target="this" hx-swap="outerHTML">
+    Remove all files
+</button>

tts/NeMo/examples/asr/asr_webapp/templates/updates/upload_files_failed.html

@@ -0,0 +1,22 @@
+<!--
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+-->
+
+{{ pre_exec }}
+<button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
+        hx-post="{{ url }}"
+        hx-target="this" hx-swap="outerHTML" hx-encoding="multipart/form-data">
+    Upload audio file(s)
+</button>

tts/NeMo/examples/asr/asr_webapp/templates/updates/upload_files_successful.html

@@ -0,0 +1,22 @@
+<!--
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+-->
+
+{{ pre_exec }}
+<button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
+        hx-post="{{ url }}"
+        hx-target="this" hx-swap="outerHTML" hx-encoding="multipart/form-data">
+    {{ msg }}
+</button>

tts/NeMo/examples/asr/asr_webapp/wsgi.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from asr_service import app
+
+if __name__ == '__main__':
+    app.run()

tts/NeMo/examples/asr/conf/citrinet/citrinet_384.yaml

@@ -0,0 +1,428 @@
+# This config contains the default values for training a Citrinet model with CTC loss and BPE-based vocabulary.
+# Default learning parameters in this config are set for effective batch size of 1k on 32 GPUs.
+# To train it with smaller batch sizes, you may need to re-tune the learning parameters or use higher accumulate_grad_batches.
+# If training for a short time, you can also reduce weight decay to 0.
+
+# Training Recipe
+# This model can be trained using the default settings in this config with FP32 precision.
+# When training under AMP, increase `warmup_steps` to 5000 for stable training.
+# In order to create Citrinet-C, find-replace `filters: 384` with `filters: C`.
+# When reducing the receptive field of these models, it is advised to reduce the amount of augmentation
+# for larger models from 10x time masking to 5x or 2x time masking.
+# For further details regarding Citrinet, visit - https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#citrinet
+
+name: &name "Citrinet-384-8x-Stride"
+
+model:
+  sample_rate: &sample_rate 16000
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    batch_size: 32
+    trim_silence: false
+    max_duration: 16.7
+    shuffle: true
+    is_tarred: false
+    tarred_audio_filepaths: null
+    use_start_end_token: false
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    batch_size: 32
+    shuffle: false
+    use_start_end_token: false
+
+  test_ds:
+    manifest_filepath: null
+    sample_rate: 16000
+    batch_size: 32
+    shuffle: false
+    use_start_end_token: false
+
+  model_defaults:
+    repeat: 5
+    dropout: 0.0
+    separable: true
+    se: true
+    se_context_size: -1
+    kernel_size_factor: 1.0
+
+  tokenizer:
+    dir: ???  # path to directory which contains either tokenizer.model (bpe) or vocab.txt (for wpe)
+    type: ???  # Can be either bpe or wpe
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    sample_rate: *sample_rate
+    normalize: "per_feature"
+    window_size: 0.025
+    window_stride: 0.01
+    window: "hann"
+    features: &n_mels 80
+    n_fft: 512
+    frame_splicing: 1
+    dither: 0.00001
+    pad_to: 16
+    stft_conv: false
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    freq_masks: 2
+    time_masks: 2
+    freq_width: 27
+    time_width: 0.05
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+      - filters: 384
+        repeat: 1
+        kernel: [5]
+        stride: [1]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [11]
+        stride: [2]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [13]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [15]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [17]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [19]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [21]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [13]
+        stride: [2]  # *stride
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [15]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [17]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [19]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [21]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [23]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [25]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [25]
+        stride: [2]  # stride
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [27]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [29]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [31]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [33]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [35]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [37]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 384
+        repeat: ${model.model_defaults.repeat}
+        kernel: [39]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: &enc_final 640
+        repeat: 1
+        kernel: [41]
+        stride: [1]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoder
+    feat_in: *enc_final
+    num_classes: -1  # filled with vocabulary size from tokenizer at runtime
+    vocabulary: []  # filled with vocabulary from tokenizer at runtime
+
+  optim:
+    name: novograd
+    lr: 0.05
+
+    # optimizer arguments
+    betas: [0.8, 0.25]
+    weight_decay: 0.001
+
+    # scheduler setup
+    sched:
+      name: CosineAnnealing
+
+      # scheduler config override
+      warmup_steps: 1000
+      warmup_ratio: null
+      min_lr: 1e-5
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 100
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  checkpoint_callback: false  # Provided by exp_manager
+  logger: false  # Provided by exp_manager
+  log_every_n_steps: 100  # Interval of logging.
+  val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  check_val_every_n_epoch: 1
+  precision: 32
+  sync_batchnorm: false
+  benchmark: false
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: true
+  create_checkpoint_callback: true
+  checkpoint_callback_params:
+    monitor: "val_wer"
+    mode: "min"
+    save_top_k: 3
+  create_wandb_logger: false
+  wandb_logger_kwargs:
+    name: null
+    project: null
+    entity: null
+  resume_if_exists: false
+  resume_ignore_no_checkpoint: false
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/asr/conf/citrinet/citrinet_512.yaml

@@ -0,0 +1,428 @@
+# This config contains the default values for training a Citrinet model with CTC loss and BPE-based vocabulary.
+# Default learning parameters in this config are set for effective batch size of 1k on 32 GPUs.
+# To train it with smaller batch sizes, you may need to re-tune the learning parameters or use higher accumulate_grad_batches.
+# If training for a short time, you can also reduce weight decay to 0.
+
+# Training Recipe
+# This model can be trained using the default settings in this config with FP32 precision.
+# When training under AMP, increase `warmup_steps` to 5000 for stable training.
+# In order to create Citrinet-C, find-replace `filters: 384` with `filters: C`.
+# When reducing the receptive field of these models, it is advised to reduce the amount of augmentation
+# for larger models from 10x time masking to 5x or 2x time masking.
+# For further details regarding Citrinet, visit - https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#citrinet
+
+name: &name "Citrinet-512-8x-Stride"
+
+model:
+  sample_rate: &sample_rate 16000
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    batch_size: 32
+    trim_silence: false
+    max_duration: 16.7
+    shuffle: true
+    is_tarred: false
+    tarred_audio_filepaths: null
+    use_start_end_token: false
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    batch_size: 32
+    shuffle: false
+    use_start_end_token: false
+
+  test_ds:
+    manifest_filepath: null
+    sample_rate: 16000
+    batch_size: 32
+    shuffle: false
+    use_start_end_token: false
+
+  model_defaults:
+    repeat: 5
+    dropout: 0.1
+    separable: true
+    se: true
+    se_context_size: -1
+    kernel_size_factor: 1.0
+
+  tokenizer:
+    dir: ???  # path to directory which contains either tokenizer.model (bpe) or vocab.txt (for wpe)
+    type: ???  # Can be either bpe or wpe
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    sample_rate: *sample_rate
+    normalize: "per_feature"
+    window_size: 0.025
+    window_stride: 0.01
+    window: "hann"
+    features: &n_mels 80
+    n_fft: 512
+    frame_splicing: 1
+    dither: 0.00001
+    pad_to: 16
+    stft_conv: false
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    freq_masks: 2
+    time_masks: 10
+    freq_width: 27
+    time_width: 0.05
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+      - filters: 512
+        repeat: 1
+        kernel: [5]
+        stride: [1]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [11]
+        stride: [2]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [13]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [15]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [17]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [19]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [21]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [13]
+        stride: [2]  # *stride
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [15]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [17]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [19]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [21]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [23]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [25]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [25]
+        stride: [2]  # stride
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [27]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [29]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [31]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [33]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [35]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [37]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: 512
+        repeat: ${model.model_defaults.repeat}
+        kernel: [39]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+      - filters: &enc_final 640
+        repeat: 1
+        kernel: [41]
+        stride: [1]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        kernel_size_factor: ${model.model_defaults.kernel_size_factor}
+
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoder
+    feat_in: *enc_final
+    num_classes: -1  # filled with vocabulary size from tokenizer at runtime
+    vocabulary: []  # filled with vocabulary from tokenizer at runtime
+
+  optim:
+    name: novograd
+    lr: 0.05
+
+    # optimizer arguments
+    betas: [0.8, 0.25]
+    weight_decay: 0.001
+
+    # scheduler setup
+    sched:
+      name: CosineAnnealing
+
+      # scheduler config override
+      warmup_steps: 1000
+      warmup_ratio: null
+      min_lr: 1e-5
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 100
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  checkpoint_callback: false  # Provided by exp_manager
+  logger: false  # Provided by exp_manager
+  log_every_n_steps: 100  # Interval of logging.
+  val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  check_val_every_n_epoch: 1
+  precision: 32
+  sync_batchnorm: false
+  benchmark: false
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: true
+  create_checkpoint_callback: true
+  checkpoint_callback_params:
+    monitor: "val_wer"
+    mode: "min"
+    save_top_k: 3
+  create_wandb_logger: false
+  wandb_logger_kwargs:
+    name: null
+    project: null
+    entity: null
+  resume_if_exists: false
+  resume_ignore_no_checkpoint: false
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/asr/conf/citrinet/config_bpe.yaml

@@ -0,0 +1,183 @@
+name: &name "ContextNet5x1"
+sample_rate: &sample_rate 16000
+repeat: &repeat 1
+dropout: &dropout 0.0
+separable: &separable true
+
+model:
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    batch_size: 32
+    trim_silence: True
+    max_duration: 16.7
+    shuffle: True
+    is_tarred: False
+    tarred_audio_filepaths: null
+    shard_strategy: "scatter"
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    batch_size: 32
+    shuffle: False
+
+  tokenizer:
+    dir: ???  # path to directory which contains either tokenizer.model (bpe) or vocab.txt (for wpe)
+    type: ???  # Can be either bpe or wpe
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    normalize: "per_feature"
+    window_size: 0.02
+    sample_rate: *sample_rate
+    window_stride: 0.01
+    window: "hann"
+    features: &n_mels 64
+    n_fft: 512
+    frame_splicing: 1
+    dither: 0.00001
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    rect_freq: 50
+    rect_masks: 5
+    rect_time: 120
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+      - filters: 128
+        repeat: 1
+        kernel: [11]
+        stride: [1]
+        dilation: [1]
+        dropout: *dropout
+        residual: true
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: 256
+        repeat: *repeat
+        kernel: [13]
+        stride: [1]
+        dilation: [1]
+        dropout: *dropout
+        residual: true
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: 256
+        repeat: *repeat
+        kernel: [15]
+        stride: [1]
+        dilation: [1]
+        dropout: *dropout
+        residual: true
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: 256
+        repeat: *repeat
+        kernel: [17]
+        stride: [1]
+        dilation: [1]
+        dropout: *dropout
+        residual: true
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: 256
+        repeat: *repeat
+        kernel: [19]
+        stride: [1]
+        dilation: [1]
+        dropout: *dropout
+        residual: true
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: 256
+        repeat: 1
+        kernel: [21]
+        stride: [1]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: &enc_feat_out 1024
+        repeat: 1
+        kernel: [1]
+        stride: [1]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoder
+    feat_in: 1024
+    num_classes: -1  # filled with vocabulary size from tokenizer at runtime
+    vocabulary: []  # filled with vocabulary from tokenizer at runtime
+
+  optim:
+    name: adam
+    # _target_: nemo.core.optim.optimizers.Adam
+    lr: .1
+
+    # optimizer arguments
+    betas: [0.9, 0.999]
+    weight_decay: 0.0001
+
+    # scheduler setup
+    sched:
+      name: CosineAnnealing
+
+      # scheduler config override
+      warmup_steps: null
+      warmup_ratio: 0.05
+      min_lr: 1e-6
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 5
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+  create_wandb_logger: False
+  wandb_logger_kwargs:
+    name: null
+    project: null
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/asr/conf/config.yaml

@@ -0,0 +1,183 @@
+name: &name "QuartzNet15x5"
+sample_rate: &sample_rate 16000
+repeat: &repeat 1
+dropout: &dropout 0.0
+separable: &separable true
+labels: &labels [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+         "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+model:
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    labels: *labels
+    batch_size: 32
+    trim_silence: True
+    max_duration: 16.7
+    shuffle: True
+    is_tarred: False
+    tarred_audio_filepaths: null
+    tarred_shard_strategy: "scatter"
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    labels: *labels
+    batch_size: 32
+    shuffle: False
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    normalize: "per_feature"
+    window_size: 0.02
+    sample_rate: *sample_rate
+    window_stride: 0.01
+    window: "hann"
+    features: &n_mels 64
+    n_fft: 512
+    frame_splicing: 1
+    dither: 0.00001
+    stft_conv: false
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    rect_freq: 50
+    rect_masks: 5
+    rect_time: 120
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+      - filters: 128
+        repeat: 1
+        kernel: [11]
+        stride: [1]
+        dilation: [1]
+        dropout: *dropout
+        residual: true
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: 256
+        repeat: *repeat
+        kernel: [13]
+        stride: [1]
+        dilation: [1]
+        dropout: *dropout
+        residual: true
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: 256
+        repeat: *repeat
+        kernel: [15]
+        stride: [1]
+        dilation: [1]
+        dropout: *dropout
+        residual: true
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: 256
+        repeat: *repeat
+        kernel: [17]
+        stride: [1]
+        dilation: [1]
+        dropout: *dropout
+        residual: true
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: 256
+        repeat: *repeat
+        kernel: [19]
+        stride: [1]
+        dilation: [1]
+        dropout: *dropout
+        residual: true
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: 256
+        repeat: 1
+        kernel: [21]
+        stride: [1]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+      - filters: &enc_feat_out 1024
+        repeat: 1
+        kernel: [1]
+        stride: [1]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: *separable
+        se: true
+        se_context_size: -1
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoder
+    feat_in: 1024
+    num_classes: 28
+    vocabulary: *labels
+
+  optim:
+    name: novograd
+    # _target_: nemo.core.optim.optimizers.Novograd
+    lr: .01
+    # optimizer arguments
+    betas: [0.8, 0.5]
+    weight_decay: 0.001
+
+    # scheduler setup
+    sched:
+      name: CosineAnnealing
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: null
+      min_lr: 0.0
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 5
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/asr/conf/conformer/conformer_ctc_bpe.yaml

@@ -0,0 +1,189 @@
+# It contains the default values for training a Conformer-CTC ASR model, large size (~120M) with CTC loss and sub-word encoding.
+
+# Architecture and training config:
+# Default learning parameters in this config are set for effective batch size of 2K. To train it with smaller effective
+# batch sizes, you may need to re-tune the learning parameters or use higher accumulate_grad_batches.
+# Here are the recommended configs for different variants of Conformer-CTC, other parameters are the same as in this config file.
+# One extra layer (compared to original paper) is added to the medium and large variants to compensate for replacing the LSTM decoder with a linear one.
+#
+#  +-------------+---------+---------+----------+------------+-----+
+#  | Model       | d_model | n_heads | n_layers | time_masks | lr  |
+#  +=============+=========+========+===========+============+=====+
+#  | Small  (13M)|   176   |    4   |    16     |     5      | 5.0 |
+#  +-------------+---------+--------+-----------+------------+-----+
+#  | Medium (30M)|   256   |    4   |    18     |     5      | 5.0 |
+#  +-------------+---------+--------+-----------+------------+-----+
+#  | Large (121M)|   512   |    8   |    18     |     10     | 2.0 |
+#  +---------------------------------------------------------------+
+#
+# If you do not want to train with AMP, you may use weight decay of 0.0 or reduce the number of time maskings to 2
+# with time_width=100. It may help when you want to train for fewer epochs and need faster convergence.
+# With weight_decay=0.0, learning rate may need to get reduced to 2.0.
+
+# You may find more info about Conformer-CTC here: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/models.html#conformer-ctc
+# Pre-trained models of Conformer-CTC can be found here: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/results.html
+# The checkpoint of the large model trained on LibriSpeech with this recipe can be found here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large_ls
+
+name: "Conformer-CTC-BPE"
+
+model:
+  sample_rate: 16000
+  log_prediction: true # enables logging sample predictions in the output during training
+  ctc_reduction: 'mean_batch'
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    batch_size: 16 # you may increase batch_size if your memory allows
+    shuffle: true
+    num_workers: 8
+    pin_memory: true
+    use_start_end_token: false
+    trim_silence: false
+    max_duration: 16.7 # it is set for LibriSpeech, you may need to update it for your dataset
+    min_duration: 0.1
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    batch_size: 16 # you may increase batch_size if your memory allows
+    shuffle: false
+    num_workers: 8
+    pin_memory: true
+    use_start_end_token: false
+
+  test_ds:
+    manifest_filepath: null
+    sample_rate: ${model.sample_rate}
+    batch_size: 16 # you may increase batch_size if your memory allows
+    shuffle: false
+    num_workers: 8
+    pin_memory: true
+    use_start_end_token: false
+
+  # recommend small vocab size of 128 or 256 when using 4x sub-sampling
+  # you may find more detail on how to train a tokenizer at: /scripts/tokenizers/process_asr_text_tokenizer.py
+  tokenizer:
+    dir: ???  # path to directory which contains either tokenizer.model (bpe) or vocab.txt (wpe)
+    type: wpe  # Can be either bpe (SentencePiece tokenizer) or wpe (WordPiece tokenizer)
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    sample_rate: ${model.sample_rate}
+    normalize: "per_feature"
+    window_size: 0.025
+    window_stride: 0.01
+    window: "hann"
+    features: 80
+    n_fft: 512
+    log: true
+    frame_splicing: 1
+    dither: 0.00001
+    pad_to: 0
+    pad_value: 0.0
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    freq_masks: 2 # set to zero to disable it
+    # you may use lower time_masks for smaller models to have a faster convergence
+    time_masks: 10 # set to zero to disable it
+    freq_width: 27
+    time_width: 0.05
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConformerEncoder
+    feat_in: ${model.preprocessor.features}
+    feat_out: -1 # you may set it if you need different output size other than the default d_model
+    n_layers: 18
+    d_model: 512
+
+    # Sub-sampling params
+    subsampling: striding # vggnet or striding, vggnet may give better results but needs more memory
+    subsampling_factor: 4 # must be power of 2
+    subsampling_conv_channels: -1 # -1 sets it to d_model
+
+    # Feed forward module's params
+    ff_expansion_factor: 4
+
+    # Multi-headed Attention Module's params
+    self_attention_model: rel_pos # rel_pos or abs_pos
+    n_heads: 8 # may need to be lower for smaller d_models
+    # [left, right] specifies the number of steps to be seen from left and right of each step in self-attention
+    att_context_size: [-1, -1] # -1 means unlimited context
+    xscaling: true # scales up the input embeddings by sqrt(d_model)
+    untie_biases: true # unties the biases of the TransformerXL layers
+    pos_emb_max_len: 5000
+
+    # Convolution module's params
+    conv_kernel_size: 31
+
+    ### regularization
+    dropout: 0.1 # The dropout used in most of the Conformer Modules
+    dropout_emb: 0.0 # The dropout used for embeddings
+    dropout_att: 0.1 # The dropout for multi-headed attention modules
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoder
+    feat_in: null
+    num_classes: -1
+    vocabulary: []
+
+  optim:
+    name: adamw
+    lr: 2.0
+    # optimizer arguments
+    betas: [0.9, 0.98]
+    # less necessity for weight_decay as we already have large augmentations with SpecAug
+    # you may need weight_decay for large models, stable AMP training, small datasets, or when lower augmentations are used
+    # weight decay of 0.0 with lr of 2.0 also works fine
+    weight_decay: 1e-3
+
+    # scheduler setup
+    sched:
+      name: NoamAnnealing
+      d_model: ${model.encoder.d_model}
+      # scheduler config override
+      warmup_steps: 10000
+      warmup_ratio: null
+      min_lr: 1e-6
+
+trainer:
+  gpus: -1 # number of GPUs, -1 would use all available GPUs
+  num_nodes: 1
+  max_epochs: 1000
+  max_steps: null # computed at runtime if not set
+  val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  gradient_clip_val: 0.0
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
+  log_every_n_steps: 10  # Interval of logging.
+  progress_bar_refresh_rate: 10
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+  num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
+  check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
+  sync_batchnorm: true
+  checkpoint_callback: false  # Provided by exp_manager
+  logger: false  # Provided by exp_manager
+
+exp_manager:
+  exp_dir: null
+  name: ${name}
+  create_tensorboard_logger: true
+  create_checkpoint_callback: true
+  checkpoint_callback_params:
+    # in case of multiple validation sets, first one is used
+    monitor: "val_wer"
+    mode: "min"
+    save_top_k: 3
+
+  # you need to set these two to True to continue the training
+  resume_if_exists: false
+  resume_ignore_no_checkpoint: false
+
+  # You may use this section to create a W&B logger
+  create_wandb_logger: false
+  wandb_logger_kwargs:
+    name: null
+    project: null

tts/NeMo/examples/asr/conf/conformer/conformer_ctc_char.yaml

@@ -0,0 +1,165 @@
+# It contains the default values for training a Conformer-CTC ASR model, large size (~120M) with CTC loss and char-based vocabulary.
+# Char-based encoding may give lower accuracy than sub-word encoding (conformer_ctc_bpe.yaml).
+
+# You may find more detail on Conformer-CTC at `examples/asr/conf/conformer/conformer_ctc_bpe.yaml`
+
+name: "Conformer-CTC-Char"
+
+model:
+  sample_rate: 16000
+  labels:  [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+            "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+  log_prediction: true # enables logging sample predictions in the output during training
+  ctc_reduction: 'mean_batch'
+
+  train_ds:
+    manifest_filepath: ???
+    labels: ${model.labels}
+    sample_rate: ${model.sample_rate}
+    batch_size: 16 # you may increase batch_size if your memory allows
+    shuffle: true
+    num_workers: 8
+    pin_memory: true
+    trim_silence: false
+    max_duration: 16.7 # it is set for LibriSpeech, you may need to update it for your dataset
+    min_duration: 0.1
+
+  validation_ds:
+    manifest_filepath: ???
+    labels: ${model.labels}
+    sample_rate: ${model.sample_rate}
+    batch_size: 16 # you may increase batch_size if your memory allows
+    shuffle: false
+    num_workers: 8
+    pin_memory: true
+
+  test_ds:
+    manifest_filepath: null
+    labels: ${model.labels}
+    sample_rate: ${model.sample_rate}
+    batch_size: 16 # you may increase batch_size if your memory allows
+    shuffle: false
+    num_workers: 8
+    pin_memory: true
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    sample_rate: ${model.sample_rate}
+    normalize: "per_feature"
+    window_size: 0.025
+    window_stride: 0.01
+    window: "hann"
+    features: 80
+    n_fft: 512
+    log: true
+    frame_splicing: 1
+    dither: 0.00001
+    pad_to: 0
+    pad_value: 0.0
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    freq_masks: 2 # set to zero to disable it
+    # you may use lower time_masks for smaller models to have a faster convergence
+    time_masks: 10 # set to zero to disable it
+    freq_width: 27
+    time_width: 0.05
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConformerEncoder
+    feat_in: ${model.preprocessor.features}
+    feat_out: -1 # you may set it if you need different output size other than the default d_model
+    n_layers: 16
+    d_model: 256
+
+    # Sub-sampling params
+    subsampling: striding # vggnet or striding, vggnet may give better results but needs more memory
+    subsampling_factor: 4 # must be power of 2
+    subsampling_conv_channels: -1 # set to -1 to make it equal to the d_model
+
+    # Feed forward module's params
+    ff_expansion_factor: 4
+
+    # Multi-headed Attention Module's params
+    self_attention_model: rel_pos # rel_pos or abs_pos
+    n_heads: 8 # may need to be lower for smaller d_models
+    # [left, right] specifies the number of steps to be seen from left and right of each step in self-attention
+    att_context_size: [-1, -1] # -1 means unlimited context
+    xscaling: true # scales up the input embeddings by sqrt(d_model)
+    untie_biases: true # unties the biases of the TransformerXL layers
+    pos_emb_max_len: 5000
+
+    # Convolution module's params
+    conv_kernel_size: 31
+
+    ### regularization
+    dropout: 0.1 # The dropout used in most of the Conformer Modules
+    dropout_emb: 0.0 # The dropout used for embeddings
+    dropout_att: 0.1 # The dropout for multi-headed attention modules
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoder
+    feat_in: null
+    num_classes: -1
+    vocabulary: ${model.labels}
+
+  optim:
+    name: adamw
+    lr: 2.0
+    # optimizer arguments
+    betas: [0.9, 0.98]
+    # less necessity for weight_decay as we already have large augmentations with SpecAug
+    # you may need weight_decay for large models, stable AMP training, small datasets, or when lower augmentations are used
+    # weight decay of 0.0 with lr of 2.0 also works fine
+    weight_decay: 1e-3
+
+    # scheduler setup
+    sched:
+      name: NoamAnnealing
+      d_model: ${model.encoder.d_model}
+      # scheduler config override
+      warmup_steps: 10000
+      warmup_ratio: null
+      min_lr: 1e-6
+
+trainer:
+  gpus: -1 # number of GPUs, -1 would use all available GPUs
+  num_nodes: 1
+  max_epochs: 1000
+  max_steps: null # computed at runtime if not set
+  val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  gradient_clip_val: 0.0
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
+  log_every_n_steps: 10  # Interval of logging.
+  progress_bar_refresh_rate: 10
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+  num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
+  check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
+  sync_batchnorm: true
+  checkpoint_callback: false  # Provided by exp_manager
+  logger: false  # Provided by exp_manager
+
+
+exp_manager:
+  exp_dir: null
+  name: ${name}
+  create_tensorboard_logger: true
+  create_checkpoint_callback: true
+  checkpoint_callback_params:
+    # in case of multiple validation sets, first one is used
+    monitor: "val_wer"
+    mode: "min"
+    save_top_k: 3
+
+  # you need to set these two to True to continue the training
+  resume_if_exists: false
+  resume_ignore_no_checkpoint: false
+
+  # You may use this section to create a W&B logger
+  create_wandb_logger: false
+  wandb_logger_kwargs:
+    name: null
+    project: null

tts/NeMo/examples/asr/conf/jasper/jasper_10x5dr.yaml

@@ -0,0 +1,208 @@
+name: &name "Jasper10x5"
+
+model:
+  sample_rate: &sample_rate 16000
+  labels: &labels [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+                   "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    labels: *labels
+    batch_size: 32
+    trim_silence: True
+    max_duration: 16.7
+    shuffle: True
+    is_tarred: False
+    tarred_audio_filepaths: null
+    tarred_shard_strategy: "scatter"
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    labels: *labels
+    batch_size: 32
+    shuffle: False
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    normalize: "per_feature"
+    window_size: 0.02
+    sample_rate: *sample_rate
+    window_stride: 0.01
+    window: "hann"
+    features: &n_mels 64
+    n_fft: 512
+    frame_splicing: 1
+    dither: 0.00001
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    rect_freq: 50
+    rect_masks: 5
+    rect_time: 120
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+    jasper:
+      - dilation: [1]
+        dropout: 0.2
+        filters: 256
+        kernel: [11]
+        repeat: 1
+        residual: false
+        stride: [2]
+      - dilation: [1]
+        dropout: 0.2
+        filters: 256
+        kernel: [11]
+        repeat: 5
+        residual: true
+        residual_dense: true
+        stride: [1]
+      - dilation: [1]
+        dropout: 0.2
+        filters: 256
+        kernel: [11]
+        repeat: 5
+        residual: true
+        residual_dense: true
+        stride: [1]
+      - dilation: [1]
+        dropout: 0.2
+        filters: 384
+        kernel: [13]
+        repeat: 5
+        residual: true
+        residual_dense: true
+        stride: [1]
+      - dilation: [1]
+        dropout: 0.2
+        filters: 384
+        kernel: [13]
+        repeat: 5
+        residual: true
+        residual_dense: true
+        stride: [1]
+      - dilation: [1]
+        dropout: 0.2
+        filters: 512
+        kernel: [17]
+        repeat: 5
+        residual: true
+        residual_dense: true
+        stride: [1]
+      - dilation: [1]
+        dropout: 0.2
+        filters: 512
+        kernel: [17]
+        repeat: 5
+        residual: true
+        residual_dense: true
+        stride: [1]
+      - dilation: [1]
+        dropout: 0.3
+        filters: 640
+        kernel: [21]
+        repeat: 5
+        residual: true
+        residual_dense: true
+        stride: [1]
+      - dilation: [1]
+        dropout: 0.3
+        filters: 640
+        kernel: [21]
+        repeat: 5
+        residual: true
+        residual_dense: true
+        stride: [1]
+      - dilation: [1]
+        dropout: 0.3
+        filters: 768
+        kernel: [25]
+        repeat: 5
+        residual: true
+        residual_dense: true
+        stride: [1]
+      - dilation: [1]
+        dropout: 0.3
+        filters: 768
+        kernel: [25]
+        repeat: 5
+        residual: true
+        residual_dense: true
+        stride: [1]
+      - dilation: [2]
+        dropout: 0.4
+        filters: 896
+        kernel: [29]
+        repeat: 1
+        residual: false
+        stride: [1]
+      - dilation: [1]
+        dropout: 0.4
+        filters: 1024
+        kernel: [1]
+        repeat: 1
+        residual: false
+        stride: [1]
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoder
+    feat_in: 1024
+    num_classes: 28
+    vocabulary: *labels
+
+  optim:
+    name: novograd
+    # _target_: nemo.core.optim.optimizers.Novograd
+    lr: .01
+    # optimizer arguments
+    betas: [0.8, 0.5]
+    weight_decay: 0.001
+
+    # scheduler setup
+    sched:
+      name: CosineAnnealing
+
+      # pytorch lightning args
+      # monitor: val_loss
+      # reduce_on_plateau: false
+
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: null
+      min_lr: 0.0
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 5
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+  create_wandb_logger: False
+  wandb_logger_kwargs:
+    name: null
+    project: null
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/asr/conf/marblenet/marblenet_3x2x64.yaml

@@ -0,0 +1,180 @@
+name: &name "MarbleNet-3x2x64"
+
+model:
+  sample_rate: 16000
+  repeat: 2
+  dropout: 0.0
+  kernel_size_factor: 1.0
+
+  labels: ['background', 'speech']
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    labels: ${model.labels}
+    batch_size: 128
+    shuffle: True
+    is_tarred: False
+    tarred_audio_filepaths: null
+    tarred_shard_strategy: "scatter"
+    augmentor:
+      shift:
+        prob: 1.0
+        min_shift_ms: -5.0
+        max_shift_ms: 5.0
+      white_noise:
+        prob: 1.0
+        min_level: -90
+        max_level: -46
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    labels: ${model.labels}
+    batch_size: 128
+    shuffle: False
+    val_loss_idx: 0
+
+  test_ds:
+    manifest_filepath: null
+    sample_rate: ${model.sample_rate}
+    labels: ${model.labels}
+    batch_size: 128
+    shuffle: False
+    test_loss_idx: 0
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMFCCPreprocessor
+    window_size: 0.025
+    window_stride: 0.01
+    window: "hann"
+    n_mels: &n_mels 64
+    n_mfcc: *n_mels
+    n_fft: 512
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    freq_masks: 2
+    time_masks: 2
+    freq_width: 15
+    time_width: 25
+    rect_masks: 5
+    rect_time: 25
+    rect_freq: 15
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+      - filters: 128
+        repeat: 1
+        kernel: [11]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: false
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 64
+        repeat: ${model.repeat}
+        kernel: [13]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: true
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 64
+        repeat: ${model.repeat}
+        kernel: [15]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: true
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 64
+        repeat: ${model.repeat}
+        kernel: [17]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: true
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 128
+        repeat: 1
+        kernel: [29]
+        stride: [1]
+        dilation: [2]
+        dropout: ${model.dropout}
+        residual: false
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: &enc_final_filters 128
+        repeat: 1
+        kernel: [1]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: false
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoderClassification
+    feat_in: *enc_final_filters
+    return_logits: true
+    pooling_type: 'avg'
+
+  optim:
+    name: sgd
+    lr: 0.01
+    # optimizer arguments
+    weight_decay: 0.001
+    momentum: 0.9
+
+    # scheduler setup
+    sched:
+      name: PolynomialHoldDecayAnnealing
+      # Scheduler params
+      power: 2.0
+      warmup_ratio: 0.05
+      hold_ratio: 0.45
+      min_lr: 0.001
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 150
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+  create_wandb_logger: False
+  wandb_logger_kwargs:
+    name: null
+    project: null
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/asr/conf/matchboxnet/matchboxnet_3x1x64_v1.yaml

@@ -0,0 +1,193 @@
+name: &name "MatchboxNet-3x1x64-v1"
+
+model:
+  sample_rate: 16000
+  timesteps: 128
+  repeat: 1
+  dropout: 0.0
+  kernel_size_factor: 1.0
+
+  labels_full: ['bed', 'bird', 'cat', 'dog', 'down', 'eight', 'five', 'four', 'go', 'happy', 'house', 'left', 'marvin',
+           'nine', 'no', 'off', 'on', 'one', 'right', 'seven', 'sheila', 'six', 'stop', 'three', 'tree', 'two', 'up',
+           'wow', 'yes', 'zero']
+
+  labels_subset: ["yes", "no", "up", "down", "left", "right", "on", "off", "stop", "go", "unknown", "silence"]
+
+  labels: ${model.labels_full}
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    labels: ${model.labels}
+    batch_size: 128
+    shuffle: True
+    is_tarred: False
+    tarred_audio_filepaths: null
+    tarred_shard_strategy: "scatter"
+    augmentor:
+      shift:
+        prob: 1.0
+        min_shift_ms: -5.0
+        max_shift_ms: 5.0
+      white_noise:
+        prob: 1.0
+        min_level: -90
+        max_level: -46
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    labels: ${model.labels}
+    batch_size: 128
+    shuffle: False
+    val_loss_idx: 0
+
+  test_ds:
+    manifest_filepath: null
+    sample_rate: ${model.sample_rate}
+    labels: ${model.labels}
+    batch_size: 128
+    shuffle: False
+    test_loss_idx: 0
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMFCCPreprocessor
+    window_size: 0.025
+    window_stride: 0.01
+    window: "hann"
+    n_mels: &n_mels 64
+    n_mfcc: *n_mels
+    n_fft: 512
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    freq_masks: 2
+    time_masks: 2
+    freq_width: 15
+    time_width: 25
+    rect_masks: 5
+    rect_time: 25
+    rect_freq: 15
+
+  crop_or_pad_augment:
+    _target_: nemo.collections.asr.modules.CropOrPadSpectrogramAugmentation
+    audio_length: ${model.timesteps}
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+      - filters: 128
+        repeat: 1
+        kernel: [11]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: false
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 64
+        repeat: ${model.repeat}
+        kernel: [13]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: true
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 64
+        repeat: ${model.repeat}
+        kernel: [15]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: true
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 64
+        repeat: ${model.repeat}
+        kernel: [17]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: true
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 128
+        repeat: 1
+        kernel: [29]
+        stride: [1]
+        dilation: [2]
+        dropout: ${model.dropout}
+        residual: false
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: &enc_final_filters 128
+        repeat: 1
+        kernel: [1]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: false
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoderClassification
+    feat_in: *enc_final_filters
+    return_logits: true
+    pooling_type: 'avg'
+
+  optim:
+    name: novograd
+    # _target_: nemo.core.optim.optimizers.Novograd
+    lr: 0.05
+    # optimizer arguments
+    betas: [0.95, 0.5]
+    weight_decay: 0.001
+
+    # scheduler setup
+    sched:
+      name: PolynomialHoldDecayAnnealing
+
+      # Scheduler params
+      power: 2.0
+      warmup_ratio: 0.05
+      hold_ratio: 0.45
+      min_lr: 0.001
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 200
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+  create_wandb_logger: False
+  wandb_logger_kwargs:
+    name: null
+    project: null
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/asr/conf/matchboxnet/matchboxnet_3x1x64_v2.yaml

@@ -0,0 +1,193 @@
+name: &name "MatchboxNet-3x1x64-v2"
+
+model:
+  sample_rate: 16000
+  timesteps: 128
+  repeat: 1
+  dropout: 0.0
+  kernel_size_factor: 1.0
+
+  labels_full: ['visual', 'wow', 'learn', 'backward', 'dog', 'two', 'left', 'happy', 'nine', 'go', 'up', 'bed', 'stop',
+           'one', 'zero', 'tree', 'seven', 'on', 'four', 'bird', 'right', 'eight', 'no', 'six', 'forward', 'house',
+           'marvin', 'sheila', 'five', 'off', 'three', 'down', 'cat', 'follow', 'yes']
+
+  labels_subset: ["yes", "no", "up", "down", "left", "right", "on", "off", "stop", "go", "unknown", "silence"]
+
+  labels: ${model.labels_full}
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    labels: ${model.labels}
+    batch_size: 128
+    shuffle: True
+    is_tarred: False
+    tarred_audio_filepaths: null
+    tarred_shard_strategy: "scatter"
+    augmentor:
+      shift:
+        prob: 1.0
+        min_shift_ms: -5.0
+        max_shift_ms: 5.0
+      white_noise:
+        prob: 1.0
+        min_level: -90
+        max_level: -46
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    labels: ${model.labels}
+    batch_size: 128
+    shuffle: False
+    val_loss_idx: 0
+
+  test_ds:
+    manifest_filepath: null
+    sample_rate: ${model.sample_rate}
+    labels: ${model.labels}
+    batch_size: 128
+    shuffle: False
+    test_loss_idx: 0
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMFCCPreprocessor
+    window_size: 0.025
+    window_stride: 0.01
+    window: "hann"
+    n_mels: &n_mels 64
+    n_mfcc: *n_mels
+    n_fft: 512
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    freq_masks: 2
+    time_masks: 2
+    freq_width: 15
+    time_width: 25
+    rect_masks: 5
+    rect_time: 25
+    rect_freq: 15
+
+  crop_or_pad_augment:
+    _target_: nemo.collections.asr.modules.CropOrPadSpectrogramAugmentation
+    audio_length: ${model.timesteps}
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+      - filters: 128
+        repeat: 1
+        kernel: [11]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: false
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 64
+        repeat: ${model.repeat}
+        kernel: [13]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: true
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 64
+        repeat: ${model.repeat}
+        kernel: [15]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: true
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 64
+        repeat: ${model.repeat}
+        kernel: [17]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: true
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: 128
+        repeat: 1
+        kernel: [29]
+        stride: [1]
+        dilation: [2]
+        dropout: ${model.dropout}
+        residual: false
+        separable: true
+        kernel_size_factor: ${model.kernel_size_factor}
+
+      - filters: &enc_final_filters 128
+        repeat: 1
+        kernel: [1]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.dropout}
+        residual: false
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoderClassification
+    feat_in: *enc_final_filters
+    return_logits: true
+    pooling_type: 'avg'
+
+  optim:
+    name: novograd
+    # _target_: nemo.core.optim.optimizers.Novograd
+    lr: 0.05
+    # optimizer arguments
+    betas: [0.95, 0.5]
+    weight_decay: 0.001
+
+    # scheduler setup
+    sched:
+      name: PolynomialHoldDecayAnnealing
+
+      # Scheduler params
+      power: 2.0
+      warmup_ratio: 0.05
+      hold_ratio: 0.45
+      min_lr: 0.001
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 200
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+  create_wandb_logger: False
+  wandb_logger_kwargs:
+    name: null
+    project: null
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/asr/conf/quartznet/quartznet_15x5.yaml

@@ -0,0 +1,281 @@
+name: &name "QuartzNet15x5"
+
+model:
+  sample_rate: &sample_rate 16000
+  repeat: &repeat 5
+  dropout: &dropout 0.0
+  separable: &separable true
+  labels: &labels [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+                   "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    labels: *labels
+    batch_size: 32
+    trim_silence: True
+    max_duration: 16.7
+    shuffle: True
+    is_tarred: False
+    tarred_audio_filepaths: null
+    tarred_shard_strategy: "scatter"
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    labels: *labels
+    batch_size: 32
+    shuffle: False
+
+  test_ds:
+    manifest_filepath: null
+    sample_rate: 16000
+    labels: *labels
+    batch_size: 32
+    shuffle: False
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    normalize: "per_feature"
+    window_size: 0.02
+    sample_rate: *sample_rate
+    window_stride: 0.01
+    window: "hann"
+    features: &n_mels 64
+    n_fft: 512
+    frame_splicing: 1
+    dither: 0.00001
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    rect_freq: 50
+    rect_masks: 5
+    rect_time: 120
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: 1
+      residual: false
+      separable: *separable
+      stride: [2]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [2]
+      dropout: *dropout
+      filters: 512
+      kernel: [87]
+      repeat: 1
+      residual: false
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: &enc_filters 1024
+      kernel: [1]
+      repeat: 1
+      residual: false
+      stride: [1]
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoder
+    feat_in: *enc_filters
+    num_classes: 28
+    vocabulary: *labels
+
+  optim:
+    name: novograd
+    # _target_: nemo.core.optim.optimizers.Novograd
+    lr: .01
+    # optimizer arguments
+    betas: [0.8, 0.5]
+    weight_decay: 0.001
+
+    # scheduler setup
+    sched:
+      name: CosineAnnealing
+
+      # pytorch lightning args
+      # monitor: val_loss
+      # reduce_on_plateau: false
+
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: null
+      min_lr: 0.0
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 5
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+  checkpoint_callback_params:
+    monitor: "val_wer"
+    mode: "min"
+  create_wandb_logger: False
+  wandb_logger_kwargs:
+    name: null
+    project: null
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/asr/conf/quartznet/quartznet_15x5_zh.yaml

@@ -0,0 +1,483 @@
+name: &name "QuartzNet15x5"
+
+model:
+  sample_rate: &sample_rate 16000
+  repeat: &repeat 5
+  dropout: &dropout 0.0
+  separable: &separable true
+  labels: &labels [' ', '''', A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P, Q, R,
+                   S, T, U, V, W, X, Y, Z, 㶧, 䶮, 一, 丁, 七, 万, 丈, 三, 上, 下, 不, 与, 丐, 丑, 专, 且, 丕,
+                   世, 丘, 丙, 业, 丛, 东, 丝, 丞, 丢, 两, 严, 丧, 个, 丫, 中, 丰, 串, 临, 丸, 丹, 为, 主, 丽, 举, 乃,
+                   久, 么, 义, 之, 乌, 乍, 乎, 乏, 乐, 乒, 乓, 乔, 乖, 乘, 乙, 九, 乞, 也, 习, 乡, 书, 买, 乱, 乳, 乾,
+                   了, 予, 争, 事, 二, 于, 亏, 云, 互, 五, 井, 亘, 亚, 些, 亟, 亡, 亢, 交, 亥, 亦, 产, 亨, 亩, 享, 京,
+                   亭, 亮, 亲, 亳, 亵, 人, 亿, 什, 仁, 仄, 仅, 仆, 仇, 今, 介, 仍, 从, 仑, 仓, 仔, 仕, 他, 仗, 付, 仙,
+                   仞, 仟, 仡, 代, 令, 以, 仨, 仪, 们, 仰, 仲, 件, 价, 任, 份, 仿, 企, 伉, 伊, 伍, 伎, 伏, 伐, 休, 众,
+                   优, 伙, 会, 伞, 伟, 传, 伢, 伤, 伦, 伪, 伫, 伯, 估, 伴, 伶, 伸, 伺, 似, 伽, 佃, 但, 位, 低, 住, 佐,
+                   佑, 体, 何, 佗, 佘, 余, 佚, 佛, 作, 佝, 佟, 你, 佣, 佩, 佬, 佯, 佰, 佳, 佶, 佻, 佼, 使, 侃, 侄, 侈,
+                   例, 侍, 侏, 侑, 侗, 供, 依, 侠, 侣, 侥, 侦, 侧, 侨, 侬, 侮, 侯, 侵, 便, 促, 俄, 俊, 俎, 俏, 俐, 俑,
+                   俗, 俘, 俚, 保, 俞, 俟, 信, 俨, 俩, 俪, 俭, 修, 俯, 俱, 俸, 俺, 俾, 倌, 倍, 倒, 倔, 倘, 候, 倚, 倜,
+                   借, 倡, 倦, 倩, 倪, 倭, 债, 值, 倾, 偃, 假, 偈, 偌, 偎, 偏, 偓, 偕, 做, 停, 健, 偶, 偷, 偻, 偿, 傀,
+                   傅, 傍, 傣, 傥, 储, 催, 傲, 傻, 像, 僚, 僧, 僮, 僵, 僻, 儋, 儒, 儡, 儿, 兀, 允, 元, 兄, 充, 兆, 先,
+                   光, 克, 免, 兑, 兔, 兖, 党, 兜, 兢, 入, 全, 八, 公, 六, 兮, 兰, 共, 关, 兴, 兵, 其, 具, 典, 兹, 养,
+                   兼, 兽, 冀, 内, 冈, 冉, 册, 再, 冒, 冕, 冗, 写, 军, 农, 冠, 冢, 冤, 冥, 冬, 冯, 冰, 冲, 决, 况, 冶,
+                   冷, 冻, 冼, 冽, 净, 凄, 准, 凇, 凉, 凋, 凌, 减, 凑, 凛, 凝, 几, 凡, 凤, 凭, 凯, 凰, 凳, 凶, 凸, 凹,
+                   出, 击, 函, 凿, 刀, 刁, 刃, 分, 切, 刊, 刍, 刎, 刑, 划, 列, 刘, 则, 刚, 创, 初, 删, 判, 刨, 利, 别,
+                   刮, 到, 制, 刷, 券, 刹, 刺, 刻, 刽, 剁, 剂, 剃, 削, 剌, 前, 剐, 剑, 剔, 剖, 剜, 剥, 剧, 剩, 剪, 副,
+                   割, 剽, 剿, 劈, 力, 劝, 办, 功, 加, 务, 劣, 动, 助, 努, 劫, 劭, 励, 劲, 劳, 劵, 劾, 势, 勃, 勇, 勉,
+                   勋, 勐, 勒, 勘, 募, 勤, 勺, 勾, 勿, 匀, 包, 匆, 匈, 匏, 匕, 化, 北, 匙, 匝, 匠, 匡, 匣, 匪, 匮, 匹,
+                   区, 医, 匾, 匿, 十, 千, 升, 午, 卉, 半, 华, 协, 卑, 卒, 卓, 单, 卖, 南, 博, 卜, 卞, 占, 卡, 卢, 卤,
+                   卦, 卧, 卫, 卯, 印, 危, 卲, 即, 却, 卵, 卷, 卸, 卿, 厂, 厄, 厅, 历, 厉, 压, 厌, 厕, 厘, 厚, 厝, 原,
+                   厢, 厥, 厦, 厨, 厩, 厮, 去, 县, 叁, 参, 又, 叉, 及, 友, 双, 反, 发, 叔, 取, 受, 变, 叙, 叛, 叠, 口,
+                   古, 句, 另, 叨, 叩, 只, 叫, 召, 叭, 叮, 可, 台, 叱, 史, 右, 叵, 叶, 号, 司, 叹, 叼, 叽, 吁, 吃, 各,
+                   吆, 合, 吉, 吊, 吋, 同, 名, 后, 吏, 吐, 向, 吒, 吓, 吕, 吖, 吗, 君, 吝, 吞, 吟, 吠, 否, 吧, 吨, 吩,
+                   含, 听, 吭, 吮, 启, 吱, 吴, 吵, 吸, 吹, 吻, 吼, 吾, 呀, 呃, 呆, 呈, 告, 呐, 呕, 呗, 员, 呛, 呜, 呢,
+                   呦, 周, 呱, 呲, 味, 呵, 呷, 呸, 呻, 呼, 命, 咀, 咂, 咄, 咆, 咋, 和, 咎, 咏, 咐, 咒, 咔, 咕, 咖, 咘,
+                   咙, 咚, 咝, 咣, 咤, 咦, 咧, 咨, 咩, 咪, 咫, 咬, 咭, 咯, 咱, 咳, 咸, 咻, 咽, 哀, 品, 哂, 哄, 哆, 哇,
+                   哈, 哉, 响, 哎, 哐, 哑, 哒, 哔, 哕, 哗, 哟, 哥, 哦, 哨, 哩, 哪, 哭, 哮, 哲, 哺, 哼, 哽, 唁, 唆, 唇,
+                   唉, 唏, 唐, 唑, 唛, 唠, 唢, 唤, 唧, 唬, 售, 唯, 唰, 唱, 唳, 唷, 唾, 啃, 啄, 商, 啊, 啕, 啖, 啜, 啡,
+                   啤, 啥, 啦, 啧, 啪, 啬, 啰, 啲, 啵, 啶, 啸, 啼, 啾, 喀, 喁, 喂, 喃, 善, 喆, 喇, 喉, 喊, 喋, 喔, 喘,
+                   喜, 喝, 喟, 喧, 喱, 喳, 喵, 喷, 喻, 喽, 嗄, 嗅, 嗑, 嗒, 嗓, 嗔, 嗖, 嗜, 嗝, 嗡, 嗣, 嗤, 嗦, 嗨, 嗪,
+                   嗫, 嗬, 嗯, 嗲, 嗷, 嗽, 嘀, 嘈, 嘉, 嘎, 嘏, 嘘, 嘛, 嘞, 嘟, 嘣, 嘭, 嘱, 嘲, 嘴, 嘶, 嘹, 嘻, 嘿, 噌,
+                   噎, 噗, 噘, 噙, 噜, 噢, 噤, 器, 噩, 噪, 噬, 噱, 噶, 噻, 噼, 嚎, 嚏, 嚓, 嚣, 嚷, 嚼, 囊, 囍, 囔, 囗,
+                   囚, 四, 回, 因, 团, 囤, 囧, 囫, 园, 囯, 困, 囱, 围, 囵, 囹, 固, 国, 图, 圃, 圄, 圆, 圈, 土, 圣, 在,
+                   圩, 圪, 圭, 地, 圳, 圹, 场, 圻, 圾, 址, 坂, 均, 坊, 坍, 坎, 坏, 坐, 坑, 块, 坚, 坛, 坝, 坞, 坟, 坠,
+                   坡, 坤, 坦, 坨, 坩, 坪, 坭, 坯, 坳, 坷, 坻, 垂, 垃, 垄, 垅, 型, 垌, 垒, 垚, 垛, 垡, 垢, 垣, 垤, 垦,
+                   垩, 垫, 垭, 垮, 埂, 埃, 埇, 埋, 城, 埔, 埕, 埚, 埝, 域, 埠, 埭, 埸, 培, 基, 堀, 堂, 堃, 堆, 堇, 堕,
+                   堡, 堤, 堪, 堰, 堵, 堺, 塌, 塍, 塑, 塔, 塘, 塞, 填, 塬, 塾, 境, 墅, 墉, 墓, 増, 墙, 增, 墟, 墨, 墩,
+                   壁, 壑, 壕, 壤, 士, 壬, 壮, 声, 壳, 壶, 壹, 处, 备, 复, 夏, 夔, 夕, 外, 夙, 多, 夜, 够, 大, 天, 太,
+                   夫, 夭, 央, 夯, 失, 头, 夷, 夸, 夹, 夺, 奁, 奂, 奄, 奇, 奈, 奉, 奋, 奎, 奏, 契, 奔, 奕, 奖, 套, 奘,
+                   奚, 奠, 奢, 奥, 女, 奴, 奶, 奸, 她, 好, 如, 妃, 妄, 妆, 妇, 妈, 妊, 妍, 妒, 妓, 妖, 妙, 妞, 妤, 妥,
+                   妨, 妩, 妪, 妫, 妮, 妯, 妲, 妹, 妻, 妾, 姆, 姊, 始, 姐, 姑, 姓, 委, 姗, 姚, 姜, 姝, 姣, 姥, 姨, 姬,
+                   姻, 姿, 威, 娃, 娄, 娅, 娆, 娇, 娈, 娉, 娌, 娓, 娘, 娜, 娟, 娠, 娣, 娥, 娩, 娱, 娲, 娴, 娶, 娼, 婀,
+                   婆, 婉, 婊, 婕, 婚, 婢, 婧, 婪, 婴, 婵, 婶, 婷, 婺, 婿, 媒, 媚, 媛, 媞, 媲, 媳, 嫁, 嫂, 嫉, 嫌, 嫒,
+                   嫔, 嫖, 嫚, 嫡, 嫣, 嫦, 嫩, 嫫, 嬅, 嬉, 嬗, 嬛, 嬴, 嬷, 孀, 子, 孑, 孔, 孕, 字, 存, 孙, 孚, 孛, 孜,
+                   孝, 孟, 孢, 季, 孤, 学, 孩, 孪, 孬, 孰, 孱, 孳, 孵, 孺, 孽, 宁, 它, 宅, 宇, 守, 安, 宋, 完, 宏, 宓,
+                   宕, 宗, 官, 宙, 定, 宛, 宜, 宝, 实, 宠, 审, 客, 宣, 室, 宥, 宦, 宪, 宫, 宰, 害, 宴, 宵, 家, 宸, 容,
+                   宽, 宾, 宿, 寂, 寄, 寅, 密, 寇, 富, 寐, 寒, 寓, 寝, 寞, 察, 寡, 寥, 寨, 寮, 寰, 寸, 对, 寺, 寻, 导,
+                   寿, 封, 射, 尅, 将, 尉, 尊, 小, 少, 尔, 尕, 尖, 尘, 尚, 尝, 尤, 尧, 尬, 就, 尴, 尸, 尹, 尺, 尼, 尽,
+                   尾, 尿, 局, 屁, 层, 居, 屈, 屉, 届, 屋, 屌, 屎, 屏, 屐, 屑, 展, 属, 屠, 屡, 履, 屯, 山, 屹, 屿, 岁,
+                   岂, 岌, 岐, 岑, 岔, 岖, 岗, 岚, 岛, 岩, 岬, 岭, 岱, 岳, 岷, 岸, 峁, 峋, 峒, 峙, 峡, 峥, 峦, 峨, 峪,
+                   峭, 峰, 峻, 崂, 崃, 崆, 崇, 崎, 崔, 崖, 崛, 崧, 崩, 崭, 崮, 崴, 崽, 嵇, 嵊, 嵋, 嵌, 嵘, 嵛, 嵩, 嵬,
+                   嶂, 嶙, 嶝, 巅, 巍, 川, 州, 巡, 巢, 工, 左, 巧, 巨, 巩, 巫, 差, 己, 已, 巳, 巴, 巷, 巾, 币, 市, 布,
+                   帅, 帆, 师, 希, 帐, 帕, 帖, 帘, 帚, 帛, 帜, 帝, 带, 帧, 席, 帮, 帷, 常, 帼, 帽, 幂, 幄, 幅, 幌, 幔,
+                   幕, 幡, 幢, 干, 平, 年, 并, 幸, 幺, 幻, 幼, 幽, 广, 庄, 庆, 庇, 床, 序, 庐, 库, 应, 底, 庖, 店, 庙,
+                   庚, 府, 庞, 废, 度, 座, 庭, 庵, 庶, 康, 庸, 庹, 庾, 廉, 廊, 廓, 廖, 延, 廷, 建, 开, 异, 弃, 弄, 弈,
+                   弊, 弋, 式, 弑, 弓, 引, 弗, 弘, 弛, 弟, 张, 弥, 弦, 弧, 弩, 弭, 弯, 弱, 弹, 强, 弼, 归, 当, 录, 彗,
+                   彝, 形, 彤, 彦, 彩, 彪, 彬, 彭, 彰, 影, 彷, 役, 彻, 彼, 往, 征, 径, 待, 徇, 很, 徉, 徊, 律, 徐, 徒,
+                   得, 徘, 徙, 徜, 御, 徨, 循, 微, 德, 徽, 心, 必, 忆, 忌, 忍, 忏, 忐, 忑, 忒, 忖, 志, 忘, 忙, 忠, 忡,
+                   忤, 忧, 忪, 快, 忱, 念, 忻, 忽, 忿, 怀, 态, 怂, 怄, 怅, 怆, 怎, 怒, 怕, 怖, 怜, 思, 怠, 怡, 急, 怦,
+                   性, 怨, 怪, 怫, 怯, 怵, 总, 怼, 怿, 恁, 恃, 恋, 恍, 恐, 恒, 恕, 恙, 恢, 恣, 恤, 恨, 恩, 恪, 恬, 恭,
+                   息, 恰, 恳, 恶, 恸, 恺, 恻, 恼, 恿, 悄, 悉, 悌, 悍, 悔, 悖, 悚, 悟, 悠, 患, 悦, 您, 悬, 悭, 悯, 悱,
+                   悲, 悴, 悸, 悻, 悼, 情, 惆, 惊, 惋, 惑, 惕, 惚, 惜, 惟, 惠, 惦, 惧, 惨, 惩, 惫, 惬, 惭, 惮, 惯, 惰,
+                   想, 惶, 惹, 惺, 愁, 愈, 愉, 意, 愕, 愚, 感, 愣, 愤, 愧, 愫, 愿, 慈, 慌, 慎, 慑, 慕, 慢, 慧, 慨, 慰,
+                   慵, 慷, 憋, 憎, 憔, 憧, 憨, 憩, 憬, 憷, 憾, 懂, 懈, 懊, 懋, 懑, 懒, 懦, 懵, 懿, 戈, 戊, 戌, 戍, 戎,
+                   戏, 成, 我, 戒, 或, 戗, 战, 戚, 戛, 戟, 截, 戬, 戮, 戳, 戴, 户, 戾, 房, 所, 扁, 扇, 扈, 扉, 手, 才,
+                   扎, 扑, 扒, 打, 扔, 托, 扛, 扞, 扣, 扦, 执, 扩, 扪, 扫, 扬, 扭, 扮, 扯, 扰, 扳, 扶, 批, 扼, 找, 承,
+                   技, 抄, 抉, 把, 抑, 抒, 抓, 投, 抖, 抗, 折, 抚, 抛, 抠, 抡, 抢, 护, 报, 抨, 披, 抬, 抱, 抵, 抹, 押,
+                   抽, 抿, 拂, 拄, 担, 拆, 拇, 拈, 拉, 拌, 拍, 拎, 拐, 拒, 拓, 拔, 拖, 拗, 拘, 拙, 拚, 招, 拜, 拟, 拢,
+                   拣, 拥, 拦, 拧, 拨, 择, 括, 拭, 拮, 拯, 拱, 拳, 拴, 拷, 拼, 拽, 拾, 拿, 持, 挂, 指, 按, 挎, 挑, 挖,
+                   挚, 挛, 挝, 挞, 挟, 挠, 挡, 挣, 挤, 挥, 挨, 挪, 挫, 振, 挺, 挽, 捂, 捅, 捆, 捉, 捋, 捌, 捍, 捎, 捏,
+                   捐, 捕, 捞, 损, 捡, 换, 捣, 捧, 据, 捶, 捷, 捺, 捻, 掀, 掂, 掇, 授, 掉, 掌, 掏, 掐, 排, 掖, 掘, 掠,
+                   探, 掣, 接, 控, 推, 掩, 措, 掬, 掮, 掰, 掳, 掴, 掷, 掸, 掺, 揄, 揉, 揍, 描, 提, 插, 握, 揣, 揩, 揪,
+                   揭, 援, 揶, 揽, 搀, 搁, 搂, 搅, 搏, 搐, 搓, 搔, 搜, 搞, 搡, 搧, 搪, 搬, 搭, 携, 搽, 摁, 摄, 摆, 摇,
+                   摈, 摊, 摒, 摔, 摘, 摞, 摧, 摩, 摸, 摹, 撂, 撅, 撇, 撑, 撒, 撕, 撞, 撤, 撩, 撬, 播, 撮, 撰, 撵, 撸,
+                   撺, 撼, 擀, 擂, 擅, 操, 擎, 擒, 擘, 擞, 擢, 擦, 攀, 攒, 攘, 攥, 攫, 支, 收, 攸, 改, 攻, 放, 政, 故,
+                   效, 敌, 敏, 救, 敕, 敖, 教, 敛, 敝, 敞, 敢, 散, 敦, 敬, 数, 敲, 整, 敷, 文, 斋, 斌, 斐, 斑, 斓, 斗,
+                   料, 斛, 斜, 斟, 斡, 斤, 斥, 斧, 斩, 断, 斯, 新, 方, 施, 旁, 旅, 旋, 旌, 族, 旖, 旗, 无, 既, 日, 旦,
+                   旧, 旨, 早, 旬, 旭, 旮, 旯, 旱, 时, 旷, 旺, 旻, 昀, 昂, 昆, 昊, 昌, 明, 昏, 易, 昔, 昕, 昙, 昝, 星,
+                   映, 春, 昧, 昨, 昭, 是, 昱, 昴, 昵, 昶, 昼, 显, 晃, 晋, 晌, 晏, 晒, 晓, 晔, 晕, 晖, 晗, 晚, 晞, 晟,
+                   晤, 晦, 晨, 普, 景, 晰, 晴, 晶, 晷, 智, 晾, 暂, 暄, 暇, 暌, 暑, 暖, 暗, 暧, 暨, 暮, 暴, 暹, 暾, 曈,
+                   曙, 曜, 曝, 曦, 曰, 曲, 曳, 更, 曹, 曼, 曾, 替, 最, 月, 有, 朋, 服, 朐, 朔, 朕, 朗, 望, 朝, 期, 朦,
+                   木, 未, 末, 本, 札, 术, 朱, 朴, 朵, 机, 朽, 杀, 杂, 权, 杆, 杈, 杉, 李, 杏, 材, 村, 杓, 杖, 杜, 杞,
+                   束, 杠, 条, 来, 杨, 杭, 杯, 杰, 杳, 杵, 杷, 松, 板, 极, 构, 枇, 枉, 枋, 析, 枕, 林, 枚, 果, 枝, 枞,
+                   枢, 枣, 枥, 枪, 枫, 枭, 枯, 枰, 枳, 架, 枷, 枸, 柃, 柄, 柏, 某, 柑, 柒, 染, 柔, 柘, 柚, 柜, 柞, 柠,
+                   查, 柩, 柬, 柯, 柱, 柳, 柴, 柿, 栀, 栅, 标, 栈, 栋, 栌, 栎, 栏, 树, 栓, 栖, 栗, 校, 栩, 株, 样, 核,
+                   根, 格, 栽, 栾, 桁, 桂, 桃, 框, 案, 桉, 桌, 桎, 桐, 桑, 桓, 桔, 桠, 桢, 档, 桥, 桦, 桨, 桩, 桴, 桶,
+                   桷, 梁, 梅, 梆, 梏, 梓, 梗, 梢, 梦, 梧, 梨, 梭, 梯, 械, 梳, 梵, 检, 棂, 棉, 棋, 棍, 棒, 棕, 棘, 棚,
+                   棠, 棣, 森, 棱, 棵, 棺, 椁, 椅, 椋, 植, 椎, 椒, 椟, 椤, 椭, 椰, 椴, 椹, 椿, 楂, 楔, 楚, 楞, 楠, 楣,
+                   楷, 楸, 楼, 概, 榄, 榆, 榈, 榉, 榔, 榕, 榛, 榜, 榨, 榫, 榭, 榴, 榷, 榻, 槃, 槌, 槎, 槐, 槛, 槟, 槭,
+                   槽, 槿, 樊, 樟, 模, 樨, 横, 樯, 樱, 樵, 樽, 樾, 橄, 橇, 橐, 橘, 橙, 橡, 橱, 檀, 檐, 檗, 檬, 欠, 次,
+                   欢, 欣, 欧, 欲, 欸, 欺, 款, 歆, 歇, 歉, 歌, 歙, 止, 正, 此, 步, 武, 歧, 歩, 歪, 歹, 死, 歼, 殁, 殃,
+                   殆, 殇, 殉, 殊, 残, 殒, 殓, 殖, 殚, 殡, 殴, 段, 殷, 殿, 毁, 毂, 毅, 毋, 母, 每, 毒, 毓, 比, 毕, 毗,
+                   毙, 毛, 毡, 毫, 毯, 毽, 氏, 民, 氓, 气, 氚, 氛, 氟, 氢, 氤, 氦, 氧, 氨, 氪, 氮, 氯, 氰, 氲, 水, 永,
+                   汀, 汁, 求, 汇, 汉, 汊, 汐, 汕, 汗, 汛, 汝, 汞, 江, 池, 污, 汤, 汨, 汩, 汪, 汰, 汲, 汴, 汶, 汹, 汽,
+                   汾, 沁, 沂, 沃, 沅, 沈, 沉, 沌, 沏, 沐, 沓, 沙, 沛, 沟, 没, 沢, 沣, 沥, 沦, 沧, 沪, 沫, 沭, 沮, 沱,
+                   河, 沸, 油, 治, 沼, 沽, 沾, 沿, 泄, 泉, 泊, 泌, 泓, 泔, 法, 泖, 泗, 泛, 泞, 泠, 泡, 波, 泣, 泥, 注,
+                   泪, 泫, 泮, 泯, 泰, 泱, 泳, 泵, 泷, 泸, 泺, 泻, 泼, 泽, 泾, 洁, 洋, 洒, 洗, 洙, 洛, 洞, 津, 洪, 洮,
+                   洱, 洲, 洵, 洹, 洺, 活, 洼, 洽, 派, 流, 浃, 浅, 浆, 浇, 浈, 浊, 测, 济, 浏, 浐, 浑, 浒, 浓, 浔, 浙,
+                   浚, 浜, 浠, 浣, 浦, 浩, 浪, 浮, 浴, 海, 浸, 涂, 涅, 消, 涉, 涌, 涎, 涑, 涓, 涕, 涛, 涝, 涞, 涟, 涠,
+                   涡, 涣, 涤, 润, 涧, 涨, 涩, 涪, 涮, 涯, 液, 涵, 涸, 涿, 淀, 淄, 淅, 淆, 淇, 淋, 淌, 淑, 淖, 淘, 淝,
+                   淞, 淡, 淤, 淦, 淫, 淬, 淮, 深, 淳, 混, 淹, 添, 淼, 清, 渊, 渌, 渍, 渎, 渐, 渑, 渔, 渗, 渚, 渝, 渠,
+                   渡, 渣, 渤, 渥, 温, 渭, 港, 渲, 渴, 游, 渺, 湃, 湄, 湉, 湍, 湎, 湖, 湘, 湛, 湫, 湾, 湿, 溃, 溅, 溆,
+                   溉, 溏, 源, 溜, 溟, 溢, 溥, 溧, 溪, 溯, 溶, 溺, 滁, 滇, 滋, 滑, 滔, 滕, 滘, 滚, 滞, 满, 滢, 滤, 滥,
+                   滦, 滨, 滩, 滴, 滹, 漂, 漆, 漉, 漏, 漓, 演, 漕, 漠, 漩, 漪, 漫, 漭, 漯, 漱, 漳, 漾, 潆, 潇, 潋, 潍,
+                   潘, 潜, 潞, 潢, 潦, 潭, 潮, 潸, 潺, 潼, 澄, 澈, 澍, 澎, 澜, 澡, 澧, 澳, 澶, 激, 濂, 濑, 濒, 濠, 濡,
+                   濮, 濯, 瀑, 瀚, 瀛, 灌, 灏, 灞, 火, 灭, 灯, 灰, 灵, 灶, 灸, 灼, 灾, 灿, 炀, 炅, 炉, 炊, 炎, 炒, 炔,
+                   炕, 炖, 炙, 炜, 炫, 炬, 炭, 炮, 炯, 炳, 炷, 炸, 点, 炼, 炽, 烀, 烁, 烂, 烃, 烈, 烊, 烘, 烙, 烛, 烟,
+                   烤, 烦, 烧, 烨, 烩, 烫, 烬, 热, 烯, 烷, 烹, 烽, 焉, 焊, 焓, 焕, 焖, 焗, 焘, 焙, 焚, 焦, 焯, 焰, 焱,
+                   然, 煊, 煌, 煎, 煜, 煞, 煤, 煦, 照, 煨, 煮, 煲, 煳, 煽, 熄, 熊, 熏, 熔, 熙, 熟, 熠, 熨, 熬, 熵, 熹,
+                   燃, 燊, 燎, 燕, 燥, 燮, 爆, 爪, 爬, 爱, 爵, 父, 爷, 爸, 爹, 爽, 片, 版, 牌, 牍, 牒, 牙, 牛, 牟, 牠,
+                   牡, 牢, 牧, 物, 牲, 牵, 特, 牺, 牾, 犀, 犁, 犄, 犇, 犊, 犒, 犟, 犬, 犯, 状, 犷, 犸, 犹, 狂, 狄, 狈,
+                   狐, 狒, 狗, 狙, 狞, 狠, 狡, 狩, 独, 狭, 狮, 狰, 狱, 狸, 狼, 猁, 猎, 猖, 猛, 猜, 猝, 猥, 猩, 猪, 猫,
+                   猬, 献, 猴, 猷, 猹, 猾, 猿, 獒, 獗, 獭, 獾, 玄, 率, 玉, 王, 玑, 玖, 玛, 玟, 玥, 玩, 玫, 玮, 环, 现,
+                   玲, 玳, 玷, 玹, 玺, 玻, 珀, 珂, 珈, 珉, 珊, 珍, 珏, 珑, 珙, 珞, 珠, 珥, 班, 珮, 珲, 珺, 球, 琅, 理,
+                   琉, 琊, 琏, 琐, 琛, 琢, 琤, 琥, 琦, 琨, 琪, 琬, 琮, 琰, 琳, 琴, 琵, 琶, 琼, 瑁, 瑄, 瑕, 瑙, 瑚, 瑛,
+                   瑜, 瑞, 瑟, 瑠, 瑭, 瑰, 瑶, 瑷, 瑾, 璀, 璃, 璇, 璋, 璐, 璞, 璟, 璧, 璨, 瓜, 瓢, 瓣, 瓦, 瓮, 瓯, 瓶,
+                   瓷, 甄, 甘, 甚, 甜, 生, 甥, 用, 甩, 甫, 甬, 甭, 田, 由, 甲, 申, 电, 男, 甸, 町, 画, 畅, 畈, 畊, 界,
+                   畏, 畔, 留, 畜, 略, 番, 畴, 畸, 畿, 疃, 疆, 疏, 疑, 疖, 疗, 疙, 疚, 疝, 疟, 疡, 疣, 疤, 疫, 疮, 疯,
+                   疱, 疲, 疴, 疵, 疸, 疹, 疼, 疽, 疾, 病, 症, 痉, 痊, 痍, 痒, 痔, 痕, 痘, 痛, 痞, 痢, 痣, 痧, 痨, 痪,
+                   痫, 痰, 痱, 痴, 痹, 痼, 瘀, 瘁, 瘙, 瘟, 瘠, 瘢, 瘤, 瘦, 瘩, 瘪, 瘫, 瘳, 瘴, 瘸, 瘾, 癌, 癖, 癜, 癞,
+                   癣, 癫, 登, 白, 百, 皂, 的, 皆, 皇, 皋, 皎, 皑, 皓, 皖, 皙, 皮, 皱, 皿, 盂, 盅, 盆, 盈, 益, 盎, 盏,
+                   盐, 监, 盒, 盔, 盖, 盗, 盘, 盛, 盟, 目, 盯, 盱, 盲, 直, 相, 盹, 盼, 盾, 省, 眈, 眉, 看, 眙, 真, 眠,
+                   眨, 眩, 眬, 眯, 眶, 眷, 眸, 眺, 眼, 着, 睁, 睇, 睐, 睑, 睛, 睡, 睢, 督, 睦, 睫, 睬, 睹, 睽, 睾, 睿,
+                   瞄, 瞅, 瞌, 瞎, 瞑, 瞒, 瞟, 瞠, 瞥, 瞧, 瞩, 瞪, 瞬, 瞭, 瞰, 瞳, 瞻, 瞿, 矍, 矗, 矛, 矜, 矢, 矣, 知,
+                   矩, 矫, 矬, 短, 矮, 石, 矶, 矸, 矽, 矾, 矿, 砀, 码, 砂, 砌, 砍, 砒, 研, 砖, 砚, 砝, 砣, 砥, 砭, 砰,
+                   破, 砷, 砸, 砺, 砼, 砾, 础, 硅, 硌, 硒, 硕, 硖, 硚, 硝, 硫, 硬, 确, 硼, 碉, 碌, 碍, 碎, 碑, 碓, 碗,
+                   碘, 碚, 碜, 碟, 碣, 碧, 碰, 碱, 碳, 碴, 碾, 磁, 磅, 磊, 磋, 磐, 磕, 磨, 磴, 磷, 磺, 礁, 示, 礼, 社,
+                   祀, 祁, 祈, 祉, 祎, 祐, 祖, 祚, 祛, 祝, 神, 祟, 祠, 祢, 祥, 票, 祭, 祯, 祷, 祸, 祺, 禀, 禁, 禄, 禅,
+                   福, 禧, 禹, 禺, 离, 禽, 禾, 秀, 私, 秃, 秆, 秉, 秋, 种, 科, 秒, 秘, 租, 秣, 秤, 秦, 秧, 秩, 积, 称,
+                   秸, 移, 秽, 稀, 程, 稍, 税, 稔, 稚, 稞, 稠, 稣, 稳, 稷, 稹, 稻, 稼, 稽, 稿, 穆, 穗, 穴, 究, 穷, 穹,
+                   空, 穿, 突, 窃, 窄, 窈, 窍, 窑, 窒, 窕, 窖, 窗, 窘, 窜, 窝, 窟, 窠, 窥, 窦, 窨, 窿, 立, 竖, 站, 竞,
+                   竟, 章, 竣, 童, 竭, 端, 竹, 竺, 竽, 竿, 笃, 笆, 笈, 笋, 笑, 笔, 笙, 笛, 笠, 符, 笨, 第, 笳, 笸, 笼,
+                   等, 筋, 筏, 筐, 筑, 筒, 答, 策, 筛, 筝, 筠, 筱, 筵, 筷, 筹, 签, 简, 箍, 箔, 箕, 算, 管, 箩, 箫, 箭,
+                   箱, 箴, 篁, 篆, 篇, 篑, 篓, 篝, 篡, 篦, 篪, 篮, 篱, 篷, 篼, 簇, 簋, 簧, 簪, 簸, 簿, 籁, 籍, 米, 类,
+                   籼, 籽, 粉, 粑, 粒, 粕, 粗, 粘, 粟, 粤, 粥, 粪, 粮, 粱, 粲, 粳, 粹, 粼, 粽, 精, 糊, 糕, 糖, 糗, 糙,
+                   糟, 糠, 糯, 系, 紊, 素, 索, 紧, 紫, 累, 絮, 綦, 繁, 纂, 纠, 纡, 红, 纣, 纤, 约, 级, 纨, 纪, 纫, 纬,
+                   纭, 纯, 纰, 纱, 纲, 纳, 纵, 纶, 纷, 纸, 纹, 纺, 纽, 纾, 线, 绀, 练, 组, 绅, 细, 织, 终, 绉, 绊, 绋,
+                   绌, 绍, 绎, 经, 绑, 绒, 结, 绔, 绕, 绘, 给, 绚, 绛, 络, 绝, 绞, 统, 绢, 绣, 绥, 继, 绩, 绪, 绫, 续,
+                   绮, 绯, 绰, 绳, 维, 绵, 绷, 绸, 绻, 综, 绽, 绿, 缀, 缄, 缅, 缆, 缇, 缉, 缎, 缓, 缔, 缕, 编, 缘, 缙,
+                   缚, 缛, 缜, 缝, 缠, 缢, 缤, 缨, 缩, 缪, 缬, 缭, 缮, 缰, 缱, 缴, 缸, 缺, 罂, 罄, 罐, 网, 罔, 罕, 罗,
+                   罚, 罡, 罢, 罩, 罪, 置, 署, 罹, 羁, 羊, 羌, 美, 羔, 羚, 羞, 羡, 群, 羧, 羯, 羲, 羸, 羹, 羽, 羿, 翁,
+                   翅, 翊, 翌, 翎, 翔, 翘, 翟, 翠, 翡, 翩, 翰, 翱, 翻, 翼, 耀, 老, 考, 耄, 者, 耆, 耋, 而, 耍, 耐, 耒,
+                   耕, 耗, 耘, 耙, 耜, 耪, 耳, 耶, 耷, 耸, 耻, 耽, 耿, 聂, 聆, 聊, 聋, 职, 联, 聘, 聚, 聪, 肃, 肆, 肇,
+                   肉, 肋, 肌, 肖, 肘, 肚, 肛, 肝, 肠, 股, 肢, 肤, 肥, 肩, 肪, 肮, 肯, 肱, 育, 肴, 肺, 肾, 肿, 胀, 胁,
+                   胃, 胆, 背, 胎, 胖, 胗, 胚, 胛, 胜, 胞, 胡, 胤, 胥, 胧, 胫, 胭, 胯, 胰, 胱, 胳, 胶, 胸, 胺, 能, 脂,
+                   脆, 脉, 脊, 脍, 脏, 脐, 脑, 脓, 脖, 脚, 脯, 脱, 脸, 脾, 腆, 腈, 腊, 腋, 腌, 腐, 腑, 腓, 腔, 腕, 腥,
+                   腩, 腭, 腮, 腰, 腱, 腴, 腹, 腺, 腻, 腼, 腾, 腿, 膀, 膈, 膊, 膏, 膑, 膛, 膜, 膝, 膨, 膳, 膺, 臀, 臂,
+                   臃, 臆, 臊, 臣, 臧, 自, 臬, 臭, 至, 致, 臻, 臼, 舀, 舅, 舆, 舌, 舍, 舐, 舒, 舔, 舛, 舜, 舞, 舟, 航,
+                   舫, 般, 舰, 舱, 舵, 舶, 舷, 舸, 船, 艇, 艋, 艘, 艮, 良, 艰, 色, 艳, 艺, 艾, 艿, 节, 芊, 芋, 芍, 芒,
+                   芗, 芙, 芜, 芝, 芥, 芦, 芩, 芪, 芬, 芭, 芮, 芯, 花, 芳, 芷, 芸, 芹, 芽, 芾, 苇, 苋, 苍, 苏, 苑, 苓,
+                   苔, 苗, 苛, 苞, 苟, 苡, 苣, 若, 苦, 苫, 苯, 英, 苷, 苹, 茁, 茂, 范, 茄, 茅, 茆, 茉, 茌, 茎, 茗, 茛,
+                   茜, 茧, 茨, 茫, 茬, 茯, 茱, 茳, 茴, 茵, 茶, 茸, 茹, 茼, 荀, 荃, 荆, 荇, 草, 荏, 荐, 荒, 荔, 荚, 荛,
+                   荞, 荟, 荠, 荡, 荣, 荤, 荧, 荨, 荫, 药, 荷, 荸, 荻, 荼, 莅, 莆, 莉, 莎, 莒, 莓, 莘, 莜, 莞, 莠, 莪,
+                   莫, 莱, 莲, 莴, 获, 莹, 莺, 莽, 菀, 菁, 菅, 菇, 菊, 菌, 菏, 菖, 菘, 菜, 菠, 菡, 菩, 菱, 菲, 萃, 萄,
+                   萋, 萌, 萍, 萎, 萝, 萤, 营, 萦, 萧, 萨, 萱, 萸, 落, 葆, 著, 葚, 葛, 葡, 董, 葩, 葫, 葬, 葱, 葳, 葵,
+                   葺, 蒂, 蒋, 蒙, 蒜, 蒯, 蒲, 蒸, 蒿, 蓁, 蓄, 蓉, 蓓, 蓝, 蓟, 蓥, 蓦, 蓬, 蓼, 蔑, 蔓, 蔗, 蔚, 蔡, 蔫,
+                   蔬, 蔷, 蔺, 蔻, 蔼, 蔽, 蕃, 蕉, 蕊, 蕙, 蕨, 蕲, 蕴, 蕾, 薄, 薇, 薏, 薛, 薪, 薯, 薰, 薷, 藁, 藉, 藏,
+                   藐, 藓, 藕, 藜, 藠, 藤, 藩, 藻, 藿, 蘑, 蘸, 虎, 虏, 虐, 虑, 虔, 虚, 虞, 虫, 虱, 虹, 虻, 虽, 虾, 蚀,
+                   蚁, 蚂, 蚊, 蚌, 蚓, 蚕, 蚝, 蚣, 蚤, 蚪, 蚬, 蚯, 蚱, 蚴, 蛀, 蛆, 蛇, 蛉, 蛊, 蛋, 蛎, 蛐, 蛔, 蛙, 蛛,
+                   蛟, 蛤, 蛮, 蛰, 蛳, 蛹, 蛾, 蜀, 蜂, 蜃, 蜇, 蜈, 蜊, 蜍, 蜒, 蜓, 蜕, 蜗, 蜘, 蜚, 蜜, 蜡, 蜢, 蜥, 蜱,
+                   蜴, 蜷, 蜻, 蜿, 蝇, 蝈, 蝉, 蝌, 蝎, 蝗, 蝙, 蝠, 蝮, 蝴, 蝶, 蝽, 螂, 螃, 螈, 融, 螨, 螳, 螺, 蟀, 蟆,
+                   蟊, 蟋, 蟑, 蟒, 蟠, 蟹, 蟾, 蠊, 蠕, 蠡, 蠢, 血, 衅, 行, 衍, 衔, 街, 衙, 衡, 衢, 衣, 补, 表, 衩, 衫,
+                   衬, 衮, 衰, 衲, 衷, 袁, 袂, 袄, 袅, 袈, 袋, 袍, 袒, 袖, 袜, 被, 袭, 袱, 裁, 裂, 装, 裆, 裔, 裕, 裘,
+                   裙, 裟, 裤, 裨, 裱, 裳, 裴, 裸, 裹, 褂, 褐, 褒, 褓, 褔, 褚, 褛, 褥, 褪, 褴, 褶, 襁, 襄, 襟, 西, 要,
+                   覃, 覆, 见, 观, 规, 觅, 视, 览, 觉, 觊, 觎, 觐, 觑, 角, 觞, 解, 觥, 触, 言, 訾, 詹, 誉, 誓, 警, 譬,
+                   计, 订, 讣, 认, 讥, 讧, 讨, 让, 讪, 训, 议, 讯, 记, 讲, 讳, 讴, 讶, 讷, 许, 讹, 论, 讼, 讽, 设, 访,
+                   诀, 证, 诃, 评, 诅, 识, 诈, 诉, 诊, 诋, 词, 诏, 译, 诓, 试, 诗, 诘, 诙, 诚, 诛, 话, 诞, 诟, 诠, 诡,
+                   询, 诣, 诤, 该, 详, 诧, 诩, 诫, 诬, 语, 误, 诱, 诲, 说, 诵, 诶, 请, 诸, 诹, 诺, 读, 诽, 课, 诿, 谀,
+                   谁, 调, 谅, 谆, 谈, 谊, 谋, 谌, 谍, 谎, 谏, 谐, 谑, 谓, 谕, 谖, 谘, 谙, 谚, 谛, 谜, 谟, 谢, 谣, 谤,
+                   谦, 谧, 谨, 谩, 谬, 谭, 谮, 谯, 谱, 谴, 谶, 谷, 豁, 豆, 豇, 豉, 豌, 豚, 象, 豢, 豪, 豫, 豹, 豺, 貂,
+                   貅, 貉, 貌, 貔, 贝, 贞, 负, 贡, 财, 责, 贤, 败, 账, 货, 质, 贩, 贪, 贫, 贬, 购, 贮, 贯, 贰, 贱, 贲,
+                   贴, 贵, 贷, 贸, 费, 贺, 贻, 贼, 贾, 贿, 赁, 赂, 赃, 资, 赅, 赈, 赉, 赊, 赋, 赌, 赎, 赏, 赐, 赓, 赔,
+                   赖, 赘, 赚, 赛, 赝, 赞, 赠, 赡, 赢, 赣, 赤, 赦, 赫, 走, 赳, 赴, 赵, 赶, 起, 趁, 超, 越, 趋, 趟, 趣,
+                   足, 趴, 趵, 趸, 趺, 趾, 跃, 跄, 跆, 跋, 跌, 跎, 跑, 跚, 跛, 距, 跟, 跤, 跨, 跪, 跬, 路, 跳, 践, 跶,
+                   跷, 跹, 跺, 跻, 踉, 踊, 踌, 踏, 踝, 踞, 踢, 踩, 踪, 踮, 踯, 踱, 踵, 踹, 踺, 蹁, 蹂, 蹄, 蹈, 蹉, 蹊,
+                   蹋, 蹒, 蹚, 蹦, 蹩, 蹬, 蹭, 蹲, 蹴, 蹶, 蹼, 蹿, 躁, 躅, 躇, 躏, 身, 躬, 躯, 躲, 躺, 车, 轧, 轨, 轩,
+                   轫, 转, 轮, 软, 轰, 轱, 轲, 轳, 轴, 轶, 轸, 轻, 轼, 载, 轿, 较, 辄, 辅, 辆, 辈, 辉, 辊, 辍, 辐, 辑,
+                   输, 辕, 辖, 辗, 辘, 辙, 辛, 辜, 辞, 辟, 辣, 辨, 辩, 辫, 辰, 辱, 边, 辽, 达, 迁, 迂, 迄, 迅, 过, 迈,
+                   迎, 运, 近, 返, 还, 这, 进, 远, 违, 连, 迟, 迢, 迥, 迦, 迩, 迪, 迫, 迭, 述, 迷, 迸, 迹, 追, 退, 送,
+                   适, 逃, 逅, 逆, 选, 逊, 逋, 逍, 透, 逐, 逑, 递, 途, 逗, 通, 逛, 逝, 逞, 速, 造, 逡, 逢, 逮, 逯, 逵,
+                   逸, 逻, 逼, 逾, 遁, 遂, 遇, 遍, 遏, 遐, 遑, 道, 遗, 遛, 遢, 遣, 遥, 遨, 遭, 遮, 遴, 遵, 避, 邀, 邂,
+                   邃, 邋, 邑, 邓, 邕, 邙, 邛, 邝, 邡, 邢, 那, 邦, 邪, 邬, 邮, 邯, 邰, 邱, 邳, 邵, 邸, 邹, 邺, 邻, 郁,
+                   郅, 郇, 郊, 郎, 郑, 郓, 郜, 郝, 郡, 郧, 部, 郫, 郭, 郯, 郴, 郸, 都, 鄂, 鄙, 鄞, 鄢, 鄱, 酉, 酊, 酋,
+                   酌, 配, 酐, 酒, 酗, 酚, 酝, 酞, 酣, 酥, 酩, 酪, 酬, 酮, 酯, 酰, 酱, 酵, 酶, 酷, 酸, 酿, 醇, 醉, 醋,
+                   醍, 醐, 醒, 醛, 醺, 采, 釉, 释, 里, 重, 野, 量, 金, 釜, 鉴, 銮, 鏖, 鑫, 钇, 针, 钉, 钊, 钎, 钏, 钐,
+                   钒, 钓, 钗, 钙, 钛, 钜, 钝, 钞, 钟, 钠, 钢, 钣, 钥, 钦, 钧, 钨, 钩, 钮, 钯, 钰, 钱, 钲, 钳, 钴, 钵,
+                   钻, 钼, 钾, 钿, 铀, 铁, 铂, 铃, 铄, 铅, 铆, 铉, 铋, 铍, 铎, 铐, 铑, 铖, 铛, 铜, 铝, 铟, 铠, 铡, 铣,
+                   铤, 铧, 铨, 铩, 铬, 铭, 铮, 铰, 铲, 银, 铷, 铸, 铺, 链, 铿, 销, 锁, 锂, 锄, 锅, 锆, 锈, 锉, 锋, 锌,
+                   锏, 锐, 锑, 锒, 错, 锚, 锟, 锡, 锢, 锣, 锤, 锥, 锦, 锨, 锭, 键, 锯, 锰, 锲, 锴, 锵, 锷, 锹, 锻, 镀,
+                   镁, 镂, 镇, 镉, 镊, 镌, 镍, 镏, 镐, 镑, 镔, 镕, 镖, 镜, 镣, 镭, 镯, 镰, 镳, 镶, 长, 门, 闩, 闪, 闫,
+                   闭, 问, 闯, 闰, 闲, 闳, 间, 闵, 闷, 闸, 闹, 闺, 闻, 闽, 闾, 阀, 阁, 阂, 阄, 阅, 阆, 阉, 阎, 阐, 阑,
+                   阔, 阕, 阖, 阙, 阚, 阜, 队, 阡, 阪, 阮, 阱, 防, 阳, 阴, 阵, 阶, 阻, 阿, 陀, 陂, 附, 际, 陆, 陇, 陈,
+                   陉, 陋, 陌, 降, 限, 陕, 陛, 陡, 院, 除, 陨, 险, 陪, 陬, 陵, 陶, 陷, 隅, 隆, 隋, 隍, 随, 隐, 隔, 隗,
+                   隘, 隙, 障, 隧, 隶, 隼, 隽, 难, 雀, 雁, 雄, 雅, 集, 雇, 雉, 雌, 雍, 雏, 雒, 雕, 雨, 雪, 雯, 雳, 零,
+                   雷, 雹, 雾, 需, 霁, 霄, 霆, 震, 霈, 霉, 霍, 霎, 霏, 霓, 霖, 霜, 霞, 霪, 露, 霸, 霹, 霾, 靑, 青, 靓,
+                   靖, 静, 靛, 非, 靠, 靡, 面, 革, 靳, 靴, 靶, 鞅, 鞋, 鞍, 鞑, 鞘, 鞠, 鞭, 韦, 韧, 韩, 韫, 韬, 韭, 音,
+                   韵, 韶, 页, 顶, 顷, 项, 顺, 须, 顽, 顾, 顿, 颀, 颁, 颂, 预, 颅, 领, 颇, 颈, 颊, 颌, 颍, 颐, 频, 颓,
+                   颖, 颗, 题, 颚, 颜, 额, 颠, 颢, 颤, 颦, 颧, 风, 飒, 飓, 飘, 飙, 飚, 飞, 食, 飧, 餍, 餐, 餮, 饕, 饥,
+                   饨, 饪, 饭, 饮, 饯, 饰, 饱, 饲, 饴, 饵, 饶, 饷, 饺, 饼, 饽, 饿, 馀, 馁, 馄, 馅, 馆, 馈, 馊, 馋, 馍,
+                   馏, 馑, 馒, 馕, 首, 馗, 香, 馥, 馨, 马, 驭, 驮, 驯, 驰, 驱, 驳, 驴, 驶, 驷, 驸, 驹, 驻, 驼, 驾, 驿,
+                   骁, 骂, 骄, 骅, 骆, 骇, 骈, 骊, 骋, 验, 骏, 骐, 骑, 骓, 骗, 骚, 骛, 骜, 骝, 骞, 骠, 骡, 骤, 骥, 骨,
+                   骰, 骷, 骸, 骺, 骼, 髂, 髅, 髋, 髌, 髓, 高, 髦, 髯, 鬃, 鬓, 鬟, 鬼, 魁, 魂, 魄, 魅, 魇, 魉, 魍, 魏,
+                   魔, 魟, 鱼, 鱿, 鲁, 鲅, 鲈, 鲍, 鲑, 鲜, 鲟, 鲠, 鲢, 鲤, 鲨, 鲫, 鲭, 鲳, 鲶, 鲷, 鲸, 鲼, 鳃, 鳄, 鳅,
+                   鳌, 鳍, 鳕, 鳖, 鳗, 鳝, 鳞, 鳟, 鸟, 鸠, 鸡, 鸢, 鸣, 鸥, 鸦, 鸩, 鸪, 鸫, 鸭, 鸯, 鸳, 鸵, 鸽, 鸾, 鸿,
+                   鹁, 鹂, 鹃, 鹅, 鹉, 鹊, 鹌, 鹏, 鹑, 鹜, 鹞, 鹤, 鹦, 鹧, 鹫, 鹭, 鹰, 鹳, 鹿, 麂, 麋, 麒, 麓, 麝, 麟,
+                   麦, 麸, 麻, 麾, 黄, 黍, 黎, 黏, 黑, 黔, 默, 黛, 黝, 黟, 黯, 鼎, 鼓, 鼠, 鼬, 鼹, 鼻, 鼾, 齐, 齿, 龃,
+                   龄, 龅, 龈, 龉, 龊, 龌, 龙, 龚, 龟, "\U0002B5AF", "\U0002B689"]
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    labels: *labels
+    batch_size: 32
+    trim_silence: True
+    normalize: False
+    max_duration: 16.7
+    shuffle: True
+    is_tarred: False
+    tarred_audio_filepaths: null
+    tarred_shard_strategy: "scatter"
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    normalize: False
+    labels: *labels
+    batch_size: 32
+    shuffle: False
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    normalize: "per_feature"
+    window_size: 0.02
+    sample_rate: *sample_rate
+    window_stride: 0.01
+    window: "hann"
+    features: &n_mels 64
+    n_fft: 512
+    frame_splicing: 1
+    dither: 0.00001
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    rect_freq: 50
+    rect_masks: 5
+    rect_time: 120
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: 1
+      residual: false
+      separable: *separable
+      stride: [2]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [2]
+      dropout: *dropout
+      filters: 512
+      kernel: [87]
+      repeat: 1
+      residual: false
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: &enc_filters 1024
+      kernel: [1]
+      repeat: 1
+      residual: false
+      stride: [1]
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoder
+    feat_in: *enc_filters
+    num_classes: 5206
+    vocabulary: *labels
+
+  optim:
+    name: novograd
+    # _target_: nemo.core.optim.optimizers.Novograd
+    lr: .01
+    # optimizer arguments
+    betas: [0.8, 0.5]
+    weight_decay: 0.001
+
+    # scheduler setup
+    sched:
+      name: CosineAnnealing
+
+      # pytorch lightning args
+      # monitor: val_loss
+      # reduce_on_plateau: false
+
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: null
+      min_lr: 0.0
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 5
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+  checkpoint_callback_params:
+    monitor: "val_wer"
+    mode: "min"
+  create_wandb_logger: False
+  wandb_logger_kwargs:
+    name: null
+    project: null
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/asr/experimental/conformer/conformer_rnnt_bpe.yaml

@@ -0,0 +1,201 @@
+# It contains the default values for training a Conformer ASR model medium size with RNNT loss and BPE-based vocabulary.
+# Default learning parameters in this config are set for effective batch size of 2K.
+# To train it with smaller batch sizes, you may need to re-tune the learning parameters or use higher accumulate_grad_batches.
+
+# You may find more info about Conformer-CTC here: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/models.html#conformer-ctc
+# Pre-trained models of Conformer-CTC can be found here: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/results.html
+
+name: "Conformer-RNNT-BPE"
+
+model:
+  sample_rate: &sample_rate 16000
+  compute_eval_loss: false # disables loss evaluation for validation and test data
+  log_prediction: true # enables logging sample predictions in the output during training
+  ctc_reduction: 'mean_batch'
+
+  model_defaults:
+    se: true
+    se_context_size: -1
+    # encoder / decoder / joint values
+    enc_hidden: ${model.encoder.d_model}
+    pred_hidden: 640
+    joint_hidden: 640
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    batch_size: 16 # you may increase batch_size if your memory allows
+    shuffle: true
+    num_workers: 8
+    pin_memory: false
+    use_start_end_token: true
+    trim_silence: false
+    max_duration: 16.7 # it is set for LibriSpeech, you may need to update it for your dataset
+    min_duration: 0.1
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    batch_size: 16
+    shuffle: false
+    num_workers: 8
+    pin_memory: false
+    use_start_end_token: true
+
+  test_ds:
+    manifest_filepath: null
+    sample_rate: ${model.sample_rate}
+    batch_size: 16
+    shuffle: false
+    num_workers: 8
+    pin_memory: false
+    use_start_end_token: true
+
+  # you may find more detail on how to train a tokenizer at: /scripts/tokenizers/process_asr_text_tokenizer.py
+  tokenizer:
+    dir: ???  # path to directory which contains either tokenizer.model (bpe) or vocab.txt (for wpe)
+    type: wpe  # Can be either bpe (SentencePiece tokenizer) or wpe (WordPiece tokenizer)
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    sample_rate: *sample_rate
+    normalize: "per_feature"
+    window_size: 0.025
+    window_stride: 0.01
+    window: "hann"
+    features: 80
+    n_fft: 512
+    frame_splicing: 1
+    dither: 0.00001
+    pad_to: 0
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    freq_masks: 2 # set to zero to disable it
+    time_masks: 10 # set to zero to disable it
+    freq_width: 27
+    time_width: 0.05
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConformerEncoder
+    feat_in: ${model.preprocessor.features}
+    feat_out: -1 # you may set it if you need different output size other than the default d_model
+    n_layers: 16
+    d_model: 256
+
+    # Sub-sampling params
+    subsampling: striding # vggnet or striding
+    subsampling_factor: 4 # must be power of 2
+    subsampling_conv_channels: -1 # set to -1 to make it equal to the d_model
+
+    # Feed forward module's params
+    ff_expansion_factor: 4
+
+    # Multi-headed Attention Module's params
+    self_attention_model: rel_pos # rel_pos or abs_pos
+    n_heads: 4
+    xscaling: true # scales up the inputs by sqrt(d_model)
+    untie_biases: true # unties the biases of the TransformerXL layers
+    pos_emb_max_len: 5000
+
+    # Convolution module's params
+    conv_kernel_size: 31
+
+    ### regularization
+    dropout: 0.1 # The dropout used in most of the Conformer Modules
+    dropout_emb: 0.0 # The dropout used for embeddings
+    dropout_att: 0.1 # The dropout for multi-headed attention modules
+
+  decoder:
+    _target_: nemo.collections.asr.modules.RNNTDecoder
+    normalization_mode: null
+    random_state_sampling: false
+    blank_as_pad: true
+
+    prednet:
+      pred_hidden: ${model.model_defaults.pred_hidden}
+      pred_rnn_layers: 1
+      t_max: null
+      dropout: 0.1
+
+  joint:
+    _target_: nemo.collections.asr.modules.RNNTJoint
+    log_softmax: null  # 'null' would set it automatically according to CPU/GPU device
+
+    # fused mode
+    experimental_fuse_loss_wer: true
+    fused_batch_size: 16
+
+    jointnet:
+      joint_hidden: ${model.model_defaults.joint_hidden}
+      activation: "relu"
+      dropout: 0.1
+
+  decoding:
+    strategy: "greedy_batch"
+
+    # greedy strategy config
+    greedy:
+      max_symbols: 30
+
+    # beam strategy config
+    beam:
+      beam_size: 2
+      score_norm: true
+      tsd_max_sym_exp: 50  # for Time Synchronous Decoding
+      alsd_max_target_len: 2.0  # for Alignment-Length Synchronous Decoding
+
+  optim:
+    name: adam
+    lr: 5.0
+    # optimizer arguments
+    betas: [0.9, 0.98]
+    weight_decay: 1e-6
+
+    # scheduler setup
+    sched:
+      name: NoamAnnealing
+      d_model: ${model.encoder.d_model}
+      # scheduler config override
+      warmup_steps: 10000
+      warmup_ratio: null
+      min_lr: 1e-6
+
+trainer:
+  gpus: -1 # number of GPUs, -1 would use all available GPUs
+  num_nodes: 1
+  max_epochs: 1000
+  max_steps: null # computed at runtime if not set
+  val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  gradient_clip_val: 0.0
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
+  log_every_n_steps: 10  # Interval of logging.
+  progress_bar_refresh_rate: 10
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+  num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
+  check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
+  sync_batchnorm: true
+  checkpoint_callback: false  # Provided by exp_manager
+  logger: false  # Provided by exp_manager
+
+
+exp_manager:
+  exp_dir: null
+  name: ${name}
+  create_tensorboard_logger: true
+  create_checkpoint_callback: true
+  checkpoint_callback_params:
+    # in case of multiple validation sets, first one is used
+    monitor: "val_wer"
+    mode: "min"
+    save_top_k: 3
+  resume_if_exists: false
+  resume_ignore_no_checkpoint: false
+
+  create_wandb_logger: false
+  wandb_logger_kwargs:
+    name: null
+    project: null

tts/NeMo/examples/asr/experimental/contextnet_rnnt/config_rnnt.yaml

@@ -0,0 +1,237 @@
+name: &name "ConvRNNT5x1"
+
+model:
+  sample_rate: 16000
+  compute_eval_loss: true
+
+  labels: [ " ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+            "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'" ]
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    batch_size: 32
+    trim_silence: true
+    max_duration: 16.7
+    labels: ${model.labels}
+    shuffle: true
+    is_tarred: false
+    tarred_audio_filepaths: null
+    tarred_shard_strategy: "scatter"
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    batch_size: 32
+    shuffle: false
+    labels: ${model.labels}
+
+  test_ds:
+    manifest_filepath: null
+    sample_rate: ${model.sample_rate}
+    batch_size: 32
+    shuffle: false
+    labels: ${model.labels}
+
+  model_defaults:
+    repeat: 5
+    dropout: 0.0
+    separable: true
+    se: true
+    se_context_size: -1
+    # encoder / decoder / joint values
+    enc_hidden: 1024
+    pred_hidden: 320
+    joint_hidden: 320
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    normalize: "per_feature"
+    window_size: 0.02
+    sample_rate: ${model.sample_rate}
+    window_stride: 0.01
+    window: "hann"
+    features: &n_mels 80
+    n_fft: 512
+    frame_splicing: 1
+    dither: 0.00001
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    rect_freq: 50
+    rect_masks: 5
+    rect_time: 120
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+      - filters: 128
+        repeat: 1
+        kernel: [11]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+
+      - filters: 256
+        repeat: ${model.model_defaults.repeat}
+        kernel: [13]
+        stride: [2]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+
+      - filters: 256
+        repeat: ${model.model_defaults.repeat}
+        kernel: [15]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+
+      - filters: 256
+        repeat: ${model.model_defaults.repeat}
+        kernel: [17]
+        stride: [2]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+
+      - filters: 256
+        repeat: ${model.model_defaults.repeat}
+        kernel: [19]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+
+      - filters: 256
+        repeat: 1
+        kernel: [21]
+        stride: [2]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+
+      - filters: ${model.model_defaults.enc_hidden}
+        repeat: 1
+        kernel: [1]
+        stride: [1]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+
+  decoder:
+    _target_: nemo.collections.asr.modules.RNNTDecoder
+    normalization_mode: null
+    random_state_sampling: false
+    blank_as_pad: true
+
+    prednet:
+      pred_hidden: ${model.model_defaults.pred_hidden}
+      pred_rnn_layers: 1
+      t_max: null
+      dropout: 0.0
+
+  joint:
+    _target_: nemo.collections.asr.modules.RNNTJoint
+    log_softmax: null  # sets it according to cpu/gpu device
+
+    # fused mode
+    experimental_fuse_loss_wer: false
+    fused_batch_size: 1
+
+    jointnet:
+      joint_hidden: ${model.model_defaults.joint_hidden}
+      activation: "relu"
+      dropout: 0.0
+
+  decoding:
+    strategy: "greedy_batch"
+
+    # greedy strategy config
+    greedy:
+      max_symbols: 30
+
+    # beam strategy config
+    beam:
+      beam_size: 2
+      score_norm: true
+      tsd_max_sym_exp: 50  # for Time Synchronous Decoding
+      alsd_max_target_len: 2.0  # for Alignment-Length Synchronous Decoding
+
+  optim:
+    name: adam
+    # _target_: nemo.core.optim.optimizers.Adam
+    lr: .1
+
+    # optimizer arguments
+    betas: [0.9, 0.999]
+    weight_decay: 0.0001
+
+    # scheduler setup
+    sched:
+      name: CosineAnnealing
+
+      # scheduler config override
+      warmup_steps: null
+      warmup_ratio: 0.05
+      min_lr: 1e-6
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 5
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  precision: 32
+  accumulate_grad_batches: 1
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+  checkpoint_callback_params:
+    monitor: "val_wer"
+    mode: "min"
+  create_wandb_logger: False
+  wandb_logger_kwargs:
+    name: null
+    project: null

tts/NeMo/examples/asr/experimental/contextnet_rnnt/config_rnnt_bpe.yaml

@@ -0,0 +1,238 @@
+name: &name "ConvRNNTBPE5x1"
+
+model:
+  sample_rate: 16000
+  compute_eval_loss: true
+
+  tokenizer:
+    dir: ???  # path to directory which contains either tokenizer.model (bpe) or vocab.txt (for wpe)
+    type: ???  # Can be either bpe or wpe
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    batch_size: 32
+    trim_silence: true
+    max_duration: 16.7
+    labels: []
+    shuffle: true
+    is_tarred: false
+    tarred_audio_filepaths: null
+    tarred_shard_strategy: "scatter"
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    batch_size: 32
+    shuffle: false
+    labels: []
+
+  test_ds:
+    manifest_filepath: null
+    sample_rate: ${model.sample_rate}
+    batch_size: 32
+    shuffle: false
+    labels: []
+
+  model_defaults:
+    repeat: 5
+    dropout: 0.0
+    separable: true
+    se: true
+    se_context_size: -1
+    # encoder / decoder / joint values
+    enc_hidden: 1024
+    pred_hidden: 320
+    joint_hidden: 320
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    normalize: "per_feature"
+    window_size: 0.02
+    sample_rate: ${model.sample_rate}
+    window_stride: 0.01
+    window: "hann"
+    features: &n_mels 80
+    n_fft: 512
+    frame_splicing: 1
+    dither: 0.00001
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    rect_freq: 50
+    rect_masks: 5
+    rect_time: 120
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+      - filters: 128
+        repeat: 1
+        kernel: [11]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+
+      - filters: 256
+        repeat: ${model.model_defaults.repeat}
+        kernel: [13]
+        stride: [2]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+
+      - filters: 256
+        repeat: ${model.model_defaults.repeat}
+        kernel: [15]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+
+      - filters: 256
+        repeat: ${model.model_defaults.repeat}
+        kernel: [17]
+        stride: [2]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+
+      - filters: 256
+        repeat: ${model.model_defaults.repeat}
+        kernel: [19]
+        stride: [1]
+        dilation: [1]
+        dropout: ${model.model_defaults.dropout}
+        residual: true
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+
+      - filters: 256
+        repeat: 1
+        kernel: [21]
+        stride: [2]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+        stride_last: true
+        residual_mode: "stride_add"
+
+      - filters: ${model.model_defaults.enc_hidden}
+        repeat: 1
+        kernel: [1]
+        stride: [1]
+        dilation: [1]
+        dropout: 0.0
+        residual: false
+        separable: ${model.model_defaults.separable}
+        se: ${model.model_defaults.se}
+        se_context_size: ${model.model_defaults.se_context_size}
+
+  decoder:
+    _target_: nemo.collections.asr.modules.RNNTDecoder
+    normalization_mode: null
+    random_state_sampling: false
+    blank_as_pad: true
+
+    prednet:
+      pred_hidden: ${model.model_defaults.pred_hidden}
+      pred_rnn_layers: 1
+      t_max: null
+      dropout: 0.0
+
+  joint:
+    _target_: nemo.collections.asr.modules.RNNTJoint
+    log_softmax: null  # sets it according to cpu/gpu device
+
+    # fused mode
+    experimental_fuse_loss_wer: false
+    fused_batch_size: 1
+
+    jointnet:
+      joint_hidden: ${model.model_defaults.joint_hidden}
+      activation: "relu"
+      dropout: 0.0
+
+  decoding:
+    strategy: "greedy_batch"
+
+    # greedy strategy config
+    greedy:
+      max_symbols: 30
+
+    # beam strategy config
+    beam:
+      beam_size: 2
+      score_norm: true
+      tsd_max_sym_exp: 50  # for Time Synchronous Decoding
+      alsd_max_target_len: 2.0  # for Alignment-Length Synchronous Decoding
+
+  optim:
+    name: adam
+    # _target_: nemo.core.optim.optimizers.Adam
+    lr: .1
+
+    # optimizer arguments
+    betas: [0.9, 0.999]
+    weight_decay: 0.0001
+
+    # scheduler setup
+    sched:
+      name: CosineAnnealing
+
+      # scheduler config override
+      warmup_steps: null
+      warmup_ratio: 0.05
+      min_lr: 1e-6
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 5
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  precision: 32
+  accumulate_grad_batches: 1
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+  checkpoint_callback_params:
+    monitor: "val_wer"
+    mode: "min"
+  create_wandb_logger: False
+  wandb_logger_kwargs:
+    name: null
+    project: null

tts/NeMo/examples/asr/experimental/quartznet/quartznet_15x5_aug.yaml

@@ -0,0 +1,290 @@
+name: &name "QuartzNet15x5"
+
+model:
+  sample_rate: &sample_rate 16000
+  repeat: &repeat 5
+  dropout: &dropout 0.0
+  separable: &separable true
+  labels: &labels [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+                   "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    labels: *labels
+    batch_size: 32
+    trim_silence: True
+    max_duration: 16.7
+    shuffle: True
+    is_tarred: False
+    tarred_audio_filepaths: null
+    tarred_shard_strategy: "scatter"
+    augmentor:
+      rir_noise_aug:
+        prob: 0.5
+        rir_manifest_path: ???
+        rir_tar_filepaths: ???
+        rir_prob: 0.5
+        noise_manifest_paths: ???
+        noise_tar_filepaths: ???
+        min_snr_db: [0,0]
+        max_snr_db: [30,30]
+        orig_sample_rate: ???
+        bg_noise_manifest_paths: ???
+        bg_noise_tar_filepaths: ???
+        bg_min_snr_db: [10,10]
+        bg_max_snr_db: [40,40]
+        bg_orig_sample_rate: ???
+      transcode_aug:
+        prob: 0.1
+
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    labels: *labels
+    batch_size: 32
+    shuffle: False
+
+  preprocessor:
+    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
+    normalize: "per_feature"
+    window_size: 0.02
+    sample_rate: *sample_rate
+    window_stride: 0.01
+    window: "hann"
+    features: &n_mels 64
+    n_fft: 512
+    frame_splicing: 1
+    dither: 0.00001
+    stft_conv: false
+
+  spec_augment:
+    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    rect_freq: 50
+    rect_masks: 5
+    rect_time: 120
+
+  encoder:
+    _target_: nemo.collections.asr.modules.ConvASREncoder
+    feat_in: *n_mels
+    activation: relu
+    conv_mask: true
+
+    jasper:
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: 1
+      residual: false
+      separable: *separable
+      stride: [2]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [2]
+      dropout: *dropout
+      filters: 512
+      kernel: [87]
+      repeat: 1
+      residual: false
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: &enc_filters 1024
+      kernel: [1]
+      repeat: 1
+      residual: false
+      stride: [1]
+
+  decoder:
+    _target_: nemo.collections.asr.modules.ConvASRDecoder
+    feat_in: *enc_filters
+    num_classes: 28
+    vocabulary: *labels
+
+  optim:
+    name: novograd
+    # _target_: nemo.core.optim.optimizers.Novograd
+    lr: .01
+    # optimizer arguments
+    betas: [0.8, 0.5]
+    weight_decay: 0.001
+
+    # scheduler setup
+    sched:
+      name: CosineAnnealing
+
+      # pytorch lightning args
+      # monitor: val_loss
+      # reduce_on_plateau: false
+
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: null
+      min_lr: 0.0
+      last_epoch: -1
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 5
+  max_steps: null # computed at runtime if not set
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+  create_wandb_logger: False
+  wandb_logger_kwargs:
+    name: null
+    project: null
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/asr/experimental/sclite/speech_to_text_sclite.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script is based on speech_to_text_infer.py and allows you to score the hypotheses
+with sclite. A local installation from https://github.com/usnistgov/SCTK is required.
+Hypotheses and references are first saved in trn format and are scored after applying a glm
+file (if provided).
+"""
+
+import errno
+import json
+import os
+import subprocess
+from argparse import ArgumentParser
+
+import torch
+
+from nemo.collections.asr.metrics.wer import WER
+from nemo.collections.asr.models import EncDecCTCModel
+from nemo.utils import logging
+
+try:
+    from torch.cuda.amp import autocast
+except ImportError:
+    from contextlib import contextmanager
+
+    @contextmanager
+    def autocast(enabled=None):
+        yield
+
+
+def score_with_sctk(sctk_dir, ref_fname, hyp_fname, out_dir, glm=""):
+    sclite_path = os.path.join(sctk_dir, "bin", "sclite")
+    if not os.path.exists(sclite_path):
+        raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), sclite_path)
+    # apply glm
+    if os.path.exists(glm):
+        rfilter_path = os.path.join(sctk_dir, "bin", "rfilter1")
+        if not os.path.exists(rfilter_path):
+            raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), rfilter_path)
+        hypglm = os.path.join(out_dir, os.path.basename(hyp_fname)) + ".glm"
+        rfilt_cmd = [rfilter_path] + [glm]
+        with open(hypglm, "w") as hypf, open(hyp_fname, "r") as hyp_in:
+            subprocess.run(rfilt_cmd, stdin=hyp_in, stdout=hypf)
+        refglm = os.path.join(out_dir, os.path.basename(ref_fname)) + ".glm"
+        with open(refglm, "w") as reff, open(ref_fname, "r") as ref_in:
+            subprocess.run(rfilt_cmd, stdin=ref_in, stdout=reff)
+    else:
+        refglm = ref_fname
+        hypglm = hyp_fname
+
+    _ = subprocess.check_output(f"{sclite_path} -h {hypglm}  -r {refglm} -i wsj -o all", shell=True)
+
+
+can_gpu = torch.cuda.is_available()
+
+
+def get_utt_info(manifest_path):
+    info_list = []
+    with open(manifest_path, "r") as utt_f:
+        for line in utt_f:
+            utt = json.loads(line)
+            info_list.append(utt)
+
+    return info_list
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--asr_model", type=str, default="QuartzNet15x5Base-En", required=False, help="Pass: 'QuartzNet15x5Base-En'",
+    )
+    parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data")
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument(
+        "--dont_normalize_text",
+        default=False,
+        action='store_true',
+        help="Turn off trasnscript normalization. Recommended for non-English.",
+    )
+    parser.add_argument("--out_dir", type=str, required=True, help="Destination dir for output files")
+    parser.add_argument("--sctk_dir", type=str, required=False, default="", help="Path to sctk root dir")
+    parser.add_argument("--glm", type=str, required=False, default="", help="Path to glm file")
+    args = parser.parse_args()
+    torch.set_grad_enabled(False)
+
+    if not os.path.exists(args.out_dir):
+        os.makedirs(args.out_dir)
+
+    use_sctk = os.path.exists(args.sctk_dir)
+
+    if args.asr_model.endswith('.nemo'):
+        logging.info(f"Using local ASR model from {args.asr_model}")
+        asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model)
+    else:
+        logging.info(f"Using NGC cloud ASR model {args.asr_model}")
+        asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model)
+    asr_model.setup_test_data(
+        test_data_config={
+            'sample_rate': 16000,
+            'manifest_filepath': args.dataset,
+            'labels': asr_model.decoder.vocabulary,
+            'batch_size': args.batch_size,
+            'normalize_transcripts': not args.dont_normalize_text,
+        }
+    )
+    if can_gpu:
+        asr_model = asr_model.cuda()
+    asr_model.eval()
+    labels_map = dict([(i, asr_model.decoder.vocabulary[i]) for i in range(len(asr_model.decoder.vocabulary))])
+
+    wer = WER(vocabulary=asr_model.decoder.vocabulary)
+    hypotheses = []
+    references = []
+    all_log_probs = []
+    for test_batch in asr_model.test_dataloader():
+        if can_gpu:
+            test_batch = [x.cuda() for x in test_batch]
+        with autocast():
+            log_probs, encoded_len, greedy_predictions = asr_model(
+                input_signal=test_batch[0], input_signal_length=test_batch[1]
+            )
+        for r in log_probs.cpu().numpy():
+            all_log_probs.append(r)
+        hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions)
+        for batch_ind in range(greedy_predictions.shape[0]):
+            reference = ''.join([labels_map[c] for c in test_batch[2][batch_ind].cpu().detach().numpy()])
+            references.append(reference)
+        del test_batch
+
+    info_list = get_utt_info(args.dataset)
+    hypfile = os.path.join(args.out_dir, "hyp.trn")
+    reffile = os.path.join(args.out_dir, "ref.trn")
+    with open(hypfile, "w") as hyp_f, open(reffile, "w") as ref_f:
+        for i in range(len(hypotheses)):
+            utt_id = os.path.splitext(os.path.basename(info_list[i]['audio_filepath']))[0]
+            # rfilter in sctk likes each transcript to have a space at the beginning
+            hyp_f.write(" " + hypotheses[i] + " (" + utt_id + ")" + "\n")
+            ref_f.write(" " + references[i] + " (" + utt_id + ")" + "\n")
+
+    if use_sctk:
+        score_with_sctk(args.sctk_dir, reffile, hypfile, args.out_dir, glm=args.glm)
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/experimental/structured/conf/quartznet_15x5.yaml

@@ -0,0 +1,235 @@
+name: &name "QuartzNet15x5"
+
+model:
+  sample_rate: &sample_rate 16000
+  repeat: &repeat 5
+  dropout: &dropout 0.0
+  separable: &separable true
+  labels: &labels [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+                   "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: *sample_rate
+    labels: *labels
+    batch_size: 32
+    max_duration: 16.7
+    is_tarred: False
+    tarred_audio_filepaths: null
+
+  validation_ds:
+    labels: *labels
+    manifest_filepath: ???
+    sample_rate: *sample_rate
+    batch_size: 32
+
+  preprocessor:
+    sample_rate: *sample_rate
+    window_size: 0.02
+    window_stride: 0.01
+    features: &n_mels 64
+
+  spec_augment:
+    rect_freq: 50
+    rect_masks: 5
+    rect_time: 120
+
+  encoder:
+    feat_in: *n_mels
+    activation: relu
+
+    jasper:
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: 1
+      residual: false
+      separable: *separable
+      stride: [2]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [33]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 256
+      kernel: [39]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [51]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [63]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: 512
+      kernel: [75]
+      repeat: *repeat
+      residual: true
+      separable: *separable
+      stride: [1]
+
+    - dilation: [2]
+      dropout: *dropout
+      filters: 512
+      kernel: [87]
+      repeat: 1
+      residual: false
+      separable: *separable
+      stride: [1]
+
+    - dilation: [1]
+      dropout: *dropout
+      filters: &enc_filters 1024
+      kernel: [1]
+      repeat: 1
+      residual: false
+      stride: [1]
+
+  decoder:
+    feat_in: *enc_filters
+    num_classes: 28
+    vocabulary: *labels
+
+  optim:
+    name: novograd
+    lr: .01
+
+    # optimizer arguments
+    betas: [0.8, 0.5]
+    weight_decay: 0.001
+
+    # scheduler setup
+    sched:
+      name: CosineAnnealing
+
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: null
+      min_lr: 0.0
+
+trainer:
+  gpus: 0 # number of gpus
+  max_epochs: 5
+  num_nodes: 1
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+
+exp_manager:
+  name: *name

tts/NeMo/examples/asr/experimental/structured/speech_to_text_hybrid.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytorch_lightning as pl
+
+from nemo.collections.asr.models import EncDecCTCModel, configs
+from nemo.core.config import hydra_runner
+from nemo.utils.config_utils import update_model_config
+from nemo.utils.exp_manager import exp_manager
+
+
+"""
+python speech_to_text_hybrid.py \
+    --config-path="conf/quartznet" \
+    --config-name="quartznet_15x5" \
+    model.train_ds.manifest_filepath="/home/smajumdar/PycharmProjects/NeMo-som/examples/asr/an4/train_manifest.json" \
+    model.validation_ds.manifest_filepath="/home/smajumdar/PycharmProjects/NeMo-som/examples/asr/an4/test_manifest.json" \
+    trainer.gpus=1
+"""
+
+
+@hydra_runner(config_path="conf/quartznet", config_name="quartznet_15x5")
+def main(cfg):
+    # Generate default asr model config
+    asr_model_config = configs.EncDecCTCModelConfig()
+
+    # Merge hydra updates with model config
+    # `drop_missing_subconfig=True` is necessary here. Without it, while the data class will instantiate and be added
+    # to the config, it contains test_ds.sample_rate = MISSING and test_ds.labels = MISSING.
+    # This will raise a OmegaConf MissingMandatoryValue error when processing the dataloaders inside
+    # model_utils.resolve_test_dataloaders(model=self) (used for multi data loader support).
+    # In general, any operation that tries to use a DictConfig with MISSING in it will fail,
+    # other than explicit update operations to change MISSING to some actual value.
+    asr_model_config = update_model_config(asr_model_config, cfg, drop_missing_subconfigs=True)
+
+    # From here on out, its a general OmegaConf DictConfig, directly usable by our code.
+    trainer = pl.Trainer(**asr_model_config.trainer)
+    exp_manager(trainer, asr_model_config.get("exp_manager", None))
+    asr_model = EncDecCTCModel(cfg=asr_model_config.model, trainer=trainer)
+
+    trainer.fit(asr_model)
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/experimental/structured/speech_to_text_structured.py

@@ -0,0 +1,146 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import asdict
+
+import pytorch_lightning as pl
+
+import nemo.collections.asr as nemo_asr
+from nemo.collections.asr.models import EncDecCTCModel, configs
+from nemo.utils.exp_manager import exp_manager
+
+"""
+python speech_to_text_structured.py
+"""
+
+# Generate default asr model config
+cfg = configs.EncDecCTCModelConfig()
+
+# set global values
+cfg.model.repeat = 5
+cfg.model.separable = True
+
+# fmt: off
+LABELS = [
+    " ", "a", "b", "c", "d", "e",
+    "f", "g", "h", "i", "j", "k",
+    "l", "m", "n", "o", "p", "q",
+    "r", "s", "t", "u", "v", "w",
+    "x", "y", "z", "'",
+]
+# fmt: on
+
+qn_15x5 = [
+    nemo_asr.modules.conv_asr.JasperEncoderConfig(
+        filters=256,
+        repeat=1,
+        kernel=[33],
+        stride=[2],
+        separable=cfg.model.separable,
+        dilation=[1],
+        dropout=cfg.model.dropout,
+        residual=False,
+    ),
+    nemo_asr.modules.conv_asr.JasperEncoderConfig(
+        filters=256,
+        repeat=1,
+        kernel=[33],
+        stride=[1],
+        separable=cfg.model.separable,
+        dilation=[1],
+        dropout=cfg.model.dropout,
+        residual=True,
+    ),
+    # ... repeat 14 more times
+    nemo_asr.modules.conv_asr.JasperEncoderConfig(
+        filters=1024, repeat=1, kernel=[1], stride=[1], dilation=[1], dropout=cfg.model.dropout, residual=False,
+    ),
+]
+
+
+def main():
+    # Update values
+    # MODEL UPDATES
+    cfg.name = "Mini QuartzNet"
+    cfg.model.labels = LABELS
+
+    # train ds
+    cfg.model.train_ds.manifest_filepath = "<path to train dataset>"
+    cfg.model.train_ds.labels = LABELS
+    cfg.model.train_ds.sample_rate = cfg.model.sample_rate
+
+    # validation ds
+    cfg.model.validation_ds.manifest_filepath = "<path to test dataset>"
+    cfg.model.validation_ds.labels = LABELS
+    cfg.model.validation_ds.sample_rate = cfg.model.sample_rate
+
+    # del `test_ds` does not work!
+    # Refer - https://stackoverflow.com/questions/58119758/how-to-remove-dataclass-attributes
+    # Hydra/OmegaConf dont allow custom .asdict() methods either
+    # For now, explicitly set parameters
+    cfg.model.test_ds.sample_rate = cfg.model.sample_rate
+    cfg.model.test_ds.labels = cfg.model.labels
+
+    # preprocessor
+    cfg.model.preprocessor.sample_rate = cfg.model.sample_rate
+
+    # spec aug
+    cfg.model.spec_augment.rect_masks = 5
+    cfg.model.spec_augment.rect_freq = 50
+    cfg.model.spec_augment.rect_time = 120
+
+    # encoder
+    cfg.model.encoder.feat_in = cfg.model.preprocessor.features
+    cfg.model.encoder.activation = 'relu'
+    cfg.model.encoder.jasper = qn_15x5
+
+    # decoder
+    cfg.model.decoder.feat_in = qn_15x5[-1].filters
+    cfg.model.decoder.num_classes = len(LABELS)
+    cfg.model.decoder.vocabulary = LABELS
+
+    # optim
+    cfg.model.optim.name = 'novograd'
+    cfg.model.optim.lr = 0.01
+
+    # `betas` dont exist inside the base config,
+    # so they cannot be added as such!
+    # Same for `weight_decay`.
+    cfg.model.optim.betas = [0.8, 0.5]
+    cfg.model.optim.weight_decay = 0.001
+
+    # sched
+    # As parameters such as warmup_steps and warmup_ratio
+    # dont exist inside the shell config, these values are not added!
+    cfg.model.optim.sched.name = "CosineAnnealing"
+    cfg.model.optim.sched.warmup_steps = None
+    cfg.model.optim.sched.warmup_ratio = 0.01
+
+    # Trainer config
+    cfg.trainer.gpus = 1
+    cfg.trainer.max_epochs = 5
+
+    # Exp Manager config
+    cfg.exp_manager.name = cfg.name
+
+    # Note usage of asdict
+    trainer = pl.Trainer(**asdict(cfg.trainer))
+    exp_manager(trainer, asdict(cfg.exp_manager))
+    asr_model = EncDecCTCModel(cfg=cfg.model, trainer=trainer)
+
+    trainer.fit(asr_model)
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/experimental/structured/speech_to_text_structured_v2.py

@@ -0,0 +1,90 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import asdict
+
+import pytorch_lightning as pl
+
+from nemo.collections.asr.models import EncDecCTCModel, configs
+from nemo.core.config import modelPT, optimizers, schedulers
+from nemo.utils.exp_manager import exp_manager
+
+"""
+python speech_to_text_structured_v2.py
+"""
+
+# fmt: off
+LABELS = [
+    " ", "a", "b", "c", "d", "e",
+    "f", "g", "h", "i", "j", "k",
+    "l", "m", "n", "o", "p", "q",
+    "r", "s", "t", "u", "v", "w",
+    "x", "y", "z", "'",
+]
+
+optim_cfg = optimizers.NovogradParams(
+    lr=0.01,
+    betas=(0.8, 0.5),
+    weight_decay=0.001
+)
+
+sched_cfg = schedulers.CosineAnnealingParams(
+    warmup_steps=None,
+    warmup_ratio=None,
+    min_lr=0.0,
+)
+# fmt: on
+
+
+def main():
+    # NeMo Model config
+    cfg = modelPT.NemoConfig(name='Custom QuartzNet')
+
+    # Generate default asr model config
+    builder = configs.EncDecCTCModelConfigBuilder(name='quartznet_15x5')
+
+    # set model global values
+    builder.set_labels(LABELS)
+    builder.set_optim(cfg=optim_cfg, sched_cfg=sched_cfg)
+
+    model_cfg = builder.build()
+
+    # set the model config to the NeMo Model
+    cfg.model = model_cfg
+
+    # Update values
+    # MODEL UPDATES
+    # train ds
+    model_cfg.train_ds.manifest_filepath = ""
+
+    # validation ds
+    model_cfg.validation_ds.manifest_filepath = ""
+
+    # Trainer config
+    cfg.trainer.gpus = 1
+    cfg.trainer.max_epochs = 5
+
+    # Exp Manager config
+    cfg.exp_manager.name = cfg.name
+
+    # Note usage of asdict
+    trainer = pl.Trainer(**asdict(cfg.trainer))
+    exp_manager(trainer, asdict(cfg.exp_manager))
+    asr_model = EncDecCTCModel(cfg=cfg.model, trainer=trainer)
+
+    trainer.fit(asr_model)
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/experimental/wav2vec/configs/wav2vec_pretrain.yaml

@@ -0,0 +1,123 @@
+name: &name "Wav2vec_Pretrain"
+# For more details on the model parameters, see nemo/collections/asr/models/wav2vec/wav2vec_config.py
+# These parameters are based off the FairSeq implementation.
+# See here: https://github.com/pytorch/fairseq/blob/master/examples/wav2vec/config/pretraining/wav2vec2_base_librispeech.yaml
+model:
+  transformer_encoder:
+    dropout: 0.1 # Dropout probability for the transformer
+    conv: # Config for convolutional model that generates positional embeddings required for attention layer
+      conv_pos: 128 # Number of filters for convolutional positional embeddings
+      conv_pos_groups: 16 # Number of groups for convolutional positional embeddings
+    encoder:
+      encoder_layers: 12 # Number of encoder layers in transformer model
+      encoder_layerdrop: 0.05 # Probability of dropping transformer layers
+      embedding_dim: 768 # Encoder embedding dim
+      ffn_embedding_dim: 3072 # Encoder embedding dim for feed forward
+      num_attention_heads: 8 # Number of encoder attention heads
+      dropout: 0.1 # Dropout probability for transformer encoder
+      activation_fn: gelu # Activation for transformer
+      layer_norm_first: False # Apply layer norm first within the transformer
+  loss:
+    prob_ppl_weight: 0.1 # Weight applied to quantized prob perplexity loss
+    feature_loss_weight: 0 # Weight applied to feature L2 Norm
+  quantizer:
+    quantize_targets: True # Quantize targets
+    quantize_input: False # Quantize inputs
+    same_quantizer: False # Use the same quantizer for inputs and targets (If both enabled)
+    latent_vars: 320 # Number of latent variables in each group of the codebook
+    latent_groups: 2 # Number of groups within the codebook
+    latent_dim: 0 # If greater than 0, use the dim for latest variables else inferred by final dim / latent groups
+    latent_temp: [2, 0.5, 0.999995] # Quantize temperature (start, stop, decay factor)
+  conv_feature_encoder: # Config for feature extractor taking raw audio waveform and converting to features
+    extractor_mode: default # Mode for feature extractor. [Default, layer_norm]
+    conv_bias: False # Include bias in convolution feature extractor model
+  masking:
+    mask_prob: 0.65 # Probability of replacing token with mask
+    mask_type: static # Used to select configuration to compute mask lengths [static, uniform, normal, poisson]
+    mask_other: 0 # Secondary mask used for complex distributions (see help in compute_mask_indices)
+    mask_length: 10 # Length of mask when masking time steps
+    no_mask_overlap: False # Whether to allow masks to overlap
+    mask_min_space: 1 # Minimum space beetween spans (if no overlap is enabled)
+    mask_channel_prob: 0 # Probability of replacing a feature with 0
+    mask_channel_type: static # Used to select configuration to compute mask channel lengths [static, uniform, normal, poisson]
+    mask_channel_other: 0 # Secondary mask argument (used for more complex distributions (see help in compute_mask_indices)
+    mask_channel_length: 10 # Length of masks for features (channels)
+    no_mask_channel_overlap: False # Whether to allow channel masks to overlap
+    mask_channel_min_space: 1 # Minimum space between spans (if no overlap is enabled)
+  dropout_input: 0.1 # Dropout applied to input raw features
+  dropout_features: 0.1 # Dropout applied to the features generator by convolutions
+  final_dim: 0 # Project final representations and targets to this dimension
+  n_negatives: 100 # Number of negatives to sample from the same audio sample
+  cross_sample_negatives: 0 # Number of negatives to sample from any sample in the batch
+  codebook_negatives: 0 # Number of negative examples in codebook
+  negatives_from_everywhere: False # Sample negatives from everywhere, not just masked states
+  logit_temp: 0.1 # Temperature to divide logits by
+  target_glu: False # Adds project and applies GLU to targets
+  feature_grad_mult: 0.1 # Multiply extracted feature gradients
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    batch_size: 32
+    shuffle: True
+    num_workers: 4
+    labels: [ ] # No vocab needed for pre-training
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    batch_size: 32
+    shuffle: False
+    num_workers: 1
+    labels: [ ]
+  test_ds:
+    manifest_filepath: null
+    sample_rate: 16000
+    batch_size: 32
+    shuffle: False
+    num_workers: 4
+    labels: [ ]
+  optim:
+    name: adamw
+    lr: 0.0005
+    eps: 1e-06
+    # optimizer arguments
+    betas: [ 0.9, 0.98 ]
+    weight_decay: 0.01
+
+    # scheduler setup
+    sched:
+      name: PolynomialDecayAnnealing
+      min_lr: 0.000
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.1
+trainer:
+  gpus: 0 # number of gpus
+  num_nodes: 1
+  max_epochs: 100
+  max_steps: null # computed at runtime if not set
+  val_check_interval: 0.5 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  gradient_clip_val: 0.0
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
+  log_every_n_steps: 10 # Interval of logging.
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+  num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
+  check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
+  sync_batchnorm: true
+
+  checkpoint_callback: false # Provided by exp_manager
+  logger: false # Provided by exp_manager
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: true
+  create_checkpoint_callback: true
+  create_wandb_logger: false
+  wandb_logger_kwargs:
+    name: null
+    project: null
+  resume_if_exists: false
+  resume_ignore_no_checkpoint: false

tts/NeMo/examples/asr/experimental/wav2vec/configs/wav2vec_pretrain_large.yaml

@@ -0,0 +1,123 @@
+name: &name "Wav2vec_large_Pretrain"
+# For more details on the model parameters, see nemo/collections/asr/models/wav2vec/wav2vec_config.py
+# These parameters are based off the FairSeq implementation.
+# See here: https://github.com/pytorch/fairseq/blob/master/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml
+model:
+  transformer_encoder:
+    dropout: 0.0 # Dropout probability for the transformer
+    conv: # Config for convolutional model that generates positional embeddings required for attention layer
+      conv_pos: 128 # Number of filters for convolutional positional embeddings
+      conv_pos_groups: 16 # Number of groups for convolutional positional embeddings
+    encoder:
+      encoder_layers: 24 # Number of encoder layers in transformer model
+      encoder_layerdrop: 0.0 # Probability of dropping transformer layers
+      embedding_dim: 1024 # Encoder embedding dim
+      ffn_embedding_dim: 4096 # Encoder embedding dim for feed forward
+      num_attention_heads: 16 # Number of encoder attention heads
+      dropout: 0.0 # Dropout probability for transformer encoder
+      activation_fn: gelu # Activation for transformer
+      layer_norm_first: True # Apply layer norm first within the transformer
+  loss:
+    prob_ppl_weight: 0.1 # Weight applied to quantized prob perplexity loss
+    feature_loss_weight: 0 # Weight applied to feature L2 Norm
+  quantizer:
+    quantize_targets: True # Quantize targets
+    quantize_input: False # Quantize inputs
+    same_quantizer: False # Use the same quantizer for inputs and targets (If both enabled)
+    latent_vars: 320 # Number of latent variables in each group of the codebook
+    latent_groups: 2 # Number of groups within the codebook
+    latent_dim: 0 # If greater than 0, use the dim for latest variables else inferred by final dim / latent groups
+    latent_temp: [2, 0.1, 0.999995] # Quantize temperature (start, stop, decay factor)
+  conv_feature_encoder: # Config for feature extractor taking raw audio waveform and converting to features
+    extractor_mode: layer_norm # Mode for feature extractor. [default, layer_norm]
+    conv_bias: True # Include bias in convolution feature extractor model
+  masking:
+    mask_prob: 0.65 # Probability of replacing token with mask
+    mask_type: static # Used to select configuration to compute mask lengths [static, uniform, normal, poisson]
+    mask_other: 0 # Secondary mask used for complex distributions (see help in compute_mask_indices)
+    mask_length: 10 # Length of mask when masking time steps
+    no_mask_overlap: False # Whether to allow masks to overlap
+    mask_min_space: 1 # Minimum space beetween spans (if no overlap is enabled)
+    mask_channel_prob: 0 # Probability of replacing a feature with 0
+    mask_channel_type: static # Used to select configuration to compute mask channel lengths [static, uniform, normal, poisson]
+    mask_channel_other: 0 # Secondary mask argument (used for more complex distributions (see help in compute_mask_indices)
+    mask_channel_length: 10 # Length of masks for features (channels)
+    no_mask_channel_overlap: False # Whether to allow channel masks to overlap
+    mask_channel_min_space: 1 # Minimum space between spans (if no overlap is enabled)
+  dropout_input: 0.0 # Dropout applied to input raw features
+  dropout_features: 0.0 # Dropout applied to the features generator by convolutions
+  final_dim: 768 # Project final representations and targets to this dimension
+  n_negatives: 100 # Number of negatives to sample from the same audio sample
+  cross_sample_negatives: 0 # Number of negatives to sample from any sample in the batch
+  codebook_negatives: 0 # Number of negative examples in codebook
+  negatives_from_everywhere: False # Sample negatives from everywhere, not just masked states
+  logit_temp: 0.1 # Temperature to divide logits by
+  target_glu: False # Adds project and applies GLU to targets
+  feature_grad_mult: 1.0 # Multiply extracted feature gradients
+
+  train_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    batch_size: 32
+    shuffle: True
+    num_workers: 4
+    labels: [ ] # No vocab needed for pre-training
+  validation_ds:
+    manifest_filepath: ???
+    sample_rate: 16000
+    batch_size: 32
+    shuffle: False
+    num_workers: 1
+    labels: [ ]
+  test_ds:
+    manifest_filepath: null
+    sample_rate: 16000
+    batch_size: 32
+    shuffle: False
+    num_workers: 4
+    labels: [ ]
+  optim:
+    name: adamw
+    lr: 0.0005
+    eps: 1e-06
+    # optimizer arguments
+    betas: [ 0.9, 0.98 ]
+    weight_decay: 0.01
+
+    # scheduler setup
+    sched:
+      name: PolynomialDecayAnnealing
+      min_lr: 0.000
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.1
+trainer:
+  gpus: 0 # number of gpus
+  num_nodes: 1
+  max_epochs: 100
+  max_steps: null # computed at runtime if not set
+  val_check_interval: 0.5 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  accelerator: ddp
+  accumulate_grad_batches: 1
+  gradient_clip_val: 0.0
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
+  log_every_n_steps: 10 # Interval of logging.
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+  num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
+  check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
+  sync_batchnorm: true
+
+  checkpoint_callback: false # Provided by exp_manager
+  logger: false # Provided by exp_manager
+exp_manager:
+  exp_dir: null
+  name: *name
+  create_tensorboard_logger: true
+  create_checkpoint_callback: true
+  create_wandb_logger: false
+  wandb_logger_kwargs:
+    name: null
+    project: null
+  resume_if_exists: false
+  resume_ignore_no_checkpoint: false

tts/NeMo/examples/asr/experimental/wav2vec/wav2vec_pretrain.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.asr.models.wav2vec.wav2vec_model import Wav2VecEncoderModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+"""
+Pre-train a wav2vec 2.0 transformer model on audio. Uses a contrastive loss function to pre-train on unlabelled audio,
+using a task similar to masked language modeling in NLP. In wav2vec, we mask portions of the audio 
+and the model is trained by minimising the distance of the ground truth for the masked section, 
+using the ground truth quantized codebook representation. Distractors are obtained from other time steps.
+See :class:`Wav2VecCriterion` for more information.
+
+Reference: https://arxiv.org/abs/2006.11477
+
+    python examples/asr/experimental/wav2vec/wav2vec_pretrain.py \
+        model.train_ds.manifest_filepath="./examples/asr/train.tsv" \
+        model.validation_ds.manifest_filepath="./examples/asr/valid.tsv" \
+        trainer.gpus=1 \
+        trainer.max_epochs=100
+        
+Basic run (on CPU for 50 epochs):
+    python examples/asr/experimental/wav2vec/wav2vec_pretrain.py \
+        model.train_ds.manifest_filepath="./examples/asr/train.tsv" \
+        model.validation_ds.manifest_filepath="./examples/asr/valid.tsv" \
+        trainer.gpus=1 \
+        trainer.max_epochs=50
+
+Using wav2vec-large with mixed precision:
+    python examples/asr/experimental/wav2vec/wav2vec_pretrain.py \
+        --config-name=wav2vec_pretrain_large \
+        model.train_ds.manifest_filepath="./examples/asr/train.tsv" \
+        model.validation_ds.manifest_filepath="./examples/asr/valid.tsv" \
+        trainer.gpus=1 \
+        trainer.max_epochs=100 \
+        trainer.precision=16
+
+Add PyTorch Lightning Trainer arguments from CLI:
+    python wav2vec.py \
+        ... \
+        +trainer.fast_dev_run=true
+
+Hydra logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/.hydra)"
+PTL logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/lightning_logs)"
+
+Override some args of optimizer:
+    python examples/asr/experimental/wav2vec/wav2vec_pretrain.py \
+        model.train_ds.manifest_filepath="./examples/asr/train.tsv" \
+        model.validation_ds.manifest_filepath="./examples/asr/valid.tsv" \
+        trainer.gpus=2 \
+        trainer.max_epochs=2 \
+        model.optim.args.params.betas=[0.8,0.5] \
+        model.optim.args.params.weight_decay=0.0001
+
+Override optimizer entirely
+    python examples/asr/experimental/wav2vec/wav2vec_pretrain.py \
+        model.train_ds.manifest_filepath="./examples/asr/train.tsv" \
+        model.validation_ds.manifest_filepath="./examples/asr/valid.tsv" \
+        trainer.gpus=2 \
+        trainer.max_epochs=2 \
+        ~model.optim.args \
+        +model.optim.args.betas=[0.8,0.5]\
+        +model.optim.args.weight_decay=0.0005
+
+"""
+
+
+@hydra_runner(config_path="configs", config_name="wav2vec_pretrain")
+def main(cfg: DictConfig):
+    logging.info("Application config\n" + OmegaConf.to_yaml(cfg))
+
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    wav2vec_encoder_model = Wav2VecEncoderModel(cfg=cfg.model, trainer=trainer)
+    trainer.fit(wav2vec_encoder_model)
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/quantization/speech_to_text_calibrate.py

@@ -0,0 +1,160 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Script for calibrating a pretrained ASR model for quantization
+"""
+
+from argparse import ArgumentParser
+
+import torch
+from omegaconf import open_dict
+
+from nemo.collections.asr.models import EncDecCTCModel
+from nemo.utils import logging
+
+try:
+    from pytorch_quantization import calib
+    from pytorch_quantization import nn as quant_nn
+    from pytorch_quantization import quant_modules
+    from pytorch_quantization.tensor_quant import QuantDescriptor
+except ImportError:
+    raise ImportError(
+        "pytorch-quantization is not installed. Install from "
+        "https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
+    )
+
+try:
+    from torch.cuda.amp import autocast
+except ImportError:
+    from contextlib import contextmanager
+
+    @contextmanager
+    def autocast(enabled=None):
+        yield
+
+
+can_gpu = torch.cuda.is_available()
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--asr_model", type=str, default="QuartzNet15x5Base-En", required=True, help="Pass: 'QuartzNet15x5Base-En'",
+    )
+    parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data")
+    parser.add_argument("--batch_size", type=int, default=256)
+    parser.add_argument(
+        "--dont_normalize_text",
+        default=False,
+        action='store_false',
+        help="Turn off trasnscript normalization. Recommended for non-English.",
+    )
+    parser.add_argument('--num_calib_batch', default=1, type=int, help="Number of batches for calibration.")
+    parser.add_argument('--calibrator', type=str, choices=["max", "histogram"], default="max")
+    parser.add_argument('--percentile', nargs='+', type=float, default=[99.9, 99.99, 99.999, 99.9999])
+    parser.add_argument("--amp", action="store_true", help="Use AMP in calibration.")
+    parser.set_defaults(amp=False)
+
+    args = parser.parse_args()
+    torch.set_grad_enabled(False)
+
+    # Initialize quantization
+    quant_desc_input = QuantDescriptor(calib_method=args.calibrator)
+    quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
+    quant_nn.QuantConvTranspose2d.set_default_quant_desc_input(quant_desc_input)
+    quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
+
+    if args.asr_model.endswith('.nemo'):
+        logging.info(f"Using local ASR model from {args.asr_model}")
+        asr_model_cfg = EncDecCTCModel.restore_from(restore_path=args.asr_model, return_config=True)
+        with open_dict(asr_model_cfg):
+            asr_model_cfg.encoder.quantize = True
+        asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model, override_config_path=asr_model_cfg)
+
+    else:
+        logging.info(f"Using NGC cloud ASR model {args.asr_model}")
+        asr_model_cfg = EncDecCTCModel.from_pretrained(model_name=args.asr_model, return_config=True)
+        with open_dict(asr_model_cfg):
+            asr_model_cfg.encoder.quantize = True
+        asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model, override_config_path=asr_model_cfg)
+
+    asr_model.setup_test_data(
+        test_data_config={
+            'sample_rate': 16000,
+            'manifest_filepath': args.dataset,
+            'labels': asr_model.decoder.vocabulary,
+            'batch_size': args.batch_size,
+            'normalize_transcripts': args.dont_normalize_text,
+            'shuffle': True,
+        }
+    )
+    asr_model.preprocessor.featurizer.dither = 0.0
+    asr_model.preprocessor.featurizer.pad_to = 0
+    if can_gpu:
+        asr_model = asr_model.cuda()
+    asr_model.eval()
+
+    # Enable calibrators
+    for name, module in asr_model.named_modules():
+        if isinstance(module, quant_nn.TensorQuantizer):
+            if module._calibrator is not None:
+                module.disable_quant()
+                module.enable_calib()
+            else:
+                module.disable()
+
+    for i, test_batch in enumerate(asr_model.test_dataloader()):
+        if can_gpu:
+            test_batch = [x.cuda() for x in test_batch]
+        if args.amp:
+            with autocast():
+                _ = asr_model(input_signal=test_batch[0], input_signal_length=test_batch[1])
+        else:
+            _ = asr_model(input_signal=test_batch[0], input_signal_length=test_batch[1])
+        if i >= args.num_calib_batch:
+            break
+
+    # Save calibrated model(s)
+    model_name = args.asr_model.replace(".nemo", "") if args.asr_model.endswith(".nemo") else args.asr_model
+    if not args.calibrator == "histogram":
+        compute_amax(asr_model, method="max")
+        asr_model.save_to(F"{model_name}-max-{args.num_calib_batch*args.batch_size}.nemo")
+    else:
+        for percentile in args.percentile:
+            print(F"{percentile} percentile calibration")
+            compute_amax(asr_model, method="percentile")
+            asr_model.save_to(F"{model_name}-percentile-{percentile}-{args.num_calib_batch*args.batch_size}.nemo")
+
+        for method in ["mse", "entropy"]:
+            print(F"{method} calibration")
+            compute_amax(asr_model, method=method)
+            asr_model.save_to(F"{model_name}-{method}-{args.num_calib_batch*args.batch_size}.nemo")
+
+
+def compute_amax(model, **kwargs):
+    for name, module in model.named_modules():
+        if isinstance(module, quant_nn.TensorQuantizer):
+            if module._calibrator is not None:
+                if isinstance(module._calibrator, calib.MaxCalibrator):
+                    module.load_calib_amax()
+                else:
+                    module.load_calib_amax(**kwargs)
+            print(F"{name:40}: {module}")
+    if can_gpu:
+        model.cuda()
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/quantization/speech_to_text_quant_infer.py

@@ -0,0 +1,216 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Script for post training quantization of ASR models
+"""
+
+import collections
+from argparse import ArgumentParser
+from pprint import pprint
+
+import torch
+from omegaconf import open_dict
+
+from nemo.collections.asr.metrics.wer import WER, word_error_rate
+from nemo.collections.asr.models import EncDecCTCModel
+from nemo.utils import logging
+
+try:
+    from pytorch_quantization import nn as quant_nn
+    from pytorch_quantization import quant_modules
+except ImportError:
+    raise ImportError(
+        "pytorch-quantization is not installed. Install from "
+        "https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
+    )
+
+
+try:
+    from torch.cuda.amp import autocast
+except ImportError:
+    from contextlib import contextmanager
+
+    @contextmanager
+    def autocast(enabled=None):
+        yield
+
+
+can_gpu = torch.cuda.is_available()
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--asr_model", type=str, default="QuartzNet15x5Base-En", required=True, help="Pass: 'QuartzNet15x5Base-En'",
+    )
+    parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data")
+    parser.add_argument("--wer_target", type=float, default=None, help="used by test")
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--wer_tolerance", type=float, default=1.0, help="used by test")
+    parser.add_argument(
+        "--dont_normalize_text",
+        default=False,
+        action='store_false',
+        help="Turn off trasnscript normalization. Recommended for non-English.",
+    )
+    parser.add_argument(
+        "--use_cer", default=False, action='store_true', help="Use Character Error Rate as the evaluation metric"
+    )
+    parser.add_argument('--sensitivity', action="store_true", help="Perform sensitivity analysis")
+    parser.add_argument('--onnx', action="store_true", help="Export to ONNX")
+    parser.add_argument('--quant-disable-keyword', type=str, nargs='+', help='disable quantizers by keyword')
+    args = parser.parse_args()
+    torch.set_grad_enabled(False)
+
+    quant_modules.initialize()
+
+    if args.asr_model.endswith('.nemo'):
+        logging.info(f"Using local ASR model from {args.asr_model}")
+        asr_model_cfg = EncDecCTCModel.restore_from(restore_path=args.asr_model, return_config=True)
+        with open_dict(asr_model_cfg):
+            asr_model_cfg.encoder.quantize = True
+        asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model, override_config_path=asr_model_cfg)
+
+    else:
+        logging.info(f"Using NGC cloud ASR model {args.asr_model}")
+        asr_model_cfg = EncDecCTCModel.from_pretrained(model_name=args.asr_model, return_config=True)
+        with open_dict(asr_model_cfg):
+            asr_model_cfg.encoder.quantize = True
+        asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model, override_config_path=asr_model_cfg)
+    asr_model.setup_test_data(
+        test_data_config={
+            'sample_rate': 16000,
+            'manifest_filepath': args.dataset,
+            'labels': asr_model.decoder.vocabulary,
+            'batch_size': args.batch_size,
+            'normalize_transcripts': args.dont_normalize_text,
+        }
+    )
+    asr_model.preprocessor.featurizer.dither = 0.0
+    asr_model.preprocessor.featurizer.pad_to = 0
+    if can_gpu:
+        asr_model = asr_model.cuda()
+    asr_model.eval()
+
+    if args.quant_disable_keyword:
+        for name, module in asr_model.named_modules():
+            if isinstance(module, quant_nn.TensorQuantizer):
+                for keyword in args.quant_disable_keyword:
+                    if keyword in name:
+                        logging.warning(F"Disable {name}")
+                        module.disable()
+
+    labels_map = dict([(i, asr_model.decoder.vocabulary[i]) for i in range(len(asr_model.decoder.vocabulary))])
+    wer = WER(vocabulary=asr_model.decoder.vocabulary, use_cer=args.use_cer)
+    wer_quant = evaluate(asr_model, labels_map, wer)
+    logging.info(f'Got WER of {wer_quant}. Tolerance was {args.wer_tolerance}')
+
+    if args.sensitivity:
+        if wer_quant < args.wer_tolerance:
+            logging.info("Tolerance is already met. Skip sensitivity analyasis.")
+            return
+        quant_layer_names = []
+        for name, module in asr_model.named_modules():
+            if isinstance(module, quant_nn.TensorQuantizer):
+                module.disable()
+                layer_name = name.replace("._input_quantizer", "").replace("._weight_quantizer", "")
+                if layer_name not in quant_layer_names:
+                    quant_layer_names.append(layer_name)
+        logging.info(F"{len(quant_layer_names)} quantized layers found.")
+
+        # Build sensitivity profile
+        quant_layer_sensitivity = {}
+        for i, quant_layer in enumerate(quant_layer_names):
+            logging.info(F"Enable {quant_layer}")
+            for name, module in asr_model.named_modules():
+                if isinstance(module, quant_nn.TensorQuantizer) and quant_layer in name:
+                    module.enable()
+                    logging.info(F"{name:40}: {module}")
+
+            # Eval the model
+            wer_value = evaluate(asr_model, labels_map, wer)
+            logging.info(F"WER: {wer_value}")
+            quant_layer_sensitivity[quant_layer] = args.wer_tolerance - wer_value
+
+            for name, module in asr_model.named_modules():
+                if isinstance(module, quant_nn.TensorQuantizer) and quant_layer in name:
+                    module.disable()
+                    logging.info(F"{name:40}: {module}")
+
+        # Skip most sensitive layers until WER target is met
+        for name, module in asr_model.named_modules():
+            if isinstance(module, quant_nn.TensorQuantizer):
+                module.enable()
+        quant_layer_sensitivity = collections.OrderedDict(sorted(quant_layer_sensitivity.items(), key=lambda x: x[1]))
+        pprint(quant_layer_sensitivity)
+        skipped_layers = []
+        for quant_layer, _ in quant_layer_sensitivity.items():
+            for name, module in asr_model.named_modules():
+                if isinstance(module, quant_nn.TensorQuantizer):
+                    if quant_layer in name:
+                        logging.info(F"Disable {name}")
+                        if not quant_layer in skipped_layers:
+                            skipped_layers.append(quant_layer)
+                        module.disable()
+            wer_value = evaluate(asr_model, labels_map, wer)
+            if wer_value <= args.wer_tolerance:
+                logging.info(
+                    F"WER tolerance {args.wer_tolerance} is met by skipping {len(skipped_layers)} sensitive layers."
+                )
+                print(skipped_layers)
+                export_onnx(args, asr_model)
+                return
+        raise ValueError(f"WER tolerance {args.wer_tolerance} can not be met with any layer quantized!")
+
+    export_onnx(args, asr_model)
+
+
+def export_onnx(args, asr_model):
+    if args.onnx:
+        if args.asr_model.endswith("nemo"):
+            onnx_name = args.asr_model.replace(".nemo", ".onnx")
+        else:
+            onnx_name = args.asr_model
+        logging.info(F"Export to {onnx_name}")
+        quant_nn.TensorQuantizer.use_fb_fake_quant = True
+        asr_model.export(onnx_name, onnx_opset_version=13)
+        quant_nn.TensorQuantizer.use_fb_fake_quant = False
+
+
+def evaluate(asr_model, labels_map, wer):
+    # Eval the model
+    hypotheses = []
+    references = []
+    for test_batch in asr_model.test_dataloader():
+        if can_gpu:
+            test_batch = [x.cuda() for x in test_batch]
+        with autocast():
+            log_probs, encoded_len, greedy_predictions = asr_model(
+                input_signal=test_batch[0], input_signal_length=test_batch[1]
+            )
+        hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions)
+        for batch_ind in range(greedy_predictions.shape[0]):
+            seq_len = test_batch[3][batch_ind].cpu().detach().numpy()
+            seq_ids = test_batch[2][batch_ind].cpu().detach().numpy()
+            reference = ''.join([labels_map[c] for c in seq_ids[0:seq_len]])
+            references.append(reference)
+        del test_batch
+    wer_value = word_error_rate(hypotheses=hypotheses, references=references, use_cer=wer.use_cer)
+
+    return wer_value
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/quantization/speech_to_text_quant_infer_trt.py

@@ -0,0 +1,226 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Script for inference ASR models using TensorRT
+"""
+
+import collections
+import os
+import time
+from argparse import ArgumentParser
+from pprint import pprint
+
+import numpy as np
+import pycuda.autoinit
+import pycuda.driver as cuda
+import tensorrt as trt
+import torch
+from omegaconf import open_dict
+
+from nemo.collections.asr.metrics.wer import WER, word_error_rate
+from nemo.collections.asr.models import EncDecCTCModel
+from nemo.utils import logging
+
+TRT_LOGGER = trt.Logger()
+
+
+can_gpu = torch.cuda.is_available()
+
+try:
+    from torch.cuda.amp import autocast
+except ImportError:
+    from contextlib import contextmanager
+
+    @contextmanager
+    def autocast(enabled=None):
+        yield
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--asr_model", type=str, default="QuartzNet15x5Base-En", required=True, help="Pass: 'QuartzNet15x5Base-En'",
+    )
+    parser.add_argument(
+        "--asr_onnx",
+        type=str,
+        default="./QuartzNet15x5Base-En-max-32.onnx",
+        help="Pass: 'QuartzNet15x5Base-En-max-32.onnx'",
+    )
+    parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data")
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument(
+        "--dont_normalize_text",
+        default=False,
+        action='store_false',
+        help="Turn off trasnscript normalization. Recommended for non-English.",
+    )
+    parser.add_argument(
+        "--use_cer", default=False, action='store_true', help="Use Character Error Rate as the evaluation metric"
+    )
+    parser.add_argument('--qat', action="store_true", help="Use onnx file exported from QAT tools")
+    args = parser.parse_args()
+    torch.set_grad_enabled(False)
+
+    if args.asr_model.endswith('.nemo'):
+        logging.info(f"Using local ASR model from {args.asr_model}")
+        asr_model_cfg = EncDecCTCModel.restore_from(restore_path=args.asr_model, return_config=True)
+        with open_dict(asr_model_cfg):
+            asr_model_cfg.encoder.quantize = True
+        asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model, override_config_path=asr_model_cfg)
+
+    else:
+        logging.info(f"Using NGC cloud ASR model {args.asr_model}")
+        asr_model_cfg = EncDecCTCModel.from_pretrained(model_name=args.asr_model, return_config=True)
+        with open_dict(asr_model_cfg):
+            asr_model_cfg.encoder.quantize = True
+        asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model, override_config_path=asr_model_cfg)
+    asr_model.setup_test_data(
+        test_data_config={
+            'sample_rate': 16000,
+            'manifest_filepath': args.dataset,
+            'labels': asr_model.decoder.vocabulary,
+            'batch_size': args.batch_size,
+            'normalize_transcripts': args.dont_normalize_text,
+        }
+    )
+    asr_model.preprocessor.featurizer.dither = 0.0
+    asr_model.preprocessor.featurizer.pad_to = 0
+    if can_gpu:
+        asr_model = asr_model.cuda()
+    asr_model.eval()
+    labels_map = dict([(i, asr_model.decoder.vocabulary[i]) for i in range(len(asr_model.decoder.vocabulary))])
+    wer = WER(vocabulary=asr_model.decoder.vocabulary, use_cer=args.use_cer)
+    wer_result = evaluate(asr_model, args.asr_onnx, labels_map, wer, args.qat)
+    logging.info(f'Got WER of {wer_result}.')
+
+
+def get_min_max_input_shape(asr_model):
+    max_shape = (1, 64, 1)
+    min_shape = (64, 64, 99999)
+    for test_batch in asr_model.test_dataloader():
+        test_batch = [x.cuda() for x in test_batch]
+        processed_signal, processed_signal_length = asr_model.preprocessor(
+            input_signal=test_batch[0], length=test_batch[1]
+        )
+        shape = processed_signal.cpu().numpy().shape
+        if shape[0] > max_shape[0]:
+            max_shape = (shape[0], *max_shape[1:])
+        if shape[0] < min_shape[0]:
+            min_shape = (shape[0], *min_shape[1:])
+        if shape[2] > max_shape[2]:
+            max_shape = (*max_shape[0:2], shape[2])
+        if shape[2] < min_shape[2]:
+            min_shape = (*min_shape[0:2], shape[2])
+    return min_shape, max_shape
+
+
+def build_trt_engine(asr_model, onnx_path, qat):
+    trt_engine_path = "{}.trt".format(onnx_path)
+    if os.path.exists(trt_engine_path):
+        return trt_engine_path
+
+    min_input_shape, max_input_shape = get_min_max_input_shape(asr_model)
+    workspace_size = 512
+    with trt.Builder(TRT_LOGGER) as builder:
+        network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+        if qat:
+            network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_PRECISION)
+        with builder.create_network(flags=network_flags) as network, trt.OnnxParser(
+            network, TRT_LOGGER
+        ) as parser, builder.create_builder_config() as builder_config:
+            parser.parse_from_file(onnx_path)
+            builder_config.max_workspace_size = workspace_size * (1024 * 1024)
+            if qat:
+                builder_config.set_flag(trt.BuilderFlag.INT8)
+
+            profile = builder.create_optimization_profile()
+            profile.set_shape("audio_signal", min=min_input_shape, opt=max_input_shape, max=max_input_shape)
+            builder_config.add_optimization_profile(profile)
+
+            engine = builder.build_engine(network, builder_config)
+            serialized_engine = engine.serialize()
+            with open(trt_engine_path, "wb") as fout:
+                fout.write(serialized_engine)
+    return trt_engine_path
+
+
+def trt_inference(stream, trt_ctx, d_input, d_output, input_signal, input_signal_length):
+    print("infer with shape: {}".format(input_signal.shape))
+
+    trt_ctx.set_binding_shape(0, input_signal.shape)
+    assert trt_ctx.all_binding_shapes_specified
+
+    h_output = cuda.pagelocked_empty(tuple(trt_ctx.get_binding_shape(1)), dtype=np.float32)
+
+    h_input_signal = cuda.register_host_memory(np.ascontiguousarray(input_signal.cpu().numpy().ravel()))
+    cuda.memcpy_htod_async(d_input, h_input_signal, stream)
+    trt_ctx.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
+    cuda.memcpy_dtoh_async(h_output, d_output, stream)
+    stream.synchronize()
+
+    greedy_predictions = torch.tensor(h_output).argmax(dim=-1, keepdim=False)
+    return greedy_predictions
+
+
+def evaluate(asr_model, asr_onnx, labels_map, wer, qat):
+    # Eval the model
+    hypotheses = []
+    references = []
+    stream = cuda.Stream()
+    vocabulary_size = len(labels_map) + 1
+    engine_file_path = build_trt_engine(asr_model, asr_onnx, qat)
+    with open(engine_file_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
+        trt_engine = runtime.deserialize_cuda_engine(f.read())
+        trt_ctx = trt_engine.create_execution_context()
+
+        profile_shape = trt_engine.get_profile_shape(profile_index=0, binding=0)
+        print("profile shape min:{}, opt:{}, max:{}".format(profile_shape[0], profile_shape[1], profile_shape[2]))
+        max_input_shape = profile_shape[2]
+        input_nbytes = trt.volume(max_input_shape) * trt.float32.itemsize
+        d_input = cuda.mem_alloc(input_nbytes)
+        max_output_shape = [max_input_shape[0], vocabulary_size, (max_input_shape[-1] + 1) // 2]
+        output_nbytes = trt.volume(max_output_shape) * trt.float32.itemsize
+        d_output = cuda.mem_alloc(output_nbytes)
+
+        for test_batch in asr_model.test_dataloader():
+            if can_gpu:
+                test_batch = [x.cuda() for x in test_batch]
+            processed_signal, processed_signal_length = asr_model.preprocessor(
+                input_signal=test_batch[0], length=test_batch[1]
+            )
+
+            greedy_predictions = trt_inference(
+                stream,
+                trt_ctx,
+                d_input,
+                d_output,
+                input_signal=processed_signal,
+                input_signal_length=processed_signal_length,
+            )
+            hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions)
+            for batch_ind in range(greedy_predictions.shape[0]):
+                seq_len = test_batch[3][batch_ind].cpu().detach().numpy()
+                seq_ids = test_batch[2][batch_ind].cpu().detach().numpy()
+                reference = ''.join([labels_map[c] for c in seq_ids[0:seq_len]])
+                references.append(reference)
+            del test_batch
+        wer_value = word_error_rate(hypotheses=hypotheses, references=references, use_cer=wer.use_cer)
+
+    return wer_value
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/speech_to_label.py

@@ -0,0 +1,146 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+# Task 1: Speech Command
+
+## Preparing the dataset
+Use the `process_speech_commands_data.py` script under <NEMO_ROOT>/scripts/dataset_processing in order to prepare the dataset.
+
+```sh
+python <NEMO_ROOT>/scripts/dataset_processing/process_speech_commands_data.py \
+    --data_root=<absolute path to where the data should be stored> \
+    --data_version=<either 1 or 2, indicating version of the dataset> \
+    --class_split=<either "all" or "sub", indicates whether all 30/35 classes should be used, or the 10+2 split should be used> \
+    --rebalance \
+    --log
+```
+
+## Train to convergence
+```sh
+python speech_to_label.py \
+    # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
+    model.train_ds.manifest_filepath="<path to train manifest>" \
+    model.validation_ds.manifest_filepath=["<path to val manifest>","<path to test manifest>"] \
+    trainer.gpus=2 \
+    trainer.accelerator="ddp" \
+    trainer.max_epochs=200 \
+    exp_manager.create_wandb_logger=True \
+    exp_manager.wandb_logger_kwargs.name="MatchboxNet-3x1x64-v1" \
+    exp_manager.wandb_logger_kwargs.project="MatchboxNet-v1" \
+    +trainer.precision=16 \
+    +trainer.amp_level=O1  # needed if using PyTorch < 1.6
+```
+
+
+# Task 2: Voice Activity Detection
+
+## Preparing the dataset
+Use the `process_vad_data.py` script under <NEMO_ROOT>/scripts/dataset_processing in order to prepare the dataset.
+
+```sh
+python process_vad_data.py \
+    --out_dir=<output path to where the generated manifest should be stored> \
+    --speech_data_root=<path where the speech data are stored> \
+    --background_data_root=<path where the background data are stored> \
+    --rebalance_method=<'under' or 'over' of 'fixed'> \
+    --log
+    (Optional --demo (for demonstration in tutorial). If you want to use your own background noise data, make sure to delete --demo)
+```
+
+## Train to convergence
+```sh
+python speech_to_label.py \
+    --config-path=<path to dir of configs e.g. "conf">
+    --config-name=<name of config without .yaml e.g. "matchboxnet_3x1x64_vad"> \
+    model.train_ds.manifest_filepath="<path to train manifest>" \
+    model.validation_ds.manifest_filepath=["<path to val manifest>","<path to test manifest>"] \
+    trainer.gpus=2 \
+    trainer.accelerator="ddp" \
+    trainer.max_epochs=200 \
+    exp_manager.create_wandb_logger=True \
+    exp_manager.wandb_logger_kwargs.name="MatchboxNet-3x1x64-vad" \
+    exp_manager.wandb_logger_kwargs.project="MatchboxNet-vad" \
+    +trainer.precision=16 \
+    +trainer.amp_level=O1  # needed if using PyTorch < 1.6
+```
+
+# Optional: Use tarred dataset to speed up data loading. Apply to both tasks.
+## Prepare tarred dataset. 
+   Prepare ONE manifest that contains all training data you would like to include. Validation should use non-tarred dataset.
+   Note that it's possible that tarred datasets impacts validation scores because it drop values in order to have same amount of files per tarfile; 
+   Scores might be off since some data is missing. 
+
+   Use the `convert_to_tarred_audio_dataset.py` script under <NEMO_ROOT>/scripts/speech_recognition in order to prepare tarred audio dataset.
+   For details, please see TarredAudioToClassificationLabelDataset in <NEMO_ROOT>/nemo/collections/asr/data/audio_to_label.py
+
+python speech_to_label.py \
+    --config-path=<path to dir of configs e.g. "conf">
+    --config-name=<name of config without .yaml e.g. "matchboxnet_3x1x64_vad"> \
+    model.train_ds.manifest_filepath=<path to train tarred_audio_manifest.json> \
+    model.train_ds.is_tarred=True \
+    model.train_ds.tarred_audio_filepaths=<path to train tarred audio dataset e.g. audio_{0..2}.tar> \
+    +model.train_ds.num_worker=<num_shards used generating tarred dataset> \
+    model.validation_ds.manifest_filepath=<path to validation audio_manifest.json>\
+    trainer.gpus=2 \
+    trainer.accelerator="ddp" \
+    trainer.max_epochs=200 \
+    exp_manager.create_wandb_logger=True \
+    exp_manager.wandb_logger_kwargs.name="MatchboxNet-3x1x64-vad" \
+    exp_manager.wandb_logger_kwargs.project="MatchboxNet-vad" \
+    +trainer.precision=16 \
+    +trainer.amp_level=O1  # needed if using PyTorch < 1.6
+
+# Fine-tune a model
+
+For documentation on fine-tuning this model, please visit -
+https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#fine-tuning-configurations
+
+# Pretrained Models
+
+For documentation on existing pretrained models, please visit -
+https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/speech_classification/results.html#
+
+"""
+import pytorch_lightning as pl
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models import EncDecClassificationModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf/matchboxnet", config_name="matchboxnet_3x1x64_v1")
+def main(cfg):
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    asr_model = EncDecClassificationModel(cfg=cfg.model, trainer=trainer)
+
+    # Initialize the weights of the model from another model, if provided via config
+    asr_model.maybe_init_from_pretrained_checkpoint(cfg)
+
+    trainer.fit(asr_model)
+
+    if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None:
+        gpu = 1 if cfg.trainer.gpus != 0 else 0
+        trainer = pl.Trainer(gpus=gpu)
+        if asr_model.prepare_test(trainer):
+            trainer.test(asr_model)
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/speech_to_text.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+# Training the model
+
+Basic run (on CPU for 50 epochs):
+    python examples/asr/speech_to_text.py \
+        # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
+        model.train_ds.manifest_filepath="<path to manifest file>" \
+        model.validation_ds.manifest_filepath="<path to manifest file>" \
+        trainer.gpus=0 \
+        trainer.max_epochs=50
+
+
+Add PyTorch Lightning Trainer arguments from CLI:
+    python speech_to_text.py \
+        ... \
+        +trainer.fast_dev_run=true
+
+Hydra logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/.hydra)"
+PTL logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/lightning_logs)"
+
+Override some args of optimizer:
+    python speech_to_text.py \
+    # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
+    model.train_ds.manifest_filepath="./an4/train_manifest.json" \
+    model.validation_ds.manifest_filepath="./an4/test_manifest.json" \
+    trainer.gpus=2 \
+    trainer.max_epochs=2 \
+    model.optim.args.betas=[0.8,0.5] \
+    model.optim.args.weight_decay=0.0001
+
+Override optimizer entirely
+    python speech_to_text.py \
+    # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
+    model.train_ds.manifest_filepath="./an4/train_manifest.json" \
+    model.validation_ds.manifest_filepath="./an4/test_manifest.json" \
+    trainer.gpus=2 \
+    trainer.max_epochs=2 \
+    model.optim.name=adamw \
+    model.optim.lr=0.001 \
+    ~model.optim.args \
+    +model.optim.args.betas=[0.8,0.5]\
+    +model.optim.args.weight_decay=0.0005
+
+# Fine-tune a model
+
+For documentation on fine-tuning this model, please visit -
+https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#fine-tuning-configurations
+
+# Pretrained Models
+
+For documentation on existing pretrained models, please visit -
+https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/results.html
+
+"""
+
+import pytorch_lightning as pl
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models import EncDecCTCModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="config")
+def main(cfg):
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    asr_model = EncDecCTCModel(cfg=cfg.model, trainer=trainer)
+
+    # Initialize the weights of the model from another model, if provided via config
+    asr_model.maybe_init_from_pretrained_checkpoint(cfg)
+
+    trainer.fit(asr_model)
+
+    if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None:
+        gpu = 1 if cfg.trainer.gpus != 0 else 0
+        test_trainer = pl.Trainer(
+            gpus=gpu,
+            precision=trainer.precision,
+            amp_level=trainer.accelerator_connector.amp_level,
+            amp_backend=cfg.trainer.get("amp_backend", "native"),
+        )
+        if asr_model.prepare_test(test_trainer):
+            test_trainer.test(asr_model)
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/speech_to_text_bpe.py

@@ -0,0 +1,101 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+# Preparing the Tokenizer for the dataset
+Use the `process_asr_text_tokenizer.py` script under <NEMO_ROOT>/scripts/tokenizers/ in order to prepare the tokenizer.
+
+```sh
+python <NEMO_ROOT>/scripts/tokenizers/process_asr_text_tokenizer.py \
+        --manifest=<path to train manifest files, seperated by commas>
+        OR
+        --data_file=<path to text data, seperated by commas> \
+        --data_root="<output directory>" \
+        --vocab_size=<number of tokens in vocabulary> \
+        --tokenizer=<"spe" or "wpe"> \
+        --no_lower_case \
+        --spe_type=<"unigram", "bpe", "char" or "word"> \
+        --spe_character_coverage=1.0 \
+        --log
+```
+
+# Training the model
+```sh
+python speech_to_text_bpe.py \
+    # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
+    model.train_ds.manifest_filepath=<path to train manifest> \
+    model.validation_ds.manifest_filepath=<path to val/test manifest> \
+    model.tokenizer.dir=<path to directory of tokenizer (not full path to the vocab file!)> \
+    model.tokenizer.type=<either bpe or wpe> \
+    trainer.gpus=-1 \
+    trainer.accelerator="ddp" \
+    trainer.max_epochs=100 \
+    model.optim.name="adamw" \
+    model.optim.lr=0.001 \
+    model.optim.betas=[0.9,0.999] \
+    model.optim.weight_decay=0.0001 \
+    model.optim.sched.warmup_steps=2000
+    exp_manager.create_wandb_logger=True \
+    exp_manager.wandb_logger_kwargs.name="<Name of experiment>" \
+    exp_manager.wandb_logger_kwargs.project="<Name of project>"
+```
+
+# Fine-tune a model
+
+For documentation on fine-tuning this model, please visit -
+https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#fine-tuning-configurations
+
+# Pretrained Models
+
+For documentation on existing pretrained models, please visit -
+https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/results.html
+
+"""
+
+import pytorch_lightning as pl
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models.ctc_bpe_models import EncDecCTCModelBPE
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf/citrinet/", config_name="config_bpe")
+def main(cfg):
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    asr_model = EncDecCTCModelBPE(cfg=cfg.model, trainer=trainer)
+
+    # Initialize the weights of the model from another model, if provided via config
+    asr_model.maybe_init_from_pretrained_checkpoint(cfg)
+
+    trainer.fit(asr_model)
+
+    if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None:
+        gpu = 1 if cfg.trainer.gpus != 0 else 0
+        test_trainer = pl.Trainer(
+            gpus=gpu,
+            precision=trainer.precision,
+            amp_level=trainer.accelerator_connector.amp_level,
+            amp_backend=cfg.trainer.get("amp_backend", "native"),
+        )
+        if asr_model.prepare_test(test_trainer):
+            test_trainer.test(asr_model)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/asr/speech_to_text_infer.py

@@ -0,0 +1,114 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script serves three goals:
+    (1) Demonstrate how to use NeMo Models outside of PytorchLightning
+    (2) Shows example of batch ASR inference
+    (3) Serves as CI test for pre-trained checkpoint
+"""
+
+from argparse import ArgumentParser
+
+import torch
+
+from nemo.collections.asr.metrics.wer import WER, word_error_rate
+from nemo.collections.asr.models import EncDecCTCModel
+from nemo.utils import logging
+
+try:
+    from torch.cuda.amp import autocast
+except ImportError:
+    from contextlib import contextmanager
+
+    @contextmanager
+    def autocast(enabled=None):
+        yield
+
+
+can_gpu = torch.cuda.is_available()
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--asr_model", type=str, default="QuartzNet15x5Base-En", required=True, help="Pass: 'QuartzNet15x5Base-En'",
+    )
+    parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data")
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--wer_tolerance", type=float, default=1.0, help="used by test")
+    parser.add_argument(
+        "--dont_normalize_text",
+        default=False,
+        action='store_true',
+        help="Turn off trasnscript normalization. Recommended for non-English.",
+    )
+    parser.add_argument(
+        "--use_cer", default=False, action='store_true', help="Use Character Error Rate as the evaluation metric"
+    )
+    args = parser.parse_args()
+    torch.set_grad_enabled(False)
+
+    if args.asr_model.endswith('.nemo'):
+        logging.info(f"Using local ASR model from {args.asr_model}")
+        asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model)
+    else:
+        logging.info(f"Using NGC cloud ASR model {args.asr_model}")
+        asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model)
+
+    asr_model.preprocessor.featurizer.pad_to = 0
+    asr_model.preprocessor.featurizer.dither = 0.0
+
+    asr_model.setup_test_data(
+        test_data_config={
+            'sample_rate': 16000,
+            'manifest_filepath': args.dataset,
+            'labels': asr_model.decoder.vocabulary,
+            'batch_size': args.batch_size,
+            'normalize_transcripts': not args.dont_normalize_text,
+        }
+    )
+    if can_gpu:
+        asr_model = asr_model.cuda()
+    asr_model.eval()
+    labels_map = dict([(i, asr_model.decoder.vocabulary[i]) for i in range(len(asr_model.decoder.vocabulary))])
+    wer = WER(vocabulary=asr_model.decoder.vocabulary)
+    hypotheses = []
+    references = []
+    for test_batch in asr_model.test_dataloader():
+        if can_gpu:
+            test_batch = [x.cuda() for x in test_batch]
+        with autocast():
+            log_probs, encoded_len, greedy_predictions = asr_model(
+                input_signal=test_batch[0], input_signal_length=test_batch[1]
+            )
+        hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions)
+        for batch_ind in range(greedy_predictions.shape[0]):
+            seq_len = test_batch[3][batch_ind].cpu().detach().numpy()
+            seq_ids = test_batch[2][batch_ind].cpu().detach().numpy()
+            reference = ''.join([labels_map[c] for c in seq_ids[0:seq_len]])
+            references.append(reference)
+        del test_batch
+
+    wer_value = word_error_rate(hypotheses=hypotheses, references=references, use_cer=args.use_cer)
+    if not args.use_cer:
+        if wer_value > args.wer_tolerance:
+            raise ValueError(f"got wer of {wer_value}. it was higher than {args.wer_tolerance}")
+        logging.info(f'Got WER of {wer_value}. Tolerance was {args.wer_tolerance}')
+    else:
+        logging.info(f'Got CER of {wer_value}')
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/speech_to_text_rnnt.py

@@ -0,0 +1,104 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+# Training the model
+
+Basic run (on CPU for 50 epochs):
+    python examples/asr/speech_to_text_rnnt.py \
+        # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
+        model.train_ds.manifest_filepath="<path to manifest file>" \
+        model.validation_ds.manifest_filepath="<path to manifest file>" \
+        trainer.gpus=0 \
+        trainer.max_epochs=50
+
+
+Add PyTorch Lightning Trainer arguments from CLI:
+    python speech_to_text_rnnt.py \
+        ... \
+        +trainer.fast_dev_run=true
+
+Hydra logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/.hydra)"
+PTL logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/lightning_logs)"
+
+Override some args of optimizer:
+    python speech_to_text_rnnt.py \
+    --config-path="experimental/contextnet_rnnt" \
+    --config-name="config_rnnt" \
+    model.train_ds.manifest_filepath="./an4/train_manifest.json" \
+    model.validation_ds.manifest_filepath="./an4/test_manifest.json" \
+    trainer.gpus=2 \
+    trainer.precision=16 \
+    trainer.max_epochs=2 \
+    model.optim.betas=[0.8,0.5] \
+    model.optim.weight_decay=0.0001
+
+Override optimizer entirely
+    python speech_to_text_rnnt.py \
+    --config-path="experimental/contextnet_rnnt" \
+    --config-name="config_rnnt" \
+    model.train_ds.manifest_filepath="./an4/train_manifest.json" \
+    model.validation_ds.manifest_filepath="./an4/test_manifest.json" \
+    trainer.gpus=2 \
+    trainer.precision=16 \
+    trainer.max_epochs=2 \
+    model.optim.name=adamw \
+    model.optim.lr=0.001 \
+    ~model.optim.args \
+    +model.optim.args.betas=[0.8,0.5]\
+    +model.optim.args.weight_decay=0.0005
+
+# Fine-tune a model
+
+For documentation on fine-tuning this model, please visit -
+https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#fine-tuning-configurations
+
+"""
+
+import pytorch_lightning as pl
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models import EncDecRNNTModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="experimental/contextnet_rnnt", config_name="config_rnnt")
+def main(cfg):
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    asr_model = EncDecRNNTModel(cfg=cfg.model, trainer=trainer)
+
+    # Initialize the weights of the model from another model, if provided via config
+    asr_model.maybe_init_from_pretrained_checkpoint(cfg)
+
+    trainer.fit(asr_model)
+
+    if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None:
+        gpu = 1 if cfg.trainer.gpus != 0 else 0
+        test_trainer = pl.Trainer(
+            gpus=gpu,
+            precision=trainer.precision,
+            amp_level=trainer.accelerator_connector.amp_level,
+            amp_backend=cfg.trainer.get("amp_backend", "native"),
+        )
+        if asr_model.prepare_test(test_trainer):
+            test_trainer.test(asr_model)
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/speech_to_text_rnnt_bpe.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+# Preparing the Tokenizer for the dataset
+Use the `process_asr_text_tokenizer.py` script under <NEMO_ROOT>/scripts/tokenizers/ in order to prepare the tokenizer.
+
+```sh
+python <NEMO_ROOT>/scripts/tokenizers/process_asr_text_tokenizer.py \
+        --manifest=<path to train manifest files, seperated by commas>
+        OR
+        --data_file=<path to text data, seperated by commas> \
+        --data_root="<output directory>" \
+        --vocab_size=<number of tokens in vocabulary> \
+        --tokenizer=<"spe" or "wpe"> \
+        --no_lower_case \
+        --spe_type=<"unigram", "bpe", "char" or "word"> \
+        --spe_character_coverage=1.0 \
+        --log
+```
+
+# Training the model
+```sh
+python speech_to_text_rnnt_bpe.py \
+    # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
+    model.train_ds.manifest_filepath=<path to train manifest> \
+    model.validation_ds.manifest_filepath=<path to val/test manifest> \
+    model.tokenizer.dir=<path to directory of tokenizer (not full path to the vocab file!)> \
+    model.tokenizer.type=<either bpe or wpe> \
+    trainer.gpus=-1 \
+    trainer.accelerator="ddp" \
+    trainer.max_epochs=100 \
+    model.optim.name="adamw" \
+    model.optim.lr=0.001 \
+    model.optim.betas=[0.9,0.999] \
+    model.optim.weight_decay=0.0001 \
+    model.optim.sched.warmup_steps=2000
+    exp_manager.create_wandb_logger=True \
+    exp_manager.wandb_logger_kwargs.name="<Name of experiment>" \
+    exp_manager.wandb_logger_kwargs.project="<Name of project>"
+```
+
+# Fine-tune a model
+
+For documentation on fine-tuning this model, please visit -
+https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#fine-tuning-configurations
+
+"""
+
+import pytorch_lightning as pl
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models import EncDecRNNTBPEModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="experimental/contextnet_rnnt", config_name="config_rnnt_bpe")
+def main(cfg):
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    asr_model = EncDecRNNTBPEModel(cfg=cfg.model, trainer=trainer)
+
+    # Initialize the weights of the model from another model, if provided via config
+    asr_model.maybe_init_from_pretrained_checkpoint(cfg)
+
+    trainer.fit(asr_model)
+
+    if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None:
+        gpu = 1 if cfg.trainer.gpus != 0 else 0
+        test_trainer = pl.Trainer(
+            gpus=gpu,
+            precision=trainer.precision,
+            amp_level=trainer.accelerator_connector.amp_level,
+            amp_backend=cfg.trainer.get("amp_backend", "native"),
+        )
+        if asr_model.prepare_test(test_trainer):
+            test_trainer.test(asr_model)
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/transcribe_speech.py

@@ -0,0 +1,167 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import glob
+import json
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+import pytorch_lightning as pl
+import torch
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.metrics.rnnt_wer import RNNTDecodingConfig
+from nemo.collections.asr.models import ASRModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging, model_utils
+
+
+"""
+# Transcribe audio
+# Arguments
+# model_path: path to .nemo ASR checkpoint
+# pretrained_name: name of pretrained ASR model (from NGC registry)
+# audio_dir: path to directory with audio files
+# dataset_manifest: path to dataset JSON manifest file (in NeMo format)
+#
+# ASR model can be specified by either "model_path" or "pretrained_name".
+# Data for transcription can be defined with either "audio_dir" or "dataset_manifest".
+# Results are returned in a JSON manifest file.
+
+python transcribe_speech.py \
+    model_path=null \
+    pretrained_name=null \
+    audio_dir="" \
+    dataset_manifest="" \
+    output_filename=""
+"""
+
+
+@dataclass
+class TranscriptionConfig:
+    # Required configs
+    model_path: Optional[str] = None  # Path to a .nemo file
+    pretrained_name: Optional[str] = None  # Name of a pretrained model
+    audio_dir: Optional[str] = None  # Path to a directory which contains audio files
+    dataset_manifest: Optional[str] = None  # Path to dataset's JSON manifest
+
+    # General configs
+    output_filename: Optional[str] = None
+    batch_size: int = 32
+    cuda: Optional[bool] = None  # will switch to cuda if available, defaults to CPU otherwise
+    amp: bool = False
+    audio_type: str = "wav"
+
+    # decoding strategy for RNNT models
+    rnnt_decoding: RNNTDecodingConfig = RNNTDecodingConfig()
+
+
+@hydra_runner(config_name="TranscriptionConfig", schema=TranscriptionConfig)
+def main(cfg: TranscriptionConfig):
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+
+    if cfg.model_path is None and cfg.pretrained_name is None:
+        raise ValueError("Both cfg.model_path and cfg.pretrained_name cannot be None!")
+    if cfg.audio_dir is None and cfg.dataset_manifest is None:
+        raise ValueError("Both cfg.audio_dir and cfg.dataset_manifest cannot be None!")
+
+    # setup GPU
+    if cfg.cuda is None:
+        cfg.cuda = torch.cuda.is_available()
+
+    if type(cfg.cuda) == int:
+        device_id = int(cfg.cuda)
+    else:
+        device_id = 0
+
+    device = torch.device(f'cuda:{device_id}' if cfg.cuda else 'cpu')
+
+    # setup model
+    if cfg.model_path is not None:
+        # restore model from .nemo file path
+        model_cfg = ASRModel.restore_from(restore_path=cfg.model_path, return_config=True)
+        classpath = model_cfg.target  # original class path
+        imported_class = model_utils.import_class_by_path(classpath)  # type: ASRModel
+        logging.info(f"Restoring model : {imported_class.__name__}")
+        asr_model = imported_class.restore_from(restore_path=cfg.model_path, map_location=device)  # type: ASRModel
+        model_name = os.path.splitext(os.path.basename(cfg.model_path))[0]
+    else:
+        # restore model by name
+        asr_model = ASRModel.from_pretrained(model_name=cfg.pretrained_name, map_location=device)  # type: ASRModel
+        model_name = cfg.pretrained_name
+
+    trainer = pl.Trainer(gpus=int(cfg.cuda))
+    asr_model.set_trainer(trainer)
+    asr_model = asr_model.eval()
+
+    # Setup decoding strategy
+    if hasattr(asr_model, 'change_decoding_strategy'):
+        asr_model.change_decoding_strategy(cfg.rnnt_decoding)
+
+    # get audio filenames
+    if cfg.audio_dir is not None:
+        filepaths = list(glob.glob(os.path.join(cfg.audio_dir, f"*.{cfg.audio_type}")))
+    else:
+        # get filenames from manifest
+        filepaths = []
+        with open(cfg.dataset_manifest, 'r') as f:
+            for line in f:
+                item = json.loads(line)
+                filepaths.append(item['audio_filepath'])
+    logging.info(f"\nTranscribing {len(filepaths)} files...\n")
+
+    # setup AMP (optional)
+    if cfg.amp and torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'):
+        logging.info("AMP enabled!\n")
+        autocast = torch.cuda.amp.autocast
+    else:
+
+        @contextlib.contextmanager
+        def autocast():
+            yield
+
+    # transcribe audio
+    with autocast():
+        with torch.no_grad():
+            transcriptions = asr_model.transcribe(filepaths, batch_size=cfg.batch_size)
+    logging.info(f"Finished transcribing {len(filepaths)} files !")
+
+    if cfg.output_filename is None:
+        # create default output filename
+        if cfg.audio_dir is not None:
+            cfg.output_filename = os.path.dirname(os.path.join(cfg.audio_dir, '.')) + '.json'
+        else:
+            cfg.output_filename = cfg.dataset_manifest.replace('.json', f'_{model_name}.json')
+
+    logging.info(f"Writing transcriptions into file: {cfg.output_filename}")
+
+    with open(cfg.output_filename, 'w', encoding='utf-8') as f:
+        if cfg.audio_dir is not None:
+            for idx, text in enumerate(transcriptions):
+                item = {'audio_filepath': filepaths[idx], 'pred_text': text}
+                f.write(json.dumps(item) + "\n")
+        else:
+            with open(cfg.dataset_manifest, 'r') as fr:
+                for idx, line in enumerate(fr):
+                    item = json.loads(line)
+                    item['pred_text'] = transcriptions[idx]
+                    f.write(json.dumps(item) + "\n")
+
+    logging.info("Finished writing predictions !")
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/asr/vad_infer.py

@@ -0,0 +1,171 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+During inference, we perform frame-level prediction by two approaches: 
+    1) shift the window of length time_length (e.g. 0.63s) by shift_length (e.g. 10ms) to generate the frame and use the prediction of the window to represent the label for the frame;
+       [this script demonstrate how to do this approach]
+    2) generate predictions with overlapping input segments. Then a smoothing filter is applied to decide the label for a frame spanned by multiple segments. 
+       [get frame level prediction by this script and use vad_overlap_posterior.py in NeMo/scripts/voice_activity_detection
+       One can also find posterior about converting frame level prediction 
+       to speech/no-speech segment in start and end times format in that script.]
+       
+       Image https://raw.githubusercontent.com/NVIDIA/NeMo/main/tutorials/asr/images/vad_post_overlap_diagram.png 
+       will help you understand this method.
+   
+Usage:
+python vad_infer.py  --vad_model="vad_marblenet" --dataset=<FULL PATH OF MANIFEST TO BE PERFORMED INFERENCE ON> --out_dir='frame/demo' --time_length=0.63
+
+"""
+
+
+import json
+import logging
+import os
+from argparse import ArgumentParser
+
+import torch
+
+from nemo.collections.asr.models import EncDecClassificationModel
+from nemo.collections.asr.parts.utils.vad_utils import get_vad_stream_status, prepare_manifest
+from nemo.utils import logging
+
+try:
+    from torch.cuda.amp import autocast
+except ImportError:
+    from contextlib import contextmanager
+
+    @contextmanager
+    def autocast(enabled=None):
+        yield
+
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--vad_model", type=str, default="MatchboxNet-VAD-3x2", required=False, help="Pass: 'MatchboxNet-VAD-3x2'"
+    )
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        required=True,
+        help="Path of json file of evaluation data. Audio files should have unique names.",
+    )
+    parser.add_argument("--out_dir", type=str, default="vad_frame", help="Dir of your vad outputs")
+    parser.add_argument("--time_length", type=float, default=0.63)
+    parser.add_argument("--shift_length", type=float, default=0.01)
+    parser.add_argument("--normalize_audio", type=bool, default=False)
+    parser.add_argument("--num_workers", type=float, default=20)
+    parser.add_argument("--split_duration", type=float, default=400)
+    parser.add_argument(
+        "--dont_auto_split",
+        default=False,
+        action='store_true',
+        help="Whether to automatically split manifest entry by split_duration to avoid potential CUDA out of memory issue.",
+    )
+
+    args = parser.parse_args()
+
+    torch.set_grad_enabled(False)
+
+    if args.vad_model.endswith('.nemo'):
+        logging.info(f"Using local VAD model from {args.vad_model}")
+        vad_model = EncDecClassificationModel.restore_from(restore_path=args.vad_model)
+    else:
+        logging.info(f"Using NGC cloud VAD model {args.vad_model}")
+        vad_model = EncDecClassificationModel.from_pretrained(model_name=args.vad_model)
+
+    if not os.path.exists(args.out_dir):
+        os.mkdir(args.out_dir)
+
+    # Prepare manifest for streaming VAD
+    manifest_vad_input = args.dataset
+    if not args.dont_auto_split:
+        logging.info("Split long audio file to avoid CUDA memory issue")
+        logging.debug("Try smaller split_duration if you still have CUDA memory issue")
+        config = {
+            'manifest_filepath': manifest_vad_input,
+            'time_length': args.time_length,
+            'split_duration': args.split_duration,
+            'num_workers': args.num_workers,
+        }
+        manifest_vad_input = prepare_manifest(config)
+    else:
+        logging.warning(
+            "If you encounter CUDA memory issue, try splitting manifest entry by split_duration to avoid it."
+        )
+
+    # setup_test_data
+    vad_model.setup_test_data(
+        test_data_config={
+            'vad_stream': True,
+            'sample_rate': 16000,
+            'manifest_filepath': manifest_vad_input,
+            'labels': ['infer',],
+            'num_workers': args.num_workers,
+            'shuffle': False,
+            'time_length': args.time_length,
+            'shift_length': args.shift_length,
+            'trim_silence': False,
+            'normalize_audio': args.normalize_audio,
+        }
+    )
+
+    vad_model = vad_model.to(device)
+    vad_model.eval()
+
+    time_unit = int(args.time_length / args.shift_length)
+    trunc = int(time_unit / 2)
+    trunc_l = time_unit - trunc
+    all_len = 0
+
+    data = []
+    for line in open(args.dataset, 'r'):
+        file = json.loads(line)['audio_filepath'].split("/")[-1]
+        data.append(file.split(".wav")[0])
+    logging.info(f"Inference on {len(data)} audio files/json lines!")
+
+    status = get_vad_stream_status(data)
+    for i, test_batch in enumerate(vad_model.test_dataloader()):
+        test_batch = [x.to(device) for x in test_batch]
+        with autocast():
+            log_probs = vad_model(input_signal=test_batch[0], input_signal_length=test_batch[1])
+            probs = torch.softmax(log_probs, dim=-1)
+            pred = probs[:, 1]
+
+            if status[i] == 'start':
+                to_save = pred[:-trunc]
+            elif status[i] == 'next':
+                to_save = pred[trunc:-trunc_l]
+            elif status[i] == 'end':
+                to_save = pred[trunc_l:]
+            else:
+                to_save = pred
+
+            all_len += len(to_save)
+            outpath = os.path.join(args.out_dir, data[i] + ".frame")
+            with open(outpath, "a") as fout:
+                for f in range(len(to_save)):
+                    fout.write('{0:0.4f}\n'.format(to_save[f]))
+        del test_batch
+        if status[i] == 'end' or status[i] == 'single':
+            logging.debug(f"Overall length of prediction of {data[i]} is {all_len}!")
+            all_len = 0
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/cv/mnist_lenet5_image_classification_pure_lightning.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+# from hydra.utils import instantiate
+import pytorch_lightning as ptl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.cv.models import MNISTLeNet5, MNISTLeNet5Config
+from nemo.core.config import Config, TrainerConfig, hydra_runner
+from nemo.utils import logging
+
+
+@dataclass
+class AppConfig(Config):
+    """
+    This is structured config for this application.
+
+    Args:
+        name: Description of the application.
+        trainer: configuration of the trainer.
+        model: configuation of the model.
+    """
+
+    name: str = "Training of a LeNet-5 Model using a pure PyTorchLightning approach - using DDP on 2 GPUs."
+    trainer: TrainerConfig = TrainerConfig(gpus=2, accelerator="ddp")
+    model: MNISTLeNet5Config = MNISTLeNet5Config()
+
+
+@hydra_runner(schema=AppConfig, config_name="AppConfig")
+def main(cfg: DictConfig):
+    # Show configuration - user can modify every parameter from command line!
+    logging.info("Application config\n" + OmegaConf.to_yaml(cfg))
+
+    # The "model" - with dataloader/dataset inside of it.
+    lenet5 = MNISTLeNet5(cfg.model)
+
+    # Setup train data loader and optimizer
+    lenet5.setup_training_data()
+
+    # Setup optimizer and scheduler
+    lenet5.setup_optimization()
+
+    # Create trainer.
+    trainer = ptl.Trainer(**(cfg.trainer))
+    # trainer = instantiate(cfg.trainer)
+
+    # Train.
+    trainer.fit(model=lenet5)
+
+
+if __name__ == "__main__":
+    main()  # TODO: No cfg in function call, and no hydra runner

tts/NeMo/examples/nlp/dialogue_state_tracking/conf/sgdqa_config.yaml

@@ -0,0 +1,115 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+pretrained_model: null # pretrained QAModel model from list_available_models()
+do_training: true # training mode, for testing change to false
+trainer:
+  gpus: 8 # number of GPUs (0 for CPU), or list of the GPUs to use e.g. [0, 1]
+  num_nodes: 1
+  max_epochs: 3
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  gradient_clip_val: 1.0
+  amp_level: O1 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2 to enable the AMP.
+  accelerator: ddp
+  log_every_n_steps: 5  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+  num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+
+model:
+  nemo_path: null # filename to save the model and associated artifacts to .nemo file
+  tokenizer:
+      tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+      vocab_file: null # path to vocab file
+      tokenizer_model: null # only used if tokenizer is sentencepiece
+      special_tokens: null
+
+  language_model:
+    pretrained_model_name: bert-base-cased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+  encoder:
+    dropout: 0.1
+
+  dataset:
+    max_seq_length: 128 # the maximum length BERT supports is 512
+    max_num_cat_slot: 6 # maximum number of different categorical slots per service in dataset 
+    max_num_noncat_slot: 12 # maximum number of different non-categorical slots per service in dataset 
+    max_value_per_cat_slot: 12 # maximum number of different categorical slot values per service in dataset 
+    max_num_intent: 4 # maximum number of different intents per service in dataset 
+    task_name: sgd_single_domain # or from [sgd_all, sgd_all_single, sgd_multi_domain, debug_sample]
+    state_tracker: nemotracker # or baseline
+    data_dir: ???
+    dialogues_example_dir: ??? # cache location to store or load preprocessed data from. To use cache, specify `use_cache`
+    subsample: false # balances negative and positive training examples for improved performance
+    use_cache: false # uses a cache to store the processed dataset, you may use it for large datasets for speed up
+    use_fuzzy_match: true # Whether to use fuzzy string matching when comparing non-categorical slot values. Should be set to False when conducting multiwoz style evaluation.
+    joint_acc_across_turn: false # Whether to compute joint goal accuracy across turn instead of across service. Should be set to True when conducting multiwoz style evaluation.
+
+  train_ds:
+    ds_item: "train"
+    batch_size: 128
+    shuffle: true
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    ds_item: ["dev"] # ["dev", "test"]
+    batch_size: 8
+    shuffle: false
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  test_ds:
+    ds_item: ["test"]
+    batch_size: 8
+    shuffle: false
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adamw
+    lr: 1e-4
+    # optimizer arguments
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+
+    # scheduler setup
+    sched:
+      name: PolynomialDecayAnnealing
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.02
+      last_epoch: -1
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
+  name: "SGDQA"  # The name of your model
+  create_tensorboard_logger: True  # Whether you want exp_manger to create a tb logger
+  create_checkpoint_callback: True  # Whether you want exp_manager to create a modelcheckpoint callback
+  resume_if_exists: false
+  resume_ignore_no_checkpoint: false

tts/NeMo/examples/nlp/dialogue_state_tracking/sgd_qa.py

@@ -0,0 +1,145 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+# Copyright 2019 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script contains an example of how to train and test the NeMo SGD-QA Model (https://arxiv.org/abs/2105.08049).
+The SGD-QA model is a fast multi-pass schema-guided state-tracking model, that is trained on the Google schema-guided state tracking dataset (https://arxiv.org/abs/1909.05855).
+The model takes dialogue as input and outputs the dialogue state, which includes slot-value pairs. 
+The model consists of two components: a neural natural language understanding model (NLU), and a rule-based state tracker.
+The NLU takes in a dialogue turn and different schema (entity) information options and outputs their match score. The state tracker takes the highest rated entities and composes
+the dialogue state across turns.
+
+***Data format***
+The SGD-QA requires a JSON schema and dialogue files for each dataset split. 
+In the following we will show an example for a service entry in the schema file.
+* service_name
+* description
+* slots
+    * name
+    * description
+    * is_categorical
+    * possible values
+* intents
+    * name
+    * description
+    * required_slots (not used)
+    * is_transactional (not used)
+    * optional_slots (not used)
+    * result_slots (not used)
+
+
+In the following we will show an example for a dialogue. More information at https://arxiv.org/abs/1909.05855
+* dialogue_id
+* services
+* turns
+    * frames
+        * actions
+            * act
+            * slot
+            * values
+        * service
+        * slots
+            * exclusive_end
+            * slot
+            * start
+        * state
+            * active_intent
+            * requeste_slots
+            * slot_values 
+    * speaker - [USER, SYSTEM]
+    * utterance
+
+
+***Downloading the dataset***
+#   git clone https://github.com/google-research-datasets/dstc8-schema-guided-dialogue.git
+
+***Setting the configs***
+The model and the PT trainer are defined in a config file that declares multiple important sections.
+The most important ones are:
+    model: All arguments that are related to the Model - language model, SGD-QA encoder and decoder, loss, optimizer,
+            schedulers, and datasets/data loaders.
+    trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs,
+            precision level, etc.
+
+This script uses the `/examples/nlp/dialogue_state_tracking/conf/sgdqa_config.yaml` config file
+by default. You may update the config file from the file directly. The other option is to set another config file via command-line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+
+***Model Training***
+# python sgd_qa.py
+    do_training=True
+    model.dataset.data_dir=<DATA_DIR_WITH_JSON_DATA>
+    model.dataset.dialogues_example_dir=<DATA_DIR_WITH_PREPROCESSED_DATA>
+    model.validation_ds.ds_item=<LIST_OF_SPLITS>
+    trainer.max_epochs=<NUM_EPOCHS>
+    trainer.gpus=[<CHANGE_TO_GPU_YOU_WANT_TO_USE>]
+
+
+***Model Evaluation***
+#   python sgd_qa.py
+    do_training=False
+    model.test_ds.ds_item=<LIST_OF_SPLITS>
+
+To load a pretrained checkpoint from the cloud prior to training (e.g. for fine-tuning) or evaluation you can set cfg.from_pretrained=<MODEL_NAME>. You can find all pretrained model names by using 
+SGDQAModel.list_available_models(). To load a local checkpoint use model.restore_from(<PATH_TO_CHECKPOINT>)
+
+"""
+
+import os
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.dialogue_state_tracking.sgdqa_model import SGDQAModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="sgdqa_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    if cfg.pretrained_model or (cfg.model.nemo_path and os.path.exists(cfg.model.nemo_path)):
+        if cfg.pretrained_model:
+            logging.info(f'Loading pretrained model {cfg.pretrained_model}')
+            model = SGDQAModel.from_pretrained(cfg.pretrained_model)
+        else:
+            logging.info(f'Restoring model from {cfg.model.nemo_path}')
+            model = SGDQAModel.restore_from(cfg.model.nemo_path)
+        if cfg.do_training:
+            model.setup_training_data(train_data_config=cfg.model.train_ds)
+            model.setup_multiple_validation_data(val_data_config=cfg.model.validation_ds)
+    else:
+        logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+        model = SGDQAModel(cfg.model, trainer=trainer)
+
+    if cfg.do_training:
+        trainer.fit(model)
+        if cfg.model.nemo_path:
+            model.save_to(cfg.model.nemo_path)
+
+    if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.ds_item is not None:
+        gpu = 1 if cfg.trainer.gpus != 0 else 0
+        trainer = pl.Trainer(gpus=gpu)
+        model.setup_multiple_test_data(test_data_config=cfg.model.test_ds)
+        if model.prepare_test(trainer):
+            trainer.test(model)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/entity_linking/build_index.py

@@ -0,0 +1,201 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import pickle as pkl
+import random
+from argparse import ArgumentParser
+
+import h5py
+import numpy as np
+import torch
+from omegaconf import DictConfig, OmegaConf
+from sklearn.decomposition import PCA
+from tqdm import tqdm
+
+from nemo.collections.nlp.models import EntityLinkingModel
+from nemo.utils import logging
+
+try:
+    import faiss
+except ModuleNotFoundError:
+    logging.warning("Faiss is required for building the index. Please install faiss-gpu")
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def build_index(cfg: DictConfig, model: object):
+    """
+    Builds faiss index from index dataset specified in the config.
+        
+    Args:
+        cfg (DictConfig): Config file specifying index parameters
+        model (object): Encoder model
+    """
+
+    # Get index dataset embeddings
+    # PCA model exists and index embeddings have already been PCAed, no need to re-extract/PCA them
+    if cfg.apply_pca and os.path.isfile(cfg.pca.pca_save_name) and os.path.isfile(cfg.pca_embeddings_save_name):
+        logging.info("Loading reduced dimensionality embeddings")
+        embeddings = h5py.File(cfg.pca_embeddings_save_name, "r")
+        embeddings = embeddings[cfg.index_ds.name][:]
+
+    elif os.path.isfile(cfg.embedding_save_name):
+        logging.info("Loading previously extracted index dataset embeddings")
+        embeddings = h5py.File(cfg.embedding_save_name, "r")
+        embeddings = embeddings[cfg.index_ds.name][:]
+
+    else:
+        logging.info("Encoding index dataset, this may take a while")
+        index_dataloader = model.setup_dataloader(cfg.index_ds, is_index_data=True)
+        embeddings, concept_ids = get_index_embeddings(cfg, index_dataloader, model)
+
+    # Create pca model to reduce dimensionality of index dataset and decrease memory footprint
+    if cfg.apply_pca:
+
+        # Need to train PCA model and apply PCA transformation with newly trained model
+        if not os.path.isfile(cfg.pca.pca_save_name):
+            logging.info("Fitting PCA model for embedding dimensionality reduction")
+            pca_train_set = random.sample(list(embeddings), k=int(len(embeddings) * cfg.pca.sample_fraction))
+            pca = PCA(n_components=cfg.pca.output_dim)
+            pca.fit(pca_train_set)
+            pkl.dump(pca, open(cfg.pca.pca_save_name, "wb"))
+            embeddings = reduce_embedding_dim(pca, embeddings, cfg)
+
+        # PCA model already trained, just need to reduce dimensionality of all embeddings
+        elif not os.path.isfile(cfg.pca_embeddings_save_name):
+            pca = pkl.load(open(cfg.pca.pca_save_name, "rb"))
+            embeddings = reduce_embedding_dim(pca, embeddings, cfg)
+
+    # Build faiss index from embeddings
+    logging.info(f"Training index with embedding dim size {cfg.dims} using {faiss.get_num_gpus()} gpus")
+    quantizer = faiss.IndexFlatL2(cfg.dims)
+    index = faiss.IndexIVFFlat(quantizer, cfg.dims, cfg.nlist)
+    index = faiss.index_cpu_to_all_gpus(index)
+    index.train(embeddings)
+
+    logging.info("Adding dataset embeddings to index")
+    for i in tqdm(range(0, embeddings.shape[0], cfg.index_batch_size)):
+        index.add(embeddings[i : i + cfg.index_batch_size])
+
+    logging.info("Saving index")
+    faiss.write_index(faiss.index_gpu_to_cpu(index), cfg.index_save_name)
+    logging.info("Index built and saved")
+
+
+def reduce_embedding_dim(pca, embeddings, cfg):
+    """Apply PCA transformation to index dataset embeddings"""
+
+    logging.info("Applying PCA transformation to entire index dataset")
+    embeddings = np.array(pca.transform(embeddings), dtype=np.float32)
+    emb_file = h5py.File(cfg.pca_embeddings_save_name, "w")
+    emb_file.create_dataset(cfg.index_ds.name, data=embeddings)
+    emb_file.close()
+
+    return embeddings
+
+
+def get_index_embeddings(cfg: DictConfig, dataloader: object, model: object):
+    """Use entity linking encoder to get embeddings for full index dataset"""
+    embeddings = []
+    concept_ids = []
+
+    with torch.no_grad():
+        for batch in tqdm(dataloader):
+            input_ids, token_type_ids, input_mask, batch_concept_ids = batch
+            input_ids = input_ids.to(device)
+            token_type_ids = token_type_ids.to(device)
+            input_mask = input_mask.to(device)
+            batch_embeddings = model.forward(
+                input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=input_mask
+            )
+
+            embeddings.extend(batch_embeddings.detach().cpu().numpy())
+            concept_ids.extend(batch_concept_ids.numpy())
+
+    emb_file = h5py.File(cfg.embedding_save_name, "w")
+    emb_file.create_dataset(cfg.index_ds.name, data=embeddings)
+    emb_file.close()
+
+    pkl.dump(concept_ids, open(cfg.concept_id_save_name, "wb"))
+
+    return embeddings, concept_ids
+
+
+def load_model(cfg: DictConfig, restore: bool):
+    """
+    Loads encoder model.
+
+    Args:
+        cfg: Config file specifying model parameters
+        restore: Whether to restore model weights trained
+                 by the user. Otherwise will load weights
+                 used before self alignment pretraining. 
+    """
+
+    if restore:
+        model = EntityLinkingModel.restore_from(cfg.nemo_path)
+    else:
+        cfg.train_ds = None
+        cfg.validation_ds = None
+        cfg.test_ds = None
+        model = EntityLinkingModel(cfg)
+
+    model = model.to(device)
+
+    return model
+
+
+def main(cfg: DictConfig, restore: bool):
+    """
+    Builds new index if one hasn't been built yet.
+
+    Args:
+        cfg: Config file specifying index parameters
+        restore: Whether to restore model weights trained
+                 by the user. Otherwise will load weights
+                 used before self alignment pretraining.
+    """
+
+    logging.info("Loading entity linking encoder model")
+    model = load_model(cfg.model, restore)
+
+    if not os.path.isfile(cfg.index.index_save_name) or (
+        cfg.apply_pca and not os.path.isfile(cfg.index.pca.pca_save_name)
+    ):
+        logging.info("Building index")
+        build_index(cfg.index, model)
+    else:
+        logging.info("Index and pca model (if required) already exists. Skipping build index step.")
+
+    if not os.path.isfile(cfg.index.idx_to_id):
+        logging.info("Mapping entity index postions to ids")
+        map_idx_to_ids(cfg.index)
+    else:
+        logging.info("Map from concept index to id already exists. Skipping mapping step.")
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--restore", action="store_true", help="Whether to restore encoder model weights from nemo path"
+    )
+    parser.add_argument("--project_dir", required=False, type=str, default=".")
+    parser.add_argument("--cfg", required=False, type=str, default="./conf/umls_medical_entity_linking_config.yaml")
+    args = parser.parse_args()
+
+    cfg = OmegaConf.load(args.cfg)
+    cfg.project_dir = args.project_dir
+
+    main(cfg, args.restore)

tts/NeMo/examples/nlp/entity_linking/conf/tiny_example_entity_linking_config.yaml

@@ -0,0 +1,90 @@
+project_dir: null
+name: SelfAlignmentPretrainingForMedicalEntityLinking
+trainer:
+  gpus: 1
+  num_nodes: 1
+  max_epochs: 2
+  max_steps: null
+  accumulate_grad_batches: 1
+  precision: 16
+  amp_level: O1
+  accelerator: ddp
+  gradient_clip_val: 0.0
+  log_every_n_steps: 1
+  val_check_interval: 2
+  checkpoint_callback: false
+  logger: false
+model:
+  nemo_path: ???
+  max_seq_length: 128
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    config_file: null
+    config: null
+    lm_checkpoint: null
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name}
+    vocab_file: null
+    tokenizer_model: null
+    do_lower_case: true
+  loss_params: null
+  train_ds:
+    data_file: ???
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 8
+    shuffle: true
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  validation_ds:
+    data_file: ??? 
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 8
+    shuffle: false
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  optim:
+    name: adam
+    lr: 3.0e-05
+    weight_decay: 0.0
+    sched:
+      name: CosineAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.1
+      min_lr: 0.0
+      last_epoch: -1
+index:
+  dims: 768
+  nlist: 2
+  top_n: 3
+  query_num_factor: 20
+  index_save_name: ???
+  index_batch_size: 10
+  index_ds:
+    name: tiny_example
+    data_file: ???
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 100
+    shuffle: false
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  idx_to_id: ${project_dir}/idx_to_id.pkl
+  id_to_string: ${project_dir}/id_to_string.pkl
+  concept_id_save_name: ${project_dir}/tiny_example_concept_ids.pkl
+  embedding_save_name: ${project_dir}/tiny_example_concept_embeddings.hdf5
+  pca_embeddings_save_name: null
+  apply_pca: false
+  pca: null
+exp_manager:
+  exp_dir: .
+  name: ${project_dir}/SelfAlignmentPretrainingTinyExample
+  create_tensorboard_logger: true
+  create_checkpoint_callback: true
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/entity_linking/conf/umls_medical_entity_linking_config.yaml

@@ -0,0 +1,95 @@
+project_dir: ???
+name: SelfAlignmentPretrainingForMedicalEntityLinking
+trainer:
+  gpus: 1
+  num_nodes: 1
+  max_epochs: 2
+  max_steps: null
+  accumulate_grad_batches: 1
+  precision: 16
+  amp_level: O1
+  accelerator: ddp
+  gradient_clip_val: 0.0
+  log_every_n_steps: 1
+  val_check_interval: 1000
+  checkpoint_callback: false
+  logger: false
+model:
+  nemo_path: ${project_dir}/sap_bert_umls.nemo
+  raw_data: ${project_dir}/data/MRCONSO.RRF
+  max_seq_length: 128
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    config_file: null
+    config: null
+    lm_checkpoint: null
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name}
+    vocab_file: null
+    tokenizer_model: null
+    do_lower_case: true
+  train_ds:
+    data_file: ${project_dir}/data/umls_train_pairs.tsv
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 128
+    shuffle: true
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  validation_ds:
+    data_file: ${project_dir}/data/umls_validation_pairs.tsv
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 128
+    shuffle: false
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  optim:
+    name: adam
+    lr: 3.0e-05
+    weight_decay: 0.0
+    sched:
+      name: CosineAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.1
+      min_lr: 0.0
+      last_epoch: -1
+index:
+  dims: 256
+  nlist: 300
+  top_n: 5
+  query_num_factor: 20
+  index_save_name: ${project_dir}/medical_entity_linking_index
+  index_batch_size: 1000
+  raw_data: ${model.raw_data}
+  index_ds:
+    name: umls
+    data_file: ${project_dir}/data/umls_index_concepts.tsv
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 128
+    shuffle: false
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  idx_to_id: ${project_dir}/data/idx_to_id.pkl
+  id_to_string: ${project_dir}/data/id_to_string.pkl
+  concept_id_save_name: ${project_dir}/data/concept_ids.pkl
+  embedding_save_name: ${project_dir}/data/medical_concept_embeddings.hdf5
+  pca_embeddings_save_name: ${project_dir}/data/medical_concept_reduced_${index.dims}dim_embeddings.hdf5
+  apply_pca: true
+  pca:
+    input_dim: 756
+    output_dim: ${index.dims}
+    sample_fraction: 0.5
+    pca_save_name: ${project_dir}/${index.pca.input_dim}_to_${index.pca.output_dim}_pca_model.pkl
+exp_manager:
+  exp_dir: ${project_dir}/medical_entity_linking_experiments
+  name: sap_bert_umls
+  create_tensorboard_logger: true
+  create_checkpoint_callback: true
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/entity_linking/data/umls_dataset_processing.py

@@ -0,0 +1,189 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import itertools
+import pickle as pkl
+import random
+from argparse import ArgumentParser
+
+import pandas as pd
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+# Info on these headers can be found here on the UMLS website https://www.ncbi.nlm.nih.gov/books/NBK9685/
+# section 3.3.4 Table 1
+HEADERS = [
+    'CUI',
+    'LAT',
+    'TS',
+    'LUI',
+    'STT',
+    'SUI',
+    'ISPREF',
+    'AUI',
+    'SAUI',
+    'SCUI',
+    'SDUI',
+    'SAB',
+    'TTY',
+    'CODE',
+    'STR',
+    'SRL',
+    'SUPPRESS',
+    'CVF',
+]
+
+
+def process_umls_training_dataset(data_path, train_save_name, val_save_name, max_pairs, train_split, headers):
+    """
+    Generates and saves UMLS self alignment pretraining train and validation data. Takes the raw .RRF UMLS 
+    data file and creates different pair combinations for entities with the same CUI. Each row in the output
+    will be formatted as 'CUI EntitySynonym1 EntitySynonym2' with each item in a row separated by tabs.
+    Saves two .tsv output files, one for the train split and one for the validation split.
+    Only data marked as English is added to the train and val splits. 
+
+    Arguments:
+        data_path (str): path to MRCONSO.RRF UMLS data file
+        train_save_name (str): path to where training data will be saved
+        val_save_name (str): path to where validation data will be saved
+        max_pairs (int): max number of pairs for any one CUI added to the train 
+                   or validation splits
+        train_split (float): precentage of raw data to be added to train set split
+        headers (list): column lables within MRCONSO.RRF
+    """
+
+    print("Loading training data file...")
+    df = pd.read_table(data_path, names=headers, index_col=False, delimiter='|')
+    train_file = open(train_save_name, 'w')
+    val_file = open(val_save_name, 'w')
+
+    cui = df["CUI"].iloc[0]
+    names = []
+    random.seed(2021)
+
+    for idx in tqdm(range(len(df))):
+        # Address incorrectly formatted data
+        if type(df["STR"].iloc[idx]) != str or "|" in df["STR"].iloc[idx]:
+            continue
+
+        # Collect all english concept strings matching the current CUI
+        if df["CUI"].iloc[idx] == cui and df["LAT"].iloc[idx] == "ENG":
+            concept_string = df["STR"].iloc[idx]
+            names.append(concept_string)
+
+        else:
+            # Pair off concept synonyms to make training and val sets
+            pairs = list(itertools.combinations(names, 2))
+
+            if len(pairs) == 0:
+                # Not enough concepts gathered to make a pair
+                cui = df["CUI"].iloc[idx]
+                names = [df["STR"].iloc[idx]]
+                continue
+
+            # Removing leading C to convert label string to int
+            cui = int(cui[1:])
+            random.shuffle(pairs)
+
+            # Keep up to max pairs number pairs for any one concept
+            for pair in pairs[:max_pairs]:
+
+                # Want concepts in train and val splits to be randomly selected and mutually exclusive
+                add_to_train = random.random()
+
+                if add_to_train <= train_split:
+                    train_file.write(f'{cui}\t{pair[0]}\t{pair[1]}\n')
+                else:
+                    val_file.write(f'{cui}\t{pair[0]}\t{pair[1]}\n')
+
+            # Switch to next concept
+            cui = df["CUI"].iloc[idx]
+            names = [df["STR"].iloc[idx]]
+
+    train_file.close()
+    val_file.close()
+    print("Finished making training and validation data")
+
+
+def process_umls_index_dataset(data_path, data_savename, id2string_savename, headers):
+    """
+    Generates data file needed to build a UMLS index and a hash table mapping each
+    CUI to one canonical concept string. Takes the raw .RRF data file and saves 
+    a .tsv indec concept file as well as the a .pkl file of cui to concept string 
+    mappings. Only data marked as English is added to the index data file. 
+
+    Arguments:
+        data_path (str): path to MRCONSO.RRF UMLS data file
+        data_savename (str): path to where .tsv index data will be saved
+        id2string_savename (str): path to where .pkl cui to string mapping will
+                                  be saved
+        headers (list): column lables within MRCONSO.RRF
+    """
+
+    print("Loading index data file...")
+    df = pd.read_table(data_path, names=headers, index_col=False, delimiter='|')
+    id2string = {}
+
+    with open(data_savename, "w") as outfile:
+        for idx, row in tqdm(df.iterrows(), total=df.shape[0]):
+            # Address incorrectly formatted data
+            if type(row["STR"]) != str or "|" in row["STR"]:
+                continue
+
+            cui = row["CUI"]
+            sent = row["STR"]
+
+            # Removing leading C to convert label string to int
+            cui = int(cui[1:])
+
+            # Only keeping english concepts
+            if row["LAT"] == "ENG":
+                outfile.write(f'{cui}\t{sent}\n')
+
+                # Matching each cui to one canonical string represention
+                if cui not in id2string and ":" not in sent:
+                    id2string[cui] = sent
+
+    outfile.close()
+    pkl.dump(id2string, open(id2string_savename, "wb"))
+    print("Finished saving index data and id to concept mapping")
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument("--index", action="store_true", help="Whether to process data for building an index")
+    parser.add_argument("--project_dir", required=False, type=str, default=".")
+    parser.add_argument("--cfg", required=False, type=str, default="conf/umls_medical_entity_linking_config.yaml")
+    parser.add_argument(
+        "--max_pairs", required=False, type=int, default=50, help="Max number of train pairs for a single concepts"
+    )
+    parser.add_argument(
+        "--train_split", required=False, type=float, default=0.99, help="Precentage of data to add to train set"
+    )
+
+    args = parser.parse_args()
+    cfg = OmegaConf.load(args.cfg)
+    cfg.project_dir = args.project_dir
+
+    if args.index:
+        process_umls_index_dataset(cfg.index.raw_data, cfg.index.index_ds.data_file, cfg.index.id_to_string, HEADERS)
+    else:
+        process_umls_training_dataset(
+            cfg.model.raw_data,
+            cfg.model.train_ds.data_file,
+            cfg.model.validation_ds.data_file,
+            args.max_pairs,
+            args.train_split,
+            HEADERS,
+        )

tts/NeMo/examples/nlp/entity_linking/query_index.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import pickle as pkl
+from argparse import ArgumentParser
+from collections import OrderedDict
+from typing import Dict
+
+import numpy as np
+import torch
+from build_index import load_model
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.utils import logging
+
+try:
+    import faiss
+except ModuleNotFoundError:
+    logging.warning("Faiss is required for building the index. Please install faiss-gpu")
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def get_query_embedding(query, model):
+    """Use entity linking encoder to get embedding for index query"""
+    model_input = model.tokenizer(
+        query,
+        add_special_tokens=True,
+        padding=True,
+        truncation=True,
+        max_length=512,
+        return_token_type_ids=True,
+        return_attention_mask=True,
+    )
+
+    query_emb = model.forward(
+        input_ids=torch.LongTensor([model_input["input_ids"]]).to(device),
+        token_type_ids=torch.LongTensor([model_input["token_type_ids"]]).to(device),
+        attention_mask=torch.LongTensor([model_input["attention_mask"]]).to(device),
+    )
+
+    return query_emb
+
+
+def query_index(
+    query: str, cfg: DictConfig, model: object, index: object, pca: object, idx2id: dict, id2string: dict,
+) -> Dict:
+
+    """
+    Query the nearest neighbor index of entities to find the 
+    concepts in the index dataset that are most similar to the 
+    query.
+
+    Args:
+        query (str): entity to look up in the index
+        cfg (DictConfig): config object to specifiy query parameters
+        model (EntityLinkingModel): entity linking encoder model
+        index (object): faiss index
+        pca (object): sklearn pca transformation to be applied to queries 
+        idx2id (dict): dictionary mapping unique concept dataset index to 
+                       its CUI
+        id2string (dict): dictionary mapping each unqiue CUI to a 
+                          representative english description of
+                          the concept
+    Returns:
+        A dictionary with the concept ids of the index's most similar 
+        entities as the keys and a tuple containing the string 
+        representation of that concept and its cosine similarity to 
+        the query as the values. 
+    """
+    query_emb = get_query_embedding(query, model).detach().cpu().numpy()
+
+    if cfg.apply_pca:
+        query_emb = pca.transform(query_emb)
+
+    dist, neighbors = index.search(query_emb.astype(np.float32), cfg.query_num_factor * cfg.top_n)
+    dist, neighbors = dist[0], neighbors[0]
+    unique_ids = OrderedDict()
+    neighbor_idx = 0
+
+    # Many of nearest neighbors could map to the same concept id, their idx is their unique identifier
+    while len(unique_ids) < cfg.top_n and neighbor_idx < len(neighbors):
+        concept_id_idx = neighbors[neighbor_idx]
+        concept_id = idx2id[concept_id_idx]
+
+        # Only want one instance of each unique concept
+        if concept_id not in unique_ids:
+            concept = id2string[concept_id]
+            unique_ids[concept_id] = (concept, 1 - dist[neighbor_idx])
+
+        neighbor_idx += 1
+
+    unique_ids = dict(unique_ids)
+
+    return unique_ids
+
+
+def main(cfg: DictConfig, restore: bool):
+    """
+    Loads faiss index and allows commandline queries 
+    to the index. Builds new index if one hasn't been built yet.
+
+    Args:
+        cfg: Config file specifying index parameters
+        restore: Whether to restore model weights trained
+                 by the user. Otherwise will load weights
+                 used before self alignment pretraining.
+    """
+
+    if not os.path.isfile(cfg.index.index_save_name) or (
+        cfg.apply_pca and not os.path.isfile(cfg.index.pca.pca_save_name) or not os.path.isfile(cfg.index.idx_to_id)
+    ):
+        logging.warning("Either no index and/or no mapping from entity idx to ids exists. Please run `build_index.py`")
+        return
+
+    logging.info("Loading entity linking encoder model")
+    model = load_model(cfg.model, restore)
+
+    logging.info("Loading index and associated files")
+    index = faiss.read_index(cfg.index.index_save_name)
+    idx2id = pkl.load(open(cfg.index.idx_to_id, "rb"))
+    id2string = pkl.load(open(cfg.index.id_to_string, "rb"))  # Should be created during dataset prep
+
+    if cfg.index.apply_pca:
+        pca = pkl.load(open(cfg.index.pca.pca_save_name, "rb"))
+
+    while True:
+        query = input("enter index query: ")
+        output = query_index(query, cfg.top_n, cfg.index, model, index, pca, idx2id, id2string)
+
+        if query == "exit":
+            break
+
+        for concept_id in output:
+            concept_details = output[concept_id]
+            concept_id = "C" + str(concept_id).zfill(7)
+            print(concept_id, concept_details)
+
+        print("----------------\n")
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--restore", action="store_true", help="Whether to restore encoder model weights from nemo path"
+    )
+    parser.add_argument("--project_dir", required=False, type=str, default=".")
+    parser.add_argument("--cfg", required=False, type=str, default="./conf/umls_medical_entity_linking_config.yaml")
+    args = parser.parse_args()
+
+    cfg = OmegaConf.load(args.cfg)
+    cfg.project_dir = args.project_dir
+
+    main(cfg, args.restore)

tts/NeMo/examples/nlp/entity_linking/self_alignment_pretraining.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# Please see tutorial at Nemo/tutorials/nlp/Entity_Linking_Medical.ipynb for
+# more information on entity linking and self alignment pretraining.
+
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+
+from nemo.collections.nlp.models import EntityLinkingModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="umls_medical_entity_linking_config.yaml")
+def main(cfg: DictConfig) -> None:
+    logging.info(f"\nConfig Params:\n{OmegaConf.to_yaml(cfg)}")
+    trainer = Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    logging.info(f"Loading weights from pretrained model {cfg.model.language_model.pretrained_model_name}")
+    model = EntityLinkingModel(cfg=cfg.model, trainer=trainer)
+    logging.info("===========================================================================================")
+    logging.info('Starting training...')
+    trainer.fit(model)
+    logging.info('Training finished!')
+    logging.info("===========================================================================================")
+
+    if cfg.model.nemo_path:
+        # '.nemo' file contains the last checkpoint and the params to initialize the model
+        model.save_to(cfg.model.nemo_path)
+        logging.info(f'Model is saved into `.nemo` file: {cfg.model.nemo_path}')
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/glue_benchmark/glue_benchmark.py

@@ -0,0 +1,73 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+## Tasks
+This script works with all GLUE Benchmark tasks, more details about the GLUE Benchmark could be found at
+https://gluebenchmark.com/
+
+More details on how to use this script could be found in tutorials/nlp/GLUE_Benchmark.ipynb
+
+## Model Training
+
+To train GLUEModel with the default config file, run:
+    python glue_benchmark.py \
+    model.dataset.data_dir=<PATH_TO_DATA_DIR>  \
+    model.task_name=TASK_NAME \
+    trainer.max_epochs=<NUM_EPOCHS> \
+    trainer.gpus="[<CHANGE_TO_GPU_YOU_WANT_TO_USE>]
+
+Supported task names:
+["cola", "sst-2", "mrpc", "sts-b", "qqp", "mnli", "qnli", "rte", "wnli"]
+Note, MNLI task includes both matched and mismatched dev sets
+"""
+
+import os
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import GLUEModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_name="glue_benchmark_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager_cfg = cfg.get("exp_manager", None)
+
+    if exp_manager_cfg:
+        exp_manager_cfg.name = cfg.model.task_name
+        logging.info(f'Setting task_name to {exp_manager_cfg.name} in exp_manager')
+    exp_manager(trainer, exp_manager_cfg)
+
+    if cfg.model.nemo_path and os.path.exists(cfg.model.nemo_path):
+        model = GLUEModel.restore_from(cfg.model.nemo_path)
+        logging.info(f'Restoring model from {cfg.model.nemo_path}')
+        model.update_data_dir(data_dir=cfg.model.dataset.data_dir)
+        model.setup_training_data()
+        model.setup_multiple_validation_data()
+        trainer.fit(model)
+    else:
+        model = GLUEModel(cfg.model, trainer=trainer)
+        trainer.fit(model)
+        if cfg.model.nemo_path:
+            model.save_to(cfg.model.nemo_path)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/glue_benchmark/glue_benchmark_config.yaml

@@ -0,0 +1,82 @@
+# GLUE Benchmark with pre-trained BERT models
+supported_tasks: &supported_tasks ['cola', 'sst-2', 'mrpc', 'sts-b', 'qqp', 'mnli', 'qnli', 'rte', 'wnli']
+
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU
+  num_nodes: 1
+  max_epochs: 3
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 16
+  accelerator: ddp
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+
+model:
+  task_name: &task_name mrpc # choose from: ["cola", "sst-2", "mrpc", "sts-b", "qqp", "mnli", "qnli", "rte", "wnli"] GLUE task name, MNLI includes both matched and mismatched dev sets
+  supported_tasks: *supported_tasks
+  output_dir: null # dir to write write predictions
+  nemo_path: null # filename to save the model and associated artifacts to .nemo file
+  dataset:
+    data_dir: ??? # /path/to/data
+    max_seq_length: 128
+    use_cache: true
+
+    # shared across dataloaders:
+    num_workers:  2
+    pin_memory: false
+    drop_last: false
+
+  train_ds:
+    ds_item: 'train.tsv'
+    shuffle: true
+    num_samples: -1
+    batch_size: 32
+
+  validation_ds:
+    ds_item: 'dev.tsv' # for MNLI 'dev_matched.tsv' and 'dev_mismatched.tsv' will de used
+    shuffle: false
+    num_samples: -1
+    batch_size: 32
+
+  tokenizer:
+      tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+      vocab_file: null # path to vocab file
+      tokenizer_model: null # only used if tokenizer is sentencepiece
+      special_tokens: null # only necessary for adding transformer/bert-specific special tokens to tokenizer if the tokenizer does not already have these inherently.
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+  optim:
+    name: adam
+    lr: 5e-5
+    weight_decay: 0.00
+
+    sched:
+      name: WarmupAnnealing
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.1
+      last_epoch: -1
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./NeMo_experiments"
+  name: *task_name # The name of your model
+  create_tensorboard_logger: True  # Whether you want exp_manger to create a tb logger
+  create_checkpoint_callback: True  # Whether you want exp_manager to create a modelcheckpoint callback
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/information_retrieval/bert_dpr.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import BertDPRModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="bert_ir_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    bert_dpr_model = BertDPRModel(cfg.model, trainer=trainer)
+    trainer.fit(bert_dpr_model)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/information_retrieval/bert_joint_ir.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import BertJointIRModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="bert_ir_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    bert_joint_ir_model = BertJointIRModel(cfg.model, trainer=trainer)
+    trainer.fit(bert_joint_ir_model)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/information_retrieval/conf/bert_ir_config.yaml

@@ -0,0 +1,99 @@
+# Fine-tuning BERT model for information retrieval
+name: &name BertIR
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices
+  num_nodes: 1
+  max_epochs: 2 # the number of training epochs
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  precision: 16 # 16 to use AMP
+  amp_level: O1 # O1 or O2 if using AMP
+  accelerator: ddp
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 0.05  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  checkpoint_callback: false  # provided by exp_manager
+  logger: false  # provided by exp_manager
+
+model:
+  nemo_path: null # exported .nemo path
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    sim_score_dropout: 0.1
+    lm_checkpoint: null
+    config:
+      attention_probs_dropout_prob: 0.1
+      hidden_act: gelu
+      hidden_dropout_prob: 0.1
+      hidden_size: 768
+      initializer_range: 0.02
+      intermediate_size: 3072
+      max_position_embeddings: 512
+      num_attention_heads: 12
+      num_hidden_layers: 12
+      type_vocab_size: 2
+      vocab_size: 30522
+    config_file: null # json file, precedence over config
+
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # tokenizer model for sentencepiece
+    special_tokens: null
+
+  train_ds:
+    passages: null # path to file with passages and their indices
+    queries: null # path to file with training queries and their indices
+    query_to_passages: null
+    # path to file with training examples which have the form of
+    # (query_id, relevant_passage_id, irrelevant_passage_1_id, ..., irrelevant_passage_n_id)
+    num_negatives: 10
+    batch_size: 6
+    psg_cache_format: npz
+    shuffle: true
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 1
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    passages: null # path to file with passages and their indices
+    queries: null # path to file with validation queries and their indices
+    query_to_passages: null # path to file with passages to re-rank for each validation query
+    num_negatives: 10
+    batch_size: 6
+    psg_cache_format: pkl
+    shuffle: false
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 1
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adam
+    lr: 1e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0
+
+    sched:
+      name: WarmupAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.05
+      last_epoch: -1
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+exp_manager:
+  exp_dir: null  # where to store logs and checkpoints
+  name: *name  # name of experiment
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/information_retrieval/construct_random_negatives.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+
+import numpy as np
+import pandas as pd
+
+
+def construct_negatives(input_file, output_file, num_passages, num_negatives):
+    qrels = pd.read_csv(input_file, delimiter="\t", header=None)
+    with open(output_file, "w") as f:
+        for i in range(len(qrels)):
+            query_id, rel_passage_id = qrels[0][i], qrels[2][i]
+            negatives = np.random.randint(num_passages, size=num_negatives)
+            output_ids = [query_id, rel_passage_id] + negatives.tolist()
+            output_str = [str(id_) for id_ in output_ids]
+            print("\t".join(output_str), file=f)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Negative passages construction")
+    parser.add_argument("--data", type=str, default="msmarco_dataset", help="path to folder with data")
+    parser.add_argument("--num_passages", type=int, default=8841823, help="total number of passages")
+    parser.add_argument("--num_negatives", type=int, default=10, help="number of negatives per positive")
+    args = parser.parse_args()
+
+    for mode in ["train", "dev"]:
+        construct_negatives(
+            input_file=f"{args.data}/qrels.{mode}.tsv",
+            output_file=f"{args.data}/query2passages.{mode}.tsv",
+            num_passages=args.num_passages,
+            num_negatives=args.num_negatives,
+        )
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/information_retrieval/get_msmarco.sh

@@ -0,0 +1,22 @@
+echo "=== Acquiring MSMARCO dataset ==="
+echo "---"
+
+mkdir -p msmarco_dataset
+cd msmarco_dataset
+
+echo "- Downloading passages"
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/collection.tar.gz
+tar -xzvf collection.tar.gz
+rm collection.tar.gz
+
+echo "- Downloading queries"
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/queries.tar.gz
+tar -xzvf queries.tar.gz
+rm queries.tar.gz
+rm queries.eval.tsv
+
+echo "- Downloading relevance labels"
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/qrels.train.tsv
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/qrels.dev.tsv
+
+echo "---"

tts/NeMo/examples/nlp/intent_slot_classification/conf/intent_slot_classification_config.yaml

@@ -0,0 +1,111 @@
+# Intent and Slot classification with pretrained BERT models
+
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU
+  num_nodes: 1
+  max_epochs: 50
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 32 # Should be set to 16 for O1 and O2 amp_level to enable the AMP.
+  accelerator: ddp
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+
+  checkpoint_callback: false  # Provided by exp_manager
+  logger: false  # Provided by exp_manager
+
+model:
+  nemo_path: null # filename to save the model and associated artifacts to .nemo file
+  data_dir: ??? # /path/to/data
+  class_labels:
+    intent_labels_file: intent_labels.csv
+    slot_labels_file: slot_labels.csv
+  class_balancing: null # or weighted_loss
+  intent_loss_weight: 0.6 # relation of intent to slot loss in total loss (between 0 to 1)
+  pad_label: -1 # if -1 not slot token will be used
+  ignore_extra_tokens: false
+  ignore_start_end: true # do not use first and last token for slot training
+
+  train_ds:
+    prefix: train
+    batch_size: 32
+    shuffle: true
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    prefix: test
+    batch_size: 32
+    shuffle: false
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  test_ds:
+    prefix: test
+    batch_size: 32
+    shuffle: false
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # only used if tokenizer is sentencepiece
+    special_tokens: null
+
+  language_model:
+    max_seq_length: 50
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+  head:
+    num_output_layers: 2
+    fc_dropout: 0.1
+
+  optim:
+    name: adam
+    lr: 2e-5
+    args:
+      name: auto
+      params:
+        weight_decay: 0.01
+
+    sched:
+      name: WarmupAnnealing
+      iters_per_batch: null # computed at runtime
+      max_steps: null # computed at runtime or explicitly set here
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+      # scheduler config override
+      args:
+        name: auto
+        params:
+          warmup_steps: null
+          warmup_ratio: 0.1
+          last_epoch: -1
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
+  name: "IntentSlot"  # The name of your model
+  create_tensorboard_logger: true  # Whether you want exp_manger to create a tb logger
+  create_checkpoint_callback: true  # Whether you want exp_manager to create a modelcheckpoint callback
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/intent_slot_classification/data/assistant_utils.py

@@ -0,0 +1,151 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import re
+import shutil
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import DATABASE_EXISTS_TMP, if_exist, write_files
+from nemo.utils import logging
+
+
+def copy_input_files(infold):
+    """ Put training files in convenient place for conversion to our format. """
+    our_infold = infold + "/dataset"
+
+    if os.path.exists(our_infold + "/trainset") and os.path.exists(our_infold + "/testset"):
+        logging.info("Input folders exists")
+        return
+
+    logging.info(f"Copying files to input folder: {our_infold}")
+    os.makedirs(infold, exist_ok=True)
+
+    old_infold = (
+        infold + '/CrossValidation/autoGeneFromRealAnno/autoGene_2018_03_22-13_01_25_169/CrossValidation/KFold_1'
+    )
+    if not os.path.exists(our_infold + "/trainset"):
+        shutil.copytree(old_infold + '/trainset', our_infold + '/trainset')
+
+    if not os.path.exists(our_infold + "/testset"):
+        shutil.copytree(old_infold + '/testset/csv', our_infold + '/testset')
+
+
+def get_intents(infold):
+    """ Get list of intents from file names. """
+    intents = [f[:-4] for f in os.listdir(infold)]
+    intents.sort()
+    print(f'Found {len(intents)} intents')
+    return intents
+
+
+def get_intent_queries(infold, intent_names, mode):
+    """ Get list of queries with their corresponding intent number. """
+    intent_queries = ['sentence\tlabel\n']
+
+    for index, intent in enumerate(intent_names):
+        queries = open(f'{infold}/{mode}set/{intent}.csv', 'r').readlines()
+        for query in queries[1:]:
+            phrases = query.split(";")
+            intent_query = phrases[4][1:-1] + "\t" + str(index)
+            intent_queries.append(intent_query)
+
+    return intent_queries
+
+
+def get_slots(infold, modes):
+    """
+    Find a lost of unique slot types in training and testing data.
+    We use a single slot type name both for starting and continuation tokes (not using B-, I- notation).
+    """
+    slots = set()
+
+    for mode in modes:
+        path = f'{infold}/{mode}set'
+        for filename in os.listdir(path):
+            lines = open(f'{path}/{filename}', 'r').readlines()
+            for line in lines[1:]:
+                query = line.split(";")[3]
+                slot_phrases = re.findall('\[.*?\]', query)
+                for slot_phrase in slot_phrases:
+                    slot = slot_phrase.split(" : ")[0][1:]
+                    slots.add(slot)
+
+    slots = sorted(slots)
+    slots.append("O")
+    print(f'Found {len(slots)} slot types')
+    return slots
+
+
+def get_slot_queries(infold, slot_dict, mode, intent_names):
+    """ Convert each word in a query to corresponding slot number. """
+    slot_queries = []
+    outside_slot = len(slot_dict) - 1
+
+    # keep the same order of files/queries as for intents
+    for intent in intent_names:
+        lines = open(f'{infold}/{mode}set/{intent}.csv', 'r').readlines()
+        for line in lines[1:]:
+            slot_query = ""
+            query = line.split(";")[3]
+            words = query.split(" ")
+            current_slot = outside_slot
+            for word in words:
+                if word[0] == "[":
+                    current_slot = slot_dict[word[1:]]
+                elif word[0] == ":":
+                    continue
+                else:
+                    slot_query += str(current_slot) + " "
+                    if word[-1] == ']':
+                        current_slot = outside_slot
+
+            slot_queries.append(slot_query.strip())
+
+    return slot_queries
+
+
+def process_assistant(infold, outfold, modes=['train', 'test']):
+    """
+    https://github.com/xliuhw/NLU-Evaluation-Data - this dataset includes
+    about 25 thousand examples with 66 various multi-domain intents and 57 entity types.
+    """
+    if if_exist(outfold, [f'{mode}_slots.tsv' for mode in modes]):
+        logging.info(DATABASE_EXISTS_TMP.format('robot', outfold))
+        return outfold
+
+    logging.info(f'Processing assistant commands dataset and store at {outfold}')
+    os.makedirs(outfold, exist_ok=True)
+
+    # copy train/test files to the convenient directory to work with
+    copy_input_files(infold)
+    infold += "/dataset"
+
+    # get list of intents from train folder (test folder supposed to be the same)
+    intent_names = get_intents(infold + "/trainset")
+    write_files(intent_names, f'{outfold}/dict.intents.csv')
+
+    # get all train and test queries with their intent
+    for mode in modes:
+        intent_queries = get_intent_queries(infold, intent_names, mode)
+        write_files(intent_queries, f'{outfold}/{mode}.tsv')
+
+    # get list of all unique slots in training and testing files
+    slot_types = get_slots(infold, modes)
+    write_files(slot_types, f'{outfold}/dict.slots.csv')
+
+    # create files of slot queries
+    slot_dict = {k: v for v, k in enumerate(slot_types)}
+    for mode in modes:
+        slot_queries = get_slot_queries(infold, slot_dict, mode, intent_names)
+        write_files(slot_queries, f'{outfold}/{mode}_slots.tsv')

tts/NeMo/examples/nlp/intent_slot_classification/data/import_datasets.py

@@ -0,0 +1,265 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+import shutil
+from os.path import exists
+
+from assistant_utils import process_assistant
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import (
+    DATABASE_EXISTS_TMP,
+    MODE_EXISTS_TMP,
+    create_dataset,
+    get_dataset,
+    get_vocab,
+    if_exist,
+)
+from nemo.utils import logging
+
+
+def ids2text(ids, vocab):
+    return ' '.join([vocab[int(id_)] for id_ in ids])
+
+
+def process_atis(infold, outfold, modes=['train', 'test'], do_lower_case=False):
+    """ MSFT's dataset, processed by Kaggle
+    https://www.kaggle.com/siddhadev/atis-dataset-from-ms-cntk
+    """
+    vocab = get_vocab(f'{infold}/atis.dict.vocab.csv')
+
+    if if_exist(outfold, [f'{mode}.tsv' for mode in modes]):
+        logging.info(DATABASE_EXISTS_TMP.format('ATIS', outfold))
+        return outfold
+    logging.info(f'Processing ATIS dataset and storing at {outfold}.')
+
+    os.makedirs(outfold, exist_ok=True)
+
+    outfiles = {}
+    for mode in modes:
+        outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'w')
+        outfiles[mode].write('sentence\tlabel\n')
+        outfiles[mode + '_slots'] = open(f'{outfold}/{mode}_slots.tsv', 'w')
+
+        queries = open(f'{infold}/atis.{mode}.query.csv', 'r').readlines()
+        intents = open(f'{infold}/atis.{mode}.intent.csv', 'r').readlines()
+        slots = open(f'{infold}/atis.{mode}.slots.csv', 'r').readlines()
+
+        for i, query in enumerate(queries):
+            sentence = ids2text(query.strip().split()[1:-1], vocab)
+            if do_lower_case:
+                sentence = sentence.lower()
+            outfiles[mode].write(f'{sentence}\t{intents[i].strip()}\n')
+            slot = ' '.join(slots[i].strip().split()[1:-1])
+            outfiles[mode + '_slots'].write(slot + '\n')
+
+    shutil.copyfile(f'{infold}/atis.dict.intent.csv', f'{outfold}/dict.intents.csv')
+    shutil.copyfile(f'{infold}/atis.dict.slots.csv', f'{outfold}/dict.slots.csv')
+    for mode in modes:
+        outfiles[mode].close()
+
+
+def process_snips(infold, outfold, do_lower_case, modes=['train', 'test'], dev_split=0.1):
+    if not os.path.exists(infold):
+        link = 'https://github.com/snipsco/spoken-language-understanding-research-datasets'
+        raise ValueError(f'Data not found at {infold}. ' f'You may request to download the SNIPS dataset from {link}.')
+
+    exist = True
+    for dataset in ['light', 'speak', 'all']:
+        if if_exist(f'{outfold}/{dataset}', [f'{mode}.tsv' for mode in modes]):
+            logging.info(DATABASE_EXISTS_TMP.format('SNIPS-' + dataset, outfold))
+        else:
+            exist = False
+    if exist:
+        return outfold
+
+    logging.info(f'Processing SNIPS dataset and storing at folders "speak", "light" and "all" under {outfold}.')
+    logging.info(
+        f'Processing and importing "smart-speaker-en-close-field" -> "speak" and "smart-speaker-en-close-field" -> "light".'
+    )
+
+    os.makedirs(outfold, exist_ok=True)
+
+    speak_dir = 'smart-speaker-en-close-field'
+    light_dir = 'smart-lights-en-close-field'
+
+    light_files = [f'{infold}/{light_dir}/dataset.json']
+    speak_files = [f'{infold}/{speak_dir}/training_dataset.json']
+    speak_files.append(f'{infold}/{speak_dir}/test_dataset.json')
+
+    light_train, light_dev, light_slots, light_intents = get_dataset(light_files, dev_split)
+    speak_train, speak_dev, speak_slots, speak_intents = get_dataset(speak_files)
+
+    create_dataset(light_train, light_dev, light_slots, light_intents, do_lower_case, f'{outfold}/light')
+    create_dataset(speak_train, speak_dev, speak_slots, speak_intents, do_lower_case, f'{outfold}/speak')
+    create_dataset(
+        light_train + speak_train,
+        light_dev + speak_dev,
+        light_slots | speak_slots,
+        light_intents | speak_intents,
+        do_lower_case,
+        f'{outfold}/all',
+    )
+
+
+def process_jarvis_datasets(
+    infold, outfold, modes=['train', 'test', 'dev'], do_lower_case=False, ignore_prev_intent=False
+):
+    """ process and convert Jarvis datasets into NeMo's BIO format
+    """
+    dataset_name = "jarvis"
+    if if_exist(outfold, ['dict.intents.csv', 'dict.slots.csv']):
+        logging.info(DATABASE_EXISTS_TMP.format(dataset_name, outfold))
+        return outfold
+
+    logging.info(f'Processing {dataset_name} dataset and storing at {outfold}')
+
+    os.makedirs(outfold, exist_ok=True)
+
+    outfiles = {}
+    intents_list = {}
+    slots_list = {}
+    slots_list_all = {}
+
+    outfiles['dict_intents'] = open(f'{outfold}/dict.intents.csv', 'w')
+    outfiles['dict_slots'] = open(f'{outfold}/dict.slots.csv', 'w')
+
+    outfiles['dict_slots'].write('O\n')
+    slots_list["O"] = 0
+    slots_list_all["O"] = 0
+
+    for mode in modes:
+        if if_exist(outfold, [f'{mode}.tsv']):
+            logging.info(MODE_EXISTS_TMP.format(mode, dataset_name, outfold, mode))
+            continue
+
+        if not if_exist(infold, [f'{mode}.tsv']):
+            logging.info(f'{mode} mode of {dataset_name}' f' is skipped as it was not found.')
+            continue
+
+        outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'w')
+        outfiles[mode].write('sentence\tlabel\n')
+        outfiles[mode + '_slots'] = open(f'{outfold}/{mode}_slots.tsv', 'w')
+
+        queries = open(f'{infold}/{mode}.tsv', 'r').readlines()
+
+        for i, query in enumerate(queries):
+            line_splits = query.strip().split("\t")
+            if len(line_splits) == 3:
+                intent_str, slot_tags_str, sentence = line_splits
+            else:
+                intent_str, sentence = line_splits
+                slot_tags_str = ""
+
+            if intent_str not in intents_list:
+                intents_list[intent_str] = len(intents_list)
+                outfiles['dict_intents'].write(f'{intent_str}\n')
+
+            if ignore_prev_intent:
+                start_token = 2
+            else:
+                start_token = 1
+
+            if do_lower_case:
+                sentence = sentence.lower()
+            sentence_cld = " ".join(sentence.strip().split()[start_token:-1])
+            outfiles[mode].write(f'{sentence_cld}\t' f'{str(intents_list[intent_str])}\n')
+
+            slot_tags_list = []
+            if slot_tags_str.strip():
+                slot_tags = slot_tags_str.strip().split(",")
+                for st in slot_tags:
+                    if not st.strip():
+                        continue
+                    [start_i, end_i, slot_name] = st.strip().split(":")
+                    slot_tags_list.append([int(start_i), int(end_i), slot_name])
+                    if slot_name not in slots_list:
+                        slots_list[slot_name] = len(slots_list)
+                        slots_list_all[f'B-{slot_name}'] = len(slots_list_all)
+                        slots_list_all[f'I-{slot_name}'] = len(slots_list_all)
+                        outfiles['dict_slots'].write(f'B-{slot_name}\n')
+                        outfiles['dict_slots'].write(f'I-{slot_name}\n')
+
+            slot_tags_list.sort(key=lambda x: x[0])
+            slots = []
+            processed_index = 0
+            for tag_start, tag_end, tag_str in slot_tags_list:
+                if tag_start > processed_index:
+                    words_list = sentence[processed_index:tag_start].strip().split()
+                    slots.extend([str(slots_list_all['O'])] * len(words_list))
+                words_list = sentence[tag_start:tag_end].strip().split()
+                slots.append(str(slots_list_all[f'B-{tag_str}']))
+                slots.extend([str(slots_list_all[f'I-{tag_str}'])] * (len(words_list) - 1))
+                processed_index = tag_end
+
+            if processed_index < len(sentence):
+                words_list = sentence[processed_index:].strip().split()
+                slots.extend([str(slots_list_all['O'])] * len(words_list))
+
+            slots = slots[1:-1]
+            slot = ' '.join(slots)
+            outfiles[mode + '_slots'].write(slot + '\n')
+
+        outfiles[mode + '_slots'].close()
+        outfiles[mode].close()
+
+    outfiles['dict_slots'].close()
+    outfiles['dict_intents'].close()
+
+    return outfold
+
+
+if __name__ == "__main__":
+    # Parse the command-line arguments.
+    parser = argparse.ArgumentParser(description="Process and convert datasets into NeMo\'s format.")
+    parser.add_argument(
+        "--dataset_name", required=True, type=str, choices=['atis', 'snips', 'jarvis', 'assistant'],
+    )
+    parser.add_argument(
+        "--source_data_dir", required=True, type=str, help='path to the folder containing the dataset files'
+    )
+    parser.add_argument("--target_data_dir", required=True, type=str, help='path to save the processed dataset')
+    parser.add_argument("--do_lower_case", action='store_true')
+    parser.add_argument(
+        "--ignore_prev_intent",
+        action='store_true',
+        help='ignores previous intent while importing datasets in jarvis\'s format',
+    )
+
+    args = parser.parse_args()
+
+    dataset_name = args.dataset_name
+    source_dir = args.source_data_dir
+    target_dir = args.target_data_dir
+
+    if not exists(source_dir):
+        raise FileNotFoundError(f"{source_dir} does not exist.")
+
+    if dataset_name == 'atis':
+        process_atis(infold=source_dir, outfold=target_dir, do_lower_case=args.do_lower_case)
+    elif dataset_name == 'snips':
+        process_snips(infold=source_dir, outfold=target_dir, do_lower_case=args.do_lower_case)
+    elif dataset_name == 'jarvis':
+        process_jarvis_datasets(
+            infold=source_dir,
+            outfold=target_dir,
+            modes=["train", "test", "dev"],
+            do_lower_case=args.do_lower_case,
+            ignore_prev_intent=args.ignore_prev_intent,
+        )
+    elif dataset_name == 'assistant':
+        process_assistant(infold=source_dir, outfold=target_dir)
+    else:
+        raise ValueError(f'Dataset {dataset_name} is not supported.')

tts/NeMo/examples/nlp/intent_slot_classification/intent_slot_classification.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import IntentSlotClassificationModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="intent_slot_classification_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config Params:\n {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    # initialize the model using the config file
+    model = IntentSlotClassificationModel(cfg.model, trainer=trainer)
+
+    # training
+    logging.info("================================================================================================")
+    logging.info('Starting training...')
+    trainer.fit(model)
+    logging.info('Training finished!')
+
+    # Stop further testing as fast_dev_run does not save checkpoints
+    if trainer.fast_dev_run:
+        return
+
+    # after model training is done, you can load the model from the saved checkpoint
+    # and evaluate it on a data file or on given queries.
+    logging.info("================================================================================================")
+    logging.info("Starting the testing of the trained model on test set...")
+    logging.info("We will load the latest model saved checkpoint from the training...")
+
+    # for evaluation and inference you can load the previously trained model saved in .nemo file
+    # like this in your code, but we will just reuse the trained model here
+    # eval_model = IntentSlotClassificationModel.restore_from(restore_path=checkpoint_path)
+    eval_model = model
+
+    # we will setup testing data reusing the same config (test section)
+    eval_model.update_data_dir_for_testing(data_dir=cfg.model.data_dir)
+    eval_model.setup_test_data(test_data_config=cfg.model.test_ds)
+
+    trainer.test(model=eval_model, ckpt_path=None, verbose=False)
+    logging.info("Testing finished!")
+
+    # run an inference on a few examples
+    logging.info("======================================================================================")
+    logging.info("Evaluate the model on the given queries...")
+
+    # this will work well if you train the model on Assistant dataset
+    # for your own dataset change the examples appropriately
+    queries = [
+        'set alarm for seven thirty am',
+        'lower volume by fifty percent',
+        'what is my schedule for tomorrow',
+    ]
+
+    pred_intents, pred_slots = eval_model.predict_from_examples(queries, cfg.model.test_ds)
+
+    logging.info('The prediction results of some sample queries with the trained model:')
+    for query, intent, slots in zip(queries, pred_intents, pred_slots):
+        logging.info(f'Query : {query}')
+        logging.info(f'Predicted Intent: {intent}')
+        logging.info(f'Predicted Slots: {slots}')
+
+    logging.info("Inference finished!")
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/language_modeling/bert_pretraining.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning.plugins import DDPPlugin
+
+from nemo.collections.nlp.models.language_modeling import BERTLMModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="bert_pretraining_from_text_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config:\n {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(plugins=[DDPPlugin(find_unused_parameters=True)], **cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    bert_model = BERTLMModel(cfg.model, trainer=trainer)
+    trainer.fit(bert_model)
+    if cfg.model.nemo_path:
+        bert_model.save_to(cfg.model.nemo_path)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_preprocessed_config.yaml

@@ -0,0 +1,80 @@
+# BERT Pretraining from Preprocessed (tokenized) data
+name: &name PretrainingBERTFromPreprocessed
+trainer:
+  gpus: 8 # the number of gpus, 0 for CPU, or list with gpu indices
+  num_nodes: 1
+  max_steps: 2285714 # precedence over max_epochs
+  num_sanity_val_steps: 0 # needed for bert pretraining from preproc
+  replace_sampler_ddp: false # needed for bert pretraining from preproc
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  precision: 16 # 16 to use AMP
+  amp_level: O1 # O1 or O2 if using AMP
+  accelerator: ddp
+  gradient_clip_val: 1.0
+  log_every_n_steps: 1
+  val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch
+  checkpoint_callback: false # provided by exp_manager
+  logger: false # provided by exp_manager
+
+model:
+  nemo_path: null # exported .nemo path
+  only_mlm_loss: true # only use masked language model without next sentence prediction
+  num_tok_classification_layers: 1 # number of token classification head output layers
+  num_seq_classification_layers: 2 # number of sequence classification head output layers
+
+
+  language_model:
+    pretrained_model_name: bert-base-uncased # huggingface model name
+    lm_checkpoint: null
+    config:
+      attention_probs_dropout_prob: 0.1
+      hidden_act: gelu
+      hidden_dropout_prob: 0.1
+      hidden_size: 768
+      initializer_range: 0.02
+      intermediate_size: 3072
+      max_position_embeddings: 512
+      num_attention_heads: 12
+      num_hidden_layers: 12
+      type_vocab_size: 2
+      vocab_size: 30522
+    config_file: null # json file, precedence over config
+
+  tokenizer: null
+
+  train_ds:
+    data_file: null # path to hdf5 file (or directory)
+    max_predictions_per_seq: 80
+    batch_size: 16
+    shuffle: true
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adamw
+    lr: 0.4375e-4
+    weight_decay: 0.01
+
+    sched:
+      name: SquareRootAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.01
+      min_lr: 0.0
+      last_epoch: -1
+
+
+exp_manager:
+  exp_dir: null # where to store logs and checkpoints
+  name: *name # name of experiment
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_text_config.yaml

@@ -0,0 +1,109 @@
+# BERT Pretraining from Text
+name: &name PretrainingBERTFromText
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices
+  num_nodes: 1
+  max_epochs: 2 # the number of training epochs
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  precision: 16 # 16 to use AMP
+  amp_level: O1 # O1 or O2 if using AMP
+  accelerator: ddp
+  gradient_clip_val: 0.0
+  log_every_n_steps: 1
+  val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch
+  checkpoint_callback: false # provided by exp_manager
+  logger: false # provided by exp_manager
+
+model:
+  nemo_path: null # exported .nemo path
+  only_mlm_loss: false # only use masked language model without next sentence prediction
+  num_tok_classification_layers: 1 # number of token classification head output layers
+  num_seq_classification_layers: 2 # number of sequence classification head output layers
+  max_seq_length: 128
+  # The maximum total input sequence length after tokenization. Sequences longer than this
+  # will be truncated, and sequences shorter than this will be padded.
+  mask_prob: 0.15
+  # Probability of masking a token in the input text during data processing.
+  short_seq_prob: 0.1
+  # Probability of having a sequence shorter than the maximum sequence length `max_seq_length` in data processing.",
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config:
+      attention_probs_dropout_prob: 0.1
+      hidden_act: gelu
+      hidden_dropout_prob: 0.1
+      hidden_size: 768
+      initializer_range: 0.02
+      intermediate_size: 3072
+      max_position_embeddings: 512
+      num_attention_heads: 12
+      num_hidden_layers: 12
+      type_vocab_size: 2
+      vocab_size: 30522
+    config_file: null # json file, precedence over config
+
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # tokenizer model for sentencepiece
+    special_tokens: # only necessary for adding transformer/bert-specific special tokens to tokenizer if the tokenizer does not already have these inherently.
+        unk_token: '[UNK]'
+        sep_token: '[SEP]'
+        pad_token: '[PAD]'
+        bos_token: '[CLS]'
+        mask_token: '[MASK]'
+        eos_token: '[SEP]'
+        cls_token: '[CLS]'
+
+  train_ds:
+    data_file: null # path to data file
+    max_seq_length: ${model.max_seq_length}
+    mask_prob: ${model.mask_prob}
+    short_seq_prob: ${model.short_seq_prob}
+    batch_size: 16 # per GPU
+    shuffle: true
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    data_file: null # path to data file
+    max_seq_length: ${model.max_seq_length}
+    mask_prob: ${model.mask_prob}
+    short_seq_prob: ${model.short_seq_prob}
+    batch_size: 16 # per GPU
+    shuffle: false
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adamw
+    lr: 3e-5
+    weight_decay: 0.0
+
+    sched:
+      name: CosineAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.1
+      min_lr: 0.0
+      last_epoch: -1
+
+
+exp_manager:
+  exp_dir: null # where to store logs and checkpoints
+  name: *name # name of experiment
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/language_modeling/conf/transformer_lm_config.yaml

@@ -0,0 +1,96 @@
+name: TransformerLanguageModel
+do_training: True # set to False if only preprocessing data
+
+model:
+  label_smoothing: 0.0
+  preproc_out_dir: null # path to store data preprocessing outputs
+  
+  train_ds:
+    use_tarred_dataset: true
+    tar_files: ??? # path to tarred files
+    metadata_file: ??? # metadata for tarred dataset
+    shard_strategy: scatter
+    tokens_in_batch: 4096
+    tar_shuffle_n: 256
+    
+  validation_ds:
+    file_name: ??? # path to file with validation data
+    tokens_in_batch: 512
+    clean: false
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+    
+  test_ds:
+    file_name: ??? # path to file with test data
+    tokens_in_batch: 512
+    clean: false
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+    
+  optim:
+    name: adam
+    lr: 0.001
+    betas:
+      - 0.9
+      - 0.98
+    weight_decay: 0.0
+    sched:
+      name: InverseSquareRootAnnealing
+      min_lr: 0.0
+      last_epoch: -1
+      warmup_ratio: 0.1
+
+  tokenizer:
+    tokenizer_name: yttm
+    tokenizer_model: ???
+    vocab_file: null
+    special_tokens: null
+    training_sample_size: null # valid for sentencepiece tokenizer
+
+  encoder:
+    library: nemo
+    model_name: null
+    pretrained: false
+    max_sequence_length: 512
+    num_token_types: 0
+    embedding_dropout: 0.1
+    learn_positional_encodings: false
+    hidden_size: 512
+    num_layers: 6
+    inner_size: 2048
+    num_attention_heads: 8
+    ffn_dropout: 0.1
+    attn_score_dropout: 0.1
+    attn_layer_dropout: 0.1
+    hidden_act: relu
+    mask_future: true
+    pre_ln: false
+
+  head:
+    num_layers: 1
+    activation: relu
+    log_softmax: true
+    dropout: 0.0
+    use_transformer_init: true
+
+trainer:
+  gpus: 4
+  num_nodes: 1
+  max_epochs: 200
+  amp_level: O2 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
+  accelerator: ddp
+  checkpoint_callback: False
+  logger: False
+  log_every_n_steps: 50  # Interval of logging.
+  check_val_every_n_epoch: 1
+
+exp_manager:
+  name: TransformerLM
+  files_to_copy: []

tts/NeMo/examples/nlp/language_modeling/convert_weights_to_nemo1.0.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Converts BERT NeMo0.* checkpoints to NeMo1.0 format.
+"""
+
+from argparse import ArgumentParser
+
+import torch
+
+parser = ArgumentParser()
+parser.add_argument("--bert_encoder", required=True, help="path to BERT encoder, e.g. /../BERT-STEP-2285714.pt")
+parser.add_argument(
+    "--bert_token_classifier",
+    required=True,
+    help="path to BERT token classifier, e.g. /../BertTokenClassifier-STEP-2285714.pt",
+)
+parser.add_argument(
+    "--bert_sequence_classifier",
+    required=False,
+    default=None,
+    help="path to BERT sequence classifier, e.g /../SequenceClassifier-STEP-2285714.pt",
+)
+parser.add_argument(
+    "--output_path", required=False, default="converted_model.pt", help="output path to newly converted model"
+)
+args = parser.parse_args()
+
+bert_in = torch.load(args.bert_encoder)
+tok_in = torch.load(args.bert_token_classifier)
+if args.bert_sequence_classifier:
+    seq_in = torch.load(args.bert_sequence_classifier)
+
+new_dict = {}
+new_model = {"state_dict": new_dict}
+for k in bert_in:
+    new_name = k.replace("bert.", "bert_model.")
+    new_dict[new_name] = bert_in[k]
+
+for k in tok_in:
+    new_name = "mlm_classifier." + k
+    new_dict[new_name] = tok_in[k]
+
+if args.bert_sequence_classifier:
+    for k in seq_in:
+        new_name = "nsp_classifier." + k
+        new_dict[new_name] = seq_in[k]
+
+torch.save(new_model, args.output_path)

tts/NeMo/examples/nlp/language_modeling/get_wkt2.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+"""
+This file is adapted from
+https://github.com/salesforce/awd-lstm-lm/blob/master/getdata.sh
+Copyright by the AWD LSTM authors.
+"""
+DATA_DIR=$1
+echo "- Downloading WikiText-2"
+
+wget --continue -P $DATA_DIR https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip
+unzip -q $DATA_DIR/wikitext-2-v1.zip -d $DATA_DIR
+cd $DATA_DIR/wikitext-2
+mv wiki.train.tokens train.txt
+sed -i -e "s/<unk>/[UNK]/g" train.txt
+mv wiki.valid.tokens valid.txt
+sed -i -e "s/<unk>/[UNK]/g" valid.txt
+mv wiki.test.tokens test.txt
+sed -i -e "s/<unk>/[UNK]/g" test.txt
+cd ..
+rm wikitext-2-v1.zip
+
+echo "- WikiText-2 saved at $DATA_DIR/wikitext-2"

tts/NeMo/examples/nlp/language_modeling/transformer_lm.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.language_modeling import TransformerLMModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="transformer_lm_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    transformer_lm = TransformerLMModel(cfg.model, trainer=trainer)
+    trainer.fit(transformer_lm)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/machine_translation/conf/aayn_base.yaml

@@ -0,0 +1,156 @@
+name: AttentionIsAllYouNeed
+do_training: True # set to False if only preprocessing data
+do_testing: False # set to True to run evaluation on test data after training
+
+model:
+  beam_size: 4
+  len_pen: 0.6
+  multilingual: False
+  max_generation_delta: 5
+  label_smoothing: 0.1
+  shared_tokenizer: True # train tokenizer model across src and tgt train data
+  preproc_out_dir: null # path to store data preprocessing outputs
+  src_language: 'en'
+  tgt_language: 'de'
+
+  train_ds:
+    src_file_name: null
+    tgt_file_name: null
+    use_tarred_dataset: False # if true tar_file_name and meta_file_name will be used (or created automatically) 
+    # config for preprocessing training data and creating a tarred datset automatically
+    tar_file_prefix: parallel # prefix for tar file names
+    tar_files: null # if data has already been preprocessed (rest of config ignored)
+    metadata_file: null # metadata for tarred dataset
+    lines_per_dataset_fragment: 1000000 # Number of lines to consider for bucketing and padding
+    num_batches_per_tarfile: 100 # Number of batches (pickle files) within each tarfile
+    tar_shuffle_n: 100 # How many samples to look ahead and load to be shuffled
+    shard_strategy: scatter # tarred dataset shard distribution strategy
+    n_preproc_jobs: -2 # number of processes to use for data preprocessing (-2 means all but 2)
+    tokens_in_batch: 512
+    clean: true
+    max_seq_length: 512
+    shuffle: true
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+    concat_sampling_technique: temperature # only used with ConcatTranslationDataset 
+    concat_sampling_temperature: 5 # only used with ConcatTranslationDataset 
+    concat_sampling_probabilities: null # only used with ConcatTranslationDataset 
+
+  validation_ds:
+    src_file_name: ???
+    tgt_file_name: ???
+    tokens_in_batch: 512
+    clean: false
+    max_seq_length: 512
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+
+  test_ds:
+    src_file_name: ???
+    tgt_file_name: ???
+    tokens_in_batch: 512
+    clean: false
+    max_seq_length: 512
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+
+  optim:
+    name: adam
+    lr: 0.001
+    betas:
+      - 0.9
+      - 0.98
+    weight_decay: 0.0
+    sched:
+      name: InverseSquareRootAnnealing
+      min_lr: 0.0
+      last_epoch: -1
+      warmup_ratio: 0.1
+
+  encoder_tokenizer:
+    library: yttm
+    tokenizer_model: null
+    vocab_size: null # vocab size for training bpe
+    bpe_dropout: null
+    vocab_file: null
+    special_tokens: null
+    training_sample_size: null # valid for sentencepiece tokenizer
+
+  decoder_tokenizer:
+    library: yttm
+    tokenizer_model: null
+    vocab_size: null # vocab size for training bpe
+    bpe_dropout: null
+    vocab_file: null
+    special_tokens: null
+    training_sample_size: null # valid for sentencepiece tokenizer
+
+  encoder:
+    library: nemo
+    model_name: null
+    pretrained: false
+    max_sequence_length: 512
+    num_token_types: 0
+    embedding_dropout: 0.1
+    learn_positional_encodings: false
+    hidden_size: 512
+    num_layers: 6
+    inner_size: 2048
+    num_attention_heads: 8
+    ffn_dropout: 0.1
+    attn_score_dropout: 0.1
+    attn_layer_dropout: 0.1
+    hidden_act: relu
+    mask_future: false
+    pre_ln: false
+    pre_ln_final_layer_norm: true
+
+  decoder:
+    library: nemo
+    model_name: null
+    pretrained: false
+    max_sequence_length: 512
+    num_token_types: 0
+    embedding_dropout: 0.1
+    learn_positional_encodings: false
+    hidden_size: 512
+    inner_size: 2048
+    num_layers: 6
+    num_attention_heads: 8
+    ffn_dropout: 0.1
+    attn_score_dropout: 0.1
+    attn_layer_dropout: 0.1
+    hidden_act: relu
+    pre_ln: false
+    pre_ln_final_layer_norm: true
+
+  head:
+    num_layers: 1
+    activation: relu
+    log_softmax: true
+    dropout: 0.0
+    use_transformer_init: true
+
+trainer:
+  gpus: 4
+  num_nodes: 1
+  max_epochs: 200
+  amp_level: O2 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
+  accelerator: ddp
+  checkpoint_callback: False
+  logger: False
+  log_every_n_steps: 50  # Interval of logging.
+  check_val_every_n_epoch: 1
+
+exp_manager:
+  name: AAYNBase
+  files_to_copy: []

tts/NeMo/examples/nlp/machine_translation/conf/huggingface.yaml

@@ -0,0 +1,132 @@
+name: HuggingFaceEncoder
+do_training: True # set to False if only preprocessing data
+do_testing: False # set to True to run evaluation on test data after training
+
+model:
+  beam_size: 4
+  len_pen: 0.6
+  max_generation_delta: 5
+  label_smoothing: 0.1
+  shared_tokenizer: false
+  preproc_out_dir: null
+  src_language: 'en'
+  tgt_language: 'de'
+
+  train_ds:
+    src_file_name: null
+    tgt_file_name: null
+    use_tarred_dataset: False # if true tar_file_name and meta_file_name will be used (or created automatically) 
+    # config for preprocessing training data and creating a tarred datset automatically
+    tar_file_prefix: parallel # prefix for tar file names
+    tar_files: null # if data has already been preprocessed (rest of config ignored)
+    metadata_file: null # metadata for tarred dataset
+    lines_per_dataset_fragment: 1000000 # Number of lines to consider for bucketing and padding
+    num_batches_per_tarfile: 100 # Number of batches (pickle files) within each tarfile
+    tar_shuffle_n: 100 # How many samples to look ahead and load to be shuffled
+    shard_strategy: scatter # tarred dataset shard distribution strategy
+    n_preproc_jobs: -2 # number of processes to use for data preprocessing (-2 means all but 2)
+    tokens_in_batch: 512
+    clean: true
+    max_seq_length: 512
+    shuffle: true
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+
+  validation_ds:
+    src_file_name: null
+    tgt_file_name: null
+    tokens_in_batch: 512
+    clean: false
+    max_seq_length: 512
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+
+  test_ds:
+    src_file_name: null
+    tgt_file_name: null
+    tokens_in_batch: 512
+    clean: false
+    max_seq_length: 512
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+
+  optim:
+    name: adam
+    lr: 0.001
+    betas:
+      - 0.9
+      - 0.98
+    weight_decay: 0.0
+    sched:
+      name: InverseSquareRootAnnealing
+      min_lr: 0.0
+      last_epoch: -1
+      warmup_ratio: 0.1
+
+  encoder_tokenizer:
+    library: huggingface
+    tokenizer_model: null
+    vocab_file: null
+    special_tokens: null
+    vocab_size: null
+
+  decoder_tokenizer:
+    library: yttm
+    tokenizer_model: null
+    vocab_file: null
+    special_tokens: null
+    vocab_size: null
+
+  encoder:
+    library: huggingface
+    model_name: bert-base-uncased
+    pretrained: false
+
+  decoder:
+    library: nemo
+    model_name: null
+    pretrained: false
+    max_sequence_length: 512
+    num_token_types: 2
+    embedding_dropout: 0.1
+    learn_positional_encodings: false
+    hidden_size: 512
+    inner_size: 2048
+    num_layers: 6
+    num_attention_heads: 8
+    ffn_dropout: 0.1
+    attn_score_dropout: 0.1
+    attn_layer_dropout: 0.1
+    hidden_act: relu
+    pre_ln: false
+
+  head:
+    num_layers: 1
+    activation: relu
+    log_softmax: true
+    dropout: 0.0
+    use_transformer_init: true
+
+trainer:
+  gpus: 4
+  num_nodes: 1
+  max_epochs: 200
+  amp_level: O2 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
+  accelerator: ddp
+  checkpoint_callback: False
+  logger: False
+  log_every_n_steps: 50  # Interval of logging.
+  check_val_every_n_epoch: 1
+
+exp_manager:
+  name: HuggingFaceEncoder
+  files_to_copy: []

tts/NeMo/examples/nlp/machine_translation/create_tarred_monolingual_dataset.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+
+from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='NMT dataset pre-processing')
+    parser.add_argument('--tokenizer_model', type=str, required=True, help='Path to tokenizer model')
+    parser.add_argument('--tokenizer_name', type=str, default='yttm', help='BPE Tokenizer Name, Options: [yttm]')
+    parser.add_argument('--bpe_droput', type=float, default=0.0, help='BPE dropout to use')
+    parser.add_argument('--clean', action="store_true", help='Whether to clean dataset based on length diff')
+    parser.add_argument('--pkl_file_prefix', type=str, default='parallel', help='Prefix for tar and pickle files')
+    parser.add_argument('--fname', type=str, required=True, help='Path to monolingual data file')
+    parser.add_argument('--out_dir', type=str, required=True, help='Path to store dataloader and tokenizer models')
+    parser.add_argument('--max_seq_length', type=int, default=512, help='Max Sequence Length')
+    parser.add_argument('--min_seq_length', type=int, default=1, help='Min Sequence Length')
+    parser.add_argument('--tokens_in_batch', type=int, default=16000, help='# Tokens per batch per GPU')
+    parser.add_argument(
+        '--lines_per_dataset_fragment',
+        type=int,
+        default=1000000,
+        help='Number of lines to consider for bucketing and padding',
+    )
+    parser.add_argument(
+        '--num_batches_per_tarfile',
+        type=int,
+        default=1000,
+        help='Number of batches (pickle files) within each tarfile',
+    )
+
+    args = parser.parse_args()
+    if not os.path.exists(args.out_dir):
+        os.mkdir(args.out_dir)
+    if not os.path.exists(args.tokenizer_model):
+        assert FileNotFoundError("Could not find tokenizer model %s" % (args.tokenizer))
+
+    tokenizer_model = MTDataPreproc.get_monolingual_tokenizer(
+        tokenizer_name=args.tokenizer_name, tokenizer_model=args.tokenizer_model, bpe_dropout=args.bpe_droput
+    )
+
+    MTDataPreproc.preprocess_monolingual_dataset(
+        clean=args.clean,
+        fname=args.fname,
+        out_dir=args.out_dir,
+        tokenizer=tokenizer_model,
+        max_seq_length=args.max_seq_length,
+        min_seq_length=args.min_seq_length,
+        tokens_in_batch=args.tokens_in_batch,
+        lines_per_dataset_fragment=args.lines_per_dataset_fragment,
+        num_batches_per_tarfile=args.num_batches_per_tarfile,
+        pkl_file_prefix=args.pkl_file_prefix,
+        global_rank=0,
+        world_size=1,
+    )

tts/NeMo/examples/nlp/machine_translation/create_tarred_parallel_dataset.py

@@ -0,0 +1,143 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+
+from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='NMT dataset pre-processing')
+    parser.add_argument('--shared_tokenizer', action="store_true", help='Whether to share encoder/decoder tokenizers')
+    parser.add_argument('--clean', action="store_true", help='Whether to clean dataset based on length diff')
+    parser.add_argument('--tar_file_prefix', type=str, default='parallel', help='Prefix for tar files')
+    parser.add_argument('--bpe_dropout', type=float, default=0.1, help='Whether to share encoder/decoder tokenizers')
+    parser.add_argument('--src_fname', type=str, required=True, help='Path to the source file')
+    parser.add_argument('--tgt_fname', type=str, required=True, help='Path to the target file')
+    parser.add_argument('--out_dir', type=str, required=True, help='Path to store dataloader and tokenizer models')
+    parser.add_argument('--encoder_model_name', type=str, default=None, help='For use with pretrained encoders')
+    parser.add_argument(
+        '--decoder_model_name', type=str, default=None, help='For use with pretrained decoders (not yet supported)'
+    )
+    parser.add_argument(
+        '--encoder_tokenizer_model', type=str, default='None', help='Path to pre-trained encoder tokenizer model'
+    )
+    parser.add_argument(
+        '--encoder_tokenizer_name', type=str, default='yttm', help='Encoder BPE Tokenizer Name, Options: [yttm]'
+    )
+    parser.add_argument('--encoder_tokenizer_vocab_size', type=int, default=32000, help='Encoder Vocab size after BPE')
+    parser.add_argument(
+        '--encoder_tokenizer_coverage', type=float, default=0.999, help='Encoder Character coverage for BPE'
+    )
+    parser.add_argument('--encoder_tokenizer_bpe_dropout', type=float, default=0.1, help='Encoder BPE dropout prob')
+    parser.add_argument(
+        '--decoder_tokenizer_model', type=str, default='None', help='Path to pre-trained decoder tokenizer model'
+    )
+    parser.add_argument(
+        '--decoder_tokenizer_name', type=str, default='yttm', help='Encoder BPE Tokenizer Name, Options: [yttm]'
+    )
+    parser.add_argument('--decoder_tokenizer_vocab_size', type=int, default=32000, help='Encoder Vocab size after BPE')
+    parser.add_argument(
+        '--decoder_tokenizer_coverage', type=float, default=0.999, help='Encoder Character coverage for BPE'
+    )
+    parser.add_argument('--decoder_tokenizer_bpe_dropout', type=float, default=0.1, help='Encoder BPE dropout prob')
+    parser.add_argument('--max_seq_length', type=int, default=512, help='Max Sequence Length')
+    parser.add_argument('--min_seq_length', type=int, default=1, help='Min Sequence Length')
+    parser.add_argument('--tokens_in_batch', type=int, default=16000, help='# Tokens per batch per GPU')
+    parser.add_argument('--coverage', type=float, default=0.999, help='BPE character coverage [0-1]')
+    parser.add_argument(
+        '--lines_per_dataset_fragment',
+        type=int,
+        default=1000000,
+        help='Number of lines to consider for bucketing and padding',
+    )
+    parser.add_argument(
+        '--num_batches_per_tarfile',
+        type=int,
+        default=1000,
+        help='Number of batches (pickle files) within each tarfile',
+    )
+    parser.add_argument(
+        '--n_preproc_jobs', type=int, default=-2, help='Number of processes to use for creating the tarred dataset.',
+    )
+
+    args = parser.parse_args()
+    if not os.path.exists(args.out_dir):
+        os.mkdir(args.out_dir)
+
+    if (
+        args.encoder_tokenizer_model != 'None'
+        and args.decoder_tokenizer_model == 'None'
+        or args.decoder_tokenizer_model != 'None'
+        and args.encoder_tokenizer_model == 'None'
+    ):
+        if args.shared_tokenizer:
+            raise ValueError(
+                '''
+                If using a pre-trained shared tokenizer,
+                both encoder and decoder tokenizers must be the same
+                '''
+            )
+        else:
+            raise ValueError('Both encoder and decoder pre-trained tokenizer models must be specified')
+
+    if args.encoder_tokenizer_model == 'None' and args.decoder_tokenizer_model == 'None':
+        encoder_tokenizer_model, decoder_tokenizer_model = MTDataPreproc.train_tokenizers(
+            out_dir=args.out_dir,
+            src_fname=args.src_fname,
+            tgt_fname=args.tgt_fname,
+            shared_tokenizer=args.shared_tokenizer,
+            encoder_tokenizer_name=args.encoder_tokenizer_name,
+            encoder_tokenizer_vocab_size=args.encoder_tokenizer_vocab_size,
+            encoder_tokenizer_coverage=args.encoder_tokenizer_coverage,
+            decoder_tokenizer_name=args.decoder_tokenizer_name,
+            decoder_tokenizer_vocab_size=args.decoder_tokenizer_vocab_size,
+            decoder_tokenizer_coverage=args.decoder_tokenizer_coverage,
+            global_rank=0,
+        )
+    else:
+        encoder_tokenizer_model, decoder_tokenizer_model = args.encoder_tokenizer_model, args.decoder_tokenizer_model
+
+    encoder_tokenizer, decoder_tokenizer = MTDataPreproc.get_enc_dec_tokenizers(
+        encoder_tokenizer_name=args.encoder_tokenizer_name,
+        encoder_tokenizer_model=encoder_tokenizer_model,
+        encoder_bpe_dropout=args.encoder_tokenizer_bpe_dropout,
+        decoder_tokenizer_name=args.decoder_tokenizer_name,
+        decoder_tokenizer_model=decoder_tokenizer_model,
+        decoder_bpe_dropout=args.decoder_tokenizer_bpe_dropout,
+    )
+
+    _, _ = MTDataPreproc.preprocess_parallel_dataset(
+        clean=args.clean,
+        src_fname=args.src_fname,
+        tgt_fname=args.tgt_fname,
+        out_dir=args.out_dir,
+        encoder_tokenizer_name=args.encoder_tokenizer_name,
+        encoder_model_name=args.encoder_model_name,
+        encoder_tokenizer_model=encoder_tokenizer_model,
+        encoder_bpe_dropout=args.encoder_tokenizer_bpe_dropout,
+        decoder_tokenizer_name=args.decoder_tokenizer_name,
+        decoder_model_name=args.decoder_model_name,
+        decoder_tokenizer_model=decoder_tokenizer_model,
+        decoder_bpe_dropout=args.decoder_tokenizer_bpe_dropout,
+        max_seq_length=args.max_seq_length,
+        min_seq_length=args.min_seq_length,
+        tokens_in_batch=args.tokens_in_batch,
+        lines_per_dataset_fragment=args.lines_per_dataset_fragment,
+        num_batches_per_tarfile=args.num_batches_per_tarfile,
+        tar_file_prefix=args.tar_file_prefix,
+        global_rank=0,
+        world_size=1,
+        n_jobs=args.n_preproc_jobs,
+    )

tts/NeMo/examples/nlp/machine_translation/enc_dec_nmt.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Optional
+
+from omegaconf import OmegaConf
+from pytorch_lightning import Trainer
+
+from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc
+from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import MTEncDecModelConfig
+from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel
+from nemo.core.config import hydra_runner
+from nemo.core.config.modelPT import NemoConfig
+from nemo.core.config.pytorch_lightning import TrainerConfig
+from nemo.utils import logging
+from nemo.utils.config_utils import update_model_config
+from nemo.utils.exp_manager import ExpManagerConfig, exp_manager
+
+
+"""
+Usage:
+ 1. If you need to start docker and install NeMo, otherwise skip this step:
+ 
+    a. ```docker run --gpus all -it --rm -v /home/okuchaiev/repos/NeMo/:/NeMo -p 6006:6006  -v /mnt:/mnt --shm-size=16g --ulimit memlock=-1 --ulimit stack=67108864 --device=/dev/snd nvcr.io/nvidia/pytorch:20.11-py3```
+    b. ```cd /NeMo```
+    c. ```./reinstall.sh```
+ 
+ 2. Train a new tokenizer (or use pre-trained one):
+    ```yttm bpe --data /mnt/D1/Data/NMT/wmt16_de_en/train.clean.en-de.shuffled.common --model tokenizer.BPE.8192.model --vocab_size 8192```
+
+(To use WANDB, optionally, do login first)
+``wandb login [YOUR WANDB login]``
+    
+ 3. Start training:
+ 
+
+ (This example for "base" model on 2 GPUs for 150000 steps with batch size of 12500 tokens per GPU)
+ 
+ python enc_dec_nmt.py \
+      --config-path=conf \
+      --config-name=aayn_base \
+      trainer.gpus=[0,1] \
+      ~trainer.max_epochs \
+      +trainer.max_steps=150000 \
+      model.beam_size=4 \
+      model.max_generation_delta=5 \
+      model.label_smoothing=0.1 \
+      model.encoder_tokenizer.tokenizer_model=tokenizer.BPE.8192.model  \
+      model.decoder_tokenizer.tokenizer_model=tokenizer.BPE.8192.model  \
+      model.encoder.num_layers=6 \
+      model.encoder.hidden_size=512 \
+      model.encoder.inner_size=2048 \
+      model.encoder.num_attention_heads=8 \
+      model.encoder.ffn_dropout=0.1 \
+      model.decoder.num_layers=6 \
+      model.decoder.hidden_size=512 \
+      model.decoder.inner_size=2048 \
+      model.decoder.num_attention_heads=8 \
+      model.decoder.ffn_dropout=0.1 \
+      model.train_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/train.clean.de.shuffled \
+      model.train_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/train.clean.en.shuffled \
+      model.train_ds.tokens_in_batch=12500 \
+      model.validation_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.ref \
+      model.validation_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.src \
+      model.validation_ds.tokens_in_batch=8192 \
+      model.test_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.ref \
+      model.test_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.src \
+      model.optim.lr=0.001  \
+      model.optim.sched.warmup_ratio=0.05 \
+      +exp_manager.create_wandb_logger=True \
+      +exp_manager.wandb_logger_kwargs.name=TEST-nmt-base \
+      +exp_manager.wandb_logger_kwargs.project=nmt-de-en \
+      +exp_manager.create_checkpoint_callback=True \
+      +exp_manager.checkpoint_callback_params.monitor=val_sacreBLEU \
+      +exp_manager.exp_dir=nmt_base \
+      +exp_manager.checkpoint_callback_params.mode=max 
+"""
+
+
+@dataclass
+class MTEncDecConfig(NemoConfig):
+    name: Optional[str] = 'MTEncDec'
+    do_training: bool = True
+    do_testing: bool = False
+    model: MTEncDecModelConfig = MTEncDecModelConfig()
+    trainer: Optional[TrainerConfig] = TrainerConfig()
+    exp_manager: Optional[ExpManagerConfig] = ExpManagerConfig(name='MTEncDec', files_to_copy=[])
+
+
+@hydra_runner(config_path="conf", config_name="aayn_base")
+def main(cfg: MTEncDecConfig) -> None:
+    # merge default config with user specified config
+    default_cfg = MTEncDecConfig()
+    cfg = update_model_config(default_cfg, cfg)
+    logging.info("\n\n************** Experiment configuration ***********")
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+
+    # training is managed by PyTorch Lightning
+    trainer = Trainer(**cfg.trainer)
+
+    # tokenizers will be trained and and tarred training data will be created if needed
+    # model config is then updated
+    if cfg.model.preproc_out_dir is not None:
+        MTDataPreproc(cfg=cfg.model, trainer=trainer)
+
+    # experiment logs, checkpoints, and auto-resume are managed by exp_manager and PyTorch Lightning
+    exp_manager(trainer, cfg.exp_manager)
+
+    # everything needed to train translation models is encapsulated in the NeMo MTEncdDecModel
+    mt_model = MTEncDecModel(cfg.model, trainer=trainer)
+
+    logging.info("\n\n************** Model parameters and their sizes ***********")
+    for name, param in mt_model.named_parameters():
+        print(name, param.size())
+    logging.info("***********************************************************\n\n")
+
+    if cfg.do_training:
+        trainer.fit(mt_model)
+
+    if cfg.do_testing:
+        trainer.test(mt_model)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/machine_translation/nmt_transformer_infer.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Given NMT model's .nemo file, this script can be used to translate text.
+USAGE Example:
+1. Obtain text file in src language. You can use sacrebleu to obtain standard test sets like so:
+    sacrebleu -t wmt14 -l de-en --echo src > wmt14-de-en.src
+2. Translate:
+    python nmt_transformer_infer.py --model=[Path to .nemo file] --srctext=wmt14-de-en.src --tgtout=wmt14-de-en.pre
+"""
+
+
+from argparse import ArgumentParser
+
+import torch
+
+import nemo.collections.nlp as nemo_nlp
+from nemo.collections.nlp.modules.common.transformer import (
+    BeamSearchSequenceGenerator,
+    BeamSearchSequenceGeneratorWithLanguageModel,
+)
+from nemo.utils import logging
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument("--model", type=str, required=True, help="")
+    parser.add_argument("--srctext", type=str, required=True, help="")
+    parser.add_argument("--tgtout", type=str, required=True, help="")
+    parser.add_argument("--batch_size", type=int, default=256, help="")
+    parser.add_argument("--beam_size", type=int, default=4, help="")
+    parser.add_argument("--len_pen", type=float, default=0.6, help="")
+    parser.add_argument("--max_delta_length", type=int, default=5, help="")
+    parser.add_argument("--target_lang", type=str, default=None, help="")
+    parser.add_argument("--source_lang", type=str, default=None, help="")
+    # shallow fusion specific parameters
+    parser.add_argument("--lm_model", type=str, default=None, help="")
+    parser.add_argument("--fusion_coef", type=float, default=0.0, help="")
+
+    args = parser.parse_args()
+    torch.set_grad_enabled(False)
+    if args.model.endswith(".nemo"):
+        logging.info("Attempting to initialize from .nemo file")
+        model = nemo_nlp.models.machine_translation.MTEncDecModel.restore_from(restore_path=args.model)
+        src_text = []
+        tgt_text = []
+    else:
+        raise NotImplemented(f"Only support .nemo files, but got: {args.model}")
+
+    if torch.cuda.is_available():
+        model = model.cuda()
+
+    if args.lm_model is not None:
+        lm_model = nemo_nlp.models.language_modeling.TransformerLMModel.restore_from(restore_path=args.lm_model).eval()
+        model.beam_search = BeamSearchSequenceGeneratorWithLanguageModel(
+            embedding=model.decoder.embedding,
+            decoder=model.decoder.decoder,
+            log_softmax=model.log_softmax,
+            bos=model.decoder_tokenizer.bos_id,
+            pad=model.decoder_tokenizer.pad_id,
+            eos=model.decoder_tokenizer.eos_id,
+            language_model=lm_model,
+            fusion_coef=args.fusion_coef,
+            max_sequence_length=model.decoder.max_sequence_length,
+            beam_size=args.beam_size,
+            len_pen=args.len_pen,
+            max_delta_length=args.max_delta_length,
+        )
+    else:
+        model.beam_search = BeamSearchSequenceGenerator(
+            embedding=model.decoder.embedding,
+            decoder=model.decoder.decoder,
+            log_softmax=model.log_softmax,
+            bos=model.decoder_tokenizer.bos_id,
+            pad=model.decoder_tokenizer.pad_id,
+            eos=model.decoder_tokenizer.eos_id,
+            max_sequence_length=model.decoder.max_sequence_length,
+            beam_size=args.beam_size,
+            len_pen=args.len_pen,
+            max_delta_length=args.max_delta_length,
+        )
+
+    logging.info(f"Translating: {args.srctext}")
+
+    count = 0
+    with open(args.srctext, 'r') as src_f:
+        for line in src_f:
+            src_text.append(line.strip())
+            if len(src_text) == args.batch_size:
+                res = model.translate(text=src_text, source_lang=args.source_lang, target_lang=args.target_lang)
+                if len(res) != len(src_text):
+                    print(len(res))
+                    print(len(src_text))
+                    print(res)
+                    print(src_text)
+                tgt_text += res
+                src_text = []
+            count += 1
+            # if count % 300 == 0:
+            #    print(f"Translated {count} sentences")
+        if len(src_text) > 0:
+            tgt_text += model.translate(text=src_text, source_lang=args.source_lang, target_lang=args.target_lang)
+
+    with open(args.tgtout, 'w') as tgt_f:
+        for line in tgt_text:
+            tgt_f.write(line + "\n")
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/README.rst

@@ -0,0 +1,7 @@
+**STEPS 2 RUN**
+===============
+0. (Make sure you have Flask installed - ``pip install -r requirements.txt``)
+1. Edit "config.json" file to only contain models you need. If model's location starts with "NGC/" - it will load this model from NVIDIA's NGC. Otherwise, specify full path to .nemo file.
+4. To run: ``python nmt_service.py``
+5. To translate: ``http://127.0.0.1:5000/translate?text=Frohe%20Weihnachten`` (here %20 means space)
+6. Run web UI: ``python -m http.server``

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/config.json

@@ -0,0 +1,12 @@
+{
+  "en-de": "NGC/nmt_en_de_transformer12x2",
+  "de-en": "NGC/nmt_de_en_transformer12x2",
+  "en-es": "NGC/nmt_en_es_transformer12x2",
+  "es-en": "NGC/nmt_es_en_transformer12x2",
+  "en-ru": "NGC/nmt_en_ru_transformer6x6",
+  "ru-en": "NGC/nmt_ru_en_transformer6x6",
+  "en-fr": "NGC/nmt_en_fr_transformer12x2",
+  "fr-en": "NGC/nmt_fr_en_transformer12x2",
+  "en-zh": "NGC/nmt_en_zh_transformer6x6",
+  "zh-en": "NGC/nmt_zh_en_transformer6x6"
+}

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/index.html

@@ -0,0 +1,140 @@
+<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="utf-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1">
+  <title>Neural Machine Translation Demo App</title>
+  <link rel="stylesheet" href="//code.jquery.com/ui/1.12.1/themes/base/jquery-ui.css">  
+  <link rel="stylesheet" href="/resources/demos/style.css">
+  <style>
+    .ui-resizable-se {
+      bottom: 17px;
+    }
+  </style>
+  <link rel="stylesheet" href="style.css">
+  <script src="https://code.jquery.com/jquery-1.12.4.js"></script>
+  <script src="https://code.jquery.com/ui/1.12.1/jquery-ui.js"></script>
+  <script>
+  $( function() {
+    $( "#tgt_text" ).resizable({
+      handles: "se"
+    });
+  } );
+  $( function() {
+    $( "#src_text" ).resizable({
+      handles: "se"
+    });
+  } );
+  $( function() {
+    $( "#translate_button" ).button();
+    $( "#translate_button" ).click( function( event ) {
+        $.ajax({
+          url: "http://10.110.40.187:5000/translate",
+          type: "GET",                    
+          data: {
+            text: $( "#src_text" ).val(),
+            langpair: $('input[type=radio][name=langpair]:checked').attr('id'),
+            do_moses: true
+          },
+          success: function (result) {            
+            $( "#tgt_text" ).val(result.translation);
+          },
+          error: function (err) {
+            alert('Sorry, can not reach translation service');            
+          },
+        })
+    } );
+  } );
+  </script>
+</head>
+<body>
+<div class="contact-us">
+  <h1>Neural Machine Translation Demo</h1>
+  <p>
+  <br/>
+  <h4>Running NeMo commit id: v1.0.2  </h4>
+  <br/>
+  </p>
+  <p>
+  <fieldset>
+    <legend>Select translation direction: </legend>
+    <table>
+      <tr>
+        <td>
+          <label for="en-de">English to German</label>
+          <input type="radio" name="langpair" id="en-de">
+        </td>
+        <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+        <td>
+          <label for="de-en">German to English</label>
+          <input type="radio" name="langpair" id="de-en">
+        </td>
+        <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+        <td>
+          <label for="es-en">Spanish to English</label>
+          <input type="radio" name="langpair" id="es-en">
+        </td>
+	<td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+	<td>
+	  <label for="en-es" selected="selected">English to Spanish</label>
+	  <input type="radio" name="langpair" id="en-es" checked>
+  </td> 
+      </tr>
+      <tr>
+         <td>
+           <label for="en-ru">English to Russian</label>
+           <input type="radio" name="langpair" id="en-ru">
+         </td>
+         <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+         <td>
+           <label for="ru-en">Russian to English</label>
+           <input type="radio" name="langpair" id="ru-en">
+         </td>
+         <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+         <td>
+           <label for="en-fr">English to French</label>
+	   <input type="radio" name="langpair" id="en-fr">
+	 </td>
+	 <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+	 <td>
+	  <label for="fr-en">French to English</label>
+          <input type="radio" name="langpair" id="fr-en">
+         </td>
+         <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+         <td>
+          <label for="en-zh">English to Chinese</label>
+          <input type="radio" name="langpair" id="en-zh">
+         </td>
+	 <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+         <td>
+          <label for="zh-en">Chinese to English</label>
+          <input type="radio" name="langpair" id="zh-en">
+         </td>
+      </tr>
+    </table>
+  </fieldset>
+  <div>
+    <table>
+      <tr>
+        <td>
+          <textarea id="src_text" rows="12" cols="60">Type something here, chose language pair and click "Translate" button.</textarea>
+        </td>
+      </tr>
+      <tr>
+        <td>
+          <button id="translate_button">Translate !</button>
+        </td>
+      </tr>
+      <tr>
+        <td>
+          <textarea id="tgt_text" rows="12" cols="60" readonly></textarea>
+        </td>
+      </tr>
+    </table>
+  </div>
+  </p>
+</div>
+ 
+ 
+</body>
+</html>

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/nmt_service.py

@@ -0,0 +1,100 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import time
+
+import flask
+import torch
+from flask import Flask, json, request
+from flask_cors import CORS
+
+import nemo.collections.nlp as nemo_nlp
+from nemo.utils import logging
+
+MODELS_DICT = {}
+
+model = None
+api = Flask(__name__)
+CORS(api)
+
+
+def initialize(config_file_path: str):
+    """
+    Loads 'language-pair to NMT model mapping'
+    """
+    __MODELS_DICT = None
+
+    logging.info("Starting NMT service")
+    logging.info(f"I will attempt to load all the models listed in {config_file_path}.")
+    logging.info(f"Edit {config_file_path} to disable models you don't need.")
+    if torch.cuda.is_available():
+        logging.info("CUDA is available. Running on GPU")
+    else:
+        logging.info("CUDA is not available. Defaulting to CPUs")
+
+    # read config
+    with open(config_file_path) as f:
+        __MODELS_DICT = json.load(f)
+
+    if __MODELS_DICT is not None:
+        for key, value in __MODELS_DICT.items():
+            logging.info(f"Loading model for {key} from file: {value}")
+            if value.startswith("NGC/"):
+                model = nemo_nlp.models.machine_translation.MTEncDecModel.from_pretrained(model_name=value[4:])
+            else:
+                model = nemo_nlp.models.machine_translation.MTEncDecModel.restore_from(restore_path=value)
+            if torch.cuda.is_available():
+                model = model.cuda()
+            MODELS_DICT[key] = model
+    else:
+        raise ValueError("Did not find the config.json or it was empty")
+    logging.info("NMT service started")
+
+
+@api.route('/translate', methods=['GET', 'POST', 'OPTIONS'])
+def get_translation():
+    try:
+        time_s = time.time()
+        langpair = request.args["langpair"]
+        src = request.args["text"]
+        do_moses = request.args.get('do_moses', False)
+        if langpair in MODELS_DICT:
+            if do_moses:
+                result = MODELS_DICT[langpair].translate(
+                    [src], source_lang=langpair.split('-')[0], target_lang=langpair.split('-')[1]
+                )
+            else:
+                result = MODELS_DICT[langpair].translate([src])
+
+            duration = time.time() - time_s
+            logging.info(
+                f"Translated in {duration}. Input was: {request.args['text']} <############> Translation was: {result[0]}"
+            )
+            res = {'translation': result[0]}
+            response = flask.jsonify(res)
+            response.headers.add('Access-Control-Allow-Origin', '*')
+            return response
+
+        else:
+            logging.error(f"Got the following langpair: {langpair} which was not found")
+    except Exception as ex:
+        res = {'translation': str(ex)}
+        response = flask.jsonify(res)
+        response.headers.add('Access-Control-Allow-Origin', '*')
+        return res
+
+
+if __name__ == '__main__':
+    initialize('config.json')
+    api.run(host='0.0.0.0')

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/requirements.txt

@@ -0,0 +1,3 @@
+nemo_toolkit[all]>=1.0.0rc1
+flask
+flask_cors

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/style.css

@@ -0,0 +1,111 @@
+@import url("https://fonts.googleapis.com/css?family=Spartan&display=swap");
+* {
+  margin: 0;
+  padding: 0;
+}
+
+body {
+  /*background: url("https://img.wallpapersafari.com/desktop/1680/1050/65/18/8AYkSt.jpg") center center;*/
+    background: url("https://images5.alphacoders.com/550/550948.jpg") center center;
+  background-size: cover;
+  width: 100vw;
+  height: 100vh;
+  display: grid;
+  align-items: center;
+  justify-items: center;
+  font-size: 12pt;
+  font-family: "Spartan";
+  color: #ffffff;
+}
+
+.contact-us {
+  /*background: #1f1d1d;*/
+  padding: 50px 100px;
+  border-top: 10px solid #229955;
+}
+
+label, input, textarea {
+  display: block;
+  width: 100%;
+  font-size: 12pt;
+  line-height: 24pt;
+  font-family: "Spartan";
+}
+
+input {
+  margin-bottom: 24pt;
+}
+
+h1 {
+  color: #ffffff;
+  font-weight: bold;
+  font-size: 22pt;
+  line-height: 24pt;
+  font-style: italic;
+  margin: 0 0 0.5em 0;
+}
+
+span {
+  font-size: 8pt;
+}
+
+em {
+  color: #229955;
+  font-weight: bold;
+}
+
+input, textarea {
+  border: none;
+  border: 1px solid rgba(0, 0, 0, 0.1);
+  border-radius: 2px;
+  background: #f8f4e5;
+  padding-left: 5px;
+  outline: none;
+}
+
+input:focus, textarea:focus {
+  border: 1px solid #6bd4b1;
+}
+
+textarea {
+  resize: none;
+}
+
+fieldset {
+      overflow: hidden
+    }
+
+button {
+  display: block;
+  float: left;
+  line-height: 24pt;
+  padding: 0 20px;
+  border: none;
+  background: #f45702;
+  color: white;
+  letter-spacing: 2px;
+  transition: 0.2s all ease-in-out;
+  border-bottom: 2px solid transparent;
+  outline: none;
+}
+button:hover {
+  background: inherit;
+  color: #ea6153;
+  border-bottom: 2px solid #f45702;
+}
+
+::selection {
+  background: #ffc7b8;
+}
+
+input:-webkit-autofill,
+input:-webkit-autofill:hover,
+input:-webkit-autofill:focus,
+textarea:-webkit-autofill,
+textarea:-webkit-autofill:hover,
+textarea:-webkit-autofill:focus {
+  border: 1px solid #6bd4b1;
+  -webkit-text-fill-color: #2A293E;
+  -webkit-box-shadow: 0 0 0px 1000px #f8f4e5 inset;
+  transition: background-color 5000s ease-in-out 0s;
+}

tts/NeMo/examples/nlp/machine_translation/translate_ddp.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+from argparse import ArgumentParser
+
+import torch
+import torch.multiprocessing as mp
+from torch.utils.data import DataLoader
+
+from nemo.collections.nlp.data.language_modeling import TarredSentenceDataset
+from nemo.collections.nlp.data.machine_translation import TarredTranslationDataset
+from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel
+from nemo.utils import logging
+
+
+def get_args():
+    parser = ArgumentParser(description='Batch translation of sentences from a pre-trained model on multiple GPUs')
+    parser.add_argument("--model", type=str, required=True, help="Path to the .nemo translation model file")
+    parser.add_argument(
+        "--text2translate", type=str, required=True, help="Path to the pre-processed tarfiles for translation"
+    )
+    parser.add_argument("--result_dir", type=str, required=True, help="Folder to write translation results")
+    parser.add_argument(
+        "--twoside", action="store_true", help="Set flag when translating the source side of a parallel dataset"
+    )
+    parser.add_argument(
+        '--metadata_path', type=str, required=True, help="Path to the JSON file that contains dataset info"
+    )
+    parser.add_argument('--topk', type=int, default=500, help="Value of k for topk sampling")
+    parser.add_argument('--src_language', type=str, required=True, help="Source lang ID for detokenization")
+    parser.add_argument('--tgt_language', type=str, required=True, help="Target lang ID for detokenization")
+    parser.add_argument(
+        '--reverse_lang_direction',
+        action="store_true",
+        help="Reverse source and target language direction for parallel dataset",
+    )
+    parser.add_argument('--n_gpus', type=int, default=-1, help="Number of GPUs to use")
+    args = parser.parse_args()
+    return args
+
+
+def translate(rank, world_size, args):
+    if args.model.endswith(".nemo"):
+        logging.info("Attempting to initialize from .nemo file")
+        model = MTEncDecModel.restore_from(restore_path=args.model, map_location=f"cuda:{rank}")
+    elif args.model.endswith(".ckpt"):
+        logging.info("Attempting to initialize from .ckpt file")
+        model = MTEncDecModel.load_from_checkpoint(checkpoint_path=args.model, map_location=f"cuda:{rank}")
+    model.replace_beam_with_sampling(topk=args.topk)
+    model.eval()
+    if args.twoside:
+        dataset = TarredTranslationDataset(
+            text_tar_filepaths=args.text2translate,
+            metadata_path=args.metadata_path,
+            encoder_tokenizer=model.encoder_tokenizer,
+            decoder_tokenizer=model.decoder_tokenizer,
+            shuffle_n=100,
+            shard_strategy="scatter",
+            world_size=world_size,
+            global_rank=rank,
+            reverse_lang_direction=args.reverse_lang_direction,
+        )
+    else:
+        dataset = TarredSentenceDataset(
+            text_tar_filepaths=args.text2translate,
+            metadata_path=args.metadata_path,
+            tokenizer=model.encoder_tokenizer,
+            shuffle_n=100,
+            shard_strategy="scatter",
+            world_size=world_size,
+            global_rank=rank,
+        )
+    loader = DataLoader(dataset, batch_size=1)
+    result_dir = os.path.join(args.result_dir, f'rank{rank}')
+    os.makedirs(result_dir, exist_ok=True)
+    originals_file_name = os.path.join(result_dir, 'originals.txt')
+    translations_file_name = os.path.join(result_dir, 'translations.txt')
+    num_translated_sentences = 0
+
+    with open(originals_file_name, 'w') as of, open(translations_file_name, 'w') as tf:
+        for batch_idx, batch in enumerate(loader):
+            for i in range(len(batch)):
+                if batch[i].ndim == 3:
+                    batch[i] = batch[i].squeeze(dim=0)
+                batch[i] = batch[i].to(rank)
+            if args.twoside:
+                src_ids, src_mask, _, _, _ = batch
+            else:
+                src_ids, src_mask = batch
+            if batch_idx % 100 == 0:
+                logging.info(
+                    f"{batch_idx} batches ({num_translated_sentences} sentences) were translated by process with "
+                    f"rank {rank}"
+                )
+            num_translated_sentences += len(src_ids)
+            inputs, translations = model.batch_translate(src=src_ids, src_mask=src_mask)
+            for src, translation in zip(inputs, translations):
+                of.write(src + '\n')
+                tf.write(translation + '\n')
+
+
+def main() -> None:
+    args = get_args()
+    world_size = torch.cuda.device_count() if args.n_gpus == -1 else args.n_gpus
+    mp.spawn(translate, args=(world_size, args), nprocs=world_size, join=True)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/question_answering/conf/question_answering_squad_config.yaml

@@ -0,0 +1,135 @@
+# Question Answering with SQUAD
+name: &name QA
+
+pretrained_model: null # pretrained QAModel model from list_available_models()
+do_training: true # training mode, for testing change to false
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices
+  num_nodes: 1
+  max_epochs: 2 # the number of training epochs
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  precision: 16 # 16 to use AMP
+  amp_level: O1 # O1 or O2 if using AMP
+  accelerator: ddp
+  gradient_clip_val: 0.0
+  val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch
+  checkpoint_callback: false # provided by exp_manager
+  logger: false # provided by exp_manager
+  num_sanity_val_steps: 0
+  log_every_n_steps: 1  # Interval of logging.
+
+model:
+  nemo_path: null # exported .nemo path
+  dataset:
+    version_2_with_negative: false
+    # If true, the examples contain some that do not have an answer.
+    doc_stride: 128
+    # When splitting up a long document into chunks,
+    # how much stride to take between chunks.
+    max_query_length: 64
+    # The maximum number of tokens for the question.
+    # Questions longer than this will be truncated to
+    # this length.
+    max_seq_length: 384
+    # The maximum total input sequence length after
+    # WordPiece tokenization. Sequences longer than this
+    # will be truncated, and sequences shorter than this
+    # will be padded.
+    max_answer_length: 30
+    # The maximum length of an answer that can be
+    # generated. This is needed because the start
+    # and end predictions are not conditioned
+    # on one another.
+    null_score_diff_threshold: 0.0
+    # If null_score - best_non_null is greater than the threshold predict null.
+    n_best_size: 20
+    # The total number of n-best predictions to generate at testing.
+    use_cache: true
+    do_lower_case: true
+    # if true does lower case
+
+    num_workers:  2
+    pin_memory: false
+    drop_last: false
+
+  train_ds:
+    file: null # .json file
+    batch_size: 24 # per GPU
+    shuffle: true
+    num_samples: -1
+    # Default values for the following params are retrieved from dataset config section, but you may override them
+    num_workers: ${model.dataset.num_workers}
+    drop_last: ${model.dataset.drop_last}
+    pin_memory: ${model.dataset.pin_memory}
+
+  validation_ds:
+    file: null # .json file
+    batch_size: 24 # per GPU
+    shuffle: false
+    num_samples: -1    
+    # Default values for the following params are retrieved from dataset config section, but you may override them
+    num_workers: ${model.dataset.num_workers}
+    drop_last: ${model.dataset.drop_last}
+    pin_memory: ${model.dataset.pin_memory}
+
+  test_ds:
+    file: null # .json file
+    batch_size: 24 # per GPU
+    shuffle: false
+    num_samples: -1
+    # Default values for the following params are retrieved from dataset config section, but you may override them
+    num_workers: ${model.dataset.num_workers}
+    drop_last: ${model.dataset.drop_last}
+    pin_memory: ${model.dataset.pin_memory}
+    
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # only used if tokenizer is sentencepiece
+    special_tokens: null # expand the following to a dictionary if special tokens need to be added.
+    #  only necessary for adding transformer/bert-specific special tokens to tokenizer if the tokenizer does not already have these inherently.
+
+  language_model:
+    pretrained_model_name: bert-base-uncased # BERT-like model name
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null # if specified initializes model from scratch
+
+  token_classifier:
+    num_layers: 1
+    dropout: 0.0
+    num_classes: 2
+    activation: relu
+    log_softmax: false
+    use_transformer_init: true
+
+
+  optim:
+    name: adamw
+    lr: 3e-5
+    weight_decay: 0.0
+    sched:
+      name: SquareRootAnnealing
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+      # scheduler config override
+      warmup_steps: null
+      warmup_ratio: 0.0
+      last_epoch: -1
+
+exp_manager:
+  exp_dir: null # where to store logs and checkpoints
+  name: *name # name of experiment
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/question_answering/get_squad.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+import urllib.request
+
+from nemo.utils import logging
+
+
+class SquadDownloader:
+    def __init__(self, save_path):
+        self.save_path = save_path + '/squad'
+
+        if not os.path.exists(self.save_path):
+            os.makedirs(self.save_path)
+
+        if not os.path.exists(self.save_path + '/v1.1'):
+            os.makedirs(self.save_path + '/v1.1')
+
+        if not os.path.exists(self.save_path + '/v2.0'):
+            os.makedirs(self.save_path + '/v2.0')
+
+        self.download_urls = {
+            'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/train-v1.1.json': 'v1.1/train-v1.1.json',
+            'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/dev-v1.1.json': 'v1.1/dev-v1.1.json',
+            'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/train-v2.0.json': 'v2.0/train-v2.0.json',
+            'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/dev-v2.0.json': 'v2.0/dev-v2.0.json',
+        }
+
+    def download(self):
+        for item in self.download_urls:
+            url = item
+            file = self.download_urls[item]
+
+            logging.info('Downloading: %s', url)
+            if os.path.isfile(self.save_path + '/' + file):
+                logging.info('** Download file already exists, skipping download')
+            else:
+                response = urllib.request.urlopen(url)
+                with open(self.save_path + '/' + file, "wb") as handle:
+                    handle.write(response.read())
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Download Squad')
+    parser.add_argument(
+        '--destDir',
+        type=str,
+        required=False,
+        help='directory to store data',
+        default=os.path.split(os.path.abspath(__file__))[0],
+    )
+    args = parser.parse_args()
+    logging.info(args.destDir)
+    squad_dl = SquadDownloader(args.destDir)
+    squad_dl.download()

tts/NeMo/examples/nlp/question_answering/question_answering_squad.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""
+This script contains an example on how to train, evaluate and perform inference with the question answering model.
+The QAModel in NeMo supports extractive question answering problems for data in the SQuAD (https://rajpurkar.github.io/SQuAD-explorer/) format.
+
+***Data format***
+The QAModel requires a JSON file for each dataset split. 
+In the following we will show example for a training file. Each title has one or multiple paragraph entries, each consisting of the text - "context", and question-answer entries. Each question-answer entry has:
+* a question
+* a globally unique id
+* a boolean flag "is_impossible" which shows if the question is answerable or not
+* in case the question is answerable one answer entry, which contains the text span and its starting character index in the context. If not answerable, the "answers" list is empty
+
+The evaluation file follows the above format except for it can provide more than one answers to the same question. 
+The inference file follows the above format except for it does not require the "answers" and "is_impossible" keywords.
+
+
+***Downloading the dataset***
+Run ./NeMo/examples/nlp/question_answering/get_squad.py to download the SQuAD dataset:
+
+#   python get_squad.py --destDir=<PATH_TO_DATA>
+
+***Setting the configs***
+The model and the PT trainer are defined in a config file which declares multiple important sections.
+The most important ones are:
+    model: All arguments that are related to the Model - language model, tokenizer, token classifier, optimizer,
+            schedulers, and datasets/data loaders.
+    trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs,
+            precision level, etc.
+
+This script uses the `/examples/nlp/question_answering/conf/question_answering_squad_config.yaml` config file
+by default. You may update the config file from the file directly. The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+
+***Model Training***
+# python question_answering_squad.py
+    model.train_ds.file=<TRAIN_JSON_FILE>
+    model.validation_ds=<VAL_JSON_FILE>
+    trainer.max_epochs=<NUM_EPOCHS>
+    trainer.gpus=[<CHANGE_TO_GPU_YOU_WANT_TO_USE>]
+
+
+***Model Evaluation***
+Set `do_training=False` in the script and run:
+
+#   python question_answering_squad.py
+    model.test_file=<TEST_JSON_FILE>
+
+To load a pretrained checkpoint from cloud prior to training (e.g. for fine-tuning) or evaluation you can set cfg.from_pretrained=<MODEL_NAME>,
+e.g. MODEL_NAME='BERTBaseUncasedSQuADv1.1'. You can find all pretrained model names by using 
+QAModel.list_available_models(). To load a local checkpoint use qa_model.restore_from(<PATH_TO_CHECKPOINT>)
+
+
+***Model Inference***
+For inference use 
+    qa_model.inference(
+        file=<INFERENCE_JSON_FILE>,
+        batch_size=<BATCH_SIZE>,
+        output_nbest_file=<OPTIONAL_OUTPUT_FILE_FOR_NBEST_LIST>,
+        output_prediction_file=<OPTIONAL_OUTPUT_FILE_FOR_PREDICTION>
+    )
+
+More details on how to use this script can be found in
+./NeMo/tutorials/nlp/Question_Answering_Squad.ipynb
+"""
+
+import os
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.question_answering.qa_model import QAModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="question_answering_squad_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+
+    if not cfg.pretrained_model:
+        logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+        model = QAModel(cfg.model, trainer=trainer)
+    else:
+        logging.info(f'Loading pretrained model {cfg.pretrained_model}')
+        model = QAModel.from_pretrained(cfg.pretrained_model)
+        if cfg.do_training:
+            model.setup_training_data(train_data_config=cfg.model.train_ds)
+            model.setup_validation_data(val_data_config=cfg.model.validation_ds)
+
+    if cfg.do_training:
+        trainer.fit(model)
+        if cfg.model.nemo_path:
+            model.save_to(cfg.model.nemo_path)
+
+    if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.file is not None:
+        trainer.test(model)
+
+    # change to path if you want results to be written to file e.g. os.path.join(exp_dir, "output_nbest_file.txt")
+    output_nbest_file = None
+    # change to path if you want results to be written to file e.g.  os.path.join(exp_dir, "output_prediction_file.txt")
+    output_prediction_file = None
+    inference_samples = 5  # for test purposes. To use entire inference dataset set to -1
+    all_preds, all_nbests = model.inference(
+        file=cfg.model.validation_ds.file,
+        batch_size=1,
+        num_samples=inference_samples,
+        output_nbest_file=output_nbest_file,
+        output_prediction_file=output_prediction_file,
+    )
+
+    for _, item in all_preds.items():
+        print(f"question: {item[0]} answer: {item[1]}")
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/text2sparql/conf/text2sparql_config.yaml

@@ -0,0 +1,105 @@
+# Text2Sparql with BART
+
+name: &name Text2Sparql
+
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices
+  num_nodes: 1
+  max_epochs: 2 # the number of training epochs
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  accelerator: ddp
+  gradient_clip_val: 0.0
+  log_every_n_steps: 1
+  val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch
+  checkpoint_callback: false # provided by exp_manager
+  logger: false # provided by exp_manager
+
+model:
+  nemo_path: null # exported .nemo path
+  max_seq_length: 150
+  batch_size: 16
+  convert_labels: true # true if Bart, false otherwise (converts pad_id to -100 for masked loss)
+  data_dir: null
+
+  language_model:
+    pretrained_model_name: facebook/bart-base # huggingface end-to-end model name
+    pretrained_encoder_model_name: null # huggingface encoder model name
+    pretrained_decoder_model_name: null # huggingface decoder model name
+    lm_checkpoint: null
+    config: null
+    config_file: null # json file, precedence over config
+
+  encoder_tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # tokenizer model for sentencepiece
+    special_tokens: null
+    add_special_tokens: true
+
+  decoder_tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # tokenizer model for sentencepiece
+    special_tokens: null
+    add_special_tokens: true
+
+  train_ds:
+    filepath: ${model.data_dir}/train.tsv # path to data file
+    shuffle: true
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    filepath: ${model.data_dir}/test_easy.tsv # path to data file
+    shuffle: false
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  test_ds:
+    filepath: ${model.data_dir}/test_hard.tsv # path to data file
+    shuffle: false
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adamw
+    lr: 4e-5
+    weight_decay: 0.0
+
+    sched:
+      name: CosineAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.06
+      min_lr: 0.0
+      last_epoch: -1
+
+  generate:
+    max_length: ${model.max_seq_length}
+    num_beams: 1
+    length_penalty: 2.0
+    early_stopping: true
+    repetition_penalty: 1.0
+    do_sample: false
+    top_k: null
+    top_p: null
+    num_return_sequences: 1
+
+exp_manager:
+  exp_dir: null # where to store logs and checkpoints
+  name: *name # name of experiment
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/text2sparql/data/import_datasets.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2020, MeetKai Inc.  All rights reserved.
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script downloads Text2Sparql data and processes it into NeMo's neural machine translation dataset format.
+
+Text2Sparql data consists of 3 files which are saved to source_data_dir:
+    - train_queries_v3.tsv
+    - test_easy_queries_v3.tsv
+    - test_hard_queries_v3.tsv
+
+After processing, the script saves them to the target_data_dir as:
+    - train.tsv
+    - test_easy.tsv
+    - test_hard.tsv
+
+
+You may run it with:
+
+python import_datasets \
+    --source_data_dir ./text2sparql_src \
+    --target_data_dir ./text2sparql_tgt
+"""
+
+import argparse
+import csv
+import os
+from urllib.request import Request, urlopen
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import MODE_EXISTS_TMP, if_exist
+from nemo.utils import logging
+
+base_url = "https://m.meetkai.com/public_datasets/knowledge/"
+prefix_map = {
+    "train_queries_v3.tsv": "train.tsv",
+    "test_easy_queries_v3.tsv": "test_easy.tsv",
+    "test_hard_queries_v3.tsv": "test_hard.tsv",
+}
+
+
+def download_text2sparql(infold: str):
+    """Downloads text2sparql train, test_easy, and test_hard data
+
+    Args:
+        infold: save directory path
+    """
+    os.makedirs(infold, exist_ok=True)
+
+    for prefix in prefix_map:
+        url = base_url + prefix
+
+        logging.info(f"Downloading: {url}")
+        if if_exist(infold, [prefix]):
+            logging.info("** Download file already exists, skipping download")
+        else:
+            req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
+            with open(os.path.join(infold, prefix), "wb") as handle:
+                handle.write(urlopen(req, timeout=20).read())
+
+
+def process_text2sparql(infold: str, outfold: str, do_lower_case: bool):
+    """ Process and convert MeetKai's text2sparql datasets to NeMo's neural machine translation format.
+
+    Args:
+        infold: directory path to raw text2sparql data containing
+            train.tsv, test_easy.tsv, test_hard.tsv
+        outfold: output directory path to save formatted data for NeuralMachineTranslationDataset
+            the first line is header (sentence [tab] label)
+            each line should be [sentence][tab][label]
+        do_lower_case: if true, convert all sentences and labels to lower
+    """
+    logging.info(f"Processing Text2Sparql dataset and storing at: {outfold}")
+
+    os.makedirs(outfold, exist_ok=True)
+
+    dataset_name = "Text2Sparql"
+    for prefix in prefix_map:
+        input_file = os.path.join(infold, prefix)
+        output_file = os.path.join(outfold, prefix_map[prefix])
+
+        if if_exist(outfold, [prefix_map[prefix]]):
+            logging.info(f"** {MODE_EXISTS_TMP.format(prefix_map[prefix], dataset_name, output_file)}")
+            continue
+
+        if not if_exist(infold, [prefix]):
+            logging.info(f"** {prefix} of {dataset_name}" f" is skipped as it was not found")
+            continue
+
+        assert input_file != output_file, "input file cannot equal output file"
+        with open(input_file, "r") as in_file:
+            with open(output_file, "w") as out_file:
+                reader = csv.reader(in_file, delimiter="\t")
+
+                # replace headers
+                out_file.write("sentence\tlabel\n")
+                next(reader)
+
+                for line in reader:
+                    sentence = line[0]
+                    label = line[1]
+                    if do_lower_case:
+                        sentence = sentence.lower()
+                        label = label.lower()
+                    out_file.write(f"{sentence}\t{label}\n")
+
+
+if __name__ == "__main__":
+    # Parse the command-line arguments.
+    parser = argparse.ArgumentParser(description="Process and convert datasets into NeMo's format")
+    parser.add_argument(
+        "--source_data_dir", required=True, type=str, help="Path to the folder containing the dataset files"
+    )
+    parser.add_argument("--target_data_dir", required=True, type=str, help="Path to save the processed dataset")
+    parser.add_argument("--do_lower_case", action="store_true")
+    args = parser.parse_args()
+
+    source_dir = args.source_data_dir
+    target_dir = args.target_data_dir
+    do_lower_case = args.do_lower_case
+
+    download_text2sparql(infold=source_dir)
+    process_text2sparql(infold=source_dir, outfold=target_dir, do_lower_case=do_lower_case)

tts/NeMo/examples/nlp/text2sparql/evaluate_text2sparql.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2020, MeetKai Inc.  All rights reserved.
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script contains an example on how to evaluate a NeuralMachineTranslationModel.
+To load the example Text2Sparql dataset, please refer to ./data/import_datasets.py.
+To train a model, please refer to text2sparql.py.
+
+
+***Setting the configs***
+This script uses the `/examples/nlp/text2sparql/conf/text2sparql_config.yaml` config file by default.
+You may update the config file from the file directly or by using the command line arguments.
+Another other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+Please refer to text2sparql.py for detailed instructions on setting the configuration.
+
+
+***How to run the script?***
+- To reload and evaluate the model, run:
+
+python evaluate_text2sparql.py \
+    model.test_ds.filepath="$TGT_DATA_DIR"/test_easy.tsv \
+    model.batch_size=16 \
+    model.nemo_path=./NeMo_logs/bart.nemo \
+    exp_manager.exp_dir=./NeMo_logs
+"""
+
+import os
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.neural_machine_translation import NeuralMachineTranslationModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+
+@hydra_runner(config_path="conf", config_name="text2sparql_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f"Config:\n {OmegaConf.to_yaml(cfg)}")
+    trainer = pl.Trainer(gpus=cfg.trainer.gpus)
+    nmt_model = NeuralMachineTranslationModel.restore_from(restore_path=cfg.model.nemo_path)
+    nmt_model.setup_test_data(cfg.model.test_ds)
+    results = trainer.test(nmt_model)
+
+    with open(cfg.model.test_ds.filepath, "r") as f:
+        lines = f.readlines()
+
+    lines[0] = lines[0].strip() + f"\tpredictions\n"
+    for i, res in enumerate(results[0]["texts"]):
+        lines[i + 1] = lines[i + 1].strip() + f"\t{res}\n"
+
+    savepath = os.path.join(cfg.exp_manager.exp_dir, os.path.basename(cfg.model.test_ds.filepath))
+    with open(savepath, "w") as f:
+        f.writelines(lines)
+        logging.info(f"Predictions saved to {savepath}")
+
+
+if __name__ == "__main__":
+    main()

tts/NeMo/examples/nlp/text2sparql/text2sparql.py

@@ -0,0 +1,112 @@
+# Copyright (c) 2020, MeetKai Inc.  All rights reserved.
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script contains an example on how to train and save a NeuralMachineTranslationModel.
+NeuralMachineTranslationModel in NeMo supports sequence to sequence problems such as language translation
+and text summarization, provided the data follows the format specified below.
+
+
+***Data format***
+NeuralMachineTranslationModel requires the data to be stored in TAB separated files (.tsv) with two columns of
+sentence and label, where the first line is a header of format:
+    sentence[TAB]label
+And each line is of the format:
+    [SENTENCE][TAB][LABEL]
+
+If your dataset is stored in another format, you need to convert it to this format to use a
+NeuralMachineTranslationModel.
+
+
+***Setting the configs***
+This script uses the `/examples/nlp/text2sparql/conf/text2sparql_config.yaml` config file by default.
+You may update the config file from the file directly or by using the command line arguments.
+Another other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+A NeuralMachineTranslationModel's config file declares multiple import sections. They are:
+    - trainer: Arguments to be passed to PyTorch Lightning.
+    - model: All arguments that relate to the Model - language_model, tokenizers, datasets, optimizer, generate.
+    - exp_manager: Arguments to be passed to NeMo's experiment manager.
+    - hydra: Arguments to be passed to Hydra.
+
+If using text2sparql_config.yaml, you must first update the following fields in the config:
+    - model.nemo_path: Model save path. Eg. [PATH]/bart.nemo
+    - model.data_dir: Path to data directory. Alternatively, you can adjust the file paths directly:
+        - model.train_ds.filepath
+        - model.validation_ds.filepath
+        - model.test_ds.filepath
+    - exp_manager.exp_dir: Directory to log results from the experiment.
+
+It is highly recommended to also adjust these parameters as necessary:
+    - trainer.gpus: Set to 0 to use CPU. Otherwise the number denotes the number of GPUs.
+    - trainer.max_epochs: Maximum number of epochs to train for.
+    - model.batch_size: 8 is sufficient to train a decent Bart model for Text2Sparql.
+    - model.max_seq_length: Maximum (tokenized) sequence length. 150 is sufficient for Text2Sparql.
+    - model.language_model.pretrained_model_name: End2end pretrained model name from huggingface.
+        - model.encoder_tokenizer.tokenizer_name: Pretrained tokenizer name from huggingface.
+        - model.decoder_tokenizer.tokenizer_name: The same as above, as the tokenizer will handle encoding and decoding.
+    - model.optim.lr: Learning rate.
+
+You can also specify an encoder and decoder rather than using an end2end model like Bart by defining these parameters:
+    - model.language_model.pretrained_encoder_model_name: Pretrained huggingface encoder model name.
+        - model.encoder_tokenizer.tokenizer_name: Pretrained huggingface encoder tokenizer name.
+    - model.language_model.pretrained_decoder_model_name: Pretrained huggingface decoder model name.
+        - model.decoder_tokenizer.tokenizer_name: Pretrained huggingface decoder tokenizer name.
+    - model.language_model.pretrained_model_name: Set this to null.
+
+
+***How to run the script?***
+- First, download the data to TGT_DATA_DIR (see ./data/import_datasets.py):
+
+SRC_DATA_DIR=./data/text2sparql_src
+TGT_DATA_DIR=./data/text2sparql_tgt
+
+python ./data/import_datasets.py \
+    --source_data_dir $SRC_DATA_DIR \
+    --target_data_dir $TGT_DATA_DIR
+
+- And run the following to train and save the model:
+
+python text2sparql.py \
+    model.train_ds.filepath="$TGT_DATA_DIR"/train.tsv \
+    model.validation_ds.filepath="$TGT_DATA_DIR"/test_easy.tsv \
+    model.test_ds.filepath="$TGT_DATA_DIR"/test_hard.tsv \
+    model.batch_size=16 \
+    model.nemo_path=./NeMo_logs/bart.nemo \
+    exp_manager.exp_dir=./NeMo_logs
+"""
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.neural_machine_translation import NeuralMachineTranslationModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="text2sparql_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f"Config:\n {OmegaConf.to_yaml(cfg)}")
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    nmt_model = NeuralMachineTranslationModel(cfg.model, trainer=trainer)
+    trainer.fit(nmt_model)
+    if cfg.model.nemo_path:
+        nmt_model.save_to(cfg.model.nemo_path)
+
+
+if __name__ == "__main__":
+    main()

tts/NeMo/examples/nlp/text_classification/conf/text_classification_config.yaml

@@ -0,0 +1,119 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Config file for text classification with pre-trained BERT models
+
+trainer:
+  gpus: 1 # number of GPUs (0 for CPU), or list of the GPUs to use e.g. [0, 1]
+  num_nodes: 1
+  max_epochs: 100
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  gradient_clip_val: 0.0
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
+  accelerator: ddp
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+  num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
+
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+
+model:
+  nemo_path: text_classification_model.nemo # filename to save the model and associated artifacts to .nemo file
+  tokenizer:
+      tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+      vocab_file: null # path to vocab file
+      tokenizer_model: null # only used if tokenizer is sentencepiece
+      special_tokens: null
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+  classifier_head:
+    num_output_layers: 2
+    fc_dropout: 0.1
+
+  class_labels:
+    class_labels_file : null # optional to specify a file containing the list of the labels
+
+  dataset:
+    num_classes: ??? # The number of classes. 0 < Label <num_classes.
+    do_lower_case: false # true for uncased models, false for cased models, will be set automatically if pre-trained tokenizer model is used
+    max_seq_length: 256 # the maximum length BERT supports is 512
+    class_balancing: null # null or 'weighted_loss'. 'weighted_loss' enables the weighted class balancing of the loss, may be used for handling unbalanced classes
+    use_cache: false # uses a cache to store the processed dataset, you may use it for large datasets for speed up
+
+  train_ds:
+    file_path: null
+    batch_size: 64
+    shuffle: true
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    file_path: null
+    batch_size: 64
+    shuffle: false
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  test_ds:
+    file_path: null
+    batch_size: 64
+    shuffle: false
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adam
+    lr: 2e-5
+    # optimizer arguments
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+
+    # scheduler setup
+    sched:
+      name: WarmupAnnealing
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.1
+      last_epoch: -1
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+  # List of some sample queries for inference after training is done
+  infer_samples: [
+    'by the end of no such thing the audience , like beatrice , has a watchful affection for the monster .',
+    'director rob marshall went out gunning to make a great one .',
+    'uneasy mishmash of styles and genres .',
+  ]
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
+  name: "TextClassification"  # The name of your model
+  create_tensorboard_logger: True  # Whether you want exp_manger to create a tb logger
+  create_checkpoint_callback: True  # Whether you want exp_manager to create a modelcheckpoint callback

tts/NeMo/examples/nlp/text_classification/data/import_datasets.py

@@ -0,0 +1,243 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script can be used to process and import IMDB, ChemProt, SST-2, and THUCnews datasets into NeMo's format.
+You may run it as the following:
+
+python import_datasets.py \
+        --dataset_name DATASET_NAME \
+        --target_data_dir TARGET_PATH \
+        --source_data_dir SOURCE_PATH
+
+The dataset should be specified by "DATASET_NAME" which can be from ["sst-2", "chemprot", "imdb", "thucnews"].
+It reads the data from "SOURCE_PATH" folder, processes and converts the data into NeMo's format.
+Then writes the results into "TARGET_PATH" folder.
+
+"""
+
+import argparse
+import csv
+import glob
+import os
+from os.path import exists
+
+import tqdm
+
+from nemo.utils import logging
+
+
+def process_imdb(infold, outfold, uncased, modes=['train', 'test']):
+    if not os.path.exists(infold):
+        link = 'https://ai.stanford.edu/~amaas/data/sentiment/'
+        raise ValueError(
+            f'Data not found at {infold}. '
+            f'Please download IMDB reviews dataset from {link} and '
+            f'extract it into the folder specified by source_data_dir argument.'
+        )
+
+    logging.info(f'Processing IMDB dataset and store at {outfold}')
+    os.makedirs(outfold, exist_ok=True)
+
+    outfiles = {}
+    for mode in modes:
+        outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'w')
+        for sent in ['neg', 'pos']:
+            if sent == 'neg':
+                label = 0
+            else:
+                label = 1
+            files = glob.glob(f'{infold}/{mode}/{sent}/*.txt')
+            for file in files:
+                with open(file, 'r') as f:
+                    review = f.read().strip()
+                if uncased:
+                    review = review.lower()
+                review = review.replace("<br />", "")
+                outfiles[mode].write(f'{review}\t{label}\n')
+
+    for mode in modes:
+        outfiles[mode].close()
+
+    class_labels_file = open(os.path.join(outfold, 'label_ids.tsv'), 'w')
+    class_labels_file.write('negative\npositive\n')
+    class_labels_file.close()
+
+
+def process_sst2(infold, outfold, uncased, splits=['train', 'dev']):
+    """Process sst2 dataset."""
+    # "test" split doesn't have labels, so it is skipped
+    if not os.path.exists(infold):
+        link = 'https://dl.fbaipublicfiles.com/glue/data/SST-2.zip'
+        raise ValueError(
+            f'Data not found at {infold}. Please download SST-2 dataset from `{link}` and '
+            f'extract it into the folder specified by `source_data_dir` argument.'
+        )
+
+    logging.info(f'Processing SST-2 dataset')
+    os.makedirs(outfold, exist_ok=True)
+
+    def _read_tsv(input_file, quotechar=None):
+        """Read a tab separated value file."""
+        with open(input_file, "r") as f:
+            reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
+            lines = []
+            for line in reader:
+                lines.append(line)
+            return lines
+
+    for split in splits:
+        # Load input file.
+        input_file = os.path.join(infold, split + '.tsv')
+        lines = _read_tsv(input_file)
+        # Create output.
+        outfile = open(os.path.join(outfold, split + '.tsv'), 'w')
+
+        # Copy lines, skip the header (line 0).
+        for line in lines[1:]:
+            text = line[0]
+            label = line[1]
+            # Lowercase when required.
+            if uncased:
+                text = text.lower()
+            # Write output.
+            outfile.write(f'{text}\t{label}\n')
+        # Close file.
+        outfile.close()
+
+    class_labels_file = open(os.path.join(outfold, 'label_ids.tsv'), 'w')
+    class_labels_file.write('negative\npositive\n')
+    class_labels_file.close()
+
+    logging.info(f'Result stored at {outfold}')
+
+
+def process_chemprot(source_dir, target_dir, uncased, modes=['train', 'test', 'dev']):
+    if not os.path.exists(source_dir):
+        link = 'https://github.com/arwhirang/recursive_chemprot/tree/master/Demo/tree_LSTM/data'
+        raise ValueError(f'Data not found at {source_dir}. ' f'Please download ChemProt from {link}.')
+
+    logging.info(f'Processing Chemprot dataset and store at {target_dir}')
+    os.makedirs(target_dir, exist_ok=True)
+
+    naming_map = {'train': 'trainingPosit_chem', 'test': 'testPosit_chem', 'dev': 'developPosit_chem'}
+
+    def _read_tsv(input_file, quotechar=None):
+        """Reads a tab separated value file."""
+        with open(input_file, "r") as f:
+            reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
+            lines = []
+            for line in reader:
+                lines.append(line)
+            return lines
+
+    outfiles = {}
+    label_mapping = {}
+    out_label_mapping = open(os.path.join(target_dir, 'label_mapping.tsv'), 'w')
+    for mode in modes:
+        outfiles[mode] = open(os.path.join(target_dir, mode + '.tsv'), 'w')
+        input_file = os.path.join(source_dir, naming_map[mode])
+        lines = _read_tsv(input_file)
+        for line in lines:
+            text = line[1]
+            label = line[2]
+            if label == "True":
+                label = line[3]
+            if uncased:
+                text = text.lower()
+            if label not in label_mapping:
+                out_label_mapping.write(f'{label}\t{len(label_mapping)}\n')
+                label_mapping[label] = len(label_mapping)
+            label = label_mapping[label]
+            outfiles[mode].write(f'{text}\t{label}\n')
+    for mode in modes:
+        outfiles[mode].close()
+    out_label_mapping.close()
+
+
+def process_thucnews(infold, outfold):
+    modes = ['train', 'test']
+    train_size = 0.8
+    if not os.path.exists(infold):
+        link = 'thuctc.thunlp.org/'
+        raise ValueError(f'Data not found at {infold}. ' f'Please download THUCNews from {link}.')
+
+    logging.info(f'Processing THUCNews dataset and store at {outfold}')
+    os.makedirs(outfold, exist_ok=True)
+
+    outfiles = {}
+    for mode in modes:
+        outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'a+', encoding='utf-8')
+    categories = ['体育', '娱乐', '家居', '彩票', '房产', '教育', '时尚', '时政', '星座', '游戏', '社会', '科技', '股票', '财经']
+    for category in categories:
+        label = categories.index(category)
+        category_files = glob.glob(f'{infold}/{category}/*.txt')
+        test_num = int(len(category_files) * (1 - train_size))
+        test_files = category_files[:test_num]
+        train_files = category_files[test_num:]
+
+        for mode in modes:
+            logging.info(f'Processing {mode} data of the category {category}')
+            if mode == 'test':
+                files = test_files
+            else:
+                files = train_files
+
+            if len(files) == 0:
+                logging.info(f'Skipping category {category} for {mode} mode')
+                continue
+
+            for file in tqdm.tqdm(files):
+                with open(file, 'r', encoding='utf-8') as f:
+                    news = f.read().strip().replace('\r', '')
+                    news = news.replace('\n', '').replace('\t', ' ')
+                    outfiles[mode].write(f'{news}\t{label}\n')
+    for mode in modes:
+        outfiles[mode].close()
+
+
+if __name__ == "__main__":
+    # Parse the command-line arguments.
+    parser = argparse.ArgumentParser(description="Process and convert datasets into NeMo\'s format.")
+    parser.add_argument("--dataset_name", required=True, type=str, choices=['imdb', 'thucnews', 'chemprot'])
+    parser.add_argument(
+        "--source_data_dir", required=True, type=str, help='The path to the folder containing the dataset files.'
+    )
+    parser.add_argument("--target_data_dir", required=True, type=str)
+    parser.add_argument("--do_lower_case", action='store_true')
+    args = parser.parse_args()
+
+    dataset_name = args.dataset_name
+    do_lower_case = args.do_lower_case
+    source_dir = args.source_data_dir
+    target_dir = args.target_data_dir
+
+    if not exists(source_dir):
+        raise FileNotFoundError(f"{source_dir} does not exist.")
+
+    if dataset_name == 'imdb':
+        process_imdb(source_dir, target_dir, do_lower_case)
+    elif dataset_name == 'thucnews':
+        process_thucnews(source_dir, target_dir)
+    elif dataset_name == "chemprot":
+        process_chemprot(source_dir, target_dir, do_lower_case)
+    elif dataset_name == "sst-2":
+        process_sst2(source_dir, target_dir, do_lower_case)
+    else:
+        raise ValueError(
+            f'Dataset {dataset_name} is not supported.'
+            + "Please make sure that you build the preprocessing process for it. "
+            + "NeMo's format assumes that a data file has a header and each line of the file follows "
+            + "the format: text [TAB] label. Label is assumed to be an integer."
+        )

tts/NeMo/examples/nlp/text_classification/model_parallel_text_classification_evaluation.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script runs model parallel text classification evaluation.
+"""
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.text_classification import TextClassificationModel
+from nemo.collections.nlp.parts.nlp_overrides import NLPDDPPlugin
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="text_classification_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'\nConfig Params:\n{OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    # TODO: can we drop strict=False
+    model = TextClassificationModel.restore_from(cfg.model.nemo_path, trainer=trainer, strict=False)
+    model.setup_test_data(test_data_config=cfg.model.test_ds)
+
+    trainer.test(model=model, ckpt_path=None)
+
+
+if __name__ == '__main__':
+    main()