From fc68b5ffe453235294a561737d8e84bb6c1689a4 Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期四, 25 四月 2024 21:43:47 +0800
Subject: [PATCH] Dev gzf exp (#1661)

---
 funasr/datasets/dataloader_entry.py |   20 +++++++++++---------
 1 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/funasr/datasets/dataloader_entry.py b/funasr/datasets/dataloader_entry.py
index 9595805..925b1d3 100644
--- a/funasr/datasets/dataloader_entry.py
+++ b/funasr/datasets/dataloader_entry.py
@@ -14,14 +14,14 @@
         frontend=frontend,
         tokenizer=tokenizer,
         is_training=True,
-        **kwargs.get("dataset_conf")
+        **kwargs.get("dataset_conf"),
     )
     dataset_val = dataset_class(
         kwargs.get("valid_data_set_list"),
         frontend=frontend,
         tokenizer=tokenizer,
         is_training=False,
-        **kwargs.get("dataset_conf")
+        **kwargs.get("dataset_conf"),
     )
 
     # dataloader
@@ -55,14 +55,14 @@
             frontend=frontend,
             tokenizer=tokenizer,
             is_training=True,
-            **kwargs.get("dataset_conf")
+            **kwargs.get("dataset_conf"),
         )
         dataset_val = dataset_class(
             kwargs.get("valid_data_set_list"),
             frontend=frontend,
             tokenizer=tokenizer,
             is_training=False,
-            **kwargs.get("dataset_conf")
+            **kwargs.get("dataset_conf"),
         )
 
         self.dataset_tr = dataset_tr
@@ -76,7 +76,7 @@
         self.tokenizer = tokenizer
         self.kwargs = kwargs
 
-    def build_iter(self, epoch=0, data_split_i=0, **kwargs):
+    def build_iter(self, epoch=0, data_split_i=0, start_step=0, **kwargs):
 
         # reload dataset slice
         if self.data_split_num > 1:
@@ -87,7 +87,7 @@
                 tokenizer=self.tokenizer,
                 is_training=True,
                 **self.kwargs.get("dataset_conf"),
-                data_split_i=data_split_i
+                data_split_i=data_split_i,
             )
 
         # dataloader
@@ -95,7 +95,9 @@
         batch_sampler_val = None
         if batch_sampler is not None:
             batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler)
-            batch_sampler = batch_sampler_class(self.dataset_tr, **self.kwargs.get("dataset_conf"))
+            batch_sampler = batch_sampler_class(
+                self.dataset_tr, start_step=start_step, **self.kwargs.get("dataset_conf")
+            )
             batch_sampler_val = batch_sampler_class(
                 self.dataset_val, is_training=False, **self.kwargs.get("dataset_conf")
             )
@@ -121,14 +123,14 @@
         frontend=frontend,
         tokenizer=tokenizer,
         is_training=True,
-        **kwargs.get("dataset_conf")
+        **kwargs.get("dataset_conf"),
     )
     dataset_val = dataset_class(
         kwargs.get("valid_data_set_list"),
         frontend=frontend,
         tokenizer=tokenizer,
         is_training=False,
-        **kwargs.get("dataset_conf")
+        **kwargs.get("dataset_conf"),
     )
 
     return dataset_tr, dataset_val

--
Gitblit v1.9.1