419 lines
19 KiB
Python
419 lines
19 KiB
Python
"""
|
||
Модуль transcriber.py содержит класс WhisperTranscriber, который использует модель Whisper от
|
||
OpenAI для транскрибации аудиофайлов в текст. Класс включает в себя методы для загрузки модели,
|
||
обработки аудио (с использованием класса AudioProcessor), и выполнения транскрибации.
|
||
Обрабатывает выбор устройства (CPU, CUDA, MPS) для выполнения вычислений и обеспечивает
|
||
возможность использования Flash Attention 2 для ускорения работы модели на поддерживаемых GPU.
|
||
"""
|
||
|
||
import os
|
||
import time
|
||
import threading
|
||
import traceback
|
||
from typing import Dict, Tuple, Union
|
||
import logging
|
||
|
||
import numpy as np
|
||
import torch
|
||
from huggingface_hub import snapshot_download
|
||
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 must be an integer, got: %s. Using default value 0", device_id)
|
||
device_id = 0
|
||
|
||
# Проверяем, доступен ли запрошенный GPU
|
||
device_count = torch.cuda.device_count()
|
||
if device_id >= device_count:
|
||
logger.warning("Requested GPU index %s is not available. Available GPU count: %s. Using GPU index 0", device_id, device_count)
|
||
device_id = 0
|
||
|
||
logger.info("Using CUDA GPU with index %s for inference", device_id)
|
||
return torch.device(f"cuda:{device_id}")
|
||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||
logger.info("Using MPS (Apple Silicon) for inference")
|
||
# Обходное решение для MPS: PyTorch проверяет is_initialized()
|
||
# при создании тензоров на MPS-устройстве, что вызывает ошибку
|
||
# в однопроцессном режиме.
|
||
# TODO: Удалить после обновления до PyTorch >= 2.5
|
||
setattr(torch.distributed, "is_initialized", lambda: False)
|
||
return torch.device("mps")
|
||
else:
|
||
logger.info("Using CPU for inference")
|
||
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("Loading model from %s", self.model_path)
|
||
self.model_path = self._prepare_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 supports Flash Attention 2 (compute capability: %d.%d)", *capability)
|
||
else:
|
||
logger.info("GPU does not support Flash Attention 2 (compute capability: %d.%d), using standard attention", *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("Using Flash Attention 2")
|
||
except Exception as e:
|
||
logger.warning("Failed to load model with 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("Model loaded successfully and is ready to use")
|
||
|
||
def _prepare_model_path(self) -> str:
|
||
"""
|
||
Ensures that the configured model path is usable.
|
||
|
||
If ``model_path`` looks like a local path and the model is missing,
|
||
this method can optionally download a fallback Whisper model into that
|
||
directory. If ``model_path`` is a Hugging Face repo ID, it is returned
|
||
unchanged and transformers will resolve it normally.
|
||
"""
|
||
model_path = self.model_path
|
||
|
||
looks_like_local_path = (
|
||
os.path.isabs(model_path)
|
||
or model_path.startswith(".")
|
||
or os.path.sep in model_path
|
||
)
|
||
if not looks_like_local_path:
|
||
return model_path
|
||
|
||
if self._is_valid_local_model_dir(model_path):
|
||
return model_path
|
||
|
||
if self.config.get("auto_download_missing_model", False):
|
||
return self._download_missing_model(model_path)
|
||
|
||
if not os.path.isdir(model_path):
|
||
raise FileNotFoundError(
|
||
"Local model directory not found: "
|
||
f"'{model_path}'. Mount the Whisper model into the container, "
|
||
"set 'model_path' to that folder, or enable "
|
||
"'auto_download_missing_model'."
|
||
)
|
||
|
||
config_file = os.path.join(model_path, "config.json")
|
||
raise FileNotFoundError(
|
||
"The configured model directory exists but does not look like a "
|
||
f"Hugging Face model folder: '{model_path}'. Missing required file "
|
||
f"'{config_file}'. Make sure the directory contains files like "
|
||
"'config.json', tokenizer files, and model weights, or enable "
|
||
"'auto_download_missing_model'."
|
||
)
|
||
|
||
def _is_valid_local_model_dir(self, model_path: str) -> bool:
|
||
"""Returns True when the directory looks like a local HF model folder."""
|
||
return os.path.isdir(model_path) and os.path.isfile(
|
||
os.path.join(model_path, "config.json")
|
||
)
|
||
|
||
def _download_missing_model(self, model_path: str) -> str:
|
||
"""
|
||
Downloads a fallback Whisper model into the configured local directory.
|
||
"""
|
||
model_repo_id = self.config.get("model_repo_id", "openai/whisper-large-v3")
|
||
model_revision = self.config.get("model_revision")
|
||
|
||
logger.warning(
|
||
"Local model was not found at %s. Downloading fallback model %s",
|
||
model_path,
|
||
model_repo_id,
|
||
)
|
||
|
||
os.makedirs(model_path, exist_ok=True)
|
||
snapshot_download(
|
||
repo_id=model_repo_id,
|
||
local_dir=model_path,
|
||
local_dir_use_symlinks=False,
|
||
revision=model_revision,
|
||
)
|
||
|
||
if not self._is_valid_local_model_dir(model_path):
|
||
raise FileNotFoundError(
|
||
"Model download completed but the target directory still does not "
|
||
f"look valid: '{model_path}'. Expected to find '{os.path.join(model_path, 'config.json')}'."
|
||
)
|
||
|
||
logger.info(
|
||
"Downloaded Whisper model %s into %s",
|
||
model_repo_id,
|
||
model_path,
|
||
)
|
||
return model_path
|
||
|
||
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("Starting transcription for file: %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("Transcription completed: produced %s text characters", 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("Timestamps were requested but not found in the transcription result")
|
||
|
||
logger.info("Timestamped transcription completed: produced %s segments", len(segments))
|
||
|
||
# Возвращаем словарь с сегментами и полным текстом
|
||
return {
|
||
"segments": segments,
|
||
"text": full_text
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error("Error while transcribing audio file '%s': %s", audio_path, e)
|
||
logger.error("Exception type: %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("Starting file processing: %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("Processing and transcription finished in %.2f seconds", elapsed_time)
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
elapsed_time = time.time() - start_time
|
||
logger.error("Error while processing file '%s' after %.2f seconds: %s", input_path, elapsed_time, e)
|
||
logger.error("Exception type: %s", type(e).__name__)
|
||
logger.error("Traceback: %s", traceback.format_exc())
|
||
raise
|
||
|
||
finally:
|
||
# Очистка временных файлов
|
||
cleanup_temp_files(temp_files)
|