import argparse
import json
from pathlib import Path

from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess


def parse_args():
    parser = argparse.ArgumentParser(description="Offline FunASR transcription")
    parser.add_argument("--model-dir", required=True, help="Complete local FunASR model directory")
    parser.add_argument("--audio", required=True, help="Audio file to transcribe")
    parser.add_argument("--language", default="auto")
    return parser.parse_args()


def main():
    args = parse_args()
    model_dir = Path(args.model_dir).resolve()
    audio_path = Path(args.audio).resolve()

    if not (model_dir / "config.yaml").is_file():
        raise FileNotFoundError(f"missing model config: {model_dir / 'config.yaml'}")
    if not (model_dir / "model.pt").is_file():
        raise FileNotFoundError(f"missing model weights: {model_dir / 'model.pt'}")
    if not audio_path.is_file():
        raise FileNotFoundError(f"missing audio: {audio_path}")

    model = AutoModel(
        model=str(model_dir),
        device="cpu",
        disable_update=True,
    )
    result = model.generate(
        input=str(audio_path),
        language=args.language,
        use_itn=False,
    )[0]
    raw_text = result.get("text", "")
    print(
        json.dumps(
            {
                "text": rich_transcription_postprocess(raw_text),
                "raw_text": raw_text,
            },
            ensure_ascii=False,
        )
    )


if __name__ == "__main__":
    main()

from PyInstaller.utils.hooks import collect_data_files


hiddenimports = [
    "funasr.frontends.wav_frontend",
    "funasr.models.sense_voice.model",
    "funasr.models.specaug.specaug",
    "funasr.tokenizer.sentencepiece_tokenizer",
]

# FunASR uses these source files while registering the classes required by
# the tested SenseVoiceSmall config.
datas = collect_data_files(
    "funasr",
    include_py_files=True,
    includes=[
        "frontends/wav_frontend.py",
        "models/sense_voice/model.py",
        "models/specaug/specaug.py",
        "tokenizer/sentencepiece_tokenizer.py",
    ],
)