whisper+whisperx ASR加对齐

忘了怎么安装了，这里记录一下整理出来的类，不过这个

from chj.comm.pic import *import json
import whisper
import whisperx
import gcclass Warp_whisper:def __init__(self, language="zh", device="cuda", compute_type="float32", model="large-v2" ):torch.backends.cudnn.enabled = Falseif not torch.cuda.is_available():device="cpu"dmodel="XXXXX/models/torch/whisper"self.asr_model=whisper.load_model(model, device, download_root=dmodel)self.txt_converter = Noneif model=="large-v2" and language=="zh":from opencc import OpenCCconverter = OpenCC('t2s')self.txt_converter = converterself.prompt=Noneelse:if language=="zh":self.prompt='以下是普通话的句子'else:self.prompt=Noneself.prompt=Noneself.language=languageself.device=deviceself.align_model, self.align_metadata = whisperx.load_align_model(language_code=language, device=device)def do_asr_algin(self, fjson, fwav):audio = whisper.load_audio(fwav)result = self.asr_model.transcribe(audio, language=self.language, initial_prompt=self.prompt)#assert result["language"] == self.languageresult_segments = result["segments"]if self.txt_converter:for e in result_segments:e['text'] = converter.convert( e['text'] )result = whisperx.align(result_segments, self.align_model, self.align_metadata, audio, self.device, return_char_alignments=False)result=result["segments"]with open(fjson, "w") as fout:json.dump(result, fout, indent=4, ensure_ascii=False)def f2_invoke():print("Doing... whisper align")basedir=sys.argv[1]din=f"{basedir}/audio_feats/wav16k"if not os.path.exists(din):print("no such dir", din)exit(1)dout=f"{basedir}/audio_feats/whisper_align"cls_asr=Warp_whisper()chj_file.mkdir(dout)for fwav in tqdm( glob.glob(f"{din}/*.wav") ):nm = chj_file.get_nm(fwav)fnm=f"{dout}/{nm}.json"if os.path.isfile(fnm): continuecls_asr.do_asr_algin(fnm,fwav)print("Finished whisper align")