148 lines
4.9 KiB
Python
148 lines
4.9 KiB
Python
import torch
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import librosa
|
|
import re
|
|
import logging
|
|
from functools import partial
|
|
from dataclasses import asdict
|
|
from typing import Optional
|
|
from transformers import (
|
|
AutoConfig,
|
|
AutoModel,
|
|
AutoTokenizer,
|
|
WhisperProcessor,
|
|
)
|
|
|
|
import transformers.models.llama.modeling_llama as _llama_mod
|
|
|
|
if not hasattr(_llama_mod, "LLAMA_ATTENTION_CLASSES"):
|
|
_llama_mod.LLAMA_ATTENTION_CLASSES = _llama_mod.ALL_ATTENTION_FUNCTIONS
|
|
|
|
from boson_multimodal.data_collator.higgs_audio_collator import HiggsAudioSampleCollator
|
|
from boson_multimodal.data_types import ChatMLSample, AudioContent, Message
|
|
from boson_multimodal.dataset.chatml_dataset import (
|
|
ChatMLDatasetSample,
|
|
prepare_chatml_sample_qwen,
|
|
)
|
|
|
|
from ..config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_model: Optional[AutoModel] = None
|
|
_tokenizer: Optional[AutoTokenizer] = None
|
|
_collator: Optional[HiggsAudioSampleCollator] = None
|
|
_config: Optional[AutoConfig] = None
|
|
|
|
|
|
def _load_model():
|
|
global _model, _tokenizer, _collator, _config
|
|
if _model is not None:
|
|
return
|
|
|
|
logger.info("Loading Higgs STT on %s", settings.stt_device)
|
|
dtype = torch.bfloat16 if settings.stt_dtype == "bfloat16" else torch.float16
|
|
|
|
_config = AutoConfig.from_pretrained(settings.stt_model_id, trust_remote_code=True)
|
|
_model = AutoModel.from_pretrained(
|
|
settings.stt_model_id,
|
|
torch_dtype=dtype,
|
|
trust_remote_code=True,
|
|
attn_implementation="eager",
|
|
device_map=settings.stt_device,
|
|
)
|
|
_model.eval()
|
|
_tokenizer = AutoTokenizer.from_pretrained(settings.stt_model_id)
|
|
_model.audio_out_bos_token_id = _tokenizer.convert_tokens_to_ids(
|
|
"<|audio_out_bos|>"
|
|
)
|
|
_model.audio_eos_token_id = _tokenizer.convert_tokens_to_ids("<|audio_eos|>")
|
|
|
|
whisper_proc = WhisperProcessor.from_pretrained("openai/whisper-large-v3")
|
|
_collator = HiggsAudioSampleCollator(
|
|
whisper_processor=whisper_proc,
|
|
audio_in_token_id=_config.audio_in_token_idx,
|
|
audio_out_token_id=_config.audio_out_token_idx,
|
|
audio_stream_bos_id=_config.audio_stream_bos_id,
|
|
audio_stream_eos_id=_config.audio_stream_eos_id,
|
|
encode_whisper_embed=_config.encode_whisper_embed,
|
|
pad_token_id=_config.pad_token_id,
|
|
return_audio_in_tokens=_config.encode_audio_in_tokens,
|
|
use_delay_pattern=_config.use_delay_pattern,
|
|
round_to=1,
|
|
audio_num_codebooks=_config.audio_num_codebooks,
|
|
chunk_size_seconds=getattr(_config, "chunk_size_seconds", 30),
|
|
encoder_padding_method=getattr(_config, "encoder_padding_method", "max_length"),
|
|
)
|
|
|
|
|
|
def transcribe_segment(audio_path: str, start_time: float, end_time: float) -> str:
|
|
_load_model()
|
|
|
|
audio_np, sr = sf.read(audio_path)
|
|
if sr != 16000:
|
|
audio_np = librosa.resample(audio_np, orig_sr=sr, target_sr=16000)
|
|
|
|
start_sample = int(start_time * 16000)
|
|
end_sample = int(end_time * 16000)
|
|
segment_audio = audio_np[start_sample:end_sample]
|
|
|
|
if len(segment_audio) == 0:
|
|
return ""
|
|
|
|
prompt = "Transcribe the speech. Output only the spoken words in lowercase with no punctuation."
|
|
messages = [
|
|
Message(role="user", content=[prompt, AudioContent(audio_url="placeholder")])
|
|
]
|
|
chatml = ChatMLSample(messages=messages)
|
|
prep_fn = partial(prepare_chatml_sample_qwen, enable_thinking=True)
|
|
input_tokens, _, _, _ = prep_fn(chatml, _tokenizer, add_generation_prompt=True)
|
|
|
|
sample = ChatMLDatasetSample(
|
|
input_ids=torch.LongTensor(input_tokens),
|
|
label_ids=None,
|
|
audio_ids_concat=None,
|
|
audio_ids_start=None,
|
|
audio_waveforms_concat=torch.tensor(segment_audio, dtype=torch.float32),
|
|
audio_waveforms_start=torch.tensor([0]),
|
|
audio_sample_rate=torch.tensor([16000]),
|
|
audio_speaker_indices=torch.tensor([0]),
|
|
)
|
|
|
|
batch = asdict(_collator([sample]))
|
|
device = next(_model.parameters()).device
|
|
batch = {
|
|
k: v.to(device).contiguous() if isinstance(v, torch.Tensor) else v
|
|
for k, v in batch.items()
|
|
}
|
|
|
|
with torch.inference_mode():
|
|
outputs = _model.generate(
|
|
**batch,
|
|
max_new_tokens=1024,
|
|
use_cache=True,
|
|
do_sample=False,
|
|
stop_strings=["<|im_end|>", "<|endoftext|>"],
|
|
tokenizer=_tokenizer,
|
|
)
|
|
|
|
output_ids = outputs[0] if isinstance(outputs, tuple) else outputs
|
|
full_text = _tokenizer.decode(output_ids[0], skip_special_tokens=False)
|
|
|
|
parts = full_text.split("assistant\n")
|
|
hyp = parts[-1] if len(parts) > 1 else full_text
|
|
hyp = re.sub(r"<think>.*?</think>", "", hyp, flags=re.DOTALL)
|
|
hyp = re.sub(r"<\|.*?\|>", "", hyp).strip()
|
|
|
|
return hyp
|
|
|
|
|
|
def unload_model():
|
|
global _model, _tokenizer, _collator, _config
|
|
_model = None
|
|
_tokenizer = None
|
|
_collator = None
|
|
_config = None
|
|
torch.cuda.empty_cache()
|