switch to beartype

from functools import partial

from typing import Optional, Union, List

from typeguard import typechecked

from beartype import beartype

import torch

from torch import nn, einsum

# the three hierarchical transformers

@typechecked

@beartype

class SemanticTransformer(nn.Module):

    def __init__(

        self,

        tokens = self.transformer(tokens, context = text_embeds, context_mask = text_mask)

        return self.to_logits(tokens)

@typechecked

@beartype

class CoarseTransformer(nn.Module):

    def __init__(

        self,

# training wrappers

@typechecked

@beartype

class SemanticTransformerWrapper(nn.Module):

    def __init__(

        self,

        return loss

@typechecked

@beartype

class CoarseTransformerWrapper(nn.Module):

    def __init__(

        self,

            coarse_loss * num_coarse_logits

        ) / (num_semantic_logits + num_coarse_logits)

@typechecked

@beartype

class FineTransformerWrapper(nn.Module):

    def __init__(

        self,

# audio LM

@typechecked

@beartype

class AudioLM(nn.Module):

    def __init__(

        self,

  ],

  install_requires=[

    'accelerate',

    'beartype',

    'einops>=0.6',

    'ema-pytorch',

    'fairseq',

    'torchaudio',

    'transformers',

    'tqdm',

    'typeguard',

    'vector-quantize-pytorch>=0.10.11'

  ],

  classifiers=[