How to use the nemo.backends.pytorch.nm.DataLayerNM function in NEMO

# Copyright (c) 2019 NVIDIA Corporation
# pylint: disable=E0401, E0602, E0611, E1101

import torch

from nemo.backends.pytorch.nm import DataLayerNM
from nemo.core.neural_types import *
from nemo.core import DeviceType
from .datasets import BertNERDataset


class BertNERDataLayer(DataLayerNM):
    @staticmethod
    def create_ports():
        input_ports = {}
        output_ports = {
            "input_ids": NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag)
            }),
            "input_type_ids": NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag)
            }),
            "input_mask": NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag)
            }),

This package contains Neural Modules responsible for ASR-related data
processing
"""
import torch
from apex import amp

from nemo.backends.pytorch.nm import DataLayerNM, TrainableNM, NonTrainableNM
from nemo.core import Optimization, DeviceType
from nemo.core.neural_types import (NeuralType, AxisType, BatchTag, TimeTag,
                                    SpectrogramSignalTag, ProcessedTimeTag)
from .parts.dataset import AudioDataset, seq_collate_fn
from .parts.features import FilterbankFeatures, WaveformFeaturizer
from .parts.spectr_augment import SpecAugment, SpecCutout


class AudioToTextDataLayer(DataLayerNM):
    """Data Layer for general ASR tasks.

    Module which reads ASR labeled data. It accepts comma-separated
    JSON manifest files describing the correspondence between wav audio files
    and their transcripts. JSON files should be of the following format::

        {"audio_filepath": path_to_wav_0, "duration": time_in_sec_0, "text": \
transcript_0}
        ...
        {"audio_filepath": path_to_wav_n, "duration": time_in_sec_n, "text": \
transcript_n}


    Args:
        manifest_filepath (str): path to JSON containing data.
        labels (list): list of characters that can be output by the ASR model.

batch_size,
            sample_rate=16000,
            int_values=False,
            eos_id=None,
            min_duration=0.1,
            max_duration=None,
            normalize_transcripts=True,
            trim_silence=False,
            load_audio=True,
            drop_last=False,
            shuffle=True,
            num_workers=0,
            # perturb_config=None,
            **kwargs
    ):
        DataLayerNM.__init__(self, **kwargs)

        self._featurizer = WaveformFeaturizer(
            sample_rate=sample_rate, int_values=int_values, augmentor=None)
        self._dataset = AudioDataset(
            manifest_filepath=manifest_filepath,
            labels=labels,
            featurizer=self._featurizer, max_duration=max_duration,
            min_duration=min_duration, normalize=normalize_transcripts,
            trim=trim_silence, logger=self._logger,
            eos_id=eos_id, load_audio=load_audio
        )

        if self._placement == DeviceType.AllGpu:
            self._logger.info('Parallelizing DATALAYER')
            sampler = torch.utils.data.distributed.DistributedSampler(
                self._dataset)

import math
import torch
from torch.utils.data import Dataset
from torchvision import transforms, datasets


import nemo
import torch

from nemo.backends.pytorch.nm import TrainableNM, NonTrainableNM, LossNM,\
    DataLayerNM
from nemo.core import NeuralType, BatchTag, ChannelTag, HeightTag, WidthTag,\
    AxisType, DeviceType, LogProbabilityTag


class MNISTDataLayer(DataLayerNM):
    """Wrapper around torchvision's MNIST dataset.

    Args:
        batch_size (int)
        root (str): Where to store the dataset
        train (bool)
        shuffle (bool)
    """

    @staticmethod
    def create_ports(input_size=(32, 32)):
        input_ports = {}
        output_ports = {
            "images": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag, 1),
                                  2: AxisType(HeightTag, input_size[1]),

# Copyright (c) 2019 NVIDIA Corporation
import torch

from nemo.backends.pytorch.nm import DataLayerNM
from nemo.core import DeviceType
from nemo.core.neural_types import *
from .parts.datasets import AudioOnlyDataset


class AudioDataLayer(DataLayerNM):
    """
    Data Layer for general speech tasks that loads only the audio.

    Module which reads speech data. It accepts comma-separated
    JSON manifest files describing the wav audio files and their metadata.
    JSON files should be of the following format::

        {"audio_filepath": path_to_wav_0, "duration": time_in_sec_0}
        ...
        {"audio_filepath": path_to_wav_n, "duration": time_in_sec_n}


    Args:
        manifest_filepath (str): path to JSON containing data.
        batch_size (int): batch sizelse.
        min_duration (float): All training files which have a duration less

return self._batch_size*2
        self._dataset = DummyDataset(batch_size)

    def __len__(self):
        return self._dataset.__len__()

    @property
    def dataset(self):
        return self._dataset

    @property
    def data_iterator(self):
        return None


class MnistGanDataLayer(DataLayerNM):
    """Wrapper around torchvision's MNIST dataset. Additionally, it returns a
    random variable to be used in the generator.

    Args:
        batch_size (int)
        root (str): Where to store the dataset
        train (bool)
        shuffle (bool)
    """

    @staticmethod
    def create_ports(input_size=(32, 32)):
        input_ports = {}
        output_ports = {
            "latent": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag, 64),

'GlueDataLayerRegression']

# from abc import abstractmethod
import sys

import torch
from torch.utils import data as pt_data

import nemo
from nemo.backends.pytorch.nm import DataLayerNM
from nemo.core.neural_types import *

from .datasets import *


class TextDataLayer(DataLayerNM):
    """
    Generic Text Data Layer NM which wraps PyTorch's dataset

    Args:
        dataset_type: type of dataset used for this datalayer
        dataset_params (dict): all the params for the dataset
    """

    def __init__(self, dataset_type, dataset_params, **kwargs):
        super().__init__(**kwargs)
        if isinstance(dataset_type, str):
            dataset_type = getattr(sys.modules[__name__], dataset_type)
        self._dataset = dataset_type(**dataset_params)

    def __len__(self):
        return len(self._dataset)

__all__ = ['TextDataLayer']

from functools import partial

import torch
from torch.utils.data import DataLoader, DistributedSampler

from nemo.backends.pytorch.common.parts import TextDataset
from nemo.backends.pytorch.nm import DataLayerNM
from nemo.core import DeviceType
from nemo.core.neural_types import *
from nemo.utils.misc import pad_to


class TextDataLayer(DataLayerNM):
    """A simple Neural Module for loading textual data

    Args:
        path: (str) Path to file with newline separate strings of text
        labels (list): List of string labels to use when to str2int translation
        eos_id (int): Label position of end of string symbol
        pad_id (int): Label position of padding symbol
        batch_size (int): Size of batches to generate in data loader
        drop_last (bool): Whether we drop last (possibly) incomplete batch.
            Defaults to False.
        num_workers (int): Number of processes to work on data loading (0 for
            just main process).
            Defaults to 0.

    """

def __init__(self, *, tokenizer, path_to_data, max_seq_length, **kwargs):
        DataLayerNM.__init__(self, **kwargs)
        self._dataset = BertTokenClassificationDataset(
            tokenizer=tokenizer,
            input_file=path_to_data,
            max_seq_length=max_seq_length)

def __init__(
            self, *,
            manifest_filepath,
            batch_size,
            min_duration=0.1,
            max_duration=None,
            trim_silence=False,
            drop_last=False,
            shuffle=True,
            num_workers=0,
            n_segments=0,
            **kwargs
    ):
        DataLayerNM.__init__(self, **kwargs)

        self._dataset = AudioOnlyDataset(
            manifest_filepath=manifest_filepath,
            max_duration=max_duration,
            min_duration=min_duration,
            trim=trim_silence,
            logger=self._logger,
            n_segments=n_segments
        )

        sampler = None
        if self._placement == DeviceType.AllGpu:
            self._logger.info('Parallelizing DATALAYER')
            sampler = torch.utils.data.distributed.DistributedSampler(
                self._dataset)