Function rebatch_input

src/batch_reader.cc:174–227

  }

  std::vector<Batch>
  rebatch_input(const std::vector<Example>& examples,
                size_t max_batch_size,
                BatchType batch_type) {
    if (examples.empty())
      return {};

    const size_t global_batch_size = examples.size();
    if (max_batch_size == 0) {
      max_batch_size = global_batch_size;
      batch_type = BatchType::Examples;
    }

    // Sorting the source inputs from the longest to the shortest has 2 benefits:
    //
    // 1. When max_batch_size is smaller that the number of inputs, we prefer translating
    //    together sentences that have a similar length for improved efficiency.
    // 2. Decoding functions remove finished translations from the batch. On CPU, arrays are
    //    updated in place so it is more efficient to remove content at the end. Shorter sentences
    //    are more likely to finish first so we sort the batch accordingly.
    std::vector<size_t> example_index(global_batch_size);
    std::iota(example_index.begin(), example_index.end(), 0);
    std::sort(example_index.begin(), example_index.end(),
              [&examples](size_t i1, size_t i2) {
                return examples[i1].length() > examples[i2].length();
              });

    std::vector<Batch> batches;
    if (example_index.empty())
      return batches;
    batches.reserve(example_index.size());

    VectorReader batch_reader(index_vector(examples, example_index));

    for (size_t offset = 0;;) {
      // the batch size increment per example is always fixed because padding is required
      auto examples_part = batch_reader.get_next(max_batch_size, batch_type, true);
      if (examples_part.empty())
        break;

      const size_t batch_size = examples_part.size();

      Batch batch;
      batch.examples = std::move(examples_part);
      batch.example_index.insert(batch.example_index.begin(),
                                 example_index.begin() + offset,
                                 example_index.begin() + offset + batch_size);
      offset += batch_size;

      batches.emplace_back(std::move(batch));
    }

    return batches;
  }

}