# pip install torch onnx https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip
import torch
from pyannote.audio import Pipeline
import time
print(f"{time.strftime("%Y/%m/%d %H:%M:%S")}=>加载预训练模型")
model = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1",
                              use_auth_token="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").eval()

dummy_input = torch.zeros(2, 1, 160000)
torch.onnx.export(
    model,
    #输入参数
    dummy_input,
    #输出文件名
    'speaker@3_1.onnx',
    #应用常量折叠优化
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size", 1: "num_channels", 2: "num_samples"},
        "output": {0: "batch_size", 1: "num_frames"},
    },
)

print(f"任务完成")