| 131 | # ----------------------------------------------------------------------------- |
| 132 | |
| 133 | def create_model(vocab_size) -> nn.Module: |
| 134 | return conformer_model.Conformer( |
| 135 | num_classes=vocab_size, |
| 136 | input_dim=80, |
| 137 | encoder_dim=144, |
| 138 | num_encoder_layers=16, |
| 139 | num_attention_heads=4, |
| 140 | feed_forward_expansion_factor=4, |
| 141 | conv_expansion_factor=2, |
| 142 | input_dropout_p=0.1, |
| 143 | feed_forward_dropout_p=0.1, |
| 144 | attention_dropout_p=0.1, |
| 145 | conv_dropout_p=0.1, |
| 146 | conv_kernel_size=31, |
| 147 | half_step_residual=True, |
| 148 | ) |
| 149 | |
| 150 | |
| 151 | # ----------------------------------------------------------------------------- |