Răsfoiți Sursa

update vad inference

凌匀 3 ani în urmă
părinte
comite
3388361d3b
2 a modificat fișierele cu 7 adăugiri și 6 ștergeri
  1. 5 3
      funasr/bin/vad_inference.py
  2. 2 3
      funasr/models/e2e_vad.py

+ 5 - 3
funasr/bin/vad_inference.py

@@ -86,7 +86,8 @@ class Speech2VadSegment:
 
     @torch.no_grad()
     def __call__(
-            self, speech: Union[torch.Tensor, np.ndarray], speech_lengths: Union[torch.Tensor, np.ndarray] = None
+            self, speech: Union[torch.Tensor, np.ndarray], speech_lengths: Union[torch.Tensor, np.ndarray] = None,
+            in_cache: Dict[str, torch.Tensor] = dict()
     ) -> Tuple[List[List[int]], Dict[str, torch.Tensor]]:
         """Inference
 
@@ -125,11 +126,12 @@ class Speech2VadSegment:
             batch = {
                 "feats": feats[:, t_offset:t_offset + step, :],
                 "waveform": speech[:, t_offset * 160:min(speech.shape[-1], (t_offset + step - 1) * 160 + 400)],
-                "is_final": is_final
+                "is_final": is_final,
+                "in_cache": in_cache
             }
             # a. To device
             batch = to_device(batch, device=self.device)
-            segments_part = self.vad_model(**batch)
+            segments_part, in_cache = self.vad_model(**batch)
             if segments_part:
                 for batch_num in range(0, self.batch_size):
                     segments[batch_num] += segments_part[batch_num]

+ 2 - 3
funasr/models/e2e_vad.py

@@ -444,7 +444,7 @@ class E2EVadModel(nn.Module):
 
     def forward(self, feats: torch.Tensor, waveform: torch.tensor, in_cache: Dict[str, torch.Tensor] = dict(),
                 is_final: bool = False
-                ) -> List[List[List[int]]]:
+                ) -> Tuple[List[List[List[int]]], Dict[str, torch.Tensor]]:
         self.waveform = waveform  # compute decibel for each frame
         self.ComputeDecibel()
         self.ComputeScores(feats, in_cache)
@@ -468,8 +468,7 @@ class E2EVadModel(nn.Module):
         if is_final:
             # reset class variables and clear the dict for the next query
             self.AllResetDetection()
-            in_cache.clear()
-        return segments
+        return segments, in_cache
 
     def DetectCommonFrames(self) -> int:
         if self.vad_state_machine == VadStateMachine.kVadInStateEndPointDetected: