get a working example of text conditioned semantic transformer training

```

## Text Conditioned Audio Synthesis

ex. Semantic Transformer

```python

import torch

from audiolm_pytorch import HubertWithKmeans, SemanticTransformer, SemanticTransformerTrainer, get_embeds, FairseqVQWav2Vec

wav2vec = HubertWithKmeans(

    checkpoint_path = './hubert/hubert_base_ls960.pt',

    kmeans_path = './hubert/hubert_base_ls960_L9_km500.bin'

)

semantic_transformer = SemanticTransformer(

    num_semantic_tokens = 500,

    dim = 1024,

    depth = 6,

    has_condition = True  # this will have to be set to True

).cuda()

# mock text video dataset (as an example)

# you will have to extend your own from `Dataset`, and return an audio tensor as well as a string (the audio description) in any order (the framework will autodetect and route it into the transformer)

from torch.utils.data import Dataset

class MockTextAudioDataset(Dataset):

    def __init__(self, length = 100, audio_length = 320 * 32):

        super().__init__()

        self.audio_length = audio_length

        self.len = length

    def __len__(self):

        return self.len

    def __getitem__(self, idx):

        mock_audio = torch.randn(self.audio_length)

        mock_caption = 'audio caption'

        return mock_caption, mock_audio

dataset = MockTextAudioDataset()

# instantiate semantic transformer trainer and train

trainer = SemanticTransformerTrainer(

    transformer = semantic_transformer,

    wav2vec = wav2vec,

    dataset = dataset,

    batch_size = 4,

    grad_accum_every = 8,

    data_max_length = 320 * 32,

    num_train_steps = 100000

)

trainer.train()

# after much training above

sample = trainer.generate(text = ['sound of rain drops on the rooftops'], batch_size = 1, max_length = 2) # (1, < 128) - may terminate early if it detects [eos]

```

## Appreciation

- <a href="https://stability.ai/">Stability.ai</a> for the generous sponsorship to work and open source cutting edge artificial intelligence research

    @eval_decorator

    @torch.no_grad()

    @beartype

    def generate(

        self,

        *,

    @eval_decorator

    @torch.no_grad()

    @beartype

    def generate(

        self,

        *,

    @eval_decorator

    @torch.no_grad()

    @beartype

    def generate(

        self,

        *,

from pathlib import Path

from functools import partial, wraps

from typing import Tuple

from beartype.door import is_bearable

import torchaudio

from torchaudio.functional import resample

        is_one_data = not isinstance(data[0], tuple)

        if is_one_data:

            return fn(data)

            data = (data,)

        outputs = []

        for datum in zip(*data):

            if is_bearable(datum, Tuple[str, ...]):

                output = list(datum)

            else:

                output = fn(datum)

            outputs.append(output)

        return tuple(map(fn, zip(*data)))

        return tuple(outputs)

    return inner

from pathlib import Path

from shutil import rmtree

from typing import Union, List, Optional

from typing import Union, List, Optional, Tuple

from typing_extensions import Annotated

from beartype import beartype

        *,

        num_train_steps,

        batch_size,

        dataset: Optional[Dataset] = None,

        data_max_length = None,

        folder,

        folder = None,

        lr = 3e-4,

        grad_accum_every = 1,

        wd = 0.,

        # create dataset

        self.ds = dataset

        if not exists(self.ds):

            assert exists(folder), 'folder must be passed in, if not passing in a custom dataset for text conditioned audio synthesis training'

            self.ds = SoundDataset(

                folder,

                max_length = data_max_length,

                seq_len_multiple_of = wav2vec.seq_len_multiple_of

            )

        self.ds_fields = None

        # split for validation

        if valid_frac > 0:

    def is_local_main(self):

        return self.accelerator.is_local_main_process

    def data_tuple_to_kwargs(self, data):

        if not exists(self.ds_fields):

            self.ds_fields = determine_types(data, DATASET_FIELD_TYPE_CONFIG)

            assert not has_duplicates(self.ds_fields), 'dataset fields must not have duplicate field names'

        return dict(zip(self.ds_fields, data))

    def train_step(self):

        device = self.device

        # update vae (generator)

        for _ in range(self.grad_accum_every):

            wave = next(self.dl_iter).to(device)

            data_kwargs = self.data_tuple_to_kwargs(next(self.dl_iter))

            loss = self.train_wrapper(raw_wave = wave, return_loss = True)

            loss = self.train_wrapper(**data_kwargs, return_loss = True)

            self.accelerator.backward(loss / self.grad_accum_every)

        # sample results every so often

        if self.is_main and not (steps % self.save_results_every):

            wave = next(self.valid_dl_iter).to(device)

            data_kwargs = self.data_tuple_to_kwargs(next(self.valid_dl_iter))

            with torch.no_grad():

                self.train_wrapper.eval()

                valid_loss = self.train_wrapper(raw_wave = wave, return_loss = True)

                valid_loss = self.train_wrapper(**data_kwargs, return_loss = True)

            self.print(f'{steps}: valid loss {valid_loss}')

        *,

        num_train_steps,

        batch_size,

        dataset: Optional[Dataset] = None,

        ds_fields: Tuple[str, ...] = ('raw_wave', 'raw_wave_for_soundstream', 'text'),

        data_max_length = None,

        folder,

        folder = None,

        lr = 3e-4,

        grad_accum_every = 1,

        wd = 0.,

        # create dataset

        self.ds = dataset

        if not exists(self.ds):

            assert exists(folder), 'folder must be passed in, if not passing in a custom dataset for text conditioned audio synthesis training'

            self.ds = SoundDataset(

                folder,

                max_length = data_max_length,

                seq_len_multiple_of = soundstream.seq_len_multiple_of

            )

        self.ds_fields = ds_fields

        # split for validation

        if valid_frac > 0:

        # update vae (generator)

        for _ in range(self.grad_accum_every):

            wave_wav2vec, wave_soundstream = tuple(map(lambda t: t.to(device), next(self.dl_iter)))

            data_kwargs = dict(zip(self.ds_fields, next(self.dl_iter)))

            loss = self.train_wrapper(

                raw_wave = wave_wav2vec,

                raw_wave_for_soundstream = wave_soundstream,

                **data_kwargs,

                return_loss = True

            )

        # sample results every so often

        if self.is_main and not (steps % self.save_results_every):

            wave_wav2vec, wave_soundstream = tuple(map(lambda t: t.to(device), next(self.valid_dl_iter)))

            data_kwargs = dict(zip(self.ds_fields, next(self.valid_dl_iter)))

            with torch.no_grad():

                self.train_wrapper.eval()

                valid_loss = self.train_wrapper(

                    raw_wave = wave_wav2vec,

                    raw_wave_for_soundstream = wave_soundstream,

                    **data_kwargs,

                    return_loss = True

                )

        *,

        num_train_steps,

        batch_size,

        dataset: Optional[Dataset] = None,

        data_max_length = None,

        folder,

        folder = None,

        lr = 3e-4,

        grad_accum_every = 1,

        wd = 0.,

        # create dataset

        self.ds = dataset

        if not exists(self.ds):

            assert exists(folder), 'folder must be passed in, if not passing in a custom dataset for text conditioned audio synthesis training'

            self.ds = SoundDataset(

                folder,

                max_length = data_max_length,

                seq_len_multiple_of = soundstream.seq_len_multiple_of

            )

        self.ds_fields = None

        # split for validation

        if valid_frac > 0:

    def is_local_main(self):

        return self.accelerator.is_local_main_process

    def data_tuple_to_kwargs(self, data):

        if not exists(self.ds_fields):

            self.ds_fields = determine_types(data, DATASET_FIELD_TYPE_CONFIG)

            assert not has_duplicates(self.ds_fields), 'dataset fields must not have duplicate field names'

        return dict(zip(self.ds_fields, data))

    def train_step(self):

        device = self.device

        # update vae (generator)

        for _ in range(self.grad_accum_every):

            wave = next(self.dl_iter).to(device)

            loss = self.train_wrapper(raw_wave = wave, return_loss = True)

            data_kwargs = self.data_tuple_to_kwargs(next(self.dl_iter))

            loss = self.train_wrapper(**data_kwargs, return_loss = True)

            self.accelerator.backward(loss / self.grad_accum_every)

        # sample results every so often

        if self.is_main and not (steps % self.save_results_every):

            wave = next(self.valid_dl_iter).to(device)

            data_kwargs = self.data_tuple_to_kwargs(next(self.valid_dl_iter))

            with torch.no_grad():

                self.train_wrapper.eval()

                valid_loss = self.train_wrapper(raw_wave = wave, return_loss = True)

                valid_loss = self.train_wrapper(**data_kwargs, return_loss = True)

            self.print(f'{steps}: valid loss {valid_loss}')

setup(

  name = 'audiolm-pytorch',

  packages = find_packages(exclude=[]),

  version = '0.1.16',

  version = '0.1.17',

  license='MIT',

  description = 'AudioLM - Language Modeling Approach to Audio Generation from Google Research - Pytorch',

  author = 'Phil Wang',