<jupyter_start><jupyter_text>FastPitch + GSTAt this assignment you should implement and train FastPitch AM with GST encoder.  AssessmentTotal score for this task is 30Categorically:* **3** - you filled in all the code gaps before the GST estimator block (and they seem sane to the reviewer)* **9** - you trained fp model* **3** - you visualized the GST space* **bonus 5** - you found some interesting (speed, speech style etc.) cluster in the corpus and tried to re-synthesize one utterance with a style of another * **3** - you filled in all the code gaps after the GST estimator block (and they seem sane to the reviewer)* **6** - you trained a GST estimator on the top of your FP model* **3** - your test files sound intelligible and corresponds to the text* **3** - your test files have normal or good intonation and sound naturalThe result of this work is the notebook filled and 16 wavfiles (15 test corpus + one about a quick brown fox).Please pack them in archive with the name {your last name}_{your first name}_hw3_fastpitch  Dependencies & Includes<jupyter_code># !pip install cudatoolkit==11.1 pytorch==1.8.1 # -- should already be at DataSphere
pip install librosa praat-parselmouth torch-optimizer tensorboardX Unidecode inflect wget tqdm scikit-learn
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import os
import sys
import argparse
import numpy as np
import re
import wget
import time
import shutil
import tqdm
from collections import defaultdict

from tensorboardX import SummaryWriter

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.nn.utils.rnn import pad_sequence

from torch_optimizer import Lamb
from torch.optim import Adam

import librosa
import soundfile as sf
import parselmouth
from scipy.io import wavfile
import IPython.display as ipd

from hparams import HParamsFastpitch
from nv_extern.tacotron2.arg_parser import parse_tacotron2_args
from nv_extern.tacotron2.model import Tacotron2
from nv_extern.waveglow.model import WaveGlow
from nv_extern.waveglow.denoiser import Denoiser
from nv_extern.common.text.symbols import get_symbols
from nv_extern.common.text.text_processing import TextProcessing
assert torch.cuda.is_available(), "Need some GPU to train the model"
device = torch.device('cuda')
print("GPU found: {}".format(torch.cuda.get_device_name(device)))
ljspeech_data_path = "LJSpeech-1.1"
prepr_data_path = "LJSpeech_prepr"

train_filelist_path = "filelists/train_filelist.txt"
val_filelist_path = "filelists/val_filelist.txt"


logs_dir = "logs_dir"
gst_logs_dir = "gst_logs_dir"

pretrained_models_path = "pretrained_models"

if not os.path.exists(pretrained_models_path):
    os.mkdir(pretrained_models_path)
    
synth_wavs_path = "synth_wavs"
_prev_bytes = 0
def bar_progress(current, total, width=80):
    global _prev_bytes
    
    pq = _prev_bytes / total
    cq = current / total
    if _prev_bytes > current or (cq - pq) * 1000 > 1:
        _prev_bytes = current
    else:
        return
    progress_message = "Downloading: %d%% [%d / %d] bytes" % (cq * 100, current, total)
    # Don't use print() as it will print in new line every time.
    sys.stdout.write("\r" + progress_message)
    sys.stdout.flush()
    
def getfile(url, dest):
    wget.download(url, dest, bar=bar_progress)
    

hparams = HParamsFastpitch({'batch_size' : 8})
def load_from_checkpoint(model, state_dict):
    if not hasattr(model, 'module'):
        state_dict = {
            name.replace('module.', ''): param
            for name, param in state_dict.items()
        }
    else:
        for key, value in state_dict.items():
            break
        if key[:6] != 'module':
            state_dict = {
                'module.' + name :  param
                for name, param in state_dict.items()
            }
            
    model.load_state_dict(state_dict)<jupyter_output><empty_output><jupyter_text>Download data<jupyter_code>ljs_arch="LJSpeech-1.1.tar.bz2"
taco2_url = "http://data.keithito.com/data/speech/" + ljs_arch
getfile(taco2_url, ljs_arch)
!tar -xf {ljs_arch}<jupyter_output><empty_output><jupyter_text>Mel, duration and pitch extraction  text<jupyter_code>os.mkdir(prepr_data_path)
def parse_filelist(fname):
    data = []
    with open(fname) as ifile:
        for l in ifile:
            fid, _, text = l.strip().split("|")
            data.append((fid, text))
    return data
fl = parse_filelist(os.path.join(ljspeech_data_path, "metadata.csv"))

tp = TextProcessing('english_basic',['english_cleaners'])
for file_id, text in tqdm.tqdm(fl):
    text_encoded = np.array(tp.encode_text(text))
    text_path = os.path.join(prepr_data_path, file_id + ".text.npy")
    np.save(text_path, text_encoded)<jupyter_output><empty_output><jupyter_text>mel<jupyter_code>mel_basis = librosa.filters.mel(hparams.sample_rate, 
                                hparams.n_fft, 
                                n_mels=hparams.n_mel_channels,
                                fmin=hparams.min_frequency, 
                                fmax=hparams.max_frequency)


def mel_spectrogram(audio, hparams):
    def linear_to_mel(s):
        return np.dot(mel_basis, s)
    
    def amp_to_db(x):
        return np.log(np.maximum(1e-5, x))
    
    spectr = librosa.stft(audio, 
                            n_fft=hparams.n_fft, 
                            hop_length=hparams.hop_length,
                            win_length=hparams.win_length,
                            window=hparams.window)
    
    spectr = np.abs(spectr)
    mel_spectr = amp_to_db(linear_to_mel(spectr))
    return mel_spectr


def get_mel(pcm_path, mel_path, hparams):
    audio, sr = sf.read(pcm_path, dtype='float32')
    audio = librosa.resample(audio, sr, hparams.sample_rate)
    np.save(mel_path, mel_spectrogram(audio, hparams))


def extract_mels(wav_path, out_dir, hparams):
    for filename in tqdm.tqdm(os.listdir(wav_path)):
        file_id = re.search(r"(.*).wav", filename).group(1)
        mel_path = os.path.join(out_dir, file_id + ".mel")
        get_mel(os.path.join(wav_path, filename), mel_path, hparams)
extract_mels(os.path.join(ljspeech_data_path, "wavs"), prepr_data_path, hparams)<jupyter_output><empty_output><jupyter_text>durations extraction As it was said on the lecture, we can extract durations with tacotron model.Firstly, we need to download one:<jupyter_code>taco2_url = "https://api.ngc.nvidia.com/v2/models/nvidia/tacotron2_pyt_ckpt_amp/versions/19.12.0/files/nvidia_tacotron2pyt_fp16.pt"
taco2_path = os.path.join(pretrained_models_path, "tacotron2")
getfile(taco2_url, taco2_path)<jupyter_output><empty_output><jupyter_text>now we should extract durations by batch-processing our records with tacotron,then, we should store them into npy<jupyter_code>class TacoDataset(torch.utils.data.Dataset):
    def __init__(self, directory):
        super(TacoDataset, self).__init__()
        self.directory = directory
        self.mel_paths = []
        self.text_paths = []
        self.file_ids = []
        
        for filename in os.listdir(directory):
            if ".text.npy" not in filename:
                continue
            file_id = re.search(r"(.*).text.npy", filename).group(1)
            self.file_ids.append(file_id)
            text_path = os.path.join(self.directory, filename)
            self.text_paths.append(text_path)
            self.mel_paths.append(text_path.replace(".text.npy", ".mel.npy"))
            
            
    def __getitem__(self, idx):
        text = torch.IntTensor(np.load(self.text_paths[idx]))
        mel = torch.FloatTensor(np.load(self.mel_paths[idx]))
        file_id = self.file_ids[idx]
        return (text, mel, file_id)
    
    def __len__(self):
        return len(self.text_paths)

    
class TacoCollate:
    def __init__(self, n_mel_channels):
        self.n_mel_channels = n_mel_channels
        
    def  __call__(self, batch):
        batch = sorted(batch, key=lambda x: x[0].size(0), reverse=True)
        
        max_text_length = max([text.size(0) for text, *_ in batch])
        max_mel_length = max([mel.size(1) for _, mel, *_ in batch])
        
        texts_padded = torch.LongTensor(len(batch), max_text_length).zero_()
        text_lengths = torch.LongTensor(len(batch)).zero_()
        mels_padded = torch.FloatTensor(len(batch), self.n_mel_channels, max_mel_length).zero_()
        mel_lengths = torch.LongTensor(len(batch)).zero_()
        
        file_ids = []
        
        for i, (text, mel, file_id) in enumerate(batch):
            texts_padded[i, :text.size(0)] = text
            text_lengths[i] = text.size(0)
            mels_padded[i, :, :mel.size(1)] = mel
            mel_lengths[i] = mel.size(1)
            file_ids.append(file_id)
            
        max_len = torch.max(text_lengths.data).item()
        return (texts_padded, text_lengths, mels_padded, max_len, mel_lengths), file_ids<jupyter_output><empty_output><jupyter_text>As it was said, you can get good alignment by dynamic programming.Here, you should write some code that will find the trajectory in attention matrix that maximizes cost:$$(j_i) = \mathsf{argmax} (Cost) = \mathsf{argmax} \prod_{i=1}^{N_{frames}} A_{i j_i}$$where:$$j_{i+1} = j\;or\;j+1$$$$j_1 = 1$$$$j_{N_{frames}} = N_{graphemes}$$<jupyter_code>def dp_alignment(attention):
    """
    Here you are given with an attention prob matrix
        with a shape of N_frames X N_graphemes
        
    You should compute the optimal way according to the formula above with DP.
    With the optimal way (sequence j_i) computed you should return durations.
    So, duration of j-th gratheme == #(i: j_i == j) -- number of frames at which
        our optimal way is 'stuck' at the grapheme.
    These durations, as an integer numpy array should be returned from the function.
    """
    
    # !!!!!!!!!!!!!!!!!!!!!!
    # INSERT YOUR CODE HERE!
    # !!!!!!!!!!!!!!!!!!!!!!

    return durations
def extract_durations(taco_checkpoint_path, data_path, hparams, batch_size=8):
    taco_checkpoint = torch.load(taco_checkpoint_path)
    taco_model = Tacotron2(**taco_checkpoint['config']).to(device)
    load_from_checkpoint(taco_model, taco_checkpoint['state_dict'])
    taco_model.eval()

    taco_dataset = TacoDataset(data_path)
    taco_collate = TacoCollate(hparams.n_mel_channels)
    taco_dataloader = DataLoader(taco_dataset, batch_size=batch_size, collate_fn=taco_collate)
    
    for batch, file_ids in tqdm.tqdm(taco_dataloader):
        with torch.no_grad():
            batch = [elem.to(device) 
                     if torch.is_tensor(elem)
                     else elem
                     for elem in batch]
            _, text_lengths, _, _, mel_lengths = batch
            
            *_, alignments = taco_model.forward(batch)
            alignments = alignments.cpu().numpy()
            
            for j, ali in enumerate(alignments):
                dur_path = os.path.join(data_path, file_ids[j] + ".dur")
                
                ali = ali[: mel_lengths[j], : text_lengths[j]]
                durations = dp_alignment(ali)
                np.save(dur_path, durations)
extract_durations(taco2_path, prepr_data_path, hparams, batch_size=32)<jupyter_output><empty_output><jupyter_text>Check yourself:Average "a" duration in the corpus should not be far from 6.782268<jupyter_code>cnt = 0
sum = 0
for i, _ in tqdm.tqdm(fl):
    text = np.load(os.path.join(prepr_data_path, i + ".text.npy"))
    dur = np.load(os.path.join(prepr_data_path, i + ".dur.npy"))
    for c, d in zip(text, dur):
        if c == 38:
            cnt += 1
            sum += d
mean_a_duration = sum/cnt
print("Mean 'a' duration in the corpus is {} frames".format(mean_a_duration))
assert np.abs(mean_a_duration - 6.7823) < 0.01<jupyter_output><empty_output><jupyter_text>pitch extractionhere we use parselmouth pitch extractor, pad and store the extracted f0 values into npy files<jupyter_code>def maybe_pad(vec, l):
    assert np.abs(vec.shape[0] - l) <= 3
    vec = vec[:l]
    if vec.shape[0] < l:
        vec = np.pad(vec, pad_width=(0, l - vec.shape[0]))
    return vec
def calculate_pitch(wav_path, durs):
    durs = durs.astype(int)
    mel_len = durs.sum().astype(int)
    durs_cum = np.cumsum(np.pad(durs, (1, 0))).astype(int)

    snd = parselmouth.Sound(wav_path)
    pitch = snd.to_pitch(time_step=snd.duration / (mel_len + 3)).selected_array["frequency"]
    assert np.abs(mel_len - pitch.shape[0]) <= 1.0

    # Average pitch over characters
    pitch_char = np.zeros((durs.shape[0],), dtype=float)
    for idx, a, b in zip(range(int(mel_len)), durs_cum[:-1], durs_cum[1:]):
        values = pitch[a:b][np.where(pitch[a:b] != 0.0)[0]]
        pitch_char[idx] = np.mean(values) if len(values) > 0 else 0.0

    pitch_char = maybe_pad(pitch_char, len(durs))

    return pitch_char
def extract_pitch(wavs_path, preproc_data_path):
    for filename in tqdm.tqdm(os.listdir(wavs_path)):
        file_id = re.search(r"(.*).wav", filename).group(1)
        wav_path = os.path.join(wavs_path, filename)
        dur_path = os.path.join(preproc_data_path, file_id + ".dur.npy")
        pitch_path = os.path.join(preproc_data_path, file_id + ".pt.npy")
      
        dur = np.load(dur_path)
        pitch = calculate_pitch(str(wav_path), dur)
        np.save(pitch_path, pitch)
extract_pitch(os.path.join(ljspeech_data_path, "wavs"), prepr_data_path)<jupyter_output><empty_output><jupyter_text>Data utilssome helper classes for training<jupyter_code>LJSPEECH_MEAN = 218.44949768191876
LJSPEECH_STD = 64.85769765940508
class FastPitchDataset(torch.utils.data.Dataset):
    """
    Provides random access to the samples
    returns data in the form of (text, mel, duration, pitch)
    """
    def __init__(self, directory, filelist_path):
        super(FastPitchDataset, self).__init__()
        self.directory = directory
        self.mel_paths = []
        self.pitch_paths = []
        self.dur_paths = []
        self.text_paths = []

        with open(filelist_path, "r") as filelist:
            file_ids = filelist.readlines()
        file_ids = [elem.strip() for elem in file_ids] 
        
        for file_id in file_ids:
            text_path = os.path.join(directory, file_id + ".text.npy")
            self.text_paths.append(text_path)
            self.pitch_paths.append(text_path.replace(".text.npy", ".pt.npy"))
            self.dur_paths.append(text_path.replace(".text.npy", ".dur.npy"))
            self.mel_paths.append(text_path.replace(".text.npy", ".mel.npy"))
        
    def __getitem__(self, idx):
        text = torch.IntTensor(np.load(self.text_paths[idx]))
        mel = torch.FloatTensor(np.load(self.mel_paths[idx]))
        dur = torch.FloatTensor(np.load(self.dur_paths[idx]))
        pitch = torch.FloatTensor(np.load(self.pitch_paths[idx]))
        
        pitch[pitch != 0] = (pitch[pitch != 0] - LJSPEECH_MEAN) / LJSPEECH_STD

        return (text, mel, dur, pitch)
    
    def __len__(self):
        return len(self.text_paths)
    
    
class FastPitchCollate:
    """
    Groups and pads the date to the batches,
    adds lengths of texts and mels tensors to the batch info
    """
    def __init__(self, n_mel_channels):
        self.n_mel_channels = n_mel_channels

    def __call__(self, batch):
        max_text_length = max([text.size(0) for text, *_ in batch])
        max_mel_length = max([mel.size(1) for _, mel, *_ in batch])
    
        texts_padded = torch.LongTensor(len(batch), max_text_length).zero_()
        text_lengths = torch.LongTensor(len(batch)).zero_()
        mels_padded = torch.FloatTensor(len(batch), self.n_mel_channels, max_mel_length).zero_()
        mel_lengths = torch.LongTensor(len(batch)).zero_()
        dur_padded = torch.zeros_like(texts_padded, dtype=batch[0][2].dtype)
        pitch_padded = torch.zeros(dur_padded.size(0), dur_padded.size(1), dtype=batch[0][3].dtype)
        
        for i, (text, mel, dur, pitch) in enumerate(batch):
            texts_padded[i, :text.size(0)] = text
            text_lengths[i] = text.size(0)
            mels_padded[i, :, :mel.size(1)] = mel
            mel_lengths[i] = mel.size(1)
            dur_padded[i, :dur.shape[0]] = dur
            pitch_padded[i, :pitch.shape[0]] = pitch
          
        return (texts_padded,
            text_lengths,
            mels_padded,
            mel_lengths,
            dur_padded,
            pitch_padded)
def batch_to_device(batch, device):
    """
    Pushes all the numeric data to GPU
    Splits input and output data of the model
    """
    (texts, 
     text_lengths, 
     mels, 
     mel_lengths, 
     dur,  
     pitch
    ) = [tensor.to(device) if tensor is not None
               else None
               for tensor in batch]
    
    x = (texts, text_lengths, mels, mel_lengths, dur, pitch)
    y = (mels, dur, text_lengths, pitch)
    return x, y
def mask_from_lens(lengths, max_len=None):
    """
    With mels' or texts' lengths given creates a 0-1 mask for losses
    """
    if max_len is None:
        max_len = torch.max(lengths).item()
    ids = torch.arange(0, max_len, device=lengths.device)
    mask = (ids < lengths.unsqueeze(1))
    return mask<jupyter_output><empty_output><jupyter_text>Loss function<jupyter_code>class FastPitchLoss(nn.Module):
    def __init__(self, hparams):
        super(FastPitchLoss, self).__init__()
        self.dur_predictor_loss_scale = hparams.dur_predictor_loss_scale
        self.pitch_predictor_loss_scale = hparams.pitch_predictor_loss_scale

    def forward(self, model_out, targets, meta_agg='mean'):
        """
        Here you get
        * model_out, which is a tuple of
            mel_output, decoding_mask, predicted_duration, log_pred_duration and predicted_pitch
        * targets, which is also a tuple:
            mel_target, duration_target, duration_lengths, pitch_target
        
        you should compute loss - a weighted (look at constructor and hparams) sum of
            (masked) l2 mels, pitch and durations losses
        also, to draw graphs you should return the components separately - in a dict meta - kv struct with keys:
            * loss
            * mel_loss
            * duration_predictor_loss
            * pitch_loss
            * duration_l1_error
        """
        mel_out, dec_mask, dur_pred, log_dur_pred, pitch_pred = model_out
        mel_tgt, dur_tgt, dur_lens, pitch_tgt = targets
        mel_tgt.requires_grad = False
        # (B,H,T) => (B,T,H)
        mel_tgt = mel_tgt.transpose(1, 2)

        dur_mask = mask_from_lens(dur_lens, max_len=dur_tgt.size(1))
        
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!

        loss = (mel_loss
                + pitch_loss * self.pitch_predictor_loss_scale
                + log_dur_loss * self.dur_predictor_loss_scale)
       
        meta = {
            'loss':                    loss.clone().detach(),
            
            'mel_loss':                mel_loss.clone().detach(),
            
            'duration_predictor_loss': log_dur_loss.clone().detach(),
            
            'pitch_loss':              pitch_loss.clone().detach(),
            
            'duration_l1_error':       (torch.abs(dur_pred - dur_tgt).sum() 
                                           / dur_mask.sum()).detach(),
        }
       
        assert meta_agg in ('sum', 'mean')
        if meta_agg == 'sum':
            # !!!!!!!!!!!!!!!!!!!!!!
            # INSERT YOUR CODE HERE!
            # !!!!!!!!!!!!!!!!!!!!!!
            
        return loss, meta<jupyter_output><empty_output><jupyter_text>Logger<jupyter_code>class FastpitchLogger:
    """
    Logger. Saves/loads model, flushes on validation, can log gradients (useful for debugging the core AM model)
    """
    
    def __init__(self, logdir, hparams):
        self.eval_interval = hparams.eval_interval
        self.global_checkpoint_coef = hparams.global_checkpoint_coef
        self.logdir = logdir
        
        self.train_logger = SummaryWriter(os.path.join(logdir, 'train'))
        self.val_logger = SummaryWriter(os.path.join(logdir, 'val'))


    def log_grads(self, step, model):
            norms = [p.grad.norm().item() for p in model.parameters()
                     if p.grad is not None]
            for stat in ('max', 'min', 'mean'):
                self.train_logger.add_scalar(f'grad_{stat}', getattr(np, stat)(norms), step)

    def log(self, logger, step, meta):
        for k, v in meta.items():
            logger.add_scalar(k, 
                            v.item() if hasattr(v, 'item') else v, 
                            step)
    
    def log_training(self, step, meta):
        print(f"train : {meta['step']}  loss {meta['loss']:.4f}")
        self.log(self.train_logger, step, meta)
            
        
    def log_validation(self, step, meta):
        print(f"val : {meta['step']}  loss {meta['loss']}")
        self.log(self.val_logger, step, meta)
            
        for w in self.train_logger.all_writers.values():
            w.flush()
        for w in self.val_logger.all_writers.values():
            w.flush()
            
    def step_to_path(self, step):
        return os.path.join(self.logdir, f"{step:07d}_snapshot.pth")
            
    @staticmethod
    def load_last_checkpoint(logdir):
        steps = [
            int(filename.replace('_snapshot.pth', ''))
            for filename in os.listdir(logdir)
            if '_snapshot.pth' in filename
        ]

        if len(steps) == 0:
            return {
                'step': 0,
                'state_dict': None
            }

        return torch.load(os.path.join(logdir, f"{max(steps):07d}_snapshot.pth"))
            
 
    def save_checkpoint(self, model, optimizer, step, hparams):
        checkpoint = {
            'step': step,
            'state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'params': hparams.__dict__
        }
        
        torch.save(checkpoint, self.step_to_path(step))

        # remove unnecessary snapshots
        if step > self.eval_interval and (step - self.eval_interval) % (self.global_checkpoint_coef * self.eval_interval):
            os.remove(self.step_to_path(step - self.eval_interval))<jupyter_output><empty_output><jupyter_text>Model It is highly recommended to go briefly through the lectureand/or FastPitch paper https://arxiv.org/abs/2006.06873 before doing the next section  FFTransformer Here is the core part of FastPitch model.You should implement self-attention block here.Note: the classes below are provided for your convenience. You are free to re-write everything in your own style if your architecture will work and you will implement in your code main idea of FastPitch (FFT encoder and decoder, duration and pitch prediction in between) Here you need to implement sin-cos pos embeddings:$$PE[p, 2 i] = sin\left(\frac{p}{10000^{2i/d_{model}}}\right)$$$$PE[p, 2 i + 1] = cos \left(\frac{p}{10000^{2i/d_{model}}}\right)$$<jupyter_code>class PositionalEmbedding(nn.Module):
    """
    Sin and cos positional embeddings
    """
    def __init__(self, d_model):
        super(PositionalEmbedding, self).__init__()
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!

    """
    * pos_seq - 1...N array
    """
    def forward(self, pos_seq):
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
class PositionwiseConvFF(nn.Module):
    """
    This position-wise convolutional layer
    It is a sequence:
    * Conv1d, d_model -> d_inner
    * ReLU
    * Conv1d, d_inner -> d_model
    * Dropout
    * Residual (+input)
    * LayerNorm
    """
    
    def __init__(self, d_model, d_inner, kernel_size, dropout):
        super(PositionwiseConvFF, self).__init__()
        self.CoreNet = nn.Sequential ...
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!

    def forward(self, inp):
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
        return output
class FFTransformer(nn.Module):
    """
    Here, you should implement the sequence:
    * embeddings of the categorical input, if embed_input is True
    * dropout(input + embeddings), if dropemb is not 0
    * n_layers of (multi-head self-attention + dropout + layernorm + pos-wise convFF)
    * all this should be masked according to the lengths
    
    Parameters:
    * n_symbols - number of symbols (in case of cat input)
    * n_layer - self-att layers
    * d_model - outer dim, n_embedding for graphemes, n_mel for audio
    * n_head - number of attention heads
    * d_head - size of each head
    * kernel size - pos-wise convolution kernel size
    * dropout - DO after the self-attention
    * dropatt - DO at self-attention probs
    * dropemb - DO after the positional embeddings addition
    * embed_input - if true, input is treated as categorical and additional nn.Embeddings should be added
    * pad_idx - only for embed_input = true: special value categorical input are padded with
    
    """
    def __init__(
        self,
        n_symbols,
        n_layer,
        n_head,
        d_model,
        d_head,
        d_inner,
        kernel_size,
        dropout,
        dropatt,
        dropemb=0.0,
        embed_input=True,
        pad_idx=0
    ):
        super(FFTransformer, self).__init__()

        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
        
    def forward(self, dec_inp, mel_lens=None):
        """
        * dec_inp - input, continuous or categorical
        * mel_lens - for embed_input = False: lengths of input, need for padding
        Return:
        * tuple: processed output and padding bitmask
        """
        if self.word_emb is None:
            inp = dec_inp
            mask = mask_from_lens(mel_lens).unsqueeze(2)
        else:
            inp = self.word_emb(dec_inp)
            mask = (dec_inp != self.pad_idx).unsqueeze(2)
        
        pos_seq = torch.arange(inp.size(1), device=inp.device, dtype=inp.dtype)
        pos_emb = self.pos_emb(pos_seq) * mask

        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
        
        return out, mask<jupyter_output><empty_output><jupyter_text>GSTthe style tokenIt is recommended for you to use already implemented FFT block as a GST encoder.For the details - look at the lecture and/or paper https://arxiv.org/pdf/1803.09017.pdf<jupyter_code>class GST(nn.Module):
    def __init__(self, hparams):
        super(GST, self).__init__()

        self.encoder = FFTransformer(
            n_symbols=None,
            n_layer=hparams.gst_n_layers,
            n_head=hparams.gst_n_heads,
            d_model=hparams.n_mel_channels,
            d_head=hparams.gst_d_head,
            d_inner=hparams.gst_conv1d_filter_size,
            kernel_size=hparams.gst_conv1d_kernel_size,
            dropout=hparams.p_gst_dropout,
            dropatt=hparams.p_gst_dropatt,
            dropemb=hparams.p_gst_dropemb,
            embed_input=False
        )
        
        # additional layer to be used in GST
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
        

    def forward(self, mels, mel_lengths):
        """
        Should take GT mels and return style embeddings and attention probs
        """
        mels_enc = self.encoder(mels, mel_lengths)
        #...
        
        """
        mel-invariant tokens should be keys in the attention here
        and encoder output -- the query
        
        
        Please note that you should prevent unlimited growth of tokens magnitude somehow
        E.g. you can shrink them with some restrictive function:
            attention_keys = tanh(tokens_embeddings)
        """
        
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
        
        return style_embed, alphas<jupyter_output><empty_output><jupyter_text>FastPitch  Temporal Predictor (predicts a single float per each temporal location)<jupyter_code>class ConvReLUNorm(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1, dropout=0.0):
        super(ConvReLUNorm, self).__init__()
        self.conv = torch.nn.Conv1d(in_channels, out_channels,
                                    kernel_size=kernel_size,
                                    padding=(kernel_size // 2))
        self.norm = torch.nn.LayerNorm(out_channels)
        self.dropout = torch.nn.Dropout(dropout)

    def forward(self, signal):
        out = F.relu(self.conv(signal))
        out = self.norm(out.transpose(1, 2)).transpose(1, 2)
        return self.dropout(out)
    

class TemporalPredictor(nn.Module):
    """
    Regressor head
    A couple of conv+reLU+layernorm blocks + projection to answer
    """
    def __init__(self, input_size, filter_size, kernel_size, dropout, n_layers=2):
        super(TemporalPredictor, self).__init__()

        self.layers = nn.Sequential(
            *[
                ConvReLUNorm(
                    input_size if i == 0 else filter_size,
                    filter_size,
                    kernel_size=kernel_size,
                    dropout=dropout,
                )
                for i in range(n_layers)
            ]
        )
        self.fc = nn.Linear(filter_size, 1, bias=True)

    def forward(self, enc_out, enc_out_mask):
        out = enc_out * enc_out_mask
        out = self.layers(out.transpose(1, 2)).transpose(1, 2)
        out = self.fc(out) * enc_out_mask
        return out.squeeze(-1)
def regulate_len(durations, enc_out):
    """
    Here you should write the part with upsample:
    Parameters:
    * durations - predicted/gt durations array
    * enc_out - encoder embeddings you need to upsample
    Returns tuple:
    * encoder upsampled states
    * frame_lengths - sums over T of durations
    
    """
    # !!!!!!!!!!!!!!!!!!!!!!
    # INSERT YOUR CODE HERE!
    # !!!!!!!!!!!!!!!!!!!!!!
    enc_repeated = [ ... ]
    enc_repeated = pad_sequence(enc_repeated, batch_first=True)
    return enc_repeated, frame_lengths
class FastPitch(nn.Module):
    """
    Here you need to construct and infer FP net
    Note that predicted f0 should be feeded into decoder through convolution with a small kernel
    """
    def __init__(self, hparams : HParamsFastpitch):
        super(FastPitch, self).__init__()

        self.encoder = FFTransformer(
            n_symbols=hparams.n_symbols,
            n_layer=hparams.in_fft_n_layers,
            n_head=hparams.in_fft_n_heads,
            d_model=hparams.symbols_embedding_dim,
            d_head=hparams.in_fft_d_head,
            d_inner=4 * hparams.symbols_embedding_dim,
            kernel_size=hparams.in_fft_conv1d_kernel_size,
            dropout=hparams.p_in_fft_dropout,
            dropatt=hparams.p_in_fft_dropatt,
            dropemb=hparams.p_in_fft_dropemb,
            embed_input=True,
            pad_idx=hparams.pad_idx)

        self.duration_predictor = TemporalPredictor(
            input_size=hparams.symbols_embedding_dim,
            filter_size=hparams.dur_predictor_filter_size,
            kernel_size=hparams.dur_predictor_kernel_size,
            dropout=hparams.p_dur_predictor_dropout,
            n_layers=hparams.dur_predictor_n_layers,
        )

        self.decoder = FFTransformer(
            n_symbols=hparams.n_symbols,
            n_layer=hparams.out_fft_n_layers,
            n_head=hparams.out_fft_n_heads,
            d_model=hparams.symbols_embedding_dim,
            d_head=hparams.out_fft_d_head,
            d_inner=4 * hparams.symbols_embedding_dim,
            kernel_size=hparams.out_fft_conv1d_kernel_size,
            dropout=hparams.p_out_fft_dropout,
            dropatt=hparams.p_out_fft_dropatt,
            dropemb=hparams.p_out_fft_dropemb,
            embed_input=False,
            pad_idx=hparams.pad_idx
        )

        self.pitch_predictor = TemporalPredictor(
            input_size=hparams.symbols_embedding_dim,
            filter_size=hparams.pitch_predictor_filter_size,
            kernel_size=hparams.pitch_predictor_kernel_size,
            dropout=hparams.p_pitch_predictor_dropout,
            n_layers=hparams.pitch_predictor_n_layers,
        )
        
        """
        Here you need to declare all additional layers you need
        Don't forget about GST module
        """
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
            
    def forward(self, inputs, use_gt_durations=True, use_gt_pitch=True, max_duration=75):
        """
        Inputs:
        * inputs - tuple of:
            * text - texts
            * text_lengths - lengths for padding
            * mel_tgt - targets for melspec
            * mel_lengths - for padding
            * dur_tgt - gt duration
            * pitch tgt - gt pitch
        * use_gt_durations and use_gt_pitch - True for teache-forcing (train and validation estimation)
        * max_duration - maximum possible duration
        """
        text, text_lengths, mel_tgt, mel_lengths, dur_tgt, pitch_tgt = inputs

        """
        Here you should wrtie the code for ecoder
        Internal state should be (before the upsampling):
            enc_state = encoded_text + gst_style + conv(pitch)
        
        """
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
        
        enc_and_style = encoder_output + gst_style
        
        pitch_pred = self.pitch_predictor(enc_and_style, enc_mask)
    
        """
        Duration is predicted in log-scale
        """
        log_dur_pred = self.duration_predictor(enc_and_style, enc_mask)
        dur_pred = torch.clamp(torch.exp(log_dur_pred) - 1, 0, max_duration)
        
        if use_gt_pitch and pitch_tgt is not None:
            pitch_conv = ...
        else:
            pitch_conv = ...
            
        enc_state = enc_and_style + pitch_conv

        upsampled, dec_lens = regulate_len(
            dur_tgt if use_gt_durations else dur_pred, enc_state)
    
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
        
        mel_out = self.projection(dec_out)
        
        return mel_out, dec_mask, dur_pred, log_dur_pred, pitch_pred
    
    def infer(self, inputs, gst_estimator):
        """
        Here you should write the code for FP inference
        Inputs:
        * inputs - text and texts_lengths
        * gst_estimator - nn.Module that predicts style from (encoded_text, text_lengths)
        Returns:
        * mel_output
        """
        text, text_lengths = inputs
        
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
        
        gst_pred = gst_estimator.forward(encoder_output, text_lengths)
        enc_and_style = encoder_output + gst_pred.unsqueeze(1)
        
        log_dur_pred = self.duration_predictor(enc_and_style, enc_mask)
        dur_pred = torch.clamp(torch.exp(log_dur_pred) - 1, min=0)

        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
        
        enc_state = enc_and_style + pitch_conv
        
        len_regulated, dec_lens = regulate_len(dur_pred, enc_state)
        
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
        
        return mel_out.transpose(1, 2)<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>def prepare_loaders(directory, train_filelist, val_filelist, hparams):
    train_dataset = FastPitchDataset(directory, train_filelist)
    valid_dataset = FastPitchDataset(directory, val_filelist)

    collate_fn = FastPitchCollate(hparams.n_mel_channels)
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=hparams.batch_size,
        num_workers=16,
        shuffle=True,
        pin_memory=True,
        drop_last=True,
        collate_fn=collate_fn)

    valid_loader = DataLoader(
        valid_dataset,
        batch_size=hparams.batch_size,
        num_workers=8,
        shuffle=False,
        pin_memory=False,
        collate_fn=collate_fn)
    
    return train_loader, valid_loader
def validate(
    model,
    criterion,
    val_loader,
    device,
    use_gt_durations=False,
):
    """
    Parameters:
    * model - FastPitch
    * criterion - loss operator
    * val_loader - loader for validation data
    * device - torch GPU id
    * use_gt_durations - if use ground truth for durations
    """
    was_training = model.training
    model.eval()
    with torch.no_grad():
        val_meta = defaultdict(float)
        for i, batch in enumerate(val_loader):
            x, y = batch_to_device(batch, device)
            
            # !!!!!!!!!!!!!!!!!!!!!!
            # INSERT YOUR CODE HERE!
            # !!!!!!!!!!!!!!!!!!!!!!
            
            loss, meta = criterion(..., meta_agg="sum")
            for k, v in meta.items():
                val_meta[k] += v
                
        val_meta = {k: v / len(val_loader.dataset) for k, v in val_meta.items()}
        val_loss = val_meta["loss"]

    if was_training:
        model.train()
    return val_loss.item(), val_meta
def adjust_learning_rate(total_iter, opt, learning_rate, warmup_iters):
    """
    For training we use one-cycle learning rate schedule
    Firstly, linear growth of scale
    Then, exponential decay
    """
    if total_iter > warmup_iters:
        scale = 1.0 / (total_iter ** 0.5)
    else:
        scale = total_iter / (warmup_iters ** 1.5)

    for param_group in opt.param_groups:
        param_group["lr"] = learning_rate * scale
def train_fastpitch(dataset_dir, train_filelist_path, val_filelist_path, logdir, hparams):
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(hparams.seed)
    torch.manual_seed(hparams.seed)
    torch.cuda.manual_seed(hparams.seed)

    model = FastPitch(hparams).to(device)
    model.train()

    optimizer = Lamb(model.parameters(),
                     lr=hparams.learning_rate,
                     betas=(0.9, 0.98),
                     eps=1e-9,
                     weight_decay=hparams.weight_decay)
    
    logger = FastpitchLogger(logdir, hparams)
    train_loader, valid_loader = prepare_loaders(dataset_dir, train_filelist_path, val_filelist_path, hparams)
    criterion = FastPitchLoss(hparams)
    
    checkpoint = FastpitchLogger.load_last_checkpoint(logdir)
    step = checkpoint['step'] or 0
    if checkpoint['state_dict']:
        load_from_checkpoint(model, checkpoint['state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = step // len(train_loader)
    
    while True:
        if hasattr(train_loader.sampler, 'set_epoch'):
            train_loader.sampler.set_epoch(epoch)
        if hasattr(train_loader.batch_sampler, 'set_epoch'):
            train_loader.batch_sampler.set_epoch(epoch)

        for batch in train_loader:

            x, y = batch_to_device(batch, device)
            
            adjust_learning_rate(step, optimizer, hparams.learning_rate, hparams.warmup_steps)
            y_pred = ...
            
            # !!!!!!!!!!!!!!!!!!!!!!
            # INSERT YOUR CODE HERE!
            # !!!!!!!!!!!!!!!!!!!!!!
            
            model.zero_grad()
            loss, meta = criterion(...)
            loss.backward()
            
            logger.log_grads(step, model)

            # !!!!!!!!!!!!!!!!!!!!!!
            # INSERT YOUR CODE HERE!
            # !!!!!!!!!!!!!!!!!!!!!!
            
            step += 1
            
            meta['step'] = step
            meta['lr'] = optimizer.param_groups[0]["lr"]
            logger.log_training(step, meta)
            
            if step % hparams.eval_interval == 0:
                loss, meta = validate(model, criterion, valid_loader, device, use_gt_durations=True)
               
                meta['step'] = step
                logger.log_validation(step, meta)
                logger.save_checkpoint(model, optimizer, step, hparams)
                    
            if step >= hparams.num_steps:
                return

        epoch += 1<jupyter_output><empty_output><jupyter_text>There are some signs of convergence after the first day of training at 1 GPU.After ~3 days the model is fully trained.<jupyter_code>train_fastpitch(prepr_data_path, train_filelist_path, val_filelist_path, logs_dir, hparams)<jupyter_output><empty_output><jupyter_text>GST SpaceHere you should visualize the latent GST space<jupyter_code>from sklearn.manifold import TSNE


"""
Please run the model GST encoder over the whole trainset or subcorpus and visualize
"""

# !!!!!!!!!!!!!!!!!!!!!!
# INSERT YOUR CODE HERE!
# !!!!!!!!!!!!!!!!!!!!!!

GSTVec = np.stack(GSTStyles)
GSTVec = GSTVec.reshape(GSTVec.shape[0], -1)
GSTEmbedded = TSNE(n_components=2).fit_transform(GSTVec)

# !!!!!!!!!!!!!!!!!!!!!!
# INSERT YOUR CODE HERE!
# !!!!!!!!!!!!!!!!!!!!!!

"""
Visualize space and try to listen to examples of some clusters you've noticed in the map
"""<jupyter_output><empty_output><jupyter_text>GST estimator<jupyter_code>class StyleEmbeddingEstimator(nn.Module):
    """
    Estimator takes encoder outputs and returns style embeddings
    You can use some aggregation layers here: GRU or FFT + FC and some dropout for regularization
    Output dimension should be hparams.symbols_embedding_dim
    """
    def __init__(self, hparams):
        super(StyleEmbeddingEstimator, self).__init__()
    
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
        
    def forward(self, x, lengths):
        # !!!!!!!!!!!!!!!!!!!!!!
        # INSERT YOUR CODE HERE!
        # !!!!!!!!!!!!!!!!!!!!!!
def load_gst_estimator_checkpoint(predictor, optimizer, filepath):
    checkpoint = torch.load(filepath, map_location="cpu")
    predictor.load_state_dict(checkpoint['state_dict'])
    optimizer.load_state_dict(checkpoint["optimizer"])
    return checkpoint['step'], checkpoint['min_eval_loss']


def save_gst_estimator_checkpoint(filepath, step, min_eval_loss, model, optimizer):
    checkpoint = {
        "step": step,
        'min_eval_loss': min_eval_loss,
        "state_dict": model.state_dict(),
        "optimizer": optimizer.state_dict()
    }
    torch.save(checkpoint, filepath)
def train_gst_estimator(dataset_dir, train_filelist_path, val_filelist_path, 
          fastpitch_checkpoint_path, logdir):
    fastpitch_checkpoint = torch.load(fastpitch_checkpoint_path)
    hparams = HParamsFastpitch(fastpitch_checkpoint['params'])
    fastpitch = FastPitch(hparams).to(device)
    load_from_checkpoint(fastpitch, fastpitch_checkpoint['state_dict'])
    fastpitch.eval()
    
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    np.random.seed(hparams.seed)
    torch.manual_seed(hparams.seed)
    torch.cuda.manual_seed(hparams.seed)

    if not os.path.exists(logdir):
        os.makedirs(logdir)
    train_writer = SummaryWriter(os.path.join(logdir, "train"))
    val_writer = SummaryWriter(os.path.join(logdir, "val"))
    
    gst_estimator = StyleEmbeddingEstimator(hparams).to(device)
    
    optimizer = torch.optim.Adam(gst_estimator.parameters(), lr=hparams.gst_estimator_lr)
    
    train_loader, eval_loader = prepare_loaders(dataset_dir, train_filelist_path, val_filelist_path, hparams)
    
    step = 0
    min_eval_loss = np.inf

    checkpoint_path = os.path.join(logdir, f"GST_estimator_best_checkpoint.pt")
    if os.path.isfile(checkpoint_path):
        print("Resume training from checkpoint: ", checkpoint_path)
        step, min_eval_loss = load_gst_estimator_checkpoint(gst_estimator, optimizer, checkpoint_path)
  
    losses = []
    gst_estimator.train()
    while True:
        for batch in train_loader:
            x, y = batch_to_device(batch, device)
            """
            First, get the GT values for styles and encoder_state:
            """
            with torch.no_grad():
                gst_true = ...
                enc_out = ... 
                
            # !!!!!!!!!!!!!!!!!!!!!!
            # INSERT YOUR CODE HERE!
            # !!!!!!!!!!!!!!!!!!!!!!
                
            optimizer.zero_grad()
            gst_pred = ...
            
            """
            It is better to train GST estimator as l1 regression
            """
            loss = ...
            
            # !!!!!!!!!!!!!!!!!!!!!!
            # INSERT YOUR CODE HERE!
            # !!!!!!!!!!!!!!!!!!!!!!
            
            losses.append(loss.item())
            loss.backward()
            optimizer.step()
            step += 1
            
            if step % hparams.gst_estimator_eval_interval == 0:
                train_writer.add_scalar('loss', np.mean(losses), step)
                print(f"train: {step:<3d} loss: {np.mean(losses):<5.4f}")
                
                losses = []
                gst_estimator.eval()
                for batch in eval_loader:
                    x, y = batch_to_device(batch, device)
                    
                    """
                    The same, but for validation:
                    """
                     
                    with torch.no_grad():
                        gst_true = ...
                        enc_out = ...
                        gst_pred = ...
                        loss = ...
                        # !!!!!!!!!!!!!!!!!!!!!!
                        # INSERT YOUR CODE HERE!
                        # !!!!!!!!!!!!!!!!!!!!!!
                        losses.append(loss.item())
                         
                val_writer.add_scalar('loss', np.mean(losses), step)
                print(f"val: {step:<3d} loss: {np.mean(losses):<5.4f}")
                
                """
                Fallback to the prev model if the new one is not better:
                """
                if np.mean(losses) < min_eval_loss:
                    min_eval_loss = np.mean(losses)
                    checkpoint_path = os.path.join(logdir, f"GST_estimator_best_checkpoint.pt")
                    save_gst_estimator_checkpoint(checkpoint_path, step, min_eval_loss, gst_estimator, optimizer)
                                
                for w in train_writer.all_writers.values():
                    w.flush()
                for w in val_writer.all_writers.values():
                    w.flush()
                 
                if step >= hparams.gst_estimator_num_steps:
                    exit()
                
                losses = []
                gst_estimator.train()
fastpitch_checkpoint_path = "logs_dir/0000200_snapshot.pth"

"""
Usually few hours is enough to train the model well
"""

train_gst_estimator(prepr_data_path, train_filelist_path, val_filelist_path, 
          fastpitch_checkpoint_path, gst_logs_dir)
gst_estimator_checkpoint_path = os.path.join(gst_logs_dir, "GST_estimator_best_checkpoint.pt")<jupyter_output><empty_output><jupyter_text>Full model inference  load waveglow checkpointhere we are using the same vocoder -- waveglow<jupyter_code>wg_url = "https://api.ngc.nvidia.com/v2/models/nvidia/waveglow_ckpt_amp_256/versions/20.01.0/zip"
wg_checkpoint_path = os.path.join(pretrained_models_path, "waveglow")
wg_zip_path = os.path.join(pretrained_models_path, "waveglow.zip")
wg_unpacked_path = os.path.join(pretrained_models_path, 'nvidia_waveglow256pyt_fp16.pt')


getfile(wg_url, wg_zip_path)
shutil.unpack_archive(wg_zip_path, pretrained_models_path)
shutil.move(wg_unpacked_path, wg_checkpoint_path)
os.remove(wg_zip_path)
def inference(wg_checkpoint_path, fastpitch_checkpoint_path, gst_estimator_checkpoint_path,
              test_texts, synth_wavs_path):
    wg_checkpoint = torch.load(wg_checkpoint_path)
    waveglow = WaveGlow(**wg_checkpoint['config']).to(device)
    load_from_checkpoint(waveglow, wg_checkpoint['state_dict'])
    waveglow = waveglow.remove_weightnorm(waveglow)
    waveglow.eval()
    
    denoiser = Denoiser(waveglow).to(device)
    
    fastpitch_checkpoint = torch.load(fastpitch_checkpoint_path)
    hparams = HParamsFastpitch(fastpitch_checkpoint['params'])
    
    fastpitch = FastPitch(hparams).to(device)
    load_from_checkpoint(fastpitch, fastpitch_checkpoint['state_dict'])
    fastpitch.eval()
    
    gst_estimator_checkpoint = torch.load(gst_estimator_checkpoint_path)
    gst_estimator = StyleEmbeddingEstimator(hparams).to(device)
    load_from_checkpoint(gst_estimator, gst_estimator_checkpoint['state_dict'])
    gst_estimator.eval()
    
    texts_encoded = [(fid, tp.encode_text(txt)) for fid, txt in test_texts.items()]

    if not os.path.exists(synth_wavs_path):
        os.mkdir(synth_wavs_path)
    
    for fid, text in tqdm.tqdm(texts_encoded):
        with torch.no_grad():
            text = torch.LongTensor(text)
            text_length = torch.LongTensor([text.size(0)]).to(device)
            text = text.unsqueeze(0).to(device)
            mel_pred = fastpitch.infer((text, text_length), gst_estimator)
            audio_pred = waveglow.infer(mel_pred, sigma=hparams.wg_sigma_infer)
            audio_pred = denoiser(audio_pred.float(), strength=hparams.wg_denoising_strength)
            audio_pred = audio_pred.squeeze(0).squeeze(0).cpu().numpy()

            wavfile.write(os.path.join(synth_wavs_path, fid + ".wav"), hparams.sample_rate, 
                          (np.clip(audio_pred, -1, 1) * 32767).astype(np.int16))<jupyter_output><empty_output><jupyter_text>Check yourself:<jupyter_code>inference(wg_checkpoint_path, fastpitch_checkpoint_path, gst_estimator_checkpoint_path, 
          {"TST": "The quick brown fox jumps over the lazy dog."}, synth_wavs_path)
ipd.display(ipd.Audio('synth_wavs/TST.wav'))<jupyter_output><empty_output><jupyter_text>The testHere is the estimation of your model:you should render records from the given file 'test_texts.txt' ,they should be intelligible and correspond to the text:<jupyter_code>with open('test_texts.txt') as ifile:
    test_data = {
        "S{:03}_test".format(i): l.strip()
        for i, l in enumerate(ifile)
    }
    
sorted(test_data.items())
inference(wg_checkpoint_path, fastpitch_checkpoint_path, gst_estimator_checkpoint_path, 
          test_data, synth_wavs_path)
for fid, text in sorted(test_data.items()):
    print(fid)
    print(text)
    wav_filename = os.path.join(synth_wavs_path, fid + ".wav")
    ipd.display(ipd.Audio(wav_filename))
    print("----------")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11992921050044243, 'fraction_numerical': 0.006705338708091633, 'mean_word_length': 2.4998279540293167, 'pattern_counts': {'":': 4, '<': 99, '<?xml version=': 0, '>': 99, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '22', 'dir': 'jupyter-structured-clean-dedup', 'id': '577326', 'license': 'MIT', 'n_tokens_mistral': 16880, 'n_tokens_neox': 16053, 'n_words': 3349, 'path': 'week_11/fastpitch_hw3.ipynb', 'repo_name': 'RunFMe/speech_course', 'size': '50855'}