LEVI_whisper_medium.en / benchmark_utils.py

set 'levi' as default WER normalizer

2a4ff3a verified 4 months ago

No virus

11.7 kB

	#%% imports
	import pandas as pd
	import time
	from tqdm import tqdm
	import torch
	from torch.cuda.amp import autocast
	import transformers
	from transformers import WhisperFeatureExtractor, WhisperTokenizer, WhisperProcessor, WhisperForConditionalGeneration, GenerationConfig
	from transformers import pipeline, AutomaticSpeechRecognitionPipeline
	from peft import PeftModel, PeftConfig
	import warnings
	import jiwer
	from jiwer.process import WordOutput
	import pandas as pd
	import numpy as np
	from pathlib import Path
	import os
	import math
	from decimal import InvalidOperation
	import contractions
	from whisper.normalizers.english import EnglishTextNormalizer
	from num2words import num2words
	import csv
	import re
	import string

	#%% define functions
	def ASRmanifest(
	manifest_csv: str,
	out_csv: str,
	corpora_root: str,
	model_path:str,
	):
	"""Run Whisper ASR on a dataset specified in a manifest
	Args:
	manifest_csv (str): path to manifest csv listing files to transcribe
	out_csv (str):path to write output csv
	corpora_root (str): root path where audio files are, inserted in place of $DATAROOT in manifest
	model_path (str): path to model directory / huggingface model name
	"""

	df = pd.read_csv(manifest_csv,keep_default_na=False)
	fieldnames = list(df.columns) + ['asr']

	asr_pipeline=prepare_pipeline(
	model_path=model_path,
	generate_opts={'max_new_tokens':448,
	'num_beams':1,#greedy
	'repetition_penalty':1,
	'do_sample':False
	}
	)

	message = "This may take a while on CPU." if asr_pipeline.device.type=="cpu" else "Using GPU"
	print(f'Running ASR for {len(df)} files. {message} ...')
	compute_time=0
	total_audio_dur=0
	# get the start time
	st = time.time()

	with open(out_csv, 'w', newline='') as csvfile:
	writer = csv.DictWriter(csvfile, fieldnames=fieldnames,delimiter=',')
	writer.writeheader()
	for i,row in tqdm(df.iterrows(), total=df.shape[0]):
	audiofile=row['wav'].replace('$DATAROOT',corpora_root)
	with torch.no_grad():
	with autocast():
	try:
	result = asr_pipeline(audiofile )
	asrtext = result['text']
	asr_pipeline.call_count = 0
	except (FileNotFoundError, ValueError) as e:
	print(f'SKIPPED: {audiofile}')
	continue
	row['asr']=asrtext
	writer.writerow( row.to_dict())
	et = time.time()
	compute_time = (et-st)
	print(f'...transcription complete in {compute_time:.1f} sec')

	def prepare_pipeline(model_path, generate_opts):
	"""Prepare a pipeline for ASR inference
	Args:
	model_path (str): path to model directory / huggingface model name
	generate_opts (dict): options to pass to pipeline
	Returns:
	pipeline: ASR pipeline
	"""
	processor = WhisperProcessor.from_pretrained(model_path)

	asr_pipeline = pipeline(
	"automatic-speech-recognition",
	model=model_path,
	tokenizer=processor.tokenizer,
	feature_extractor=processor.feature_extractor,
	generate_kwargs=generate_opts,
	model_kwargs={"load_in_8bit": False},
	device_map='auto')
	return asr_pipeline

	#%% WER evaluation functions
	def get_normalizer(text_norm_method='isat'):
	if text_norm_method=='whisper':
	normalizer=whisper_norm_text_for_wer
	elif text_norm_method=='whisper_keep_tags':
	normalizer=EnglishTextNormalizer()
	elif text_norm_method=='isat':
	normalizer = norm_text_for_wer
	elif text_norm_method=='levi':
	normalizer = levi_norm_text_for_wer
	else:
	raise NotImplementedError(f'unrecognized normalizer method: {text_norm_method}')
	return normalizer

	def strip_punct(instr, keep_math=False):
	newstr = ''
	for word in instr.split():
	if keep_math:
	word=word.strip('!"#$&\',.:;<=>?@[\\]^_`{\|}~')
	else:
	# delete punct from start and end of word
	word = word.strip(string.punctuation)
	# delete commas inside numbers
	m = re.match(r'(\d*),(\d)', word)
	if m != None:
	word = word.replace(',', '')
	# commas inside words become space
	word = re.sub(",", " ", word)
	# hyphens inside words become space
	if keep_math:
	pass
	else:
	word = re.sub("-", " ", word)
	word = word.strip()
	newstr += ' ' + word
	newstr = newstr.strip()
	return newstr

	def remove_in_brackets(text):
	# removes any clause in brackets or parens, and the brackets themselves
	return re.sub("[$\[\<].*?[$\]\>]+", " ", text)

	def caught_num2words(text):
	# first do currency replacements #TODO: plurals vs singular
	if '$' in text:
	text = re.sub('\$([0-9]+)', '\g<1> dollars', text)
	if '€' in text:
	text = re.sub('\$([0-9]+)', '\g<1> euro', text)
	if '£' in text:
	text = re.sub('\$([0-9]+)', '\g<1> pounds', text)
	if '%' in text:
	text = re.sub('([0-9]+)\%', '\g<1> percent', text)

	# strip punctuation
	text=strip_punct(text, keep_math=True)
	text=text.strip('*=/')
	# catch strings that might be converted to infinity or NaN and return as is...
	naughty_words = ['INF','Inf','inf','NAN','NaN', 'nan', 'NONE','None','none','Infinity','infinity']
	if text in naughty_words:
	return text
	try:
	if len(text.split()) > 1:
	return ' '.join([caught_num2words(word) for word in text.split()])
	else:
	return num2words(text)
	except (InvalidOperation, ValueError) as error:
	return text

	def spell_math(text):
	# spell out mathematical expressions
	# numerals preceded by hyphen become negative
	text = re.sub('\-(\d+)', 'minus \g<1>', text)
	text = re.sub('(\d+\s?)\-(\s?\d?)', '\g<1> minus \g<2>', text)
	text = re.sub('(\w+\s+)\-(\s?\w+)', '\g<1> minus \g<2>', text) # need to be more careful with - as this could be a hyphenated word not minus
	text = re.sub('(\w+\s?)\+(\s?\w+)', '\g<1> plus \g<2>', text)
	text = re.sub('(\w+\s?)\*(\s?\w+)', '\g<1> times \g<2>', text)
	text = re.sub('(\d+\s?)x(\s?\d)', '\g<1> times \g<2>', text) # need to be more careful with x as this could be a variable not times
	text = re.sub('(\w+\s?)\/(\s?\w+)', '\g<1> divided by \g<2>', text)
	text = re.sub('(\w+\s?)\=(\s?\w+)', '\g<1> equals \g<2>', text)
	return text

	def expand_contractions(str):
	expanded_words = []
	for wrd in str.split():
	expanded_words.append(contractions.fix(wrd))
	str = ' '.join(expanded_words)
	return str

	def norm_text_for_wer(text):
	# function to format text or lists of text (e.g. asr, transcript) for wer computation.
	# Converts from list to a single string and apply some text normalization operations
	# note that the clean_REV_transcript function should be applied first to remove REV-specific keywords
	# and extract text from docx format tables

	if isinstance(text,list):
	text = ' '.join(text)
	text=str(text)
	text = text.replace('\n',' ') # replace newline with space
	text = remove_in_brackets(text) # removes non-spoken annotations such as [inaudible]
	text = re.sub('%\w+','', text) # remove %HESITATION etc
	text = ' '.join([caught_num2words(str) for str in text.split(' ')]) # spell out numbers
	text = expand_contractions(text)
	text = strip_punct(text)
	text = text.lower()
	text = re.sub('\s+',' ',text) # replace multiple space with single
	return text

	def levi_norm_text_for_wer(text):
	# function to format text or lists of text (e.g. asr, transcript) for wer computation.
	# specialized for math language

	if isinstance(text,list):
	text = ' '.join(text)
	text=str(text)
	text = text.replace('\n',' ') # replace newline with space
	text = remove_in_brackets(text) # removes non-spoken annotations such as [inaudible]
	text = re.sub('%\w+','', text) # remove %HESITATION etc
	text = spell_math(text)
	text = ' '.join([caught_num2words(str) for str in text.split(' ')]) # spell out numbers
	text = expand_contractions(text)
	text = strip_punct(text, keep_math=True)
	text = text.lower()
	text = re.sub('\s+',' ',text) # replace multiple space with single
	return text

	def whisper_norm_text_for_wer(text):
	# function to format text for wer computation.
	# uses Whisper normalizer after stripping corpus-specific special tags

	if isinstance(text,list):
	text = ' '.join(text)
	text=str(text)
	text = text.replace('\n',' ') # replace newline with space
	text = re.sub('%\w+','', text) # remove %HESITATION etc
	text = remove_in_brackets(text) # removes non-spoken annotations such as [inaudible]
	normalizer = EnglishTextNormalizer()
	text = normalizer(text)
	return text

	def wer_from_df(
	df,
	refcol='ref',
	hypcol='hyp',
	return_alignments=False,
	normalise = True,
	text_norm_method='levi',
	printout=True):
	"""Compute WER from a dataframe containing a ref col and a hyp col
	WER is computed on the edit operation counts over the whole df,
	not averaged over single utterances.

	Args:
	df (pandas DataFrame): containing rows per utterance
	refcol (str, optional): column name containing reference transcript. Defaults to 'ref'.
	hypcol (str, optional): column name containing hypothesis transcript. Defaults to 'hyp'.
	return_alignments (bool, optional): Return full word-level alignments. Defaults to False.
	normalise (bool, optional): Apply text normalisatin to ref and hyp (see norm_text_for_wer). Defaults to True.
	printout (bool, optional): Print WER metrics. Defaults to True.
	"""
	normalizer=get_normalizer(text_norm_method)

	refs=df[refcol].astype(str)
	hyps = df[hypcol].astype(str)
	if normalise:
	refs=refs.apply(normalizer)
	hyps=hyps.apply(normalizer)

	#ID,ref,hyp,ref_norm,hyp_norm
	if any(s == '' for s in list(refs)):
	nonempty=refs.str.len()>0
	refs=refs[nonempty]
	hyps=hyps[nonempty]
	# print(f'{sum(~nonempty)} empty references removed (after normalisation if applied)')
	wer_meas = jiwer.compute_measures(list(refs), list(hyps))

	if not return_alignments:
	# remove alignments
	del wer_meas['ops']
	del wer_meas['truth']
	del wer_meas['hypothesis']
	wer_meas['word_count'] = wer_meas['substitutions']+wer_meas['deletions']+wer_meas['hits']
	wer_meas['sub_rate'] = wer_meas['substitutions']/wer_meas['word_count']
	wer_meas['del_rate'] = wer_meas['deletions']/wer_meas['word_count']
	wer_meas['ins_rate'] = wer_meas['insertions']/wer_meas['word_count']

	if printout:
	for key in ['wer','sub_rate','del_rate','ins_rate']:
	print((f"{key}={100*wer_meas[key]:.1f}" ))
	print(f"word_count={int(wer_meas['word_count'])}")
	return wer_meas


	def wer_from_csv(
	csv_path,
	refcol='ref',
	hypcol='hyp',
	return_alignments=False,
	normalise = True,
	text_norm_method='levi' ,
	printout=True):

	res = pd.read_csv(csv_path).astype(str)

	wer_meas=wer_from_df(res,
	refcol=refcol,
	hypcol=hypcol,
	return_alignments=return_alignments,
	normalise = normalise,
	text_norm_method=text_norm_method,
	printout=printout)
	return wer_meas