Files
wehub-resource-sync eec33d25b2
pre-commit / pre-commit (push) Failing after 1s
Build Wheel / build (3.11) (push) Failing after 1s
Build Wheel / build (3.12) (push) Failing after 0s
chore: import upstream snapshot with attribution
2026-07-13 12:29:08 +08:00

715 lines
23 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import base64
import io
import os
import random
import re
from collections.abc import Callable
import numpy as np
import torch
import torchaudio
from process_speechdata import InputSegment, StreamingInputSegment
from torchaudio.transforms import MelSpectrogram
from vllm.multimodal.media.audio import load_audio
speech_zeroemb_idx = 151667
empty_token = "<|empty|>"
mimo_audio_tokenizer = None
device = "cpu"
asr_zh_templates = [
"请将这段语音转换为文字",
"帮我识别这个音频文件中的内容",
"把这段录音转成文本",
"请转录这段语音",
"将音频内容转换成文字格式",
"识别并转写这段语音",
"把语音内容写成文字",
"转录这个音频片段",
"将这段对话转换为文本",
"麻烦帮我把这段录音整理成详细的文字记录",
]
asr_en_templates = [
"Please transcribe this audio file",
"Convert this speech recording to text",
"Transcribe the following voice message",
"Turn this audio into readable text",
"Please convert the recording to written format",
"Transcribe what you hear in this audio",
"Convert this spoken content to text",
"Please write down what is said in this recording",
"Transcribe this voice recording",
"Could you please help me transcribe this important recording?",
"Would you mind converting this voice message into a readable text format?",
"I'd really appreciate it if you could turn this audio file into a written document",
]
tts_zh_templates = [
"请将这段文字转换为语音",
"帮我把这个文本读出来",
"将这些文字生成音频",
"请朗读这段内容",
"把这段话转换成语音文件",
"生成这段文字的语音版本",
"请用语音播报这些内容",
"将文本转换为可听的音频",
"帮我朗读这段文字",
"把这些内容念出来",
]
tts_en_templates = [
"Please convert this text to speech",
"Turn this writing into audio",
"Generate speech from this text",
"Read this content out loud",
"Convert these words to voice",
"Create an audio version of this text",
"Please vocalize this content",
"Turn this text into audible format",
"Help me convert this writing to speech",
"Make this text into spoken audio",
]
def detect_language(text):
if re.search(r"[\u4e00-\u9fff]", text):
return "zh"
else:
return "en"
# ============================================
# Common Helper Functions - InputSegment Creation
# ============================================
def create_segment(text: str = "", audio=None) -> InputSegment:
"""Create a standard InputSegment with default zeroemb parameters"""
return InputSegment(
text=text,
audio=audio,
speech_zeroemb_idx=speech_zeroemb_idx,
text_zeroemb_idx=empty_token,
)
def create_streaming_segment(
text: str, audio, tokenizer, group_size: int, audio_channels: int
) -> StreamingInputSegment:
"""Create a StreamingInputSegment"""
return StreamingInputSegment(
text=text,
audio=audio,
tokenizer=tokenizer,
group_size=group_size,
audio_channels=audio_channels,
speech_zeroemb_idx=speech_zeroemb_idx,
text_zeroemb_idx=empty_token,
)
# ============================================
# Common Helper Functions - Common Token Segment Creation
# ============================================
def create_user_start() -> InputSegment:
"""Create user role start token"""
return create_segment(text="<|im_start|>user\n")
def create_user_end() -> InputSegment:
"""Create role end token"""
return create_segment(text="<|im_end|>\n")
def create_assistant_start() -> InputSegment:
"""Create assistant role start token"""
return create_segment(text="<|im_start|>assistant\n")
def create_system_start() -> InputSegment:
"""Create system role start token"""
return create_segment(text="<|im_start|>system\n")
def create_thinking_segment(thinking: bool = False) -> InputSegment:
"""Create thinking token, closed or open based on thinking parameter"""
if thinking:
return create_segment(text="<think>\n")
else:
return create_segment(text="<think>\n\n</think>\n")
def create_sostm_segment() -> InputSegment:
"""Create streaming output start token"""
return create_segment(text="<|sostm|>")
def create_assistant_start_with_sostm() -> InputSegment:
"""Create assistant start token with sostm"""
return create_segment(text="<|im_start|>assistant\n<|sostm|>")
def create_assistant_start_with_think() -> InputSegment:
"""Create assistant start token with think opening"""
return create_segment(text="<|im_start|>assistant\n<think>\n")
# ============================================
# Common Helper Functions - Composite Segment Creation
# ============================================
def create_user_turn_with_audio(audio_tokenized, extra_text: str = None) -> list[InputSegment]:
"""Create a user turn containing audio"""
segments = [
create_user_start(),
create_segment(audio=audio_tokenized),
]
if extra_text:
segments.append(create_segment(text=extra_text))
segments.append(create_user_end())
return segments
def create_user_turn_with_text(text: str) -> list[InputSegment]:
"""Create a text-only user turn"""
return [
create_user_start(),
create_segment(text=text),
create_user_end(),
]
def create_system_turn_with_voice_prompt(prompt_text: str, audio_token) -> list[InputSegment]:
"""Create a system turn with voice prompt"""
return [
create_system_start(),
create_segment(text=prompt_text),
create_segment(text="", audio=audio_token),
create_user_end(),
]
def create_system_turn_text_only(system_text: str) -> list[InputSegment]:
"""Create a text-only system turn"""
return [
create_system_start(),
create_segment(text=system_text),
create_user_end(),
]
# ============================================
# Common Helper Functions - Multi-turn Dialogue Processing
# ============================================
def process_multiturn_messages(
message_list: list[dict],
user_processor: Callable[[dict], list[InputSegment]],
assistant_processor: Callable[[dict], list[InputSegment]],
) -> list[InputSegment]:
"""
Generic multi-turn dialogue message processing function
Args:
message_list: List of messages, each containing 'role' and 'content'
user_processor: Function to process user messages
assistant_processor: Function to process assistant messages
Returns:
Processed list of InputSegments
"""
lm_prompt = []
for message in message_list:
role = message["role"]
if role == "user":
lm_prompt.extend(user_processor(message))
elif role == "assistant":
lm_prompt.extend(assistant_processor(message))
else:
raise ValueError(f"Invalid role: {role}")
return lm_prompt
def create_text_user_message(message: dict) -> list[InputSegment]:
"""Process a text-only user message"""
return [
create_user_start(),
create_segment(text=message["content"]),
create_user_end(),
]
def create_text_assistant_message(message: dict) -> list[InputSegment]:
"""Process a text-only assistant message"""
return [
create_assistant_start(),
create_segment(text=message["content"]),
create_user_end(),
]
def create_audio_user_message(message: dict) -> list[InputSegment]:
"""Process an audio user message"""
return [
create_user_start(),
create_segment(audio=preprocess_input(message["content"])),
create_user_end(),
]
def append_assistant_ending(
lm_prompt: list[InputSegment], thinking: bool = False, use_sostm: bool = False
) -> list[InputSegment]:
"""
Append assistant ending to the prompt
Args:
lm_prompt: Existing prompt list
thinking: Whether to use open thinking token
use_sostm: Whether to use sostm token (for speech output)
"""
if use_sostm:
lm_prompt.append(create_assistant_start_with_sostm())
else:
lm_prompt.append(create_assistant_start())
lm_prompt.append(create_thinking_segment(thinking))
return lm_prompt
def get_asr_sft_prompt(
input: None | str = None,
):
"""Build prompt for ASR (Automatic Speech Recognition) task"""
audio_tokenized = preprocess_input(input)
template = random.choice(asr_zh_templates + asr_en_templates)
lm_prompt = create_user_turn_with_audio(audio_tokenized, extra_text=template)
lm_prompt = append_assistant_ending(lm_prompt, thinking=False)
return lm_prompt
def resample_audio_if_needed(wav_tensor: torch.Tensor, original_sr: int):
target_sr = 24000
if original_sr != target_sr:
wav_tensor = torchaudio.functional.resample(wav_tensor, original_sr, target_sr)
return wav_tensor
def wav2mel(wav, device="cpu"):
mel_transform = MelSpectrogram(
sample_rate=mimo_audio_tokenizer.config.sampling_rate,
n_fft=mimo_audio_tokenizer.config.nfft,
hop_length=mimo_audio_tokenizer.config.hop_length,
win_length=mimo_audio_tokenizer.config.window_size,
f_min=mimo_audio_tokenizer.config.fmin,
f_max=mimo_audio_tokenizer.config.fmax,
n_mels=mimo_audio_tokenizer.config.n_mels,
power=1.0,
center=True,
).to(device)
spec = mel_transform(wav[None, :])
return torch.log(torch.clip(spec, min=1e-7)).squeeze()
def group_by_length(features: torch.Tensor, lengths: torch.Tensor, max_length: int):
if features.size(0) != lengths.sum().item():
raise ValueError(f"Feature size mismatch: {features.size(0)} vs {lengths.sum().item()}")
split_points = []
current_sum = 0
for i, seq_len in enumerate(lengths):
if current_sum + seq_len > max_length and current_sum > 0:
split_points.append(i)
current_sum = seq_len.item()
else:
current_sum += seq_len.item()
# Convert split points to group sizes
group_sizes = []
prev = 0
for point in split_points:
group_sizes.append(point - prev)
prev = point
if prev < len(lengths):
group_sizes.append(len(lengths) - prev)
len_groups = torch.split(lengths, group_sizes)
feature_sizes = [group.sum().item() for group in len_groups]
feature_groups = torch.split(features, feature_sizes)
return feature_groups, len_groups
def encode_batch(input_features: torch.Tensor, input_lens: torch.Tensor, max_length: int = 256000):
feature_groups, len_groups = group_by_length(input_features, input_lens, max_length)
encoded_parts = []
for features, lengths in zip(feature_groups, len_groups):
with torch.no_grad():
codes, _ = mimo_audio_tokenizer.encoder.encode(
input_features=features.to(device), input_lens=lengths.to(device), return_codes_only=True
)
encoded_parts.append(codes)
return torch.cat(encoded_parts, dim=-1)
def preprocess_input(input: None | str | torch.Tensor = None, device="cpu", audio_channels=4, group_size=8):
if isinstance(input, torch.Tensor) or (isinstance(input, str) and os.path.isfile(input)):
return "<|sosp|><|empty|><|eosp|>"
else:
text = input
if (
text.isupper() or text.islower()
): # If the text only contains upper-case or lower-case letters, capitalize it.
text = text.capitalize()
return text
def _build_tts_system_prompt(has_voice_prompt: bool, voice_audio_token=None) -> list[InputSegment]:
"""Build system prompt for TTS task"""
if has_voice_prompt and voice_audio_token is not None:
return [
create_system_start(),
create_segment(
# text="You need to generate a speech with the same timbre as the speech prompt, based on the specified style instructions and text content. Your timbre should be: "
text="你需要根据指定的风格指令和文本内容来生成和语音prompt具有相同音色的语音。你的音色应该是:"
),
create_segment(text="", audio=voice_audio_token),
create_user_end(),
]
else:
return create_system_turn_text_only("你需要根据指定的风格指令和文本内容来生成语音。")
# return create_system_turn_text_only(
# "You need to generate speech based on the specified style instructions and text content."
# )
def _build_tts_system_prompt_no_instruct(has_voice_prompt: bool, voice_audio_token=None) -> list[InputSegment]:
"""Build system prompt for TTS task"""
if has_voice_prompt and voice_audio_token is not None:
return [
create_system_start(),
create_segment(
# text="You need to generate a speech with the same timbre as the speech prompt, based on the specified style instructions and text content. Your timbre should be:"
text="你需要根据指定的风格指令和文本内容来生成和语音prompt具有相同音色的语音。你的音色应该是:"
),
create_segment(text="", audio=voice_audio_token),
create_user_end(),
]
else:
return []
def get_tts_sft_prompt(
input: None | str = None,
instruct=None,
read_text_only=True,
prompt_speech=None,
):
"""
Build prompt for TTS (Text-to-Speech) task
Args:
input: Input text
instruct: Style instruction (e.g., "speak happily in a child's voice")
read_text_only: Whether to read only plain text (False means text contains template)
prompt_speech: Reference audio (for voice cloning)
"""
assistant_prompt_audio_token = preprocess_input(prompt_speech) if prompt_speech is not None else None
if not read_text_only:
# Not just reading text, text contains template (template:text format)
text = preprocess_input(input)
lm_prompt = _build_tts_system_prompt(
has_voice_prompt=assistant_prompt_audio_token is not None,
voice_audio_token=assistant_prompt_audio_token,
)
lm_prompt.append(create_segment(text=f"<|im_start|>user\n{text}<|im_end|>\n"))
lm_prompt.append(create_assistant_start_with_think())
else:
# Plain text (no instruction inside)
language = detect_language(input)
template = random.choice(tts_zh_templates if language == "zh" else tts_en_templates)
text = preprocess_input(input)
if instruct is None:
# No instruct instruction
lm_prompt = _build_tts_system_prompt_no_instruct(
has_voice_prompt=assistant_prompt_audio_token is not None,
voice_audio_token=assistant_prompt_audio_token,
)
lm_prompt.extend(
[
create_segment(text=f"<|im_start|>user\n{template}: {text}<|im_end|>\n"),
create_assistant_start_with_sostm(),
]
)
else:
# Has instruct instruction
lm_prompt = _build_tts_system_prompt(
has_voice_prompt=assistant_prompt_audio_token is not None,
voice_audio_token=assistant_prompt_audio_token,
)
lm_prompt.append(create_segment(text=f"<|im_start|>user\n{template}: {text}({instruct})<|im_end|>\n"))
lm_prompt.append(create_assistant_start_with_think())
return lm_prompt
def get_audio_understanding_sft_prompt(
input_speech,
input_text,
thinking=False,
use_sostm=False,
):
"""Build prompt for audio understanding task"""
audio_tokenized = preprocess_input(input_speech)
lm_prompt = create_user_turn_with_audio(audio_tokenized, extra_text=input_text)
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking, use_sostm=use_sostm)
return lm_prompt
def _build_voice_prompt_system(prompt_speech) -> list[InputSegment]:
"""Build system prompt with voice prompt"""
return create_system_turn_with_voice_prompt(
prompt_text="Your voice should be", audio_token=preprocess_input(prompt_speech)
)
def get_spoken_dialogue_sft_prompt(
input_speech,
system_prompt=None,
prompt_speech=None,
add_history=False,
):
"""
Build prompt for spoken dialogue task
Args:
input_speech: Input speech
system_prompt: System prompt text
prompt_speech: Reference audio (for voice cloning)
add_history: Whether to add history (Note: history variable is undefined in original code)
"""
audio_tokenized = preprocess_input(input_speech)
lm_prompt = []
# Note: history variable is undefined in original code, this branch may never execute
# To use history feature, history should be passed as a parameter
if add_history:
# Simplified form of adding history
lm_prompt = create_user_turn_with_audio(audio_tokenized)
lm_prompt.append(create_assistant_start_with_sostm())
else:
# Add voice prompt (if available)
if prompt_speech:
lm_prompt.extend(_build_voice_prompt_system(prompt_speech))
# Add user turn
lm_prompt.append(create_user_start())
if system_prompt:
lm_prompt.append(create_segment(text=system_prompt))
lm_prompt.append(create_segment(audio=audio_tokenized))
lm_prompt.append(create_user_end())
lm_prompt.append(create_assistant_start_with_sostm())
return lm_prompt
def get_spoken_dialogue_sft_multiturn_prompt(
message_list,
system_prompt=None,
prompt_speech=None,
tokenizer=None,
group_size=8,
audio_channels=4,
):
"""
Build prompt for multi-turn spoken dialogue task
Args:
message_list: List of messages containing role and content
system_prompt: System prompt text
prompt_speech: Reference audio (for voice cloning)
tokenizer: Tokenizer
group_size: Group size
audio_channels: Number of audio channels
"""
lm_prompt = []
# Add voice prompt (if available)
if prompt_speech:
lm_prompt.extend(
create_system_turn_with_voice_prompt(
prompt_text="Your voice should be:", audio_token=preprocess_input(prompt_speech)
)
)
# Define message processors
def user_processor(msg):
segments = [create_user_start()]
if system_prompt:
segments.append(create_segment(text=system_prompt))
segments.append(create_segment(audio=preprocess_input(msg["content"])))
segments.append(create_user_end())
return segments
def assistant_processor(msg):
return [
create_assistant_start(),
create_streaming_segment(
text=msg["content"]["text"],
audio=preprocess_input(msg["content"]["audio"]),
tokenizer=tokenizer,
group_size=group_size,
audio_channels=audio_channels,
),
create_user_end(),
]
# Process message list
lm_prompt.extend(process_multiturn_messages(message_list, user_processor, assistant_processor))
lm_prompt.append(create_assistant_start_with_sostm())
return lm_prompt
def get_s2t_dialogue_sft_prompt(
input_speech,
thinking=False,
):
"""Build prompt for speech-to-text dialogue task"""
audio_tokenized = preprocess_input(input_speech)
lm_prompt = create_user_turn_with_audio(audio_tokenized)
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
return lm_prompt
def get_s2t_dialogue_sft_multiturn_prompt(message_list, thinking=False):
"""Build prompt for multi-turn speech-to-text dialogue task"""
lm_prompt = process_multiturn_messages(
message_list, user_processor=create_audio_user_message, assistant_processor=create_text_assistant_message
)
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
return lm_prompt
def get_text_dialogue_sft_prompt(
input_text,
thinking=False,
):
"""Build prompt for text-only dialogue task"""
lm_prompt = create_user_turn_with_text(input_text)
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
return lm_prompt
def get_text_dialogue_sft_multiturn_prompt(
message_list,
thinking=False,
):
"""Build prompt for multi-turn text-only dialogue task"""
lm_prompt = process_multiturn_messages(
message_list, user_processor=create_text_user_message, assistant_processor=create_text_assistant_message
)
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
return lm_prompt
def get_in_context_learning_s2s_prompt(
instruction,
prompt_examples,
audio,
tokenizer=None,
group_size=8,
audio_channels=4,
):
"""
Build prompt for In-Context Learning speech-to-speech task
Args:
instruction: Instruction text
prompt_examples: List of examples, each containing input_audio, output_transcription, output_audio
audio: Input audio to be processed
tokenizer: Tokenizer
group_size: Group size
audio_channels: Number of audio channels
"""
prompt = [create_segment(text=f"[Int]:{instruction}\n")]
# Add examples
for example in prompt_examples:
prompt.extend(
[
create_segment(audio=preprocess_input(example["input_audio"])),
create_segment(text="\n"),
create_streaming_segment(
text=example["output_transcription"],
audio=preprocess_input(example["output_audio"]),
tokenizer=tokenizer,
group_size=group_size,
audio_channels=audio_channels,
),
create_segment(text=" \n\n"),
]
)
# Add input audio to be processed
prompt.extend(
[
create_segment(audio=preprocess_input(audio)),
create_segment(text="\n"),
create_sostm_segment(),
]
)
return prompt
def get_audio_data(audio_url):
if audio_url.startswith("data:"):
header, b64_data = audio_url.split(",", 1)
audio_bytes = base64.b64decode(b64_data.strip())
audio_file = io.BytesIO(audio_bytes)
else:
# File path
audio_file = audio_url
audio_signal, sr = load_audio(audio_file, sr=24000)
audio_data = (audio_signal.astype(np.float32), sr)
return audio_data
def to_prompt(input_segs):
out_put = []
for input_seg in input_segs:
if isinstance(input_seg, StreamingInputSegment) and input_seg.text:
out_put.append("<|sostm|>")
if input_seg.audio is not None and isinstance(input_seg.audio, str):
out_put.append(input_seg.text)
out_put.append("<|eot|>")
out_put.append("<|empty|>")
else:
out_put.append(input_seg.text)
out_put.append("<|eot|>")
out_put.append("<|eostm|>")
else:
out_put.append(input_seg.text)
if input_seg.audio is not None:
out_put.append(input_seg.audio)
prompt = "".join(out_put)
return prompt