File size: 2,731 Bytes
27d1541 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
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
# @title inference function
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
) |