transcription/run.py
2026-05-29 10:53:16 +03:00

141 lines
5.9 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""CLI entrypoint для транскрибации совещаний."""
import argparse
import os
import sys
from pathlib import Path
from src.audio_utils import check_ffmpeg, is_audio_file, is_video_file
from src.config import get_profile, load_config
from src.document import build_document
from src.pipeline import run_pipeline
def resolve_device(preferred: str) -> str:
"""Определяет доступное устройство."""
import torch
if preferred == "cuda" and torch.cuda.is_available():
return "cuda"
if preferred == "mps" and torch.backends.mps.is_available():
return "mps"
return "cpu"
def print_first_run_info():
"""Показывает информацию о первом запуске и скачивании моделей."""
print("=" * 60)
print("ПЕРВЫЙ ЗАПУСК: скачивание моделей")
print("=" * 60)
print("При первом запуске будут скачаны модели ИИ (~45 GB):")
print(" • Whisper large-v3 ~3.0 GB (распознавание речи)")
print(" • Pyannote диаризация ~0.4 GB (разделение спикеров)")
print(" • Wav2Vec2 alignment ~1.0 GB (точные таймкоды)")
print("")
print("Это займёт время в зависимости от скорости интернета.")
print("При последующих запусках модели загружаются с диска.")
print("=" * 60)
print("")
def main():
parser = argparse.ArgumentParser(
description="Транскрибация совещаний с диаризацией и таймкодами. "
"Поддерживает аудио (wav, mp3, m4a, ogg, flac) и видео (mp4, avi, mkv, mov, etc.)."
)
parser.add_argument("--input", "-i", required=True, help="Путь к аудио или видео файлу")
parser.add_argument("--output", "-o", default=None, help="Путь к выходному файлу (docx/md/txt)")
parser.add_argument("--profile", "-p", default=None, help="Профиль конфигурации (mac_m4, gpu_8gb, cpu_best)")
parser.add_argument("--config", "-c", default=None, help="Путь к config.yaml")
parser.add_argument("--device", "-d", default=None, help="Принудительно: cpu/cuda/mps")
parser.add_argument("--model", "-m", default=None, help="Принудительно: tiny/base/small/medium/large-v3")
parser.add_argument("--language", "-l", default=None, help="Язык (ru, en, ...)")
parser.add_argument("--format", "-f", default=None, help="Формат выхода: docx, md, txt")
args = parser.parse_args()
input_path = Path(args.input)
if not input_path.exists():
print(f"Ошибка: файл не найден: {args.input}", file=sys.stderr)
sys.exit(1)
# Проверка ffmpeg для видео
if is_video_file(args.input):
if not check_ffmpeg():
print(
"Ошибка: для обработки видео нужен ffmpeg.\n"
"Установите ffmpeg:\n"
" Mac: brew install ffmpeg\n"
" Linux: sudo apt-get install ffmpeg\n"
" Windows: https://ffmpeg.org/download.html",
file=sys.stderr,
)
sys.exit(1)
print(f"[Info] Обнаружен видео файл: {input_path.suffix}")
print("[Info] Аудио будет извлечено автоматически.")
elif not is_audio_file(args.input):
print(f"[Warning] Неизвестный формат: {input_path.suffix}. Попытка обработать...")
# Загрузка конфига
config = load_config(args.config)
profile = get_profile(config, args.profile)
# Переопределения из CLI
if args.device:
profile["device"] = resolve_device(args.device)
else:
profile["device"] = resolve_device(profile.get("device", "cpu"))
if args.model:
profile["model"] = args.model
if args.language:
profile["language"] = args.language
output_cfg = config.get("output", {})
fmt = args.format or output_cfg.get("format", "docx")
if args.output:
output_path = args.output
else:
stem = input_path.stem
output_dir = Path(config.get("paths", {}).get("output_dir", "./output"))
output_path = str(output_dir / f"{stem}.{fmt}")
# Проверка HF токена
hf_token = os.environ.get("HF_TOKEN") or config.get("hf_token")
if profile.get("diarize", True) and not hf_token:
print(
"Ошибка: для диаризации нужен HuggingFace токен.\n"
"Установите env HF_TOKEN или укажите hf_token в config.yaml\n"
"Инструкция: https://huggingface.co/docs/hub/security-tokens",
file=sys.stderr,
)
sys.exit(1)
# Информация о первом запуске
print_first_run_info()
print(f"Профиль: {args.profile or config.get('active_profile')}")
print(f"Устройство: {profile['device']}")
print(f"Модель: {profile['model']}")
print(f"Язык: {profile['language']}")
print(f"Вход: {args.input}")
print(f"Выход: {output_path}")
print("-" * 40)
# Запуск пайплайна
result = run_pipeline(
input_path=args.input,
profile_name=args.profile,
config_path=args.config,
output_path=output_path,
)
# Генерация документа
build_document(result["segments"], output_path, config)
print("-" * 40)
print(f"Готово! Сохранено: {output_path}")
if __name__ == "__main__":
main()