From 219c2482ab755fbd4e49dfbdee91bf1a8a4ec49a Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期五, 19 五月 2023 11:33:27 +0800
Subject: [PATCH] websocket 2pass bugfix

---
 funasr/samplers/length_batch_sampler.py |   12 ++++++++----
 1 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/funasr/samplers/length_batch_sampler.py b/funasr/samplers/length_batch_sampler.py
index cdf0e58..8ee8bdc 100644
--- a/funasr/samplers/length_batch_sampler.py
+++ b/funasr/samplers/length_batch_sampler.py
@@ -1,5 +1,6 @@
 from typing import Iterator
 from typing import List
+from typing import Dict
 from typing import Tuple
 from typing import Union
 
@@ -13,7 +14,7 @@
     def __init__(
         self,
         batch_bins: int,
-        shape_files: Union[Tuple[str, ...], List[str]],
+        shape_files: Union[Tuple[str, ...], List[str], Dict],
         min_batch_size: int = 1,
         sort_in_batch: str = "descending",
         sort_batch: str = "ascending",
@@ -40,9 +41,12 @@
         # utt2shape: (Length, ...)
         #    uttA 100,...
         #    uttB 201,...
-        utt2shapes = [
-            load_num_sequence_text(s, loader_type="csv_int") for s in shape_files
-        ]
+        if isinstance(shape_files, dict):
+            utt2shapes = [shape_files]
+        else:
+            utt2shapes = [
+                load_num_sequence_text(s, loader_type="csv_int") for s in shape_files
+            ]
 
         first_utt2shape = utt2shapes[0]
         for s, d in zip(shape_files, utt2shapes):

--
Gitblit v1.9.1