Spaces:
Paused
Paused
File size: 8,771 Bytes
ee6e328 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 |
# Parts of the code are adapted from the snippets provided in the TorchAudio Wav2Vec forced alignment tutorial.
# The full tutorial can be found here: https://pytorch.org/audio/stable/tutorials/forced_alignment_tutorial.html
import argparse
import os
from dataclasses import dataclass
import torch
import torchaudio
from tqdm import tqdm
from transformers import AutoConfig, AutoModelForCTC, AutoProcessor
class Wav2Vec2Aligner:
def __init__(self, model_name, input_wavs_sr, cuda):
self.cuda = cuda
self.config = AutoConfig.from_pretrained(model_name)
self.model = AutoModelForCTC.from_pretrained(model_name)
self.model.eval()
if self.cuda:
self.model.to(device="cuda")
self.processor = AutoProcessor.from_pretrained(model_name)
self.resampler = torchaudio.transforms.Resample(input_wavs_sr, 16_000)
blank_id = 0
vocab = list(self.processor.tokenizer.get_vocab().keys())
for i in range(len(vocab)):
if vocab[i] == "[PAD]" or vocab[i] == "<pad>":
blank_id = i
print("Blank Token id [PAD]/<pad>", blank_id)
self.blank_id = blank_id
def speech_file_to_array_fn(self, wav_path):
speech_array, sampling_rate = torchaudio.load(wav_path)
speech = self.resampler(speech_array).squeeze().numpy()
return speech
def align_single_sample(self, item):
blank_id = self.blank_id
transcript = "|".join(item["sent"].split(" "))
if not os.path.isfile(item["wav_path"]):
print(item["wav_path"], "not found in wavs directory")
speech_array = self.speech_file_to_array_fn(item["wav_path"])
inputs = self.processor(speech_array, sampling_rate=16_000, return_tensors="pt", padding=True)
if self.cuda:
inputs = inputs.to(device="cuda")
with torch.no_grad():
logits = self.model(inputs.input_values).logits
# get the emission probability at frame level
emissions = torch.log_softmax(logits, dim=-1)
emission = emissions[0].cpu().detach()
# get labels from vocab
labels = ([""] + list(self.processor.tokenizer.get_vocab().keys()))[
:-1
] # logits don't align with the tokenizer's vocab
dictionary = {c: i for i, c in enumerate(labels)}
tokens = []
for c in transcript:
if c in dictionary:
tokens.append(dictionary[c])
def get_trellis(emission, tokens, blank_id=0):
"""
Build a trellis matrix of shape (num_frames + 1, num_tokens + 1)
that represents the probabilities of each source token being at a certain time step
"""
num_frames = emission.size(0)
num_tokens = len(tokens)
# Trellis has extra diemsions for both time axis and tokens.
# The extra dim for tokens represents <SoS> (start-of-sentence)
# The extra dim for time axis is for simplification of the code.
trellis = torch.full((num_frames + 1, num_tokens + 1), -float("inf"))
trellis[:, 0] = 0
for t in range(num_frames):
trellis[t + 1, 1:] = torch.maximum(
# Score for staying at the same token
trellis[t, 1:] + emission[t, blank_id],
# Score for changing to the next token
trellis[t, :-1] + emission[t, tokens],
)
return trellis
trellis = get_trellis(emission, tokens, blank_id)
@dataclass
class Point:
token_index: int
time_index: int
score: float
def backtrack(trellis, emission, tokens, blank_id=0):
"""
Walk backwards from the last (sentence_token, time_step) pair to build the optimal sequence alignment path
"""
# Note:
# j and t are indices for trellis, which has extra dimensions
# for time and tokens at the beginning.
# When referring to time frame index `T` in trellis,
# the corresponding index in emission is `T-1`.
# Similarly, when referring to token index `J` in trellis,
# the corresponding index in transcript is `J-1`.
j = trellis.size(1) - 1
t_start = torch.argmax(trellis[:, j]).item()
path = []
for t in range(t_start, 0, -1):
# 1. Figure out if the current position was stay or change
# Note (again):
# `emission[J-1]` is the emission at time frame `J` of trellis dimension.
# Score for token staying the same from time frame J-1 to T.
stayed = trellis[t - 1, j] + emission[t - 1, blank_id]
# Score for token changing from C-1 at T-1 to J at T.
changed = trellis[t - 1, j - 1] + emission[t - 1, tokens[j - 1]]
# 2. Store the path with frame-wise probability.
prob = emission[t - 1, tokens[j - 1] if changed > stayed else 0].exp().item()
# Return token index and time index in non-trellis coordinate.
path.append(Point(j - 1, t - 1, prob))
# 3. Update the token
if changed > stayed:
j -= 1
if j == 0:
break
else:
raise ValueError("Failed to align")
return path[::-1]
path = backtrack(trellis, emission, tokens, blank_id)
@dataclass
class Segment:
label: str
start: int
end: int
score: float
def __repr__(self):
return f"{self.label}\t{self.score:4.2f}\t{self.start*20:5d}\t{self.end*20:5d}"
@property
def length(self):
return self.end - self.start
def merge_repeats(path):
"""
Merge repeated tokens into a single segment. Note: this shouldn't affect repeated characters from the
original sentences (e.g. `ll` in `hello`)
"""
i1, i2 = 0, 0
segments = []
while i1 < len(path):
while i2 < len(path) and path[i1].token_index == path[i2].token_index:
i2 += 1
score = sum(path[k].score for k in range(i1, i2)) / (i2 - i1)
segments.append(
Segment(
transcript[path[i1].token_index],
path[i1].time_index,
path[i2 - 1].time_index + 1,
score,
)
)
i1 = i2
return segments
segments = merge_repeats(path)
with open(item["out_path"], "w") as out_align:
for seg in segments:
out_align.write(str(seg) + "\n")
def align_data(self, wav_dir, text_file, output_dir):
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# load text file
lines = open(text_file, encoding="utf8").readlines()
items = []
for line in lines:
if len(line.strip().split("\t")) != 2:
print("Script must be in format: 00001 this is my sentence")
exit()
wav_name, sentence = line.strip().split("\t")
wav_path = os.path.join(wav_dir, wav_name + ".wav")
out_path = os.path.join(output_dir, wav_name + ".txt")
items.append({"sent": sentence, "wav_path": wav_path, "out_path": out_path})
print("Number of samples found in script file", len(items))
for item in tqdm(items):
self.align_single_sample(item)
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_name", type=str, default="arijitx/wav2vec2-xls-r-300m-bengali", help="wav2vec model name"
)
parser.add_argument("--wav_dir", type=str, default="./wavs", help="directory containing wavs")
parser.add_argument("--text_file", type=str, default="script.txt", help="file containing text")
parser.add_argument("--input_wavs_sr", type=int, default=16000, help="sampling rate of input audios")
parser.add_argument(
"--output_dir", type=str, default="./out_alignment", help="output directory containing the alignment files"
)
parser.add_argument("--cuda", action="store_true")
args = parser.parse_args()
aligner = Wav2Vec2Aligner(args.model_name, args.input_wavs_sr, args.cuda)
aligner.align_data(args.wav_dir, args.text_file, args.output_dir)
if __name__ == "__main__":
main()
|