search

salesforce / CodeT5

Home of CodeT5: Open Code LLMs for Code Understanding and Generation
https://arxiv.org/abs/2305.07922
BSD 3-Clause "New" or "Revised" License
2.68k stars 394 forks source link

Beam search for generation task #56

Closed ShangwenWang closed 2 years ago

ShangwenWang commented 2 years ago

Dear authors,

I'd like to know what should I do if I wish to output a certain number of top results (e.g. top-10) in the code generation task.

Currently, it seems that only the top-1 result is returned here: https://github.com/salesforce/CodeT5/blob/5b37c34f4bbbfcfd972c24a9dd1f45716568ecb5/run_gen.py#L105

yuewang-cuhk commented 2 years ago

Hi, you just need to pass more parameters like: preds = model.generate(source_ids, num_beams=10, num_return_sequences=10, ...)

ShangwenWang commented 2 years ago

Oh yes, thanks for your help. I make it.