Initialize OpenAIWhisperEncoderWarp. Args: dropout_rate: TODO. whisper_model: Whisper Model instance. download_dir: TODO. use_specaug: TODO. use_padmask: TODO. specaug_conf: Configura
(
self,
dropout_rate: float = 0.0,
whisper_model: str = "small",
download_dir: str = None,
use_specaug: bool = False,
use_padmask: bool = False,
specaug_conf: Union[dict, None] = None,
)
| 22 | """ |
| 23 | |
| 24 | def __init__( |
| 25 | self, |
| 26 | dropout_rate: float = 0.0, |
| 27 | whisper_model: str = "small", |
| 28 | download_dir: str = None, |
| 29 | use_specaug: bool = False, |
| 30 | use_padmask: bool = False, |
| 31 | specaug_conf: Union[dict, None] = None, |
| 32 | ): |
| 33 | """Initialize OpenAIWhisperEncoderWarp. |
| 34 | |
| 35 | Args: |
| 36 | dropout_rate: TODO. |
| 37 | whisper_model: Whisper Model instance. |
| 38 | download_dir: TODO. |
| 39 | use_specaug: TODO. |
| 40 | use_padmask: TODO. |
| 41 | specaug_conf: Configuration dict for specaug. |
| 42 | """ |
| 43 | super().__init__() |
| 44 | |
| 45 | # note that originally Whisper doesn't use dropouts |
| 46 | self.dropout = torch.nn.Dropout(dropout_rate) |
| 47 | |
| 48 | assert whisper_model in whisper.available_models() |
| 49 | _model = whisper.load_model(whisper_model, download_root=download_dir, device="cpu") |
| 50 | self.encoders = copy.deepcopy(_model.encoder) |
| 51 | self.encoders.train() |
| 52 | |
| 53 | del _model |
| 54 | |
| 55 | if use_specaug: |
| 56 | self.specaug = SpecAug(**specaug_conf) |
| 57 | else: |
| 58 | self.specaug = None |
| 59 | self.use_padmask = use_padmask |
| 60 | |
| 61 | def whisper_encode( |
| 62 | self, |