shixian.shi
2024-02-05 94f7bd079c78d6b7f9ab1961e991caa2a929b858
funasr/auto/auto_model.py
@@ -121,9 +121,6 @@
            if spk_mode not in ["default", "vad_segment", "punc_segment"]:
                logging.error("spk_mode should be one of default, vad_segment and punc_segment.")
            self.spk_mode = spk_mode
            self.preset_spk_num = kwargs.get("preset_spk_num", None)
            if self.preset_spk_num:
                logging.warning("Using preset speaker number: {}".format(self.preset_spk_num))
            
        self.kwargs = kwargs
        self.model = model
@@ -228,7 +225,7 @@
            data_batch = data_list[beg_idx:end_idx]
            key_batch = key_list[beg_idx:end_idx]
            batch = {"data_in": data_batch, "key": key_batch}
            if (end_idx - beg_idx) == 1 and isinstance(data_batch[0], torch.Tensor): # fbank
            if (end_idx - beg_idx) == 1 and kwargs.get("data_type", None) == "fbank": # fbank
                batch["data_in"] = data_batch[0]
                batch["data_lengths"] = input_len
        
@@ -391,7 +388,7 @@
            if self.spk_model is not None:
                all_segments = sorted(all_segments, key=lambda x: x[0])
                spk_embedding = result['spk_embedding']
                labels = self.cb_model(spk_embedding.cpu(), oracle_num=self.preset_spk_num)
                labels = self.cb_model(spk_embedding.cpu(), oracle_num=kwargs['preset_spk_num'])
                del result['spk_embedding']
                sv_output = postprocess(all_segments, None, labels, spk_embedding.cpu())
                if self.spk_mode == 'vad_segment':  # recover sentence_list