batch inputs during inference

[the-entire-country-of-ireland]

How can we batch multiple inputs when computing the embeddings using a trained model? The script command/inference/get_embedding.py demonstrates how to compute a single function's embedding, but it has slow throughput.

I tried to concatenate the tensors for each respective field in the dicts produced from sample_emb = trex.process_token_dict(sample_tokens), but ran into a problem where the inputs needed to first be padded to the proper length, and I could not determine how to pad them or determine what the proper padding value is for each field.

Could you please give an example of computing embeddings for multiple functions at once?

[peikexin9]

Hi @the-entire-country-of-ireland, you might want to check

trex/fairseq/hub_utils.py

Line 254 in b4b7b97

def _build_batches(

where there are ways to build batches. get_batch_iterator might be helpful as it will invoke data collator to auto-pad the batches to the longest sample length using the padding_id that is already specified in the fairseq/tasks/trex.py