funasr2

游雁

2023-12-06 27f31cd42bb4e20dc19de0034fc0d80b449f1db1

 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)