import logging import math import time import sys import gc import torch import torchaudio from tqdm import tqdm from transformers import WhisperForConditionalGeneration, WhisperProcessor from transcription.device_configuration import DeviceConfiguration from ui.ui_log_handler import UILogHandler # TODO: implement transcription with shift class AudioTranscription: """ Class for automatical audio transcription using Whisper. Provides audio file loading, resampling, conversion to mono, splitting into chunks and batches, running the model, and compiling the final text. Attributes: model_name (str): Whisper model name in HuggingFace. filepath (str): Input filepath. waveform (torch.Tensor): Audiosignal in tensor form. sampling_rate (int): Input file's sampling frequency. chunks (list): Audio's chunks list. batches (list): Batches combined from chunks. chunk_size (int): Chunk size in samples. custom_chunk_length (int): Custom chunk length (in seconds). custom_batch_length (int): Custom batch length (in chunks). device (str): Inference device ("cuda", "cpu", "mps"). processor (WhisperProcessor | None): Tokenizator/preprocessor. model (WhisperForConditionalGeneration | None): Whisper model. logger (logging.Logger): Logger. torch_dtype (torch.dtype): Data type for calculations (fp16/fp32). language (str): Transcription language (default "ru"). all_transcription (list): List of strings with transcription. Args: filepath (str): Input filepath. device_configuration (DeviceConfiguration): Device configuration (GPU/CPU/MPS, model, chunk length, batch etc.). logger (logging.Logger): Logger. language (str, optional): Transcription language. Default "ru". Example: >>> from transcription.device_configuration import DeviceConfiguration >>> import logging >>> config = DeviceConfiguration(device="cuda", model_name="openai/whisper-large-v3-turbo") # recheck this, not true i think >>> logger = logging.getLogger("transcription") >>> transcriber = AudioTranscription("audio.wav", config, logger, language="en") >>> text = transcriber.transcribe_audio() >>> print(text) "This is a test audio transcription." Methods: transcribe_audio() -> str: Starts full transcription pipeline: model loading, file loading and preprocessing, splitting into chunks/batches, inference and model unloading. Returns the final transcription. """ model_name = "openai/whisper-large-v3-turbo" filepath: str waveform: torch.Tensor sampling_rate: int file_format: str chunks: list = [] batches: list = [] chunk_size: int custom_chunk_length: int custom_batch_length: int device = "cuda" processor: WhisperProcessor | None model: WhisperForConditionalGeneration | None logger: logging.Logger torch_dtype: torch.dtype language = "ru" all_transcription: list = [] def __init__( self, filepath: str, device_configuration: DeviceConfiguration, logger: logging.Logger, language: str = "ru", ) -> None: # TODO: add pretty docs here self.filepath = filepath self.language = language self.logger = logger # setting file extension self.file_format = filepath.split(".")[-1] # extracting configuration self.device = device_configuration.device self.model_name = device_configuration.model_name self.custom_chunk_length = device_configuration.chunk_length_s self.custom_batch_length = device_configuration.batch_size self.torch_dtype = device_configuration.torch_dtype self.chunks: list = [] self.batches: list = [] self.all_transcription: list = [] def _load_model(self) -> None: self.logger.info("Loading model WhisperProcessor...") try: start_time = time.time() self.processor = WhisperProcessor.from_pretrained(self.model_name) self.model = WhisperForConditionalGeneration.from_pretrained( self.model_name, torch_dtype=self.torch_dtype ).to(self.device) end_time = time.time() self.logger.info( f"Model loaded successfully in {end_time - start_time:.2f} seconds." ) except Exception as e: self.logger.error(f"Error while loading model: {e}") def _unload_model(self) -> None: self.logger.info("Unloading model...") self.model = None self.processor = None if self.device == "cuda": torch.cuda.empty_cache() # TODO: maybe do something here for MPS self.logger.info("Model unloaded successfully.") def _load_file(self) -> None: self.logger.info(f"Loading file {self.filepath}") try: # self.waveform, self.sampling_rate = torchaudio.load( # self.filepath, format=self.file_format, backend="ffmpeg" # ) self.waveform, self.sampling_rate = torchaudio.load(self.filepath) self.logger.info(f"Successfully loaded file {self.filepath}.") except Exception as e: self.logger.error(f"Unable to load file {self.filepath}: {e}") # raise RuntimeError(f"Failed to load file {self.filepath}: {e}") def _resample(self) -> None: self.waveform = torchaudio.functional.resample( self.waveform, self.sampling_rate, 16000 ) def _to_mono(self): if self.waveform.shape[0] > 1: self.waveform = self.waveform.mean(dim=0, keepdim=True) self.waveform = self.waveform.squeeze(0) def _split_to_chunks(self, shift: bool = False) -> None: self.logger.info(f"Splitting audio on chunks...") self.chunk_size = self.custom_chunk_length * 16000 # 16kHz after resampling total_samples = self.waveform.shape[0] chunks_count = (total_samples + self.chunk_size - 1) // self.chunk_size self.logger.info( f"File length - {total_samples / 16000:.1f} seconds, splitting on {chunks_count} chunks by {self.custom_chunk_length} seconds per chunk." ) self.chunks = [] for idx in tqdm(range(chunks_count)): start = idx * self.chunk_size end = min((idx + 1) * self.chunk_size, total_samples) chunk = self.waveform[start:end].cpu().numpy().astype("float32") self.chunks.append(chunk) def _resplit_chunks_to_batches(self) -> None: self.logger.info(f"Splitting chunks into batches...") self.batches = [] for i in range(0, len(self.chunks), self.custom_batch_length): batch = self.chunks[i : i + self.custom_batch_length] self.batches.append(batch) self.logger.info(f"Total: {len(self.batches)} batches, weight = {sys.getsizeof(self.batches)}") def _process_all_batches(self) -> None: start_time = time.time() try: assert self.processor is not None assert self.model is not None self.all_transcription = [] for idx in tqdm(range(len(self.batches))): inputs = self.processor( self.batches[idx], sampling_rate=16000, return_tensors="pt", padding=True, ) input_features = inputs.input_features.to(self.device).to( self.torch_dtype ) with torch.no_grad(): predicted_ids = self.model.generate( input_features, language=self.language, task="transcribe", temperature=0.0, ) texts = self.processor.batch_decode( predicted_ids, skip_special_tokens=True ) self.all_transcription.extend(texts) inputs = None input_features = None predicted_ids = None gc.collect() if self.device.startswith("cuda"): torch.cuda.empty_cache() torch.cuda.ipc_collect() end_time = time.time() self.logger.info( f"Transcription completed in {end_time - start_time:.2f} seconds" ) except Exception as e: self.logger.error(f"Errors occured while processing chunks: {e}") def transcribe_audio(self) -> str: self._load_model() self._load_file() self._resample() self._to_mono() self._split_to_chunks() self._resplit_chunks_to_batches() self._process_all_batches() self._unload_model() return " ".join(self.all_transcription)