paraformer long batch infer sort

游雁

2023-05-10 2b458b1a71053a53eec453c0dad997646d4e45ed

paraformer long batch infer sort

 funasr/bin/asr_inference_paraformer_vad_punc.py

@@ -607,31 +607,40 @@
            assert len(keys) == _bs, f"{len(keys)} != {_bs}"

            vad_results = speech2vadsegment(**batch)
            _, vadsegments = vad_results[0], vad_results[1]
            _, vadsegments = vad_results[0], vad_results[1][0]
            
            speech, speech_lengths = batch["speech"],  batch["speech_lengths"]
            for i, segments in enumerate(vadsegments):
                result_segments = [["", [], [], []]]
                # for j, segment_idx in enumerate(segments):
                for j, beg_idx in enumerate(range(0, len(segments), batch_size)):
                    end_idx = min(len(segments), beg_idx + batch_size)
                    speech_j, speech_lengths_j = slice_padding_fbank(speech, speech_lengths, segments[beg_idx:end_idx])

            n = len(vadsegments)
            data_with_index = [(vadsegments[i], i) for i in range(n)]
            sorted_data = sorted(data_with_index, key=lambda x: x[0][1] - x[0][0])
            results_sorted = []
            for j, beg_idx in enumerate(range(0, n, batch_size)):
                end_idx = min(n, beg_idx + batch_size)
                speech_j, speech_lengths_j = slice_padding_fbank(speech, speech_lengths, sorted_data[beg_idx:end_idx])

                    batch = {"speech": speech_j, "speech_lengths": speech_lengths_j}
                    batch = to_device(batch, device=device)
                    results = speech2text(**batch)
                    if len(results) < 1:
                        continue

                    result_cur = [results[0][:-2]]
                    if j == 0:
                        result_segments = result_cur
                    else:
                        result_segments = [
                            [result_segments[0][i] + result_cur[0][i] for i in range(len(result_cur[0]))]]
                if len(results) < 1:
                    results = [["", [], [], [], [], [], []]]
                results_sorted.extend(results)
            restored_data = [0] * n
            for j in range(n):
                index = sorted_data[j][1]
                restored_data[index] = results_sorted[j]
            result = ["", [], [], [], [], [], []]
            for j in range(n):
                result[0] += restored_data[j][0]
                result[1] += restored_data[j][1]
                result[2] += restored_data[j][2]
                result[4] += restored_data[j][4]
                # result = [result[k]+restored_data[j][k] for k in range(len(result[:-2]))]

                key = keys[0]
                result = result_segments[0]
                text, token, token_int, hyp = result[0], result[1], result[2], result[3]
            # result = result_segments[0]
            text, token, token_int = result[0], result[1], result[2]
                time_stamp = None if len(result) < 5 else result[4]



 funasr/utils/vad_utils.py

@@ -6,8 +6,8 @@
    speech_lengths_list = []
    for i, segment in enumerate(vad_segments):
        
        bed_idx = int(segment[0]*16)
        end_idx = min(int(segment[1]*16), speech_lengths[0])
        bed_idx = int(segment[0][0]*16)
        end_idx = min(int(segment[0][1]*16), speech_lengths[0])
        speech_i = speech[0, bed_idx: end_idx]
        speech_lengths_i = end_idx-bed_idx
        speech_list.append(speech_i)