""" Модуль transcriber.py содержит класс WhisperTranscriber, который использует модель Whisper от OpenAI для транскрибации аудиофайлов в текст. Класс включает в себя методы для загрузки модели, обработки аудио (с использованием класса AudioProcessor), и выполнения транскрибации. Обрабатывает выбор устройства (CPU, CUDA, MPS) для выполнения вычислений и обеспечивает возможность использования Flash Attention 2 для ускорения работы модели на поддерживаемых GPU. """ import time import traceback from typing import Dict, Tuple, Union import logging import numpy as np import torch from transformers import ( WhisperForConditionalGeneration, WhisperProcessor, pipeline, ) from ..audio.processor import AudioProcessor from ..audio.utils import load_audio from ..infrastructure.storage import cleanup_temp_files logger = logging.getLogger('app.transcriber') class WhisperTranscriber: """ Класс для распознавания речи с помощью модели Whisper. Attributes: config (Dict): Словарь с параметрами конфигурации. model_path (str): Путь к модели Whisper. language (str): Язык распознавания. chunk_length_s (int): Длина аудиочанка в секундах. batch_size (int): Размер пакета для обработки. max_new_tokens (int): Максимальное количество новых токенов для генерации. return_timestamps (bool): Флаг возврата временных меток. temperature (float): Параметр температуры для генерации. torch_dtype (torch.dtype): Оптимальный тип данных для тензоров. audio_processor (AudioProcessor): Объект для обработки аудио. device (torch.device): Устройство для вычислений. model (WhisperForConditionalGeneration): Загруженная модель Whisper. processor (WhisperProcessor): Процессор для модели Whisper. asr_pipeline (pipeline): Пайплайн для автоматического распознавания речи. """ def __init__(self, config: Dict): """ Инициализация транскрайбера. Args: config: Словарь с параметрами конфигурации. """ self.config = config self.model_path = config["model_path"] self.language = config["language"] self.chunk_length_s = config["chunk_length_s"] self.batch_size = config["batch_size"] self.max_new_tokens = config["max_new_tokens"] self.return_timestamps = config["return_timestamps"] self.temperature = config["temperature"] # Оптимальный тип для тензоров self.torch_dtype = torch.bfloat16 # Создаем объект для обработки аудио self.audio_processor = AudioProcessor(config) # Определяем устройство для вычислений self.device = self._get_device() # Загружаем модель при инициализации self._load_model() def _get_device(self) -> torch.device: """ Определение доступного устройства для вычислений. Returns: Объект устройства PyTorch. """ if torch.cuda.is_available(): # Получаем device_id из конфигурации, по умолчанию 0 device_id = self.config.get("device_id", 0) # Проверяем, что device_id является целым числом if not isinstance(device_id, int): logger.warning(f"device_id должен быть целым числом, получено: {device_id}. Используем значение по умолчанию 0") device_id = 0 # Проверяем, доступен ли запрошенный GPU device_count = torch.cuda.device_count() if device_id >= device_count: logger.warning(f"Запрошенный GPU с индексом {device_id} недоступен. Доступно GPU: {device_count}. Используем GPU с индексом 0") device_id = 0 logger.info(f"Используется CUDA GPU с индексом {device_id} для вычислений") return torch.device(f"cuda:{device_id}") elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): logger.info("Используется MPS (Apple Silicon) для вычислений") # Обходное решение для MPS setattr(torch.distributed, "is_initialized", lambda: False) return torch.device("mps") else: logger.info("Используется CPU для вычислений") return torch.device("cpu") def _load_model(self) -> None: """ Загрузка модели и процессора. Raises: Exception: Если не удалось загрузить модель. """ logger.info(f"Загрузка модели из {self.model_path}") model_kwargs = dict( torch_dtype=self.torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, ) try: if self.device.type == "cuda": model_kwargs["attn_implementation"] = "flash_attention_2" self.model = WhisperForConditionalGeneration.from_pretrained( self.model_path, **model_kwargs ).to(self.device) if self.device.type == "cuda": logger.info("Используется Flash Attention 2") except Exception as e: logger.warning(f"Не удалось загрузить модель с Flash Attention: {e}") model_kwargs.pop("attn_implementation", None) self.model = WhisperForConditionalGeneration.from_pretrained( self.model_path, **model_kwargs ).to(self.device) self.processor = WhisperProcessor.from_pretrained(self.model_path) self.asr_pipeline = pipeline( "automatic-speech-recognition", model=self.model, tokenizer=self.processor.tokenizer, feature_extractor=self.processor.feature_extractor, chunk_length_s=self.chunk_length_s, batch_size=self.batch_size, return_timestamps=self.return_timestamps, torch_dtype=self.torch_dtype, device=self.device, ) logger.info("Модель успешно загружена и готова к использованию") def transcribe(self, audio_path: str, return_timestamps: bool = None) -> Union[str, Dict]: """ Транскрибация аудиофайла. Args: audio_path: Путь к обработанному аудиофайлу. return_timestamps: Флаг возврата временных меток. Если None — берётся из конфига. Returns: В зависимости от параметра return_timestamps: - Если return_timestamps=False: строка с распознанным текстом - Если return_timestamps=True: словарь с ключами "segments" (список словарей с ключами start_time_ms, end_time_ms, text) и "text" (полный текст) """ if return_timestamps is None: return_timestamps = self.return_timestamps logger.info(f"Начало транскрибации файла: {audio_path}") try: # Загрузка аудио в формате numpy array audio_array, sampling_rate = load_audio(audio_path, sr=16000) # Транскрибация с корректным форматом данных result = self.asr_pipeline( {"raw": audio_array, "sampling_rate": sampling_rate}, generate_kwargs={ "language": self.language, "max_new_tokens": self.max_new_tokens, "temperature": self.temperature }, return_timestamps=return_timestamps ) # Если временные метки не запрошены, возвращаем только текст if not return_timestamps: transcribed_text = result.get("text", "") logger.info(f"Транскрибация завершена: получено {len(transcribed_text)} символов текста") return transcribed_text # Если временные метки запрошены, обрабатываем и форматируем результат segments = [] full_text = result.get("text", "") if "chunks" in result: # Для новых версий модели Whisper for chunk in result["chunks"]: start_time = chunk.get("timestamp", [0, 0])[0] end_time = chunk.get("timestamp", [0, 0])[1] text = chunk.get("text", "").strip() segments.append({ "start_time_ms": int(start_time * 1000), "end_time_ms": int(end_time * 1000), "text": text }) elif hasattr(result, "get") and "segments" in result: # Для старых версий модели Whisper for segment in result["segments"]: start_time = segment.get("start", 0) end_time = segment.get("end", 0) text = segment.get("text", "").strip() segments.append({ "start_time_ms": int(start_time * 1000), "end_time_ms": int(end_time * 1000), "text": text }) else: logger.warning("Временные метки запрошены, но не найдены в результате транскрибации") logger.info(f"Транскрибация с временными метками завершена: получено {len(segments)} сегментов") # Возвращаем словарь с сегментами и полным текстом return { "segments": segments, "text": full_text } except Exception as e: logger.error(f"Ошибка в процессе транскрибации аудиофайла '{audio_path}': {str(e)}") logger.error(f"Тип исключения: {type(e).__name__}") logger.error(f"Traceback: {traceback.format_exc()}") raise def process_file(self, input_path: str, return_timestamps: bool = None) -> Union[str, Dict]: """ Полный процесс обработки и транскрибации аудиофайла. Args: input_path: Путь к исходному аудиофайлу. return_timestamps: Флаг возврата временных меток. Если None — берётся из конфига. Returns: В зависимости от параметра return_timestamps: - Если return_timestamps=False: строка с распознанным текстом - Если return_timestamps=True: словарь с ключами "segments" и "text" """ start_time = time.time() logger.info(f"Начало обработки файла: {input_path}") temp_files = [] try: # Обработка аудио (конвертация, нормализация, добавление тишины) processed_path, temp_files = self.audio_processor.process_audio(input_path) # Транскрибация result = self.transcribe(processed_path, return_timestamps=return_timestamps) elapsed_time = time.time() - start_time logger.info(f"Обработка и транскрибация завершены за {elapsed_time:.2f} секунд") return result except Exception as e: elapsed_time = time.time() - start_time logger.error(f"Ошибка при обработке файла '{input_path}' через {elapsed_time:.2f} секунд: {str(e)}") logger.error(f"Тип исключения: {type(e).__name__}") logger.error(f"Traceback: {traceback.format_exc()}") raise finally: # Очистка временных файлов cleanup_temp_files(temp_files)