Upload model code for multitask model

wav2vecasr/models.py

@@ -0,0 +1,167 @@
+import random
+import re
+import torch
+import torch.nn as nn
+import torchaudio
+from torch.utils.data import Dataset
+
+
+class DataCollator:
+    def __init__(self, processor, padding, device, augment):
+        self.processor = processor
+        self.padding = padding
+        self.device = device
+        self.sampling_rate = 16000
+        self.augment = augment
+
+        atempos = (0.8, 1.0, 1.25)  # audio tempo atempo=tempo
+        audio_effects = (
+            ("highpass=frequency=1500",),
+            (
+                "vibrato=f=5:d=0.4",
+                "volume=1.5",
+            ),
+            (
+                "aecho=0.8:0.88:30:0.3",
+                "volume=1.5",
+            ),
+        )
+
+        self.effectors = [None]
+        for atempo in atempos:
+            for audio_effect in audio_effects:
+                effect = f"atempo={atempo}," + ",".join(audio_effect)
+                self.effectors.append(torchaudio.io.AudioEffector(effect=effect))
+
+    def __call__(self, data):
+        waveforms, lm_labels, accent_labels, gender_labels = zip(*data)
+        accent_labels = torch.tensor(accent_labels, device=self.device)
+        gender_labels = torch.tensor(gender_labels, device=self.device)
+
+        input_features = [
+            {"input_values": self.random_augment(waveform).squeeze()}
+            for waveform in waveforms
+        ]
+        label_features = [{"input_ids": lm_label} for lm_label in lm_labels]
+
+        padded_waveforms = self.processor.pad(
+            input_features,
+            padding=True,
+            return_tensors="pt",
+        )["input_values"]
+        padded_waveforms = padded_waveforms.to(self.device)
+
+        with self.processor.as_target_processor():
+            padded_lm_labels = self.processor.pad(
+                label_features,
+                padding=True,
+                return_tensors="pt",
+            )
+
+        # replace padding with -100 to ignore loss correctly
+        padded_lm_labels = padded_lm_labels["input_ids"].masked_fill(
+            padded_lm_labels.attention_mask.ne(1), -100
+        )
+        padded_lm_labels = padded_lm_labels.to(self.device)
+
+        return padded_waveforms, padded_lm_labels, accent_labels, gender_labels
+
+    def random_augment(self, waveform):
+        if not self.augment:
+            return waveform
+
+        waveform = torch.tensor(waveform)
+        waveform = torch.transpose(waveform, 0, 1)
+        effector = random.choice(self.effectors)
+        if effector is None:
+            return waveform
+
+        augmented_waveform = effector.apply(waveform, self.sampling_rate)
+        if augmented_waveform.isnan().any() | augmented_waveform.isinf().any():
+            return waveform
+
+        return augmented_waveform
+
+
+class L2ArcticDataset(Dataset):
+    def __init__(self, processor, audio_paths, lm_labels, accent_labels, gender_labels):
+        orig_sampling_rate = 44100
+        new_sampling_rate = 16000
+        resample_transform = torchaudio.transforms.Resample(
+            orig_sampling_rate, new_sampling_rate
+        )
+
+        self.waveforms = []
+        self.lm_labels = []
+        self.accent_labels = accent_labels
+        self.gender_labels = gender_labels
+
+        for audio_path in audio_paths:
+            waveform, _ = torchaudio.load(audio_path)
+            waveform = resample_transform(waveform)
+            self.waveforms.append(
+                processor(waveform, sampling_rate=new_sampling_rate).input_values[0]
+            )
+
+        with processor.as_target_processor():
+            for lm_label in lm_labels:
+                self.lm_labels.append(processor(lm_label).input_ids)
+
+    def __getitem__(self, index):
+        return (
+            self.waveforms[index],
+            self.lm_labels[index],
+            self.accent_labels[index],
+            self.gender_labels[index],
+        )
+
+    def __len__(self):
+        return len(self.waveforms)
+
+
+class MultiTaskWav2Vec2(nn.Module):
+    def __init__(
+        self,
+        wav2vec2_backbone,
+        backbone_hidden_size,
+        projection_hidden_size,
+        num_accent_class,
+    ):
+        super().__init__()
+        self.wav2vec2 = wav2vec2_backbone
+        self.accent_projector = nn.Linear(backbone_hidden_size, projection_hidden_size)
+        self.accent_classifier = nn.Linear(projection_hidden_size, num_accent_class)
+        self.gender_projector = nn.Linear(backbone_hidden_size, projection_hidden_size)
+        self.gender_classifier = nn.Linear(projection_hidden_size, 2)
+
+    def forward(self, waveform, lm_labels=None):
+        if lm_labels is not None:
+            # use hugging face wav2vecc2
+            wav2vec2_output = self.wav2vec2(input_values=waveform, labels=lm_labels)
+
+            # get partial loss based (lm_head loss or the ctc loss)
+            ctc_loss = wav2vec2_output.loss
+
+        else:
+            # use hugging face wav2vecc2
+            wav2vec2_output = self.wav2vec2(input_values=waveform)
+            ctc_loss = None
+
+        # get features from wav2vec2
+        features = wav2vec2_output.hidden_states[-1]
+
+        # get output lm logits
+        lm_logits = wav2vec2_output.logits
+
+        # get output accent logits
+        accent_projected = self.accent_projector(features)
+        accent_projected = accent_projected.mean(dim=1)
+        accent_logits = self.accent_classifier(accent_projected)
+
+        # get output gender logits
+        gender_projected = self.gender_projector(features)
+        gender_projected = gender_projected.mean(dim=1)
+        gender_logits = self.gender_classifier(gender_projected)
+
+        return ctc_loss, lm_logits, accent_logits, gender_logits
+