labhamlet/wavjepa-base

1from typing import TypedDict2from torch import nn3 4 5class TransformerLayerCFG(TypedDict):6    d_model : int7    nhead : int8    batch_first : bool9    norm_first : bool10    bias : bool11    dim_feedforward : int12    dropout : float13    layer_norm_eps : float14 15    @classmethod16    def create(cls, 17    d_model : int = 768,18    nhead : int = 12,19    batch_first : bool = True,20    norm_first : bool = False,21    bias : bool = True,22    mlp_ratio : float = 4.0,23    dropout : float = 0.0,24    layer_norm_eps : float = 1e-6) -> 'TransformerLayerCFG':25        return TransformerLayerCFG(d_model = d_model, 26                                   nhead = nhead, 27                                   batch_first = batch_first,28                                   norm_first = norm_first, 29                                   bias = bias, 30                                   dim_feedforward = int(d_model * mlp_ratio),31                                   dropout = dropout, 32                                   layer_norm_eps = layer_norm_eps)33 34 35# Norm needs to be defined by the user!36class TransformerEncoderCFG(TypedDict):37    num_layers : int38    enable_nested_tensor: bool39    mask_check: bool40 41    @classmethod42    def create(cls, 43    num_layers : int = 12, 44    enable_nested_tensor: bool = False,45    mask_check: bool = True) -> 'TransformerEncoderCFG':46        return TransformerEncoderCFG(num_layers=num_layers, 47                enable_nested_tensor = enable_nested_tensor,48                mask_check = mask_check)