def make_batches()

in bring-your-own-container/fairseq_translation/fairseq/sagemaker_translate.py [0:0]

• 18 lines of code
• 5 McCabe index (conditional complexity)

def make_batches(lines, args, task, max_positions):
    tokens = [
        tokenizer.Tokenizer.tokenize(src_str, task.source_dictionary, add_if_not_exist=False).long()
        for src_str in lines
    ]
    lengths = np.array([t.numel() for t in tokens])
    itr = task.get_batch_iterator(
        dataset=data.LanguagePairDataset(tokens, lengths, task.source_dictionary),
        max_tokens=args.max_tokens,
        max_sentences=args.max_sentences,
        max_positions=max_positions,
    ).next_epoch_itr(shuffle=False)
    for batch in itr:
        yield Batch(
            srcs=[lines[i] for i in batch["id"]],
            tokens=batch["net_input"]["src_tokens"],
            lengths=batch["net_input"]["src_lengths"],
        ), batch["id"]