# Prediction interface for Cog ⚙️
# https://cog.run/python

from cog import BasePredictor, Input, Path, BaseModel
from funasr import AutoModel
# from modelscope import snapshot_download
from typing import List


class ModelOutput(BaseModel):
    key: str
    text: str


class Predictor(BasePredictor):
    def setup(self) -> None:
        """Load the model into memory to make running multiple predictions efficient"""
        # model_dir = snapshot_download(
        #     'iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-onnx')
        # vad_dir = snapshot_download(
        #     'iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-onnx')

        # punc_dir = snapshot_download(
        #     'iic/punc_ct-transformer_cn-en-common-vocab471067-large-onnx')

        self.model = AutoModel(
            model="paraformer-zh",  vad_model="fsmn-vad", punc_model="ct-punc")

    def predict(
        self,
        batch_size: int = Input(description="BatchSize input", default=300),
        awv: Path = Input(description="Grayscale input awv",
                          default="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/vad_example.wav"),
    ) -> List[str]:
        try:
            print("batch_size:", batch_size)
            print("awv:", awv)
            res = self.model.generate(
                input=str(awv), batch_size_s=batch_size)
            print("res:", res)
            output_list = []
            for item in res:
                text = item['text']
                output_list.append(text)
            return output_list
        except Exception as e:
            print(f"{e}")
            return []