Spaces:
Running on Zero
Running on Zero
| import os | |
| from typing import Any | |
| import torch | |
| import torch.nn as nn | |
| from transformers import PreTrainedModel | |
| from transformers.modeling_outputs import CausalLMOutput, Wav2Vec2BaseModelOutput | |
| from .configuration_zipformer import ZipformerConfig | |
| from .scaling import ScheduledFloat | |
| from .utils import GradMultiply | |
| from .wav2vec2_module import ConvFeatureExtractionModel | |
| from .zipformer import Zipformer2 | |
| class ZipformerPreTrainedModel(PreTrainedModel): | |
| """ | |
| An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained | |
| models. | |
| """ | |
| config_class = ZipformerConfig | |
| base_model_prefix = "zipformer" | |
| supports_gradient_checkpointing = False | |
| def _init_weights(self, module: nn.Module) -> None: | |
| """Initialize the weights""" | |
| if isinstance(module, (nn.Linear, nn.Embedding)): | |
| # Slightly different from the TF version which uses truncated_normal for initialization | |
| # cf https://github.com/pytorch/pytorch/pull/5617 | |
| module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) | |
| elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): | |
| module.bias.data.zero_() | |
| module.weight.data.fill_(1.0) | |
| elif isinstance(module, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): | |
| if hasattr(module, "parametrizations"): | |
| nn.init.kaiming_normal_(module.parametrizations.weight.original0.data) | |
| nn.init.kaiming_normal_(module.parametrizations.weight.original1.data) | |
| nn.init.kaiming_normal_(module.weight.data) | |
| if isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d)) and module.bias is not None: | |
| module.bias.data.zero_() | |
| class ZipformerModel(ZipformerPreTrainedModel): | |
| def __init__(self, config: ZipformerConfig) -> None: | |
| super().__init__(config) | |
| feature_enc_layers = config.conv_feature_layers | |
| self.embed = feature_enc_layers[-1][0] | |
| self.feature_extractor = ConvFeatureExtractionModel( | |
| conv_layers=feature_enc_layers, | |
| dropout=0.0, | |
| mode=config.extractor_mode, | |
| conv_bias=config.conv_bias, | |
| ) | |
| encoder_input_dim = config.encoder_dim[0] | |
| self.post_extract_proj = nn.Linear(self.embed, encoder_input_dim) if self.embed != encoder_input_dim else None | |
| self.dropout_input = nn.Dropout(config.dropout_input) | |
| self.dropout_features = nn.Dropout(config.dropout_features) | |
| self.feature_grad_mult = config.feature_grad_mult | |
| self.mask_emb = nn.Parameter(torch.FloatTensor(encoder_input_dim).uniform_()) | |
| self.encoder = Zipformer2( | |
| output_downsampling_factor=1, | |
| downsampling_factor=config.downsampling_factor, | |
| num_encoder_layers=config.num_encoder_layers, | |
| encoder_dim=config.encoder_dim, | |
| encoder_unmasked_dim=config.encoder_unmasked_dim, | |
| query_head_dim=config.query_head_dim, | |
| pos_head_dim=config.pos_head_dim, | |
| value_head_dim=config.value_head_dim, | |
| pos_dim=config.pos_dim, | |
| num_heads=config.num_heads, | |
| feedforward_dim=config.feedforward_dim, | |
| cnn_module_kernel=config.cnn_module_kernel, | |
| dropout=ScheduledFloat((0.0, 0.3), (20000.0, 0.1)), | |
| warmup_batches=4000.0, | |
| ) | |
| self.layer_norm = nn.LayerNorm(self.embed) | |
| self.post_init() | |
| def from_icefall_checkpoint( | |
| cls, | |
| checkpoint_path: os.PathLike[str], | |
| **kwargs: Any, | |
| ) -> "ZipformerModel": | |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) | |
| config = ZipformerConfig.from_icefall_checkpoint(checkpoint_path, **kwargs) | |
| model = cls(config) | |
| state_dict = checkpoint["model"] | |
| state_dict.pop("final_proj.weight") | |
| state_dict.pop("final_proj.bias") | |
| model.load_state_dict(state_dict, strict=True) | |
| return model | |
| def forward_features( | |
| self, | |
| source: torch.Tensor, | |
| padding_mask: torch.BoolTensor | None = None, | |
| ) -> torch.Tensor: | |
| if self.feature_grad_mult > 0: | |
| features = self.feature_extractor(source, padding_mask) | |
| if self.feature_grad_mult != 1.0: | |
| features = GradMultiply.apply(features, self.feature_grad_mult) | |
| else: | |
| with torch.no_grad(): | |
| features = self.feature_extractor(source, padding_mask) | |
| return features | |
| def forward_padding_mask( | |
| self, | |
| features: torch.Tensor, | |
| padding_mask: torch.Tensor, | |
| ) -> torch.Tensor: | |
| output_lengths = (~padding_mask).sum(dim=-1) | |
| for _, kernel_size, stride in self.config.conv_feature_layers: | |
| output_lengths = ( | |
| torch.div( | |
| output_lengths - kernel_size, | |
| stride, | |
| rounding_mode="floor", | |
| ) | |
| + 1 | |
| ) | |
| positions = torch.arange(features.size(1), device=features.device) | |
| return positions.unsqueeze(0) >= output_lengths.unsqueeze(1) | |
| def forward( | |
| self, | |
| input_values: torch.FloatTensor, | |
| padding_mask: torch.BoolTensor | None = None, | |
| ) -> Wav2Vec2BaseModelOutput: | |
| """output layer is 1-based""" | |
| features = self.forward_features(input_values, padding_mask) | |
| features = features.transpose(1, 2) | |
| features = self.layer_norm(features) | |
| if padding_mask is None: | |
| padding_mask = torch.zeros( | |
| features.shape[:2], | |
| dtype=torch.bool, | |
| device=features.device, | |
| ) | |
| else: | |
| padding_mask = self.forward_padding_mask(features, padding_mask) | |
| if self.post_extract_proj is not None: | |
| features = self.post_extract_proj(features) | |
| features = self.dropout_input(features) | |
| x = features | |
| # feature: (B, T, D), float | |
| # target: (B, T), long | |
| # x: (B, T, D), float -> (T, B, D), float | |
| # padding_mask: (B, T), bool | |
| # mask_indices: (B, T), bool | |
| x = x.transpose(0, 1) | |
| x, _ = self.encoder( | |
| x, | |
| (~padding_mask).sum(dim=-1), | |
| src_key_padding_mask=padding_mask, | |
| ) | |
| x = x.transpose(0, 1) | |
| return Wav2Vec2BaseModelOutput( | |
| last_hidden_state=x, | |
| extract_features=features, | |
| ) | |
| class ZipformerForCTC(ZipformerPreTrainedModel): | |
| def __init__(self, config: ZipformerConfig) -> None: | |
| super().__init__(config) | |
| self.encoder = ZipformerModel(config) | |
| self.ctc_output = nn.Sequential( | |
| nn.Dropout(p=0.1), | |
| nn.Linear(max(config.encoder_dim), config.vocab_size), | |
| ) | |
| def from_icefall_checkpoint( | |
| cls, | |
| checkpoint_path: os.PathLike[str], | |
| **kwargs: Any, | |
| ) -> "ZipformerForCTC": | |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) | |
| config = ZipformerConfig.from_icefall_checkpoint(checkpoint_path, **kwargs) | |
| model = cls(config) | |
| state_dict = checkpoint["model"] | |
| state_dict.pop("encoder.final_proj.weight") | |
| state_dict.pop("encoder.final_proj.bias") | |
| model.load_state_dict(state_dict, strict=True) | |
| return model | |
| def forward( | |
| self, | |
| input_values: torch.FloatTensor, | |
| padding_mask: torch.BoolTensor | None = None, | |
| **kwargs: Any, | |
| ) -> CausalLMOutput: | |
| if padding_mask is None: | |
| padding_mask = torch.zeros_like(input_values, dtype=torch.bool) | |
| encoder_out = self.encoder(input_values, padding_mask) | |
| ctc_output = self.ctc_output(encoder_out.last_hidden_state) | |
| return CausalLMOutput(loss=None, logits=ctc_output, hidden_states=None, attentions=None) | |