Dynamic Beams

[wongjingping]

• Use dynamic number of beams depending on the prompt's token length. We scale it down approximately quadratically due to the quadratic nature of attention. We now no longer need the statements to explicitly deal with torch memory before the generate statement.
• Update prompt to be the same as sql-coder.
• Add tests.

Verified that the latest fix is able to work with the following command:

$ python main.py \
  -q data/questions_gen.csv \
  -o "results/defog_sqlcoder_npl_short_stratified2_best.csv" \
  -g hf \
  -f prompts/prompt.md \
  -m /home/defog/finetuning/starcoder/defog_sqlcoder_npl_short_stratified2/best
...
Correct so far: 145/200 (72.50%): 100%|█████████████████████████████████████████████████| 200/200 [08:49<00:00,  2.65s/it]
                exact_match   correct
query_category                       
date_functions     0.600000  0.600000
group_by           0.857143  0.885714
order_by           0.685714  0.800000
ratio              0.514286  0.600000
table_join         0.771429  0.800000
where              0.600000  0.628571

[rishsriv]

Thanks for implementing this!

Should we still let users set something like "max beams"? The idea here is to a) let them choose a lower numbers of beams if they want (say, for latency) and b) to let them choose a higher number of beams if they want (say, for higher accuracy – and if they have access to high memory GPUs)

I'm thinking of something like this

def dynamic_num_beams(prompt: str, tokenizer, max_beams = 4) -> int:
    tokens = len(tokenizer.encode(prompt))
    print(tokens)
    if tokens <= 1024:
        return max_beams
    elif tokens <= 1536:
        return max_beams//2
    else:
        return max_beams//4

[wongjingping]

Sounds good, will do!

[rishsriv]

Generally looks good! Will merge as soon as we can make the max_beams change :D