| from dataclasses import dataclass |
| from typing import Optional |
|
|
| import torch |
| from transformers import WhisperConfig |
| from transformers.modeling_outputs import Seq2SeqLMOutput, BaseModelOutput, Seq2SeqModelOutput |
|
|
|
|
| @dataclass |
| class Seq2SeqLMOutputLosses(Seq2SeqLMOutput): |
| enc_loss: Optional[torch.FloatTensor] = None |
| dec_loss: Optional[torch.FloatTensor] = None |
| encoder_logits: Optional[torch.FloatTensor] = None |
|
|
|
|
| @dataclass |
| class BaseModelOutputLogit(BaseModelOutput): |
| logits: Optional[torch.FloatTensor] = None |
|
|
|
|
| @dataclass |
| class Seq2SeqModelOutputLogit(Seq2SeqModelOutput): |
| encoder_logits: Optional[torch.FloatTensor] = None |
|
|
|
|
| class DiCoWConfig(WhisperConfig): |
| """This is a modified version of the `WhisperEncoder` model from the `transformers` library. |
| The model has been modified to support CTC loss computation in the forward pass.""" |
| model_type = "DiCoW" |
|
|
| def __init__( |
| self, |
| ctc_loss_reduction: str = "mean", |
| final_dropout: float = 0.0, |
| ctc_zero_infinity: bool = False, |
| ctc_weight: float = 0.0, |
| blank_token_id: Optional[int] = None, |
| additional_layer: bool = False, |
| additional_self_attention_layer: bool = False, |
| sub_sample: bool = False, |
| use_fddt: bool = True, |
| fddt_is_diagonal: bool = True, |
| fddt_bias_only: bool = False, |
| fddt_use_silence: bool = True, |
| fddt_use_target: bool = True, |
| fddt_use_overlap: bool = True, |
| fddt_use_non_target: bool = True, |
| remove_timestamps_from_ctc: bool = False, |
| apply_fddt_to_n_layers: int = -1, |
| fddt_init: str = 'non-disturbing', |
| n_soft_prompts: int = 16, |
| mt_num_speakers: int = 1, |
| is_mt: bool = False, |
| non_target_fddt_value: float = 0.0, |
| use_initial_fddt: bool = False, |
| scb_method: str = None, |
| scb_layers: int = -1, |
| contrastive_loss_weight: float = 0.0, |
| use_enrollment_network: bool = False, |
| spk_embedding_extraction_layer: int = -1, |
| num_speakers: int = -1, |
| sid_loss_weight: float = 0.0, |
| attend_to_enrollment: bool = False, |
| uses_enrollments: bool = False, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.ctc_loss_reduction = ctc_loss_reduction |
| self.final_dropout = final_dropout |
| self.ctc_zero_infinity = ctc_zero_infinity |
| self.ctc_weight = ctc_weight |
| self.blank_token_id = blank_token_id |
| self.additional_layer = additional_layer |
| self.additional_self_attention_layer = additional_self_attention_layer |
| self.sub_sample = sub_sample |
| self.use_fddt = use_fddt |
| self.fddt_is_diagonal = fddt_is_diagonal |
| self.fddt_bias_only = fddt_bias_only |
| self.fddt_use_silence = fddt_use_silence |
| self.fddt_use_target = fddt_use_target |
| self.fddt_use_overlap = fddt_use_overlap |
| self.fddt_use_non_target = fddt_use_non_target |
| self.remove_timestamps_from_ctc = remove_timestamps_from_ctc |
| self.apply_fddt_to_n_layers = apply_fddt_to_n_layers |
| self.fddt_init = fddt_init |
| self.n_soft_prompts = n_soft_prompts |
| self.mt_num_speakers = mt_num_speakers |
| self.non_target_fddt_value = non_target_fddt_value |
| self.use_initial_fddt = use_initial_fddt |
| self.scb_method = scb_method |
| self.scb_layers = scb_layers |
| self.contrastive_loss_weight = contrastive_loss_weight |
| self.is_mt = is_mt |
| self.use_enrollment_network = use_enrollment_network |
| self.spk_embedding_extraction_layer = spk_embedding_extraction_layer |
| self.num_speakers = num_speakers |
| self.sid_loss_weight = sid_loss_weight |
| self.attend_to_enrollment = attend_to_enrollment |
| self.use_enrollment_network = use_enrollment_network |
| self.uses_enrollments = uses_enrollments |
|
|
|
|
| _HIDDEN_STATES_START_POSITION = 2 |
|
|
