data.py

import sys
import torch
import numpy as np
import pandas as pd
from tqdm import tqdm
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch.utils.data.sampler import Sampler
from torch.nn.utils.rnn import pad_sequence


class Labels(object):

    def __init__(self):
        self.chars = ['<BLANK>', ' ', 'А', 'Б', 'В', 'Г', 'Д', 'Е', 'Ё', 'Ж', 'З', 'И', 'Й', 'К', 'Л', 'М', 'Н', 'О', 'П', 'Р', 'С', 'Т', 'У', 'Ф', 'Х', 'Ц', 'Ч', 'Ш', 'Щ', 'Ъ', 'Ы', 'Ь', 'Э', 'Ю', 'Я']
        self.index = {c: i for i, c in enumerate(self.chars)}

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

    def __call__(self, sentence):
        targets = []
        for c in sentence.strip().upper():
            targets.append(self.index[c])
        return targets

    def space(self):
        return self.index[' ']

    def blank(self):
        return self.index['<BLANK>']

    def is_accepted(self, sentence):
        sentence = sentence.strip().upper()
        if len(sentence) == 0:
            return False
        for c in sentence:
            if c not in self.index:
                return False
        return True

    def required_frames(self, sentence):
        targets = self(sentence)
        frames = len(targets)
        for t1, t2 in zip(targets[:-1], targets[1:]):
            if t1 == t2:
                frames += 1
        return frames


def load_data(source):
    if isinstance(source, list):
        data = [pd.read_csv(p, index_col='path', encoding='utf8') for p in source]
        return pd.concat(data)
    return pd.read_csv(source, index_col='path', encoding='utf8')


class TextDataset(Dataset):

    def __init__(self, source, labels, batch_size):
        data = load_data(source)
        b = labels.blank()
        self.utterances = [[b] + labels(t) for t in data['text'].values]
        self.data = []
        self.batch_size = batch_size

    def shuffle(self, epoch):
        np.random.RandomState(epoch).shuffle(self.utterances)
        data = np.concatenate(self.utterances)
        data = torch.tensor(data, dtype=torch.long)
        n = data.numel() // self.batch_size
        data = data.narrow(0, 0, n * self.batch_size)
        self.data = data.view(self.batch_size, -1).t()

    def __len__(self):
        return len(self.data) - 1

    def __getitem__(self, i):
        return self.data[i], self.data[i+1]


class AudioDataset(Dataset):

    def __init__(self, root, source, labels, model=None, length=0):
        self.root = root
        self.data = load_data(source)
        self.labels = labels
        if model is not None:
            self.filter_by_model(model)
        if length > 0:
            self.filter_by_length(length)

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

    def __getitem__(self, i):
        file = self.data.iloc[i]
        features = torch.tensor(np.load(self.root + file.name), dtype=torch.float32)
        targets = torch.tensor(self.labels(file['text']), dtype=torch.int)
        return features, targets

    def filter_by_length(self, length):
        self.filter(self.data[self.data['frames'] < length])

    def filter_by_model(self, model):
        frames = self.data['frames'].values
        frames = model.output_time(frames)
        index = []
        for i, text in enumerate(self.data['text'].values):
            if self.labels.required_frames(text) <= frames[i]:
                index.append(i)
        self.filter(self.data.iloc[index])

    def filter(self, subset):
        diff = len(self.data) - len(subset)
        ratio = diff / len(self.data) * 100
        self.data = subset.sort_values('frames')
        print('filter %7d %7.2f%%' % (diff, ratio))


class BucketingSampler(Sampler):

    def __init__(self, data, size=1, limit=sys.maxsize):
        super().__init__(data)
        index = list(range(len(data)))
        self.bins = [index[i:i + size] for i in range(0, len(index), size)]
        self.limit = limit

    def __iter__(self):
        for batch in self.bins[:self.limit]:
            yield batch

    def __len__(self):
        return len(self.bins[:self.limit])

    def shuffle(self, epoch):
        np.random.RandomState(epoch).shuffle(self.bins)


def collate_audio(batch):

    batch = sorted(batch, key=lambda b: b[0].shape[0], reverse=True)

    n = len(batch)

    xs = []
    ys = []
    xn = torch.empty(n, dtype=torch.int)
    yn = torch.empty(n, dtype=torch.int)

    for i, (x, y) in enumerate(batch):
        xs.append(x)
        ys.append(y)
        xn[i] = len(x)
        yn[i] = len(y)

    # N x 1 x D x T
    xs = pad_sequence(xs, batch_first=True)
    xs = xs.unsqueeze(dim=1).transpose(2, 3)

    # N x S
    ys = pad_sequence(ys, batch_first=True)

    return xs, ys, xn, yn


class DataLoaderCuda(DataLoader):

    def __init__(self, *args, **kwargs):
        kwargs['num_workers'] = 4
        kwargs['pin_memory'] = True
        super().__init__(*args, **kwargs)

    def __iter__(self):
        self.progress = tqdm(super().__iter__())
        for cpu in self.progress:
            gpu = []
            for values in cpu:
                gpu.append(values.cuda(non_blocking=True))
            yield gpu

    def shuffle(self, epoch):
        self.batch_sampler.shuffle(epoch)

    def set_description(self, desc):
        self.progress.set_description(desc)

    def close(self):
        self.progress.close()


def split_train_dev_test(root, labels, model, batch_size=32):

    train = [
        root + 'asr_public_phone_calls_1.csv',
        root + 'asr_public_phone_calls_2_aa.csv',
        root + 'asr_public_phone_calls_2_ab.csv',
        root + 'public_youtube1120_hq.csv',
        root + 'public_youtube700_aa.csv',
        root + 'public_youtube700_ab.csv'
    ]

    dev = [
        root + 'asr_public_phone_calls_1.csv',
        root + 'public_youtube1120_hq.csv',
    ]

    test = [
        root + 'asr_calls_2_val.csv',
        root + 'buriy_audiobooks_2_val.csv',
        root + 'public_youtube700_val.csv'
    ]

    train = AudioDataset(root, train, labels, model, 400)
    dev = AudioDataset(root, dev, labels, model, 1000)
    test = AudioDataset(root, test, labels, model, 1000)

    sampler1 = BucketingSampler(train, size=batch_size)
    sampler2 = BucketingSampler(dev, size=1, limit=1000)

    sampler2.shuffle(0)

    train = DataLoaderCuda(train, collate_fn=collate_audio, batch_sampler=sampler1)
    dev = DataLoaderCuda(dev, collate_fn=collate_audio, batch_sampler=sampler2)
    test = DataLoaderCuda(test, collate_fn=collate_audio, batch_size=16)

    return train, dev, test


if __name__ == '__main__':

    labels = Labels()

    dataset = AudioDataset('data/', 'data/public_youtube700_val.csv', labels)

    loader = DataLoader(dataset, batch_size=32, collate_fn=collate_audio)

    x_lengths = []
    y_lengths = []

    for _, _, xn, yn in tqdm(loader):
        x_lengths.extend(list(xn.numpy()))
        y_lengths.extend(list(yn.numpy()))

    import seaborn as sns
    import matplotlib.pyplot as plt

    plt.title('x lengths')
    sns.distplot(x_lengths)
    plt.show()

    plt.title('y lengths')
    sns.distplot(y_lengths)
    plt.show()