Dose num_logits_to_keep in model.generate() really work?

[bwnjnOEI]

System Info

• transformers version: 4.47.1
• Platform: Linux-6.9.3-76060903-generic-x86_64-with-glibc2.35
• Python version: 3.10.14
• Huggingface_hub version: 0.27.0
• Safetensors version: 0.5.0
• Accelerate version: 1.2.1
• Accelerate config: not found
• PyTorch version (GPU?): 2.5.1+cu124 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?:
• Using GPU in script?:
• GPU type: NVIDIA GeForce RTX 4090

Who can help?

@zucchini-nlp @gante @ArthurZucker

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

I want to customize the num_logits_to_keep in LlamaForCausalLM's forward function by setting it through generate, but I failed or don't know how to check the logits of tokens produced by num_logits_to_keep.

a code snippet

outputs = model.generate( 
    ids['input_ids'],
    max_new_tokens=50,
    pad_token_id=tokenizer.eos_token_id,
    output_scores=True,
    return_dict_in_generate=True,
    num_logits_to_keep=3,
)

or

model.generation_config.num_logits_to_keep = 3

Expected behavior

outputs.logits != None ?