Views
No views yet
1import torch
2import torchaudio
3from datasets import load_dataset
4from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
5ENCODER = {
6 "ia ": "iê ",
7 "ìa ": "iề ",
8 "ía ": "iế ",
9 "ỉa ": "iể ",
10 "ĩa ": "iễ ",
11 "ịa ": "iệ ",
12 "ya ": "yê ",
13 "ỳa ": "yề ",
14 "ýa ": "yế ",
15 "ỷa ": "yể ",
16 "ỹa ": "yễ ",
17 "ỵa ": "yệ ",
18 "ua ": "uô ",
19 "ùa ": "uồ ",
20 "úa ": "uố ",
21 "ủa ": "uổ ",
22 "ũa ": "uỗ ",
23 "ụa ": "uộ ",
24 "ưa ": "ươ ",
25 "ừa ": "ườ ",
26 "ứa ": "ướ ",
27 "ửa ": "ưở ",
28 "ữa ": "ưỡ ",
29 "ựa ": "ượ ",
30 "ke": "ce",
31 "kè": "cè",
32 "ké": "cé",
33 "kẻ": "cẻ",
34 "kẽ": "cẽ",
35 "kẹ": "cẹ",
36 "kê": "cê",
37 "kề": "cề",
38 "kế": "cế",
39 "kể": "cể",
40 "kễ": "cễ",
41 "kệ": "cệ",
42 "ki": "ci",
43 "kì": "cì",
44 "kí": "cí",
45 "kỉ": "cỉ",
46 "kĩ": "cĩ",
47 "kị": "cị",
48 "ky": "cy",
49 "kỳ": "cỳ",
50 "ký": "cý",
51 "kỷ": "cỷ",
52 "kỹ": "cỹ",
53 "kỵ": "cỵ",
54 "ghe": "ge",
55 "ghè": "gè",
56 "ghé": "gé",
57 "ghẻ": "gẻ",
58 "ghẽ": "gẽ",
59 "ghẹ": "gẹ",
60 "ghê": "gê",
61 "ghề": "gề",
62 "ghế": "gế",
63 "ghể": "gể",
64 "ghễ": "gễ",
65 "ghệ": "gệ",
66 "ngh": "\x80",
67 "uyê": "\x96",
68 "uyề": "\x97",
69 "uyế": "\x98",
70 "uyể": "\x99",
71 "uyễ": "\x9a",
72 "uyệ": "\x9b",
73 "ng": "\x81",
74 "ch": "\x82",
75 "gh": "\x83",
76 "nh": "\x84",
77 "gi": "\x85",
78 "ph": "\x86",
79 "kh": "\x87",
80 "th": "\x88",
81 "tr": "\x89",
82 "uy": "\x8a",
83 "uỳ": "\x8b",
84 "uý": "\x8c",
85 "uỷ": "\x8d",
86 "uỹ": "\x8e",
87 "uỵ": "\x8f",
88 "iê": "\x90",
89 "iề": "\x91",
90 "iế": "\x92",
91 "iể": "\x93",
92 "iễ": "\x94",
93 "iệ": "\x95",
94 "uô": "\x9c",
95 "uồ": "\x9d",
96 "uố": "\x9e",
97 "uổ": "\x9f",
98 "uỗ": "\xa0",
99 "uộ": "\xa1",
100 "ươ": "\xa2",
101 "ườ": "\xa3",
102 "ướ": "\xa4",
103 "ưở": "\xa5",
104 "ưỡ": "\xa6",
105 "ượ": "\xa7",
106}
107
108def decode_string(x):
109 for k, v in list(reversed(list(ENCODER.items()))):
110 x = x.replace(v, k)
111 return x
112test_dataset = load_dataset("common_voice", "vi", split="test[:2%]")
113processor = Wav2Vec2Processor.from_pretrained("Nhut/wav2vec2-large-xlsr-vietnamese")
114model = Wav2Vec2ForCTC.from_pretrained("Nhut/wav2vec2-large-xlsr-vietnamese")
115resampler = torchaudio.transforms.Resample(48_000, 16_000)
116# Preprocessing the datasets.
117# We need to read the aduio files as arrays
118def speech_file_to_array_fn(batch):
119 speech_array, sampling_rate = torchaudio.load(batch["path"])
120 batch["speech"] = resampler(speech_array).squeeze().numpy()
121 return batch
122test_dataset = test_dataset.map(speech_file_to_array_fn)
123inputs = processor(test_dataset["speech"][:2], sampling_rate=16_000, return_tensors="pt", padding=True)
124with torch.no_grad():
125 logits = model(inputs.input_values, attention_mask=inputs.attention_mask).logits
126predicted_ids = torch.argmax(logits, dim=-1)
127print("Prediction:", [decode_string(x) for x in processor.batch_decode(predicted_ids)])
128print("Reference:", test_dataset["sentence"][:2])1import torch
2import torchaudio
3from datasets import load_dataset, load_metric
4from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
5import re
6
7ENCODER = {
8 "ia ": "iê ",
9 "ìa ": "iề ",
10 "ía ": "iế ",
11 "ỉa ": "iể ",
12 "ĩa ": "iễ ",
13 "ịa ": "iệ ",
14 "ya ": "yê ",
15 "ỳa ": "yề ",
16 "ýa ": "yế ",
17 "ỷa ": "yể ",
18 "ỹa ": "yễ ",
19 "ỵa ": "yệ ",
20 "ua ": "uô ",
21 "ùa ": "uồ ",
22 "úa ": "uố ",
23 "ủa ": "uổ ",
24 "ũa ": "uỗ ",
25 "ụa ": "uộ ",
26 "ưa ": "ươ ",
27 "ừa ": "ườ ",
28 "ứa ": "ướ ",
29 "ửa ": "ưở ",
30 "ữa ": "ưỡ ",
31 "ựa ": "ượ ",
32 "ke": "ce",
33 "kè": "cè",
34 "ké": "cé",
35 "kẻ": "cẻ",
36 "kẽ": "cẽ",
37 "kẹ": "cẹ",
38 "kê": "cê",
39 "kề": "cề",
40 "kế": "cế",
41 "kể": "cể",
42 "kễ": "cễ",
43 "kệ": "cệ",
44 "ki": "ci",
45 "kì": "cì",
46 "kí": "cí",
47 "kỉ": "cỉ",
48 "kĩ": "cĩ",
49 "kị": "cị",
50 "ky": "cy",
51 "kỳ": "cỳ",
52 "ký": "cý",
53 "kỷ": "cỷ",
54 "kỹ": "cỹ",
55 "kỵ": "cỵ",
56 "ghe": "ge",
57 "ghè": "gè",
58 "ghé": "gé",
59 "ghẻ": "gẻ",
60 "ghẽ": "gẽ",
61 "ghẹ": "gẹ",
62 "ghê": "gê",
63 "ghề": "gề",
64 "ghế": "gế",
65 "ghể": "gể",
66 "ghễ": "gễ",
67 "ghệ": "gệ",
68 "ngh": "\x80",
69 "uyê": "\x96",
70 "uyề": "\x97",
71 "uyế": "\x98",
72 "uyể": "\x99",
73 "uyễ": "\x9a",
74 "uyệ": "\x9b",
75 "ng": "\x81",
76 "ch": "\x82",
77 "gh": "\x83",
78 "nh": "\x84",
79 "gi": "\x85",
80 "ph": "\x86",
81 "kh": "\x87",
82 "th": "\x88",
83 "tr": "\x89",
84 "uy": "\x8a",
85 "uỳ": "\x8b",
86 "uý": "\x8c",
87 "uỷ": "\x8d",
88 "uỹ": "\x8e",
89 "uỵ": "\x8f",
90 "iê": "\x90",
91 "iề": "\x91",
92 "iế": "\x92",
93 "iể": "\x93",
94 "iễ": "\x94",
95 "iệ": "\x95",
96 "uô": "\x9c",
97 "uồ": "\x9d",
98 "uố": "\x9e",
99 "uổ": "\x9f",
100 "uỗ": "\xa0",
101 "uộ": "\xa1",
102 "ươ": "\xa2",
103 "ườ": "\xa3",
104 "ướ": "\xa4",
105 "ưở": "\xa5",
106 "ưỡ": "\xa6",
107 "ượ": "\xa7",
108}
109
110def decode_string(x):
111 for k, v in list(reversed(list(ENCODER.items()))):
112 x = x.replace(v, k)
113 return x
114
115test_dataset = load_dataset("common_voice", "vi", split="test")
116wer = load_metric("wer")
117processor = Wav2Vec2Processor.from_pretrained("Nhut/wav2vec2-large-xlsr-vietnamese")
118model = Wav2Vec2ForCTC.from_pretrained("Nhut/wav2vec2-large-xlsr-vietnamese")
119model.to("cuda")
120
121chars_to_ignore_regex = '[\\\+\@\ǀ\,\?\.\!\-\;\:\"\“\%\‘\”\�]'
122resampler = torchaudio.transforms.Resample(48_000, 16_000)
123
124# Preprocessing the datasets.
125# We need to read the aduio files as arrays
126def speech_file_to_array_fn(batch):
127 batch["sentence"] = re.sub(chars_to_ignore_regex, '', batch["sentence"]).lower()
128 speech_array, sampling_rate = torchaudio.load(batch["path"])
129 batch["speech"] = resampler(speech_array).squeeze().numpy()
130 return batch
131
132test_dataset = test_dataset.map(speech_file_to_array_fn)
133# Preprocessing the datasets.
134# We need to read the aduio files as arrays
135def evaluate(batch):
136 inputs = processor(batch["speech"], sampling_rate=16_000, return_tensors="pt", padding=True)
137 with torch.no_grad():
138 logits = model(inputs.input_values.to("cuda"), attention_mask=inputs.attention_mask.to("cuda")).logits
139 pred_ids = torch.argmax(logits, dim=-1)
140 batch["pred_strings"] = processor.batch_decode(pred_ids)
141 # decode_string: We replace the encoded letter with the initial letters
142 batch["pred_strings"] = [decode_string(x) for x in batch["pred_strings"]]
143 return batch
144
145result = test_dataset.map(evaluate, batched=True, batch_size=8)
146print("WER: {:2f}".format(100 * wer.compute(predictions=result["pred_strings"], references=result["sentence"])))train, validation and FOSD datasets and VIVOS datasets were used for training as well.
The script used for training can be found here