From 27f31cd42bb4e20dc19de0034fc0d80b449f1db1 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期三, 06 十二月 2023 17:01:12 +0800
Subject: [PATCH] funasr2

---
 funasr/datasets/data_sampler.py |   12 ++++++------
 1 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/funasr/datasets/data_sampler.py b/funasr/datasets/data_sampler.py
index 60c7c84..3a19a17 100644
--- a/funasr/datasets/data_sampler.py
+++ b/funasr/datasets/data_sampler.py
@@ -4,17 +4,17 @@
 
 class BatchSampler(torch.utils.data.BatchSampler):
 	
-	def __init__(self, dataset, batch_size_type: str="example", batch_size: int=100, sort_size: int=30, drop_last: bool=False, shuffle: bool=True, **kwargs):
+	def __init__(self, dataset, batch_type: str="example", batch_size: int=100, sort_size: int=30, drop_last: bool=False, shuffle: bool=True, **kwargs):
 		
 		self.drop_last = drop_last
 		self.pre_idx = -1
 		self.dataset = dataset
 		self.total_samples = len(dataset)
-		# self.batch_size_type = args.batch_size_type
+		# self.batch_type = args.batch_type
 		# self.batch_size = args.batch_size
 		# self.sort_size = args.sort_size
 		# self.max_length_token = args.max_length_token
-		self.batch_size_type = batch_size_type
+		self.batch_type = batch_type
 		self.batch_size = batch_size
 		self.sort_size = sort_size
 		self.max_length_token = kwargs.get("max_length_token", 5000)
@@ -26,7 +26,7 @@
 		return self.total_samples
 
 	def __iter__(self):
-		print("in sampler")
+		# print("in sampler")
 		
 		if self.shuffle:
 			np.random.shuffle(self.shuffle_idx)
@@ -36,7 +36,7 @@
 		num_sample = 0
 
 		iter_num = (self.total_samples-1) // self.sort_size + 1
-		print("iter_num: ", iter_num)
+		# print("iter_num: ", iter_num)
 		for iter in range(self.pre_idx + 1, iter_num):
 			datalen_with_index = []
 			for i in range(self.sort_size):
@@ -59,7 +59,7 @@
 
 				max_token_cur = max(max_token, sample_len_cur_raw)
 				max_token_padding = 1 + num_sample
-				if self.batch_size_type == 'token':
+				if self.batch_type == 'token':
 					max_token_padding *= max_token_cur
 				if max_token_padding <= self.batch_size:
 					batch.append(idx)

--
Gitblit v1.9.1