Upload eval_model.py

eval_model.py

@@ -0,0 +1,153 @@
+from transformers import AutoTokenizer, GPT2LMHeadModel
+from datasets import load_dataset, Dataset, DatasetDict
+import random
+import string
+import torch
+
+from torchmetrics.text import WordErrorRate, CharErrorRate
+
+wer = WordErrorRate()
+cer = CharErrorRate()
+
+def process(text):
+
+    # Lower case every letter
+    text = text.lower()
+
+    # Remove punctuation
+    punctuation_to_remove = string.punctuation.replace("'", "")
+    translation_table = str.maketrans('', '', punctuation_to_remove)
+    text = text.translate(translation_table)
+
+    # Remove whitespaces from front and behind
+    while text[0] == ' ' or text[-1] == ' ':
+        if text[0] == ' ':
+            text = text[1:]
+        if text[-1] == ' ':
+            text = text[:-1]
+    
+    return text
+
+import jiwer
+from edit_distance import SequenceMatcher
+def correct_text(text):
+    transforms = jiwer.Compose(
+       [
+           jiwer.ExpandCommonEnglishContractions(),
+           jiwer.ToLowerCase(),
+           jiwer.RemoveMultipleSpaces(),
+           jiwer.Strip(),
+           jiwer.RemovePunctuation(),
+           jiwer.ReduceToListOfListOfWords(),
+       ]
+    )
+    return transforms(text)
+
+def align_gt_asr(gt, asr):
+    sm = SequenceMatcher(a=gt, b=asr)
+    best_path = []
+    opcodes = sm.get_opcodes()
+    for tag, i1, i2, j1, j2 in opcodes:
+        if tag == "delete":
+            for i in range(i1, i2):
+                best_path.append([gt[i], ""])
+        if tag == "replace" or tag == "equal":
+            for i, j in zip(range(i1, i2), range(j1, j2)):
+                best_path.append([gt[i], asr[j]])
+        if tag == "insert":
+            for j in range(j1, j2):
+                best_path.append(["", asr[j]])
+    return best_path
+
+dtype = torch.float16
+
+dataset_name = "./../libripseech_tokenized"
+dataset = DatasetDict.load_from_disk(dataset_name)
+
+with open("./../prompting/blist/all_rare_words.txt") as fin:
+    rarewords = [process(word.strip()) for word in fin]
+
+tokenizer = AutoTokenizer.from_pretrained("./../tokenizer")
+tokenizer.pad_token_id = 0
+tokenizer.pad_token = "<|padding|>"
+tokenizer.padding_side = "left"
+
+# Adding new tokens for introducing prompts
+tokenizer.add_tokens(["<|startofprompt|>", "<|sepofprompt|>", "<|endofprompt|>"])
+sot_token = tokenizer.encode("<|startoftranscript|>")[0]
+eot_token = tokenizer.encode("<|endoftranscript|>")[0]
+
+from math import ceil
+from tqdm import tqdm
+
+val_bs = 32
+n_bwords = 25
+context_length = 2048
+
+def prepare(element):
+    
+    # Add audio
+    audio_tkns = element["audio_tokens"]
+    data = "".join([f"<|audio:{tkn}|>" for tkn in audio_tkns])
+    
+    # sample context words and mix with the biasing list
+    b_words = element["b_words"]
+    if n_bwords > len(b_words):
+        context = b_words + random.sample(rarewords, n_bwords - len(b_words))
+    else:
+        context = random.sample(b_words, n_bwords)
+    random.shuffle(context)
+    
+    # add the context words
+    data += "<|startofprompt|>" + "<|sepofprompt|>".join(context) + "<|endofprompt|>"
+    
+    # Add text
+    data += "<|startoftranscript|>"
+    
+    return {"data": data, "context": context}
+
+@torch.no_grad()
+def evaluate_model(model):
+
+    transcripts = []
+    
+    processed_data = dataset["test.clean"].map(prepare)
+    data = processed_data["data"]
+
+    for idx in tqdm(range(ceil(len(data)/val_bs))):
+
+        outputs = tokenizer(data[idx * val_bs: (idx + 1) * val_bs], truncation=False, max_length=None, padding=True, return_tensors="pt").to(model.device)
+        input_ids = outputs["input_ids"]
+        par = input_ids.shape[-1]
+
+        generations = model.generate(
+            input_ids,
+            max_new_tokens=context_length - par - 1,
+            eos_token_id = eot_token
+        )
+        transcripts += tokenizer.batch_decode(generations[:, par:], skip_special_tokens=True)
+        
+    bias_word_cnt = 0
+    normal_word_cnt = 0
+    u_wer = 0.0
+    b_wer = 0.0
+    pred_list = correct_text(transcripts)
+    text_list = correct_text(processed_data["text"])
+    prompt_list = processed_data["context"]
+    for a, b, c in zip(pred_list, text_list, prompt_list):
+        aligned_pair = align_gt_asr(b, a)
+        for gt_word, asr_word in aligned_pair:
+            if gt_word in c or asr_word in c:
+                if gt_word != asr_word:
+                    b_wer += 1.0
+                if gt_word in c:
+                    bias_word_cnt += 1
+            else:
+                if gt_word != asr_word:
+                    u_wer += 1.0
+                if gt_word != "":
+                    normal_word_cnt += 1
+    u_wer = u_wer / normal_word_cnt * 100
+    b_wer = b_wer / bias_word_cnt * 100
+    
+    return wer(transcripts, processed_data["text"]).item() * 100, cer(transcripts, processed_data["text"]).item() * 100, b_wer, u_wer