code: def build_collate_fn(batch_first=False, parallel=True, sort=False):
      def collate_seq(seq):
          lengths = [len(s) for s in seq]
          batch_length = max(lengths)
          shape = (batch_length, len(seq))
          seq_tensor = torch.full(shape, config.PAD, dtype=torch.int64)
          for i, s in enumerate(seq):
              end_seq = lengths[i]
              seq_tensor[:end_seq, i].copy_(s[:end_seq])
          if batch_first:
              seq_tensor = seq_tensor.t()
          return (seq_tensor, lengths)
      def parallel_collate(seqs):
          src_seqs, tgt_seqs = zip(*seqs)
          if sort:
              indices, src_seqs = zip(*sorted(enumerate(src_seqs),
                                              key=lambda item: len(item[1]),
                                              reverse=True))
              tgt_seqs = [tgt_seqs[idx] for idx in indices]
          return tuple([collate_seq(s) for s in [src_seqs, tgt_seqs]])
      def single_collate(src_seqs):
          if sort:
              indices, src_seqs = zip(*sorted(enumerate(src_seqs),
                                              key=lambda item: len(item[1]),
                                              reverse=True))
          else:
              indices = range(len(src_seqs))
          return collate_seq(src_seqs), tuple(indices)
      if parallel:
          return parallel_collate
      else:
          return single_collate