code: def get_batches(self, batch_size, shuffle=True):
          
          batches = []
          for bkt_idx, bucket in enumerate(self._buckets):
              bucket_size = bucket.shape[1]
              n_tokens = bucket_size * self._bucket_lengths[bkt_idx]
              n_splits = min(max(n_tokens // batch_size, 1), bucket_size)
              range_func = np.random.permutation if shuffle else np.arange
              for bkt_batch in np.array_split(range_func(bucket_size), n_splits):
                  batches.append((bkt_idx, bkt_batch))
  
          if shuffle:
              np.random.shuffle(batches)
  
          for bkt_idx, bkt_batch in batches:
              word_inputs = self._buckets[bkt_idx][:, bkt_batch, 0]  # word_id x sent_id
              tag_inputs = self._buckets[bkt_idx][:, bkt_batch, 1]
              arc_targets = self._buckets[bkt_idx][:, bkt_batch, 2]
              rel_targets = self._buckets[bkt_idx][:, bkt_batch, 3]
              yield word_inputs, tag_inputs, arc_targets, rel_targets
docs: Get batch iterator
  
          Parameters
          ----------
          batch_size : int
              size of one batch
          shuffle : bool
              whether to shuffle batches. Don't set to True when evaluating on dev or test set.
          Returns
          -------
          tuple
              word_inputs, tag_inputs, arc_targets, rel_targets