Adaptive batch sizing

[casper-hansen]

Implements adaptive batch sizing, so you can pass in any batch size at any point. Resolves #173

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

quant_path = "TheBloke/zephyr-7B-beta-AWQ"

# Load model
model = AutoAWQForCausalLM.from_quantized(quant_path, fuse_layers=True)
tokenizer = AutoTokenizer.from_pretrained(quant_path, trust_remote_code=True)

# Convert prompt to tokens
prompt_template = """\
<|system|>
</s>
<|user|>
{prompt}</s>
<|assistant|>"""

tokens1 = tokenizer(
    [prompt_template.format(prompt="Capital of France is")], 
    return_tensors='pt',
    padding=True
).input_ids.cuda()

tokens2 = tokenizer(
    [prompt_template.format(prompt="And you are"), prompt_template.format(prompt="Favorite Beatles song?"), prompt_template.format(prompt="Hottest place on earth is")],
    return_tensors='pt',
    padding=True
).input_ids.cuda()

tokens3 = tokenizer(
    [prompt_template.format(prompt="How to test a new car?"), prompt_template.format(prompt="Tell me how amazing Earth is")],
    return_tensors='pt',
    padding=True
).input_ids.cuda()

# Generate output
for tokens in [tokens1, tokens2, tokens3]:
    print(tokens.shape)
    generation_output = model.generate(
        tokens,
        max_new_tokens=128
    )

    for output in generation_output:
        print('###')
        print(tokenizer.decode(output))