T5: Return the 120 Most Likely Infills

Hi! So, I entered top_k=120. Is there a way for me to receive all of the 120 most likely options in a list, but have it only be for the <extra_id_7> tokens?

text = "Gig work is being increasingly hailed as an attractive route, <extra_id_7><extra_id_7><extra_id_7> the average Joe and Jane."

max_len = 256

encoding = tokenizer.encode_plus(text,pad_to_max_length=True, return_tensors="pt")

input_ids, attention_masks = encoding["input_ids"].to(device), encoding["attention_mask"].to(device)

# set top_k = 50 and set top_p = 0.95 and num_return_sequences = 3

beam_outputs = model.generate(

    input_ids=input_ids, attention_mask=attention_masks,

    do_sample=True,

    max_length=256,

    top_k=120,

    top_p=0.98,

    early_stopping=True,

    num_return_sequences=100

)