From 993a83546ae8bec4225503afc2545dfabbb06f42 Mon Sep 17 00:00:00 2001 From: Quentin Fuxa Date: Mon, 16 Jun 2025 16:13:57 +0200 Subject: [PATCH] core refactoring --- setup.py | 2 +- whisperlivekit/__init__.py | 7 +- whisperlivekit/audio_processor.py | 9 +- whisperlivekit/basic_server.py | 28 ++-- whisperlivekit/core.py | 197 ++++++---------------------- whisperlivekit/parse_args.py | 141 ++++++++++++++++++++ whisperlivekit/web/web_interface.py | 13 ++ 7 files changed, 220 insertions(+), 177 deletions(-) create mode 100644 whisperlivekit/parse_args.py create mode 100644 whisperlivekit/web/web_interface.py diff --git a/setup.py b/setup.py index a1c6e68..222e454 100644 --- a/setup.py +++ b/setup.py @@ -1,7 +1,7 @@ from setuptools import setup, find_packages setup( name="whisperlivekit", - version="0.1.7", + version="0.1.8", description="Real-time, Fully Local Whisper's Speech-to-Text and Speaker Diarization", long_description=open("README.md", "r", encoding="utf-8").read(), long_description_content_type="text/markdown", diff --git a/whisperlivekit/__init__.py b/whisperlivekit/__init__.py index 7319521..fdbd82f 100644 --- a/whisperlivekit/__init__.py +++ b/whisperlivekit/__init__.py @@ -1,4 +1,5 @@ -from .core import WhisperLiveKit, parse_args +from .core import TranscriptionEngine from .audio_processor import AudioProcessor - -__all__ = ['WhisperLiveKit', 'AudioProcessor', 'parse_args'] \ No newline at end of file +from .web.web_interface import get_web_interface_html +from .parse_args import parse_args +__all__ = ['TranscriptionEngine', 'AudioProcessor', 'get_web_interface_html', 'parse_args'] \ No newline at end of file diff --git a/whisperlivekit/audio_processor.py b/whisperlivekit/audio_processor.py index 37a0580..08f2b29 100644 --- a/whisperlivekit/audio_processor.py +++ b/whisperlivekit/audio_processor.py @@ -8,7 +8,7 @@ import traceback from datetime import timedelta from whisperlivekit.timed_objects import ASRToken from whisperlivekit.whisper_streaming_custom.whisper_online import online_factory -from whisperlivekit.core import WhisperLiveKit +from whisperlivekit.core import TranscriptionEngine # Set up logging once logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") @@ -27,10 +27,13 @@ class AudioProcessor: Handles audio processing, state management, and result formatting. """ - def __init__(self): + def __init__(self, **kwargs): """Initialize the audio processor with configuration, models, and state.""" - models = WhisperLiveKit() + if 'transcription_engine' in kwargs and isinstance(kwargs['transcription_engine'], TranscriptionEngine): + models = kwargs['transcription_engine'] + else: + models = TranscriptionEngine(**kwargs) # Audio processing settings self.args = models.args diff --git a/whisperlivekit/basic_server.py b/whisperlivekit/basic_server.py index 1940b0f..afe62ab 100644 --- a/whisperlivekit/basic_server.py +++ b/whisperlivekit/basic_server.py @@ -2,26 +2,24 @@ from contextlib import asynccontextmanager from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.responses import HTMLResponse from fastapi.middleware.cors import CORSMiddleware - -from whisperlivekit import WhisperLiveKit, parse_args -from whisperlivekit.audio_processor import AudioProcessor - +from whisperlivekit import TranscriptionEngine, AudioProcessor, get_web_interface_html, parse_args import asyncio import logging -import os, sys -import argparse logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logging.getLogger().setLevel(logging.WARNING) logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) -kit = None +args = parse_args() +transcription_engine = None @asynccontextmanager async def lifespan(app: FastAPI): - global kit - kit = WhisperLiveKit() + global transcription_engine + transcription_engine = TranscriptionEngine( + **vars(args), + ) yield app = FastAPI(lifespan=lifespan) @@ -33,10 +31,9 @@ app.add_middleware( allow_headers=["*"], ) - @app.get("/") async def get(): - return HTMLResponse(kit.web_interface()) + return HTMLResponse(get_web_interface_html()) async def handle_websocket_results(websocket, results_generator): @@ -55,8 +52,10 @@ async def handle_websocket_results(websocket, results_generator): @app.websocket("/asr") async def websocket_endpoint(websocket: WebSocket): - audio_processor = AudioProcessor() - + global transcription_engine + audio_processor = AudioProcessor( + transcription_engine=transcription_engine, + ) await websocket.accept() logger.info("WebSocket connection opened.") @@ -94,8 +93,6 @@ def main(): """Entry point for the CLI command.""" import uvicorn - args = parse_args() - uvicorn_kwargs = { "app": "whisperlivekit.basic_server:app", "host":args.host, @@ -114,7 +111,6 @@ def main(): "ssl_keyfile": args.ssl_keyfile } - if ssl_kwargs: uvicorn_kwargs = {**uvicorn_kwargs, **ssl_kwargs} diff --git a/whisperlivekit/core.py b/whisperlivekit/core.py index d6f90cf..45467b9 100644 --- a/whisperlivekit/core.py +++ b/whisperlivekit/core.py @@ -2,148 +2,10 @@ try: from whisperlivekit.whisper_streaming_custom.whisper_online import backend_factory, warmup_asr except ImportError: from .whisper_streaming_custom.whisper_online import backend_factory, warmup_asr -from argparse import Namespace, ArgumentParser - -def parse_args(): - parser = ArgumentParser(description="Whisper FastAPI Online Server") - parser.add_argument( - "--host", - type=str, - default="localhost", - help="The host address to bind the server to.", - ) - parser.add_argument( - "--port", type=int, default=8000, help="The port number to bind the server to." - ) - parser.add_argument( - "--warmup-file", - type=str, - default=None, - dest="warmup_file", - help=""" - The path to a speech audio wav file to warm up Whisper so that the very first chunk processing is fast. - If not set, uses https://github.com/ggerganov/whisper.cpp/raw/master/samples/jfk.wav. - If False, no warmup is performed. - """, - ) - - parser.add_argument( - "--confidence-validation", - action="store_true", - help="Accelerates validation of tokens using confidence scores. Transcription will be faster but punctuation might be less accurate.", - ) - - parser.add_argument( - "--diarization", - action="store_true", - default=False, - help="Enable speaker diarization.", - ) - - parser.add_argument( - "--no-transcription", - action="store_true", - help="Disable transcription to only see live diarization results.", - ) - - parser.add_argument( - "--min-chunk-size", - type=float, - default=0.5, - help="Minimum audio chunk size in seconds. It waits up to this time to do processing. If the processing takes shorter time, it waits, otherwise it processes the whole segment that was received by this time.", - ) - - parser.add_argument( - "--model", - type=str, - default="tiny", - help="Name size of the Whisper model to use (default: tiny). Suggested values: tiny.en,tiny,base.en,base,small.en,small,medium.en,medium,large-v1,large-v2,large-v3,large,large-v3-turbo. The model is automatically downloaded from the model hub if not present in model cache dir.", - ) - - parser.add_argument( - "--model_cache_dir", - type=str, - default=None, - help="Overriding the default model cache dir where models downloaded from the hub are saved", - ) - parser.add_argument( - "--model_dir", - type=str, - default=None, - help="Dir where Whisper model.bin and other files are saved. This option overrides --model and --model_cache_dir parameter.", - ) - parser.add_argument( - "--lan", - "--language", - type=str, - default="auto", - help="Source language code, e.g. en,de,cs, or 'auto' for language detection.", - ) - parser.add_argument( - "--task", - type=str, - default="transcribe", - choices=["transcribe", "translate"], - help="Transcribe or translate.", - ) - parser.add_argument( - "--backend", - type=str, - default="faster-whisper", - choices=["faster-whisper", "whisper_timestamped", "mlx-whisper", "openai-api"], - help="Load only this backend for Whisper processing.", - ) - parser.add_argument( - "--vac", - action="store_true", - default=False, - help="Use VAC = voice activity controller. Recommended. Requires torch.", - ) - parser.add_argument( - "--vac-chunk-size", type=float, default=0.04, help="VAC sample size in seconds." - ) - - parser.add_argument( - "--no-vad", - action="store_true", - help="Disable VAD (voice activity detection).", - ) - - parser.add_argument( - "--buffer_trimming", - type=str, - default="segment", - choices=["sentence", "segment"], - help='Buffer trimming strategy -- trim completed sentences marked with punctuation mark and detected by sentence segmenter, or the completed segments returned by Whisper. Sentence segmenter must be installed for "sentence" option.', - ) - parser.add_argument( - "--buffer_trimming_sec", - type=float, - default=15, - help="Buffer trimming length threshold in seconds. If buffer length is longer, trimming sentence/segment is triggered.", - ) - parser.add_argument( - "-l", - "--log-level", - dest="log_level", - choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], - help="Set the log level", - default="DEBUG", - ) - parser.add_argument("--ssl-certfile", type=str, help="Path to the SSL certificate file.", default=None) - parser.add_argument("--ssl-keyfile", type=str, help="Path to the SSL private key file.", default=None) +from argparse import Namespace - args = parser.parse_args() - - args.transcription = not args.no_transcription - args.vad = not args.no_vad - delattr(args, 'no_transcription') - delattr(args, 'no_vad') - - return args - -class WhisperLiveKit: +class TranscriptionEngine: _instance = None _initialized = False @@ -153,14 +15,48 @@ class WhisperLiveKit: return cls._instance def __init__(self, **kwargs): - if WhisperLiveKit._initialized: + if TranscriptionEngine._initialized: return - - default_args = vars(parse_args()) + + defaults = { + "host": "localhost", + "port": 8000, + "warmup_file": None, + "confidence_validation": False, + "diarization": False, + "min_chunk_size": 0.5, + "model": "tiny", + "model_cache_dir": None, + "model_dir": None, + "lan": "auto", + "task": "transcribe", + "backend": "faster-whisper", + "vac": False, + "vac_chunk_size": 0.04, + "buffer_trimming": "segment", + "buffer_trimming_sec": 15, + "log_level": "DEBUG", + "ssl_certfile": None, + "ssl_keyfile": None, + "transcription": True, + "vad": True, + } + + config_dict = {**defaults, **kwargs} + + if 'no_transcription' in kwargs: + config_dict['transcription'] = not kwargs['no_transcription'] + if 'no_vad' in kwargs: + config_dict['vad'] = not kwargs['no_vad'] - merged_args = {**default_args, **kwargs} - - self.args = Namespace(**merged_args) + config_dict.pop('no_transcription', None) + config_dict.pop('no_vad', None) + + if 'language' in kwargs: + config_dict['lan'] = kwargs['language'] + config_dict.pop('language', None) + + self.args = Namespace(**config_dict) self.asr = None self.tokenizer = None @@ -174,11 +70,4 @@ class WhisperLiveKit: from whisperlivekit.diarization.diarization_online import DiartDiarization self.diarization = DiartDiarization() - WhisperLiveKit._initialized = True - - def web_interface(self): - import pkg_resources - html_path = pkg_resources.resource_filename('whisperlivekit', 'web/live_transcription.html') - with open(html_path, "r", encoding="utf-8") as f: - html = f.read() - return html \ No newline at end of file + TranscriptionEngine._initialized = True diff --git a/whisperlivekit/parse_args.py b/whisperlivekit/parse_args.py new file mode 100644 index 0000000..f97215a --- /dev/null +++ b/whisperlivekit/parse_args.py @@ -0,0 +1,141 @@ + +from argparse import ArgumentParser + +def parse_args(): + parser = ArgumentParser(description="Whisper FastAPI Online Server") + parser.add_argument( + "--host", + type=str, + default="localhost", + help="The host address to bind the server to.", + ) + parser.add_argument( + "--port", type=int, default=8000, help="The port number to bind the server to." + ) + parser.add_argument( + "--warmup-file", + type=str, + default=None, + dest="warmup_file", + help=""" + The path to a speech audio wav file to warm up Whisper so that the very first chunk processing is fast. + If not set, uses https://github.com/ggerganov/whisper.cpp/raw/master/samples/jfk.wav. + If False, no warmup is performed. + """, + ) + + parser.add_argument( + "--confidence-validation", + action="store_true", + help="Accelerates validation of tokens using confidence scores. Transcription will be faster but punctuation might be less accurate.", + ) + + parser.add_argument( + "--diarization", + action="store_true", + default=False, + help="Enable speaker diarization.", + ) + + parser.add_argument( + "--no-transcription", + action="store_true", + help="Disable transcription to only see live diarization results.", + ) + + parser.add_argument( + "--min-chunk-size", + type=float, + default=0.5, + help="Minimum audio chunk size in seconds. It waits up to this time to do processing. If the processing takes shorter time, it waits, otherwise it processes the whole segment that was received by this time.", + ) + + parser.add_argument( + "--model", + type=str, + default="tiny", + help="Name size of the Whisper model to use (default: tiny). Suggested values: tiny.en,tiny,base.en,base,small.en,small,medium.en,medium,large-v1,large-v2,large-v3,large,large-v3-turbo. The model is automatically downloaded from the model hub if not present in model cache dir.", + ) + + parser.add_argument( + "--model_cache_dir", + type=str, + default=None, + help="Overriding the default model cache dir where models downloaded from the hub are saved", + ) + parser.add_argument( + "--model_dir", + type=str, + default=None, + help="Dir where Whisper model.bin and other files are saved. This option overrides --model and --model_cache_dir parameter.", + ) + parser.add_argument( + "--lan", + "--language", + type=str, + default="auto", + help="Source language code, e.g. en,de,cs, or 'auto' for language detection.", + ) + parser.add_argument( + "--task", + type=str, + default="transcribe", + choices=["transcribe", "translate"], + help="Transcribe or translate.", + ) + parser.add_argument( + "--backend", + type=str, + default="faster-whisper", + choices=["faster-whisper", "whisper_timestamped", "mlx-whisper", "openai-api"], + help="Load only this backend for Whisper processing.", + ) + parser.add_argument( + "--vac", + action="store_true", + default=False, + help="Use VAC = voice activity controller. Recommended. Requires torch.", + ) + parser.add_argument( + "--vac-chunk-size", type=float, default=0.04, help="VAC sample size in seconds." + ) + + parser.add_argument( + "--no-vad", + action="store_true", + help="Disable VAD (voice activity detection).", + ) + + parser.add_argument( + "--buffer_trimming", + type=str, + default="segment", + choices=["sentence", "segment"], + help='Buffer trimming strategy -- trim completed sentences marked with punctuation mark and detected by sentence segmenter, or the completed segments returned by Whisper. Sentence segmenter must be installed for "sentence" option.', + ) + parser.add_argument( + "--buffer_trimming_sec", + type=float, + default=15, + help="Buffer trimming length threshold in seconds. If buffer length is longer, trimming sentence/segment is triggered.", + ) + parser.add_argument( + "-l", + "--log-level", + dest="log_level", + choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], + help="Set the log level", + default="DEBUG", + ) + parser.add_argument("--ssl-certfile", type=str, help="Path to the SSL certificate file.", default=None) + parser.add_argument("--ssl-keyfile", type=str, help="Path to the SSL private key file.", default=None) + + + args = parser.parse_args() + + args.transcription = not args.no_transcription + args.vad = not args.no_vad + delattr(args, 'no_transcription') + delattr(args, 'no_vad') + + return args \ No newline at end of file diff --git a/whisperlivekit/web/web_interface.py b/whisperlivekit/web/web_interface.py new file mode 100644 index 0000000..340d744 --- /dev/null +++ b/whisperlivekit/web/web_interface.py @@ -0,0 +1,13 @@ +import logging +import importlib.resources as resources + +logger = logging.getLogger(__name__) + +def get_web_interface_html(): + """Loads the HTML for the web interface using importlib.resources.""" + try: + with resources.files('whisperlivekit.web').joinpath('live_transcription.html').open('r', encoding='utf-8') as f: + return f.read() + except Exception as e: + logger.error(f"Error loading web interface HTML: {e}") + return "

Error loading interface

" \ No newline at end of file