utils/evaluation.py

import argparse
import opencc
import csv
import editdistance
from tqdm import tqdm
from collections import defaultdict
from functools import partial
from pypinyin import pinyin, lazy_pinyin, Style
from g2p_en import G2p # too slow, should use lexicon instead
import edit_distance
lexicon_fpath = '/root/distil-whisper/utils/lexicon.lst'

def cal_complete_mer(ref_data, hyp_data):
    S, D, I, N = (0, 0, 0, 0)
    count = 0
    for ref, hyp in zip(ref_data, hyp_data):
        _S, _D, _I, _N = cal_single_complete_mer(ref, hyp)
        S += _S
        D += _D
        I += _I
        N += _N
        count += 1
    return S, D, I, N, count

def cal_single_complete_mer(ref, hyp):
    sm = edit_distance.SequenceMatcher(a=ref, b=hyp)
    opcodes = sm.get_opcodes()
    
    # Substitution
    s = sum([(max(x[2] - x[1], x[4] - x[3]) if x[0] == 'replace' else 0) for x in opcodes])
    # Deletion
    d = sum([(max(x[2] - x[1], x[4] - x[3]) if x[0] == 'delete' else 0) for x in opcodes])
    # Insertion
    i = sum([(max(x[2] - x[1], x[4] - x[3]) if x[0] == 'insert' else 0) for x in opcodes])
    n = len(ref)
    return s, d, i, n

class MixErrorRate(object):
    def __init__(
        self, 
        to_simplified_chinese=True, 
        to_traditional_chinese=False, 
        phonemize=False, 
        separate_language=False, 
        test_only=False,
        count_repetitive_hallucination=False,
        calculate_complete_mer=False
    ):
        self.converter = None
        if to_simplified_chinese and to_traditional_chinese:
            raise ValueError("Can't convert to both simplified and traditional chinese at the same time.")
        if to_simplified_chinese:
            print("Convert to simplified chinese")
            self.converter = opencc.OpenCC('t2s.json')
        elif to_traditional_chinese:
            print("Convert to traditional chinese")
            self.converter = opencc.OpenCC('s2t.json')
        else:
            print("No chinese conversion")
        if phonemize:
            if separate_language:
                raise NotImplementedError("Can't separate language and phonemize at the same time.")
            print("Phonemize chinese and english words")
            print("Force traditional to simplified conversion")
            self.converter = opencc.OpenCC('t2s.json')
            self.zh_phonemizer = partial(lazy_pinyin, style=Style.BOPOMOFO, errors='ignore')
            self.zh_bopomofo_stress_marks = ['ˊ', 'ˇ', 'ˋ', '˙']
            # self.en_phonemizer = G2p()
            # self.en_valid_phonemes = [p for p in self.en_phonemizer.phonemes]
            # for p in self.en_phonemizer.phonemes:
            #     if p[-1].isnumeric():
            #         self.en_valid_phonemes.append(p[:-1])
            # use lexicon instead
            self.en_wrd2phn = defaultdict(lambda: [])
            with open(lexicon_fpath, 'r', encoding='utf-8') as f:
                for line in f:
                    word, phonemes = line.strip().split('\t')
                    self.en_wrd2phn[word] = phonemes.split()
        self.phonemize = phonemize
        self.test_only = test_only
        self.separate_language = separate_language
        self.count_repetitive_hallucination = count_repetitive_hallucination
        self.calculate_complete_mer = calculate_complete_mer
        if self.count_repetitive_hallucination:
            print("Count repetitive hallucination (6gram-5repeat)")
    
    def _from_str_to_list(self, cs_string):
        cs_list = []
        cur_en_word = ''
        for s in cs_string:
            # if is space, skip it
            if s in [' ', '\t', '\n', '\r', ',', '.', '!', '?', '。', '，', '！', '？', '、', '；', '：', '「', '」', '『', '』', '（', '）', '(', ')', '\[', '\]', '{', '}', '<', '>', '《', '》', '“', '”', '‘', '’', '…', '—', '～', '·', '•']:
                if cur_en_word != '':
                    cs_list.append(cur_en_word)
                    cur_en_word = ''
                continue
            # if it chinese character, add it to list
            if u'\u4e00' <= s <= u'\u9fff':
                if cur_en_word != '':
                    cs_list.append(cur_en_word)
                    cur_en_word = ''
                if self.converter is not None:
                    s = self.converter.convert(s)
                cs_list.append(s)
            # check character, if it is english character, add it to current word
            elif s.isalnum() or s in ["'", "-"]:
                cur_en_word += s
            else:
                print(f"Unknown character during conversion: {s}")
        if cur_en_word != '':
            cs_list.append(cur_en_word)
        return cs_list
    
    def _unit_is_en(self, token):
        if u'\u4e00' <= token[0] <= u'\u9fff':
            return False
        return True
    
    def _unit_is_zh(self, token):
        if u'\u4e00' <= token[0] <= u'\u9fff':
            return True
        return False

    def _phonemized_cs_list(self, cs_list):
        cur_zh_chars = []
        phonemes = []
        for unit in cs_list:
            if u'\u4e00' <= unit[0] <= u'\u9fff':
                cur_zh_chars.append(unit)
            else:
                if cur_zh_chars:
                    zh_phns = ''.join(self.zh_phonemizer(''.join(cur_zh_chars)))
                    phonemes.extend(filter(lambda p: p not in self.zh_bopomofo_stress_marks, zh_phns))
                    cur_zh_chars = []
                phonemes.extend(self.en_wrd2phn[unit])
        if cur_zh_chars:
            zh_phns = ''.join(self.zh_phonemizer(''.join(cur_zh_chars)))
            phonemes.extend(filter(lambda p: p not in self.zh_bopomofo_stress_marks, zh_phns))
            cur_zh_chars = []
        return phonemes

    def _count_repetitive_hallucination(self, cs_str, n=6, repeat=5, reset_len=100):
        count = 0
        ngram_counts = defaultdict(lambda: 0)
        if len(cs_str) < n:
            return 0
        prev_reset_idx = 0
        for i in range(len(cs_str) - n + 1):
            ngram = cs_str[i:i+n]
            if '|>' in ngram or '<|' in ngram: continue
            ngram_counts[ngram] += 1
            if ngram_counts[ngram] >= repeat:
                count += 1
                # reset for next round calculation
                ngram_counts = defaultdict(lambda: 0)
            if i - prev_reset_idx >= reset_len:
                ngram_counts = defaultdict(lambda: 0)
                prev_reset_idx = i
        return count

    def compute(self, predictions=None, references=None, show_progress=True, empty_error_rate=1.0, **kwargs):
        total_err = 0
        total_ref_len = 0
        total_en_err = 0
        total_en_ref_len = 0
        total_zh_err = 0
        total_zh_ref_len = 0
        repetitive_hallucination_count = 0
        ref_repetitive_hallucination_count = 0
        if self.test_only:
            predictions = predictions[:10]
            references = references[:10]
        if self.calculate_complete_mer:
            S, D, I, N = 0, 0, 0, 0
        iterator = tqdm(enumerate(zip(predictions, references)), total=len(predictions), desc="Computing Mix Error Rate...") if show_progress and len(predictions) > 20 else enumerate(zip(predictions, references))
        for i, (pred, ref) in iterator:
            # if english use word error rate, if chinese use character error rate
            # generate list for editdistance computation
            if self.count_repetitive_hallucination:
                repetitive_hallucination_count += self._count_repetitive_hallucination(pred)
                ref_repetitive_hallucination_count += self._count_repetitive_hallucination(ref)
            pred_list = self._from_str_to_list(pred)
            ref_list = self._from_str_to_list(ref)
            if self.test_only:
                print(f"Prediction List First 20@{i}: {pred_list[:20]}")
                print(f"Reference List First 20@{i}: {ref_list[:20]}")
            if self.phonemize:
                pred_list = self._phonemized_cs_list(pred_list)
                ref_list = self._phonemized_cs_list(ref_list)
            if self.calculate_complete_mer:
                _S, _D, _I, _N = cal_single_complete_mer(ref_list, pred_list)
                S += _S
                D += _D
                I += _I
                N += _N
            # compute edit distance
            if self.separate_language:
                en_pred_list = list(filter(self._unit_is_en, pred_list))
                en_ref_list = list(filter(self._unit_is_en, ref_list))
                zh_pred_list = list(filter(self._unit_is_zh, pred_list))
                zh_ref_list = list(filter(self._unit_is_zh, ref_list))
                en_err = editdistance.eval(en_pred_list, en_ref_list)
                total_en_err += en_err
                total_en_ref_len += len(en_ref_list)
                zh_err = editdistance.eval(zh_pred_list, zh_ref_list)
                total_zh_err += zh_err
                total_zh_ref_len += len(zh_ref_list)
            err = editdistance.eval(pred_list, ref_list)
            total_err += err
            total_ref_len += len(ref_list)
            if self.test_only:
                local_mer = {
                    "MER": err / len(ref_list),
                    "EN WER": en_err / len(en_ref_list) if len(en_ref_list) != 0 else 0,
                    "ZH CER": zh_err / len(zh_ref_list) if len(zh_ref_list) != 0 else 0
                }
                print(f"Local MER@{i}: {local_mer}")
        if total_ref_len == 0:
            print(f"No reference found, return {empty_error_rate*100}% error rate instead")
            return empty_error_rate # if no reference, return 100% error rate instead
        mer = total_err / total_ref_len
        if self.separate_language or self.count_repetitive_hallucination:
            result = {
                "MER": mer,
            }
            if self.separate_language:
                en_wer = total_en_err / total_en_ref_len if total_en_ref_len != 0 else 0
                zh_cer = total_zh_err / total_zh_ref_len if total_zh_ref_len != 0 else 0
                result["EN WER"] = en_wer
                result["ZH CER"] = zh_cer
            if self.count_repetitive_hallucination:
                result["Hyp Repetitive Hallucination Count"] = repetitive_hallucination_count
                result["Ref Repetitive Hallucination Count"] = ref_repetitive_hallucination_count
            return result
        if self.calculate_complete_mer:
            print(f"SUB={S/N}, DEL={D/N}, INS={I/N}, (S, D, I, N)={(S, D, I, N)}, total_len={total_ref_len}")
        return mer # mer
    
    

def load_output_csv(fpath, skip_header=True, hyp_col=1, ref_col=2, delimiter='\t'):
    with open(fpath, 'r', encoding='utf-8') as f:
        reader = csv.reader(f, delimiter=delimiter)
        predictions = []
        references = []
        if skip_header:
            columns = next(reader)
            print(f"Columns: {columns}")
        for i, row in enumerate(reader):
            predictions.append(row[hyp_col])
            references.append(row[ref_col])
        return predictions, references

def main(args):
    print(args)
    mer = MixErrorRate(
        to_simplified_chinese=args.to_simplified_chinese, 
        to_traditional_chinese=args.to_traditional_chinese, 
        separate_language=args.separate_language,
        test_only=args.test_only,
        count_repetitive_hallucination=args.count_repetitive_hallucination, 
        calculate_complete_mer=args.calculate_complete_mer
    )
    predictions, references = load_output_csv(args.csv_fpath)
    mer_value = mer.compute(predictions=predictions, references=references)
    print(f"Mix Error Rate: {mer_value}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Compute Mix Error Rate")
    parser.add_argument("--csv_fpath", type=str, required=True, help="Path to the csv file")
    parser.add_argument("--to_simplified_chinese", action="store_true", help="Convert chinese to simplified chinese")
    parser.add_argument("--to_traditional_chinese", action="store_true", help="Convert chinese to traditional chinese")
    parser.add_argument("--separate_language", action="store_true", help="Compute MER separately for chinese and english")
    parser.add_argument("--test_only", action="store_true", help="Run test and give some cs_list examples")
    parser.add_argument("--count_repetitive_hallucination", action="store_true", help="Count repetitive hallucination")
    parser.add_argument("--calculate_complete_mer", action="store_true", help="Calculate complete MER")
    args = parser.parse_args()
    main(args)