715 lines
23 KiB
Python
715 lines
23 KiB
Python
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
|