transcribe/whisper_llm_serve.py

import asyncio
import json
import queue
import threading
import time
from logging import getLogger
from typing import List, Optional, Iterator, Tuple, Any
import asyncio
import numpy as np
import config
import collections
from api_model import TransResult, Message, DebugResult

from .utils import log_block, save_to_wave, TestDataWriter, filter_words
from .translatepipes import TranslatePipes
from .strategy import (
    TranscriptStabilityAnalyzer, TranscriptToken)
from transcribe.helpers.vadprocessor import VadProcessor
# from transcribe.helpers.vad_dynamic import VadProcessor
# from transcribe.helpers.vadprocessor import VadProcessor
from transcribe.pipelines import MetaItem

logger = getLogger("TranscriptionService")


class WhisperTranscriptionService:
    """
    Whisper语音转录服务类，处理音频流转录和翻译
    """

    SERVER_READY = "SERVER_READY"
    DISCONNECT = "DISCONNECT"

    def __init__(self, websocket, pipe: TranslatePipes, language=None, dst_lang=None, client_uid=None):
        print('>>>>>>>>>>>>>>>> init service >>>>>>>>>>>>>>>>>>>>>>')
        print('src_lang:', language)
        self.source_language = language  # 源语言
        self.target_language = dst_lang  # 目标翻译语言
        self.client_uid = client_uid
        # 转录结果稳定性管理
        self.websocket = websocket
        self._translate_pipe = pipe

        # 音频处理相关
        self.sample_rate = 16000

        self.lock = threading.Lock()


        # 文本分隔符，根据语言设置
        self.text_separator = self._get_text_separator(language)
        self.loop = asyncio.get_event_loop()
        # 发送就绪状态
        #  原始音频队列
        self._frame_queue = queue.Queue()
        #  音频队列缓冲区
        self.frames_np = None
        #  完整音频队列
        self.segments_queue = collections.deque()
        self._temp_string = ""

        self._transcrible_analysis = None
        # 启动处理线程
        self._translate_thread_stop = threading.Event()
        self._frame_processing_thread_stop = threading.Event()

        self.translate_thread = self._start_thread(self._transcription_processing_loop)
        self.frame_processing_thread = self._start_thread(self._frame_processing_loop)
        # if language == "zh":
        #     self._vad = VadProcessor(prob_threshold=0.8, silence_s=0.2, cache_s=0.15)
        # else:
        #     self._vad = VadProcessor(prob_threshold=0.7, silence_s=0.2, cache_s=0.15)
        self.row_number = 0
        # for test
        self._transcrible_time_cost = 0.
        self._translate_time_cost = 0.

        if config.SAVE_DATA_SAVE:
            self._save_task_stop = threading.Event()
            self._save_queue = queue.Queue()
            self._save_thread = self._start_thread(self.save_data_loop)

        # self._c = 0

    def save_data_loop(self):
        writer = TestDataWriter()
        while not self._save_task_stop.is_set():
            test_data = self._save_queue.get()
            writer.write(test_data)  # Save test_data to CSV


    def _start_thread(self, target_function) -> threading.Thread:
        """启动守护线程执行指定函数"""
        thread = threading.Thread(target=target_function)
        thread.daemon = True
        thread.start()
        return thread

    def _get_text_separator(self, language: str) -> str:
        """根据语言返回适当的文本分隔符"""
        return "" if language == "zh" else " "

    async def send_ready_state(self) -> None:
        """发送服务就绪状态消息"""
        await self.websocket.send(json.dumps({
            "uid": self.client_uid,
            "message": self.SERVER_READY,
            "backend": "whisper_transcription"
        }))

    def set_language(self, source_lang: str, target_lang: str) -> None:
        """设置源语言和目标语言"""
        self.source_language = source_lang
        self.target_language = target_lang
        self.text_separator = self._get_text_separator(source_lang)
        # self._transcrible_analysis = TranscriptStabilityAnalyzer(self.source_language, self.text_separator)

    def add_frames(self, frame_np: np.ndarray) -> None:
        """添加音频帧到处理队列"""
        self._frame_queue.put(frame_np)

    def _apply_voice_activity_detection(self, frame_np:np.array):
        """应用语音活动检测来优化音频缓冲区"""
        processed_audio = self._translate_pipe.voice_detect(frame_np.tobytes())
        speech_audio =  np.frombuffer(processed_audio.audio, dtype=np.float32)
        speech_status = processed_audio.speech_status
        return speech_audio, speech_status

    def _frame_processing_loop(self) -> None:
        """从队列获取音频帧并合并到缓冲区"""
        while not self._frame_processing_thread_stop.is_set():
            try:
                frame_np = self._frame_queue.get(timeout=0.1)
                frame_np, speech_status = self._apply_voice_activity_detection(frame_np)
                if frame_np is None:
                    continue
                with self.lock:
                    if self.frames_np is None:
                        self.frames_np = frame_np.copy()
                    else:
                        self.frames_np = np.append(self.frames_np, frame_np)
                    if speech_status == "END" and len(frame_np) > 0:
                        self.segments_queue.appendleft(self.frames_np.copy())
                        self.frames_np = np.array([], dtype=np.float32)
            except queue.Empty:
                pass

    def _process_transcription_results_2(self, seg_text:str,partial):

        item =  TransResult(
                seg_id=self.row_number,
                context=seg_text,
                from_=self.source_language,
                to=self.target_language,
                tran_content=self._translate_text_large(seg_text),
                partial=partial
            )
        if partial == False:
            self.row_number += 1
        return item

    def _transcription_processing_loop(self) -> None:
        """主转录处理循环"""
        frame_epoch = 1
        while not self._translate_thread_stop.is_set():

            if self.frames_np is None:
                time.sleep(0.2)
                continue

            with self.lock:
                if len(self.segments_queue) >0:
                    audio_buffer = self.segments_queue.pop()
                    partial = False
                else:
                    audio_buffer = self.frames_np[:int(frame_epoch * 1.5 * self.sample_rate)]# 获取 1.5s * epoch 个音频长度
                    partial = True

            if len(audio_buffer) ==0:
                time.sleep(0.2)
                continue

            if len(audio_buffer) < int(self.sample_rate):
                silence_audio = np.zeros(self.sample_rate, dtype=np.float32)
                silence_audio[-len(audio_buffer):] = audio_buffer
                audio_buffer = silence_audio


            logger.debug(f"audio buffer size: {len(audio_buffer) / self.sample_rate:.2f}s")
            # try:
            meta_item = self._transcribe_audio(audio_buffer)
            segments = meta_item.segments
            logger.debug(f"Segments: {segments}")
            segments = filter_words(segments)
            if len(segments):
                seg_text = self.text_separator.join(seg.text for seg in segments)
                if self._temp_string:
                    seg_text = self._temp_string + seg_text

                if partial == False:
                    if len(seg_text) < config.TEXT_THREHOLD:
                        partial = True
                        self._temp_string = seg_text
                    else:
                        self._temp_string = ""


                result = self._process_transcription_results_2(seg_text, partial)
                self._send_result_to_client(result)
                time.sleep(0.1)

                if partial == False:
                    frame_epoch = 1
                else:
                    frame_epoch += 1
            # 处理转录结果并发送到客户端
            # for result in self._process_transcription_results(segments, audio_buffer):
            #     self._send_result_to_client(result)

            # except Exception as e:
            #     logger.error(f"Error processing audio: {e}")


    def _transcribe_audio(self, audio_buffer: np.ndarray)->MetaItem:
        """转录音频并返回转录片段"""
        log_block("Audio buffer length", f"{audio_buffer.shape[0]/self.sample_rate:.2f}", "s")
        start_time = time.perf_counter()

        result = self._translate_pipe.transcrible(audio_buffer.tobytes(), self.source_language)
        segments = result.segments
        time_diff = (time.perf_counter() - start_time)
        logger.debug(f"📝 Transcrible Segments: {segments} ")
        # logger.debug(f"📝 Transcrible: {self.text_separator.join(seg.text for seg in segments)} ")
        log_block("📝 Transcrible output", f"{self.text_separator.join(seg.text for seg in segments)}", "")
        log_block("📝 Transcrible time", f"{time_diff:.3f}", "s")
        self._transcrible_time_cost = round(time_diff, 3)
        return result

    def _translate_text(self, text: str) -> str:
        """将文本翻译为目标语言"""
        if not text.strip():
            return ""

        log_block("🐧 Translation input ", f"{text}")
        start_time = time.perf_counter()

        result = self._translate_pipe.translate(text, self.source_language, self.target_language)
        translated_text = result.translate_content
        time_diff = (time.perf_counter() - start_time)
        log_block("🐧 Translation time ", f"{time_diff:.3f}", "s")
        log_block("🐧 Translation out ", f"{translated_text}")
        self._translate_time_cost = round(time_diff, 3)
        return translated_text

    def _translate_text_large(self, text: str) -> str:
        """将文本翻译为目标语言"""
        if not text.strip():
            return ""

        log_block("Translation input", f"{text}")
        start_time = time.perf_counter()

        result = self._translate_pipe.translate_large(text, self.source_language, self.target_language)
        translated_text = result.translate_content
        time_diff = (time.perf_counter() - start_time)
        log_block("Translation large model time ", f"{time_diff:.3f}", "s")
        log_block("Translation large model output", f"{translated_text}")
        self._translate_time_cost = round(time_diff, 3)
        return translated_text



    def _process_transcription_results(self, segments: List[TranscriptToken], audio_buffer: np.ndarray) -> Iterator[TransResult]:
        """
        处理转录结果，生成翻译结果

        Returns:
            TransResult对象的迭代器
        """

        if not segments:
            return
        start_time = time.perf_counter()
        for ana_result in self._transcrible_analysis.analysis(segments, len(audio_buffer)/self.sample_rate):
            if (cut_index :=ana_result.cut_index)>0:
                # 更新音频缓冲区，移除已处理部分
                self._update_audio_buffer(cut_index)
            if ana_result.partial():
                translated_context = self._translate_text(ana_result.context)
            else:
                translated_context = self._translate_text_large(ana_result.context)

            yield TransResult(
                seg_id=ana_result.seg_id,
                context=ana_result.context,
                from_=self.source_language,
                to=self.target_language,
                tran_content=translated_context,
                partial=ana_result.partial()
            )
            current_time = time.perf_counter()
            time_diff = current_time - start_time
            if config.SAVE_DATA_SAVE:
                self._save_queue.put(DebugResult(
                    seg_id=ana_result.seg_id,
                    transcrible_time=self._transcrible_time_cost,
                    translate_time=self._translate_time_cost,
                    context=ana_result.context,
                    from_=self.source_language,
                    to=self.target_language,
                    tran_content=translated_context,
                    partial=ana_result.partial()
                ))
            log_block("🚦 Traffic times diff", round(time_diff, 2), 's')


    def _send_result_to_client(self, result: TransResult) -> None:
        """发送翻译结果到客户端"""
        try:
            message = Message(result=result, request_id=self.client_uid).model_dump_json(by_alias=True)
            coro = self.websocket.send_text(message)
            future = asyncio.run_coroutine_threadsafe(coro, self.loop)
            future.add_done_callback(lambda fut: fut.exception() and self.stop())
        except RuntimeError:
            self.stop()
        except Exception as e:
            logger.error(f"Error sending result to client: {e}")

    def stop(self) -> None:
        """停止所有处理线程并清理资源"""
        self._translate_thread_stop.set()
        self._frame_processing_thread_stop.set()
        if config.SAVE_DATA_SAVE:
            self._save_task_stop.set()
        logger.info(f"Stopping transcription service for client: {self.client_uid}")