src/main.py - master - sarus/audio-agent

src/main.py
146 строк5 KB
otcheskiy
key conclusion
15 июн 2026, 23:01
15 июн 2026, 23:012fe452a
import argparse
import json
import sys
import time
from datetime import datetime
from pathlib import Path

from rich.console import Console

sys.path.insert(0, str(Path(__file__).resolve().parent))

from config import ASSISTED_SUFFIX, COMPUTE_TYPE, DEVICE, LANGUAGE, MODEL_SIZE, OUTPUT_DIR
from transcriber import transcribe_audio

console = Console()

SUPPORTED_MODELS = ("tiny", "base", "small", "medium", "large-v2", "large-v3")


def format_duration(seconds: float) -> str:
    total_seconds = int(seconds)
    hours = total_seconds // 3600
    minutes = (total_seconds % 3600) // 60
    secs = total_seconds % 60
    return f"{hours:02d}:{minutes:02d}:{secs:02d}"


def format_timestamp(dt: datetime) -> str:
    return dt.strftime("%Y-%m-%d %H:%M:%S")


def format_filename_timestamp(dt: datetime) -> str:
    return dt.strftime("%Y%m%d_%H%M%S")


def parse_args() -> tuple[Path, str, bool]:
    parser = argparse.ArgumentParser(description="Транскрибация аудио в текст")
    parser.add_argument("audio_path", help="Путь к аудиофайлу")
    parser.add_argument(
        "--model",
        choices=SUPPORTED_MODELS,
        default=None,
        metavar="MODEL",
        help=(
            f"Модель Whisper: {', '.join(SUPPORTED_MODELS)}. "
            f"По умолчанию: {MODEL_SIZE} (из config.py)"
        ),
    )
    parser.add_argument(
        "--clean",
        action="store_true",
        help=(
            "[experimental] LLM cleanup через локальный сервер (LM Studio / Ollama). "
            "По умолчанию выключен."
        ),
    )
    args = parser.parse_args()
    model_size = args.model if args.model is not None else MODEL_SIZE
    return Path(args.audio_path), model_size, args.clean


def main() -> None:
    audio_path, model_size, run_clean = parse_args()

    if not audio_path.is_file():
        console.print(f"[red]Файл не найден:[/red] {audio_path}")
        sys.exit(1)

    output_dir = Path(OUTPUT_DIR)
    output_dir.mkdir(exist_ok=True)

    console.print(f"[cyan]Транскрибация:[/cyan] {audio_path}")
    console.print(f"[dim]Модель:[/dim] {model_size}")
    console.print(f"[dim]Язык:[/dim] {LANGUAGE}")
    console.print(f"[dim]Устройство:[/dim] {DEVICE} ({COMPUTE_TYPE})")
    console.print("[dim]Загрузка модели и распознавание...[/dim]")

    start_time = time.perf_counter()
    started_at = datetime.now()
    text, segments = transcribe_audio(str(audio_path), model_size)
    end_time = time.perf_counter()
    finished_at = datetime.now()
    processing_seconds = end_time - start_time

    base_name = audio_path.stem
    time_suffix = format_filename_timestamp(finished_at)
    txt_file = output_dir / f"{base_name}_{model_size}_{time_suffix}.txt"
    json_file = output_dir / f"{base_name}_{model_size}_{time_suffix}.json"

    output_data = {
        "metadata": {
            "input_file": str(audio_path).replace("\\", "/"),
            "model": model_size,
            "language": LANGUAGE,
            "device": DEVICE,
            "compute_type": COMPUTE_TYPE,
            "started_at": format_timestamp(started_at),
            "finished_at": format_timestamp(finished_at),
            "processing_time_seconds": round(processing_seconds, 2),
            "processing_time_hms": format_duration(processing_seconds),
            "segments_count": len(segments),
        },
        "segments": segments,
    }

    txt_header = (
        f"Начало: {format_timestamp(started_at)}\n"
        f"Конец: {format_timestamp(finished_at)}\n"
        f"---\n"
    )
    txt_file.write_text(txt_header + text, encoding="utf-8")
    json_file.write_text(
        json.dumps(output_data, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )

    assisted_txt_file = None
    if run_clean:
        from cleaner import CleanupError, clean_transcript

        assisted_txt_file = (
            output_dir / f"{base_name}_{model_size}_{time_suffix}_{ASSISTED_SUFFIX}.txt"
        )
        console.print("[dim]LLM cleanup [experimental] (локальный сервер)...[/dim]")
        try:
            cleaned_text = clean_transcript(text)
        except CleanupError as exc:
            console.print(f"[yellow]LLM cleanup не выполнен:[/yellow] {exc}")
            sys.exit(1)
        assisted_txt_file.write_text(txt_header + cleaned_text, encoding="utf-8")

    console.print()
    console.print(f"Модель: {model_size}")
    console.print(f"Файл: {audio_path}")
    console.print(f"Начало: {format_timestamp(started_at)}")
    console.print(f"Конец: {format_timestamp(finished_at)}")
    console.print(f"Время обработки: {format_duration(processing_seconds)}")
    console.print(f"Сегментов: {len(segments)}")
    console.print(f"TXT: {txt_file}")
    console.print(f"JSON: {json_file}")
    if assisted_txt_file is not None:
        console.print(f"TXT (assisted): {assisted_txt_file}")


if __name__ == "__main__":
    main()