add run_evaluation.py

eval_kenlm.py

@@ -106,6 +106,7 @@ def main(args):
     set_seed(42)  # set the random seed to have reproducible result.
     processor = Wav2Vec2Processor.from_pretrained(args.model_id)
     model = Wav2Vec2ForCTC.from_pretrained(args.model_id)
+    model.to(args.device)
     kenlm = KenLM(processor.tokenizer, "language_model/5gram.bin", unigrams="language_model/unigrams.txt")
 
     # map function to decode audio

run_evaluation.py

@@ -0,0 +1,120 @@
+import torch
+import torchaudio
+from datasets import load_dataset, load_metric
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, set_seed
+import argparse
+from pyctcdecode import build_ctcdecoder
+from multiprocessing import Pool
+import re
+
+"""
+This is the script to run the prediction on test set of Turkish speech dataset.
+Usage:
+python run_evaluation.y -m <wav2vec2 model_name> -d <Zindi dataset directory> -o <output file name> \
+    -b <optional batch size, default=16>
+"""
+
+
+class KenLM:
+    def __init__(self, tokenizer, model_name, num_workers=8, beam_width=128):
+        self.num_workers = num_workers
+        self.beam_width = beam_width
+        vocab_dict = tokenizer.get_vocab()
+        self.vocabulary = [x[0] for x in sorted(vocab_dict.items(), key=lambda x: x[1], reverse=False)]
+        self.vocabulary = self.vocabulary[:-2]
+        self.decoder = build_ctcdecoder(self.vocabulary, model_name)
+
+    @staticmethod
+    def lm_postprocess(text):
+        return ' '.join([x if len(x) > 1 else "" for x in text.split()]).strip()
+
+    def decode(self, logits):
+        probs = logits.cpu().numpy()
+        # probs = logits.numpy()
+        with Pool(self.num_workers) as pool:
+            text = self.decoder.decode_batch(pool, probs)
+            text = [KenLM.lm_postprocess(x) for x in text]
+        return text
+
+chars_to_ignore_regex = '[\,\?\.\!\-\;\:\"\“\‘\”\'\`…\’»«]'
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-m", "--model_name", type=str, required=True,
+                        help="The wav2vec2 model name")
+    parser.add_argument("-n", "--name", type=str, required=True,
+                        help="The name of dataset")
+    parser.add_argument("-c", "--config_name", type=str, required=True,
+                        help="The config name of the dataset")
+    parser.add_argument("-d", "--data_dir", type=str, required=False, default=None,
+                        help="The directory contains the dataset")
+    parser.add_argument("-b", "--batch_size", type=int, required=False, default=16,
+                        help="Batch size")
+    parser.add_argument("-k", "--kenlm", type=str, required=False, default=False,
+                        help="Path to KenLM model")
+    parser.add_argument("--num_workers", type=int, required=False, default=8,
+                        help="KenLM's number of workers")
+    parser.add_argument("-w", "--beam_width", type=int, required=False, default=128,
+                        help="KenLM's beam width")
+    parser.add_argument("--test_pct", type=int, required=False, default=100,
+                        help="Percentage of the test set")
+    parser.add_argument("--cpu", required=False, action='store_true',
+                        help="Force to use CPU")
+    args = parser.parse_args()
+
+    set_seed(42)  # set the random seed to have reproducible result.
+    processor = Wav2Vec2Processor.from_pretrained(args.model_name)
+    model = Wav2Vec2ForCTC.from_pretrained(args.model_name)
+    kenlm = None
+    if args.kenlm:
+        kenlm = KenLM(processor.tokenizer, args.kenlm)
+
+    # Preprocessing the datasets.
+    # We need to read the audio files as arrays
+
+    def speech_file_to_array_fn(batch):
+        if "audio" in batch:
+            speech_array = torch.tensor(batch["audio"]["array"])
+            resampler = torchaudio.transforms.Resample(48_000, 16_000)
+        else:
+            speech_array, sampling_rate = torchaudio.load(batch["path"])
+            resampler = torchaudio.transforms.Resample(sampling_rate, 16_000)
+        batch["speech"] = resampler(speech_array).squeeze().numpy()
+        return batch
+
+    def remove_special_characters(batch):
+        batch["norm_text"] = re.sub(chars_to_ignore_regex, "", batch["sentence"]).lower().strip()
+        return batch
+
+    lg_test = load_dataset(args.name, args.config_name, data_dir=args.data_dir,
+                           split=f"test[:{args.test_pct}%]", use_auth_token=True)
+    if args.cpu:
+        device = "cpu"
+    else:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    lg_test = lg_test.map(speech_file_to_array_fn)
+    lg_test = lg_test.map(remove_special_characters)
+    model = model.to(device)
+    wer = load_metric("wer")
+
+    batch_size = args.batch_size
+    def evaluate(batch):
+        inputs = processor(batch["speech"], sampling_rate=16_000, return_tensors="pt", padding=True)
+
+        with torch.no_grad():
+            logits = model(inputs.input_values.to(device)).logits
+
+        if args.kenlm:
+            batch["pred_strings"] = kenlm.decode(logits)
+        else:
+            predicted_ids = torch.argmax(logits, dim=-1)
+            batch["pred_strings"] = processor.batch_decode(predicted_ids)
+        return batch
+
+    result = lg_test.map(evaluate, batched=True, batch_size=batch_size)
+    WER = 100 * wer.compute(predictions=result["pred_strings"], references=result["norm_text"])
+    print(f"WER: {WER:.2f}%")
+
+
+if __name__ == "__main__":
+    main()