| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
| from transformers import HubertConfig, HubertModel, PreTrainedModel |
| from transformers.utils import ModelOutput |
|
|
| from .configuration_mhubert_ipa_ctc_ft import MHuBERTIPACTCFTConfig |
|
|
|
|
| @dataclass |
| class MHuBERTIPACTCFTOutput(ModelOutput): |
| logits: torch.Tensor = None |
| hidden_states: torch.Tensor = None |
|
|
|
|
| class MHuBERTIPACTCFTModel(PreTrainedModel): |
| config_class = MHuBERTIPACTCFTConfig |
| base_model_prefix = "mhubert_ipa_ctc_ft" |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| arch = config.architecture |
| backbone_cfg = HubertConfig.from_dict(config.backbone_config) |
| self.blank_id = int(arch["blank_id"]) |
| self.backbone = HubertModel(backbone_cfg) |
| self.proj = nn.Linear(arch["input_dim"], arch["proj_dim"]) |
| self.lstm = nn.LSTM( |
| arch["proj_dim"], |
| arch["lstm_hidden"], |
| num_layers=arch["lstm_layers"], |
| bidirectional=arch["lstm_bidirectional"], |
| batch_first=True, |
| dropout=arch["dropout"] if arch["lstm_layers"] > 1 else 0.0, |
| ) |
| self.drop = nn.Dropout(arch["dropout"]) |
| out_dim = arch["lstm_hidden"] * (2 if arch["lstm_bidirectional"] else 1) |
| self.head = nn.Linear(out_dim, arch["output_dim"]) |
| self.post_init() |
|
|
| def forward(self, input_values, attention_mask=None, **kwargs): |
| backbone_out = self.backbone(input_values=input_values, attention_mask=attention_mask, **kwargs) |
| x = self.proj(backbone_out.last_hidden_state) |
| out, _ = self.lstm(x) |
| logits = self.head(self.drop(out)) |
| return MHuBERTIPACTCFTOutput(logits=logits, hidden_states=backbone_out.last_hidden_state) |
|
|
|
|
