from transformers.models.wav2vec2.configuration_wav2vec2 import Wav2Vec2Config


class SpeakerEmbeddingConfig(Wav2Vec2Config):

    def __init__(
        self,
        embedding_size: int = 768,
        train_batch_speakers: int = 64,
        train_batch_per_device: int = 64,
        train_batch_samples_per_speaker: int = 10,
        disable_positional_embeddings: bool = False,
        loss_margin: float = 0.2,
        loss_scale: float = 30.0,
        use_layer_weights: bool = True,
        n_projection_layers: int = 3,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.embedding_size = embedding_size
        self.train_batch_per_device = train_batch_per_device
        self.train_batch_speakers = train_batch_speakers
        self.train_batch_samples_per_speaker = train_batch_samples_per_speaker
        self.disable_positional_embeddings = disable_positional_embeddings
        self.loss_margin = loss_margin
        self.loss_scale = loss_scale
        self.use_layer_weights = use_layer_weights
        self.n_projection_layers = n_projection_layers
        #self.conv_stride = (7, 2, 2, 2, 2, 2, 2)
        #self.conv_kernel = (15, 5, 3, 3, 3, 2, 2)
        #self.conv_dim = (512, 512, 512, 512, 512, 512, 512)