AiMeetingNotes/setup_mega_asr.py at master · RejectKid/AiMeetingNotes · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
"""Install the pinned Mega-ASR source and download its model checkpoint."""

from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path
from typing import Sequence


PROJECT_ROOT = Path(__file__).resolve().parent
MEGA_ASR_URL = "https://github.com/xzf-thu/Mega-ASR.git"
MEGA_ASR_REVISION = "bd75877b29b695d852234169fd3738bb5080046b"
MODEL_REPO_ID = "zhifeixie/Mega-ASR"
MODEL_REVISION = "8a10bfd0cccb0c8b3be2852cae1e3ee73d6ffbf1"


def run_command(command: Sequence[str], cwd: Path | None = None) -> None:
    subprocess.run(command, cwd=cwd, check=True)


def install_source(source_dir: Path, revision: str, update: bool) -> None:
    expected_file = source_dir / "src" / "MegaASR" / "model" / "megaASR.py"
    if not (source_dir / ".git").is_dir():
        if source_dir.exists() and any(source_dir.iterdir()):
            raise RuntimeError(
                f"Source directory is not empty or a Git checkout: {source_dir}"
            )
        source_dir.parent.mkdir(parents=True, exist_ok=True)
        run_command(
            ["git", "clone", "--filter=blob:none", MEGA_ASR_URL, str(source_dir)]
        )
    elif update:
        run_command(["git", "fetch", "origin", "main"], cwd=source_dir)

    run_command(["git", "checkout", "--detach", revision], cwd=source_dir)
    if not expected_file.is_file():
        raise RuntimeError(f"Mega-ASR source is incomplete: {expected_file} is missing")


def download_checkpoint(checkpoint_dir: Path, revision: str) -> None:
    try:
        from huggingface_hub import snapshot_download
    except ImportError as error:
        raise RuntimeError(
            "huggingface-hub is missing. Run "
            "'python -m pip install -r requirements.txt' first."
        ) from error

    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    snapshot_download(
        repo_id=MODEL_REPO_ID,
        revision=revision,
        local_dir=str(checkpoint_dir),
    )


def verify_installation(source_dir: Path, checkpoint_dir: Path) -> None:
    required = [
        source_dir / "src" / "MegaASR" / "model" / "megaASR.py",
        checkpoint_dir / "Qwen3-ASR-1.7B" / "config.json",
        checkpoint_dir / "mega-asr-merged" / "adapter_config.json",
        checkpoint_dir
        / "audio_quality_router"
        / "best_acc_model.safetensors",
    ]
    missing = [path for path in required if not path.is_file()]
    if missing:
        raise RuntimeError(
            "Mega-ASR setup is incomplete; missing: "
            + ", ".join(str(path) for path in missing)
        )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Install Mega-ASR source and model weights for AiMeetingNotes."
    )
    parser.add_argument(
        "--source-dir",
        type=Path,
        default=PROJECT_ROOT / "third_party" / "Mega-ASR",
    )
    parser.add_argument(
        "--checkpoint-dir",
        type=Path,
        default=PROJECT_ROOT / "models" / "Mega-ASR",
    )
    parser.add_argument("--source-revision", default=MEGA_ASR_REVISION)
    parser.add_argument("--model-revision", default=MODEL_REVISION)
    parser.add_argument("--update-source", action="store_true")
    parser.add_argument("--skip-model-download", action="store_true")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    if sys.version_info[:2] != (3, 13):
        print(
            "AiMeetingNotes requires Python 3.13; current interpreter is "
            f"{sys.version_info.major}.{sys.version_info.minor}.",
            file=sys.stderr,
        )
        return 2

    source_dir = args.source_dir.expanduser().resolve()
    checkpoint_dir = args.checkpoint_dir.expanduser().resolve()
    try:
        install_source(source_dir, args.source_revision, args.update_source)
        if not args.skip_model_download:
            download_checkpoint(checkpoint_dir, args.model_revision)
        verify_installation(source_dir, checkpoint_dir)
    except (OSError, RuntimeError, subprocess.CalledProcessError) as error:
        print(f"Setup failed: {error}", file=sys.stderr)
        return 1

    print(f"Mega-ASR source: {source_dir}")
    print(f"Mega-ASR checkpoint: {checkpoint_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())