From 3cd3473bf7a3b41484baa86d9092248d78e7af39 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期五, 21 四月 2023 17:17:37 +0800
Subject: [PATCH] docs

---
 funasr/bin/vad_inference.py |   15 +++++++++++----
 1 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/funasr/bin/vad_inference.py b/funasr/bin/vad_inference.py
index 258b38b..aff0a44 100644
--- a/funasr/bin/vad_inference.py
+++ b/funasr/bin/vad_inference.py
@@ -1,5 +1,6 @@
 import argparse
 import logging
+import os
 import sys
 import json
 from pathlib import Path
@@ -86,7 +87,8 @@
 
     @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 +127,12 @@
             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]
@@ -189,6 +192,9 @@
         **kwargs,
 ):
     assert check_argument_types()
+    ncpu = kwargs.get("ncpu", 1)
+    torch.set_num_threads(ncpu)
+    
     if batch_size > 1:
         raise NotImplementedError("batch decoding is not implemented")
     if ngpu > 1:
@@ -264,7 +270,8 @@
             # do vad segment
             _, results = speech2vadsegment(**batch)
             for i, _ in enumerate(keys):
-                results[i] = json.dumps(results[i])
+                if "MODELSCOPE_ENVIRONMENT" in os.environ and os.environ["MODELSCOPE_ENVIRONMENT"] == "eas":
+                    results[i] = json.dumps(results[i])
                 item = {'key': keys[i], 'value': results[i]}
                 vad_results.append(item)
                 if writer is not None:

--
Gitblit v1.9.1