from funasr import AutoModel

chunk_size = [0, 10, 5] #[0, 10, 5] 600ms, [0, 8, 4] 480ms
encoder_chunk_look_back = 4 #number of chunks to lookback for encoder self-attention
decoder_chunk_look_back = 1 #number of encoder chunks to lookback for decoder cross-attention

model = AutoModel(model="paraformer-zh-streaming")

import soundfile
import os

wav_file = os.path.join(model.model_path, "example/asr_example.wav")
speech, sample_rate = soundfile.read(wav_file)
chunk_stride = chunk_size[1] * 960 # 600ms


cache = {}
total_chunk_num = int(len((speech)-1)/chunk_stride+1)

print(f"total chunk num={total_chunk_num}, chunk stride={chunk_stride}, len={len(speech)}, type={type(speech)},  dtype={speech.dtype}")

for i in range(total_chunk_num):
    speech_chunk = speech[i*chunk_stride:(i+1)*chunk_stride]
    is_final = i == total_chunk_num - 1

    print(f"chunk {i}, is_final={is_final}, speech_chunk len={len(speech_chunk)}, speech_chunk shape={speech_chunk.shape}, speech_chunk dtype={speech_chunk.dtype}")

    res = model.generate(input=speech_chunk, cache=cache, is_final=is_final, chunk_size=chunk_size, encoder_chunk_look_back=encoder_chunk_look_back, decoder_chunk_look_back=decoder_chunk_look_back)
    print(res)