333 lines
17 KiB
Python
333 lines
17 KiB
Python
"""
|
||
Модуль transcriber.py содержит класс WhisperTranscriber, который использует модель Whisper от
|
||
OpenAI для транскрибации аудиофайлов в текст. Класс включает в себя методы для загрузки модели,
|
||
обработки аудио (с использованием класса AudioProcessor), и выполнения транскрибации.
|
||
Обрабатывает выбор устройства (CPU, CUDA, MPS) для выполнения вычислений и обеспечивает
|
||
возможность использования Flash Attention 2 для ускорения работы модели на поддерживаемых GPU.
|
||
"""
|
||
|
||
import time
|
||
import threading
|
||
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"]
|
||
|
||
# Lock для потокобезопасного доступа к модели —
|
||
# Waitress обслуживает запросы в нескольких потоках,
|
||
# а HuggingFace pipeline не является thread-safe
|
||
self._inference_lock = threading.Lock()
|
||
|
||
# Создаем объект для обработки аудио
|
||
self.audio_processor = AudioProcessor(config)
|
||
|
||
# Определяем устройство для вычислений
|
||
self.device = self._get_device()
|
||
|
||
# Оптимальный тип для тензоров (зависит от устройства)
|
||
self.torch_dtype = self._get_torch_dtype()
|
||
|
||
# Загружаем модель при инициализации
|
||
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("device_id должен быть целым числом, получено: %s. Используем значение по умолчанию 0", device_id)
|
||
device_id = 0
|
||
|
||
# Проверяем, доступен ли запрошенный GPU
|
||
device_count = torch.cuda.device_count()
|
||
if device_id >= device_count:
|
||
logger.warning("Запрошенный GPU с индексом %s недоступен. Доступно GPU: %s. Используем GPU с индексом 0", device_id, device_count)
|
||
device_id = 0
|
||
|
||
logger.info("Используется CUDA GPU с индексом %s для вычислений", 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: PyTorch проверяет is_initialized()
|
||
# при создании тензоров на MPS-устройстве, что вызывает ошибку
|
||
# в однопроцессном режиме.
|
||
# TODO: Удалить после обновления до PyTorch >= 2.5
|
||
setattr(torch.distributed, "is_initialized", lambda: False)
|
||
return torch.device("mps")
|
||
else:
|
||
logger.info("Используется CPU для вычислений")
|
||
return torch.device("cpu")
|
||
|
||
def _get_torch_dtype(self) -> torch.dtype:
|
||
"""
|
||
Определяет оптимальный тип данных для тензоров в зависимости от устройства.
|
||
|
||
Returns:
|
||
torch.float32 для CPU, torch.bfloat16 для GPU/MPS.
|
||
"""
|
||
if self.device.type == "cpu":
|
||
return torch.float32
|
||
return torch.bfloat16
|
||
|
||
def _load_model(self) -> None:
|
||
"""
|
||
Загрузка модели и процессора.
|
||
|
||
Raises:
|
||
Exception: Если не удалось загрузить модель.
|
||
"""
|
||
logger.info("Загрузка модели из %s", self.model_path)
|
||
|
||
model_kwargs = dict(
|
||
torch_dtype=self.torch_dtype,
|
||
low_cpu_mem_usage=True,
|
||
use_safetensors=True,
|
||
)
|
||
|
||
use_flash_attn = False
|
||
if self.device.type == "cuda":
|
||
# Flash Attention 2 требует архитектуру Ampere или новее (compute capability >= 8.0)
|
||
capability = torch.cuda.get_device_capability(self.device.index)
|
||
if capability[0] >= 8:
|
||
use_flash_attn = True
|
||
logger.info("GPU поддерживает Flash Attention 2 (compute capability: %d.%d)", *capability)
|
||
else:
|
||
logger.info("GPU не поддерживает Flash Attention 2 (compute capability: %d.%d), используется стандартный режим", *capability)
|
||
|
||
try:
|
||
if use_flash_attn:
|
||
model_kwargs["attn_implementation"] = "flash_attention_2"
|
||
self.model = WhisperForConditionalGeneration.from_pretrained(
|
||
self.model_path, **model_kwargs
|
||
).to(self.device)
|
||
if use_flash_attn:
|
||
logger.info("Используется Flash Attention 2")
|
||
except Exception as e:
|
||
logger.warning("Не удалось загрузить модель с Flash Attention: %s", 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,
|
||
language: str = None, temperature: float = None,
|
||
prompt: str = None) -> Union[str, Dict]:
|
||
"""
|
||
Транскрибация аудиофайла.
|
||
|
||
Args:
|
||
audio_path: Путь к обработанному аудиофайлу.
|
||
return_timestamps: Флаг возврата временных меток. Если None — берётся из конфига.
|
||
language: Язык распознавания. Если None — берётся из конфига.
|
||
temperature: Параметр температуры для генерации. Если None — берётся из конфига.
|
||
prompt: Текстовая подсказка для модели (имена, термины, контекст).
|
||
|
||
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
|
||
if language is None:
|
||
language = self.language
|
||
if temperature is None:
|
||
temperature = self.temperature
|
||
|
||
logger.info("Начало транскрибации файла: %s", audio_path)
|
||
|
||
try:
|
||
# Загрузка аудио в формате numpy array
|
||
audio_array, sampling_rate = load_audio(audio_path, sr=16000)
|
||
|
||
# Транскрибация с корректным форматом данных
|
||
generate_kwargs = {
|
||
"language": language,
|
||
"max_new_tokens": self.max_new_tokens,
|
||
"temperature": temperature,
|
||
}
|
||
if prompt:
|
||
generate_kwargs["prompt_ids"] = self.processor.get_prompt_ids(
|
||
prompt, return_tensors="pt"
|
||
).to(self.device)
|
||
|
||
with self._inference_lock:
|
||
result = self.asr_pipeline(
|
||
{"raw": audio_array, "sampling_rate": sampling_rate},
|
||
generate_kwargs=generate_kwargs,
|
||
return_timestamps=return_timestamps,
|
||
)
|
||
|
||
# Если временные метки не запрошены, возвращаем только текст
|
||
if not return_timestamps:
|
||
transcribed_text = result.get("text", "")
|
||
logger.info("Транскрибация завершена: получено %s символов текста", 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("Транскрибация с временными метками завершена: получено %s сегментов", len(segments))
|
||
|
||
# Возвращаем словарь с сегментами и полным текстом
|
||
return {
|
||
"segments": segments,
|
||
"text": full_text
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error("Ошибка в процессе транскрибации аудиофайла '%s': %s", audio_path, e)
|
||
logger.error("Тип исключения: %s", type(e).__name__)
|
||
logger.error("Traceback: %s", traceback.format_exc())
|
||
raise
|
||
|
||
def process_file(self, input_path: str, return_timestamps: bool = None,
|
||
language: str = None, temperature: float = None,
|
||
prompt: str = None) -> Union[str, Dict]:
|
||
"""
|
||
Полный процесс обработки и транскрибации аудиофайла.
|
||
|
||
Args:
|
||
input_path: Путь к исходному аудиофайлу.
|
||
return_timestamps: Флаг возврата временных меток. Если None — берётся из конфига.
|
||
language: Язык распознавания. Если None — берётся из конфига.
|
||
temperature: Параметр температуры для генерации. Если None — берётся из конфига.
|
||
prompt: Текстовая подсказка для модели (имена, термины, контекст).
|
||
|
||
Returns:
|
||
В зависимости от параметра return_timestamps:
|
||
- Если return_timestamps=False: строка с распознанным текстом
|
||
- Если return_timestamps=True: словарь с ключами "segments" и "text"
|
||
"""
|
||
start_time = time.time()
|
||
logger.info("Начало обработки файла: %s", 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,
|
||
language=language, temperature=temperature,
|
||
prompt=prompt)
|
||
|
||
elapsed_time = time.time() - start_time
|
||
logger.info("Обработка и транскрибация завершены за %.2f секунд", elapsed_time)
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
elapsed_time = time.time() - start_time
|
||
logger.error("Ошибка при обработке файла '%s' через %.2f секунд: %s", input_path, elapsed_time, e)
|
||
logger.error("Тип исключения: %s", type(e).__name__)
|
||
logger.error("Traceback: %s", traceback.format_exc())
|
||
raise
|
||
|
||
finally:
|
||
# Очистка временных файлов
|
||
cleanup_temp_files(temp_files) |