create_scripted_model.py

import argparse

import numpy as np
import torch
from huggingface_hub import HfApi, create_repo

from data import LibrispeechDataset
from src.scoreq_original import SCOREQScore
from src.scoreq_pytorch import SCOREQScoreTorch
from utils import SCRIPT_DATA_DIR, TARGET_SR


def create_scripted_model(args, data_domain, mode, dataset):

    data = dataset[0]
    if data_domain == "natural":
        test_wav = data["codec_audio"]
    else:
        test_wav = data["tts_audio"]
    ref_wav = data["audio"]

    device = "cpu"
    orig_scoreq = SCOREQScore(
        data_domain=data_domain, mode=mode, device=device, use_onnx=False
    )
    torch_scoreq = SCOREQScoreTorch(
        data_domain=data_domain, mode=mode, device=device, verbose=True
    )

    # # create TorchScript version
    script_scoreq = torch.jit.script(torch_scoreq)

    orig_score = np.array([orig_scoreq.predict(test_wav, ref_wav)])
    print(f"Orig. SCOREQ {data_domain} {mode}: {orig_score}")

    torch_score = torch_scoreq.score(test_wav, ref_wav)
    torch_score = torch_score.detach().cpu().numpy()
    print(f"Torch SCOREQ {data_domain} {mode}: {torch_score}")

    # there is no .score for a scripted module
    # use plain forward
    script_score = script_scoreq(test_wav, ref_wav)
    script_score = script_score.detach().cpu().numpy()
    print(f"Script SCOREQ {data_domain} {mode}: {script_score}")

    np.testing.assert_allclose(orig_score, torch_score, rtol=1e-6, atol=1e-6)
    np.testing.assert_allclose(orig_score, script_score, rtol=1e-6, atol=1e-6)

    scripted_checkpoint = f"scoreq_{data_domain}_{mode}_scripted.pt"
    save_path = SCRIPT_DATA_DIR / scripted_checkpoint
    script_scoreq.save(save_path)

    print(f"Saved scripted SCOREQ {data_domain} {mode} to: {save_path}")

    create_repo(
        repo_id=args.repo_id,
        repo_type="model",
        private=args.private,
        exist_ok=True,
    )

    api = HfApi()

    api.upload_file(
        path_or_fileobj=str(save_path),
        path_in_repo=scripted_checkpoint,
        repo_id=args.repo_id,
        repo_type="model",
        commit_message=f"Upload scripted SCOREQ {data_domain} {mode} checkpoint",
    )

    return save_path


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Create a scripted version of PyTorch SCOREQ"
    )
    parser.add_argument(
        "--repo-id",
        type=str,
        default="Blinorot/SCOREQ-PyTorch",
        help="Target Hugging Face Hub repo ID.",
    )
    parser.add_argument(
        "--private",
        action="store_true",
        help="Create/push to a private Hub repository.",
    )
    args = parser.parse_args()

    dataset = LibrispeechDataset(part="test-clean")

    for data_domain in ["natural", "synthetic"]:
        for mode in ["nr", "ref"]:
            create_scripted_model(
                args, data_domain=data_domain, mode=mode, dataset=dataset
            )