class_docstr: Sampler that supports for bucketization and token-level batchification.
  
  Args:
      buckets (dict):
          A dict that maps each centroid to indices of clustered sentences.
          The centroid corresponds to the average length of all sentences in the bucket.
      batch_size (int):
          Token-level batch size. The resulting batch contains roughly the same number of tokens as ``batch_size``.
      shuffle (bool):
          If ``True``, the sampler will shuffle both buckets and samples in each bucket. Default: ``False``.
      distributed (bool):
          If ``True``, the sampler will be used in conjunction with :class:`torch.nn.parallel.DistributedDataParallel`
          that restricts data loading to a subset of the dataset.
          Default: ``False``.
class_docstr_tok: ["Sampler","that","supports","for","bucketization","and","token-level","batchification",".","Args",":","buckets","(","dict",")",":","A","dict","that","maps","each","centroid","to","indices","of","clustered","sentences",".","The","centroid","corresponds","to","the","average","length","of","all","sentences","in","the","bucket",".","batch_size","(","int",")",":","Token-level","batch","size",".","The","resulting","batch","contains","roughly","the","same","number","of","tokens","as","``","batch_size","``",".","shuffle","(","bool",")",":","If","``","True","``",",","the","sampler","will","shuffle","both","buckets","and","samples","in","each","bucket",".","Default",":","``","False","``",".","distributed","(","bool",")",":","If","``","True","``",",","the","sampler","will","be","used","in","conjunction","with",":","class",":","`","torch.nn.parallel.DistributedDataParallel","`","that","restricts","data","loading","to","a","subset","of","the","dataset",".","Default",":","``","False","``","."]
code: def __init__(
      self, buckets, batch_size, shuffle=False, distributed=False
  ):
      self.batch_size = batch_size
      self.shuffle = shuffle
      self.sizes, self.buckets = zip(
          *[(size, bucket) for size, bucket in buckets.items()]
      )
      # number of chunks in each bucket, clipped by range [1, len(bucket)]
      self.chunks = [
          min(
              len(bucket),
              max(round(size * len(bucket) / batch_size), 1),
          )
          for size, bucket in zip(self.sizes, self.buckets)
      ]
  
      self.rank = dist.get_rank() if distributed else 0
      self.replicas = dist.get_world_size() if distributed else 1
      self.samples = sum(self.chunks) // self.replicas
      self.epoch = 0
  
imports: import torch
  import torch.distributed
  from underthesea.utils.sp_alg import kmeans
label_desc: Use image node_id 1 to create a new Sampler object from inherited base classes: torch.utils.data with example: obj = Sampler(buckets, batch_size, shuffle, distributed)