|
|
from qwen_omni_utils import process_mm_info |
|
|
import torch |
|
|
from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor |
|
|
import librosa |
|
|
import os |
|
|
from io import BytesIO |
|
|
from urllib.request import urlopen |
|
|
import argparse |
|
|
|
|
|
def inference(audio_path,model,processor,prompt, sys_prompt): |
|
|
messages = [ |
|
|
{"role": "system", "content": [{"type": "text", "text": sys_prompt}]}, |
|
|
{"role": "user", "content": [ |
|
|
{"type": "audio", "audio": audio_path}, |
|
|
{"type": "text", "text": prompt}, |
|
|
] |
|
|
}, |
|
|
] |
|
|
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
|
|
audios, images, videos = process_mm_info(messages, use_audio_in_video=True) |
|
|
inputs = processor(text=text, audio=audios, images=images, videos=videos, return_tensors="pt", padding=True, use_audio_in_video=True) |
|
|
inputs = inputs.to(model.device).to(model.dtype) |
|
|
|
|
|
output = model.generate(**inputs, use_audio_in_video=True, return_audio=False, thinker_max_new_tokens=256, thinker_do_sample=False) |
|
|
|
|
|
text = processor.batch_decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=False) |
|
|
return text |
|
|
|
|
|
def transcribe(wavs_path, out_path, gpu_id, model): |
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) |
|
|
model_path = model |
|
|
model = Qwen2_5OmniForConditionalGeneration.from_pretrained( |
|
|
model_path, |
|
|
torch_dtype=torch.bfloat16, |
|
|
device_map="auto", |
|
|
) |
|
|
prompt = "请将这段中文语音转换为纯文本，去掉标点符号。" |
|
|
processor = Qwen2_5OmniProcessor.from_pretrained(model_path) |
|
|
with open(wavs_path, "r") as f_in, open(out_path, "w") as f_out: |
|
|
for line in f_in: |
|
|
utt, path = line.strip().split(" ", maxsplit=1) |
|
|
try: |
|
|
response=inference(path,model,processor, prompt=prompt, sys_prompt="You are a speech recognition model.") |
|
|
except Exception as e: |
|
|
print(f"Inference failed: {str(e)}") |
|
|
response="None" |
|
|
text = response[0].strip() |
|
|
lines = text.strip().splitlines() |
|
|
text = lines[-1] |
|
|
print(f"[{utt}] >>> {text}") |
|
|
f_out.write(f"{utt} {text}\n") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument("--wavs_path", type=str) |
|
|
parser.add_argument("--out_path", type=str) |
|
|
parser.add_argument("--gpu", type=int, default=0) |
|
|
parser.add_argument("--model", type=str) |
|
|
args = parser.parse_args() |
|
|
transcribe( |
|
|
wavs_path=args.wavs_path, |
|
|
out_path=args.out_path, |
|
|
gpu_id=args.gpu, |
|
|
model=args.model |
|
|
) |
