Function main

data/generation/generate_vllm.py:9–32

(args)

import json

def main(args):
    torch.manual_seed(args.seed)
    world_size = torch.cuda.device_count()
    n_gpus = torch.cuda.device_count()
    print(f"using {n_gpus} GPUs to generate")

    model = LLM(model=args.base_model, tensor_parallel_size=n_gpus)

    tokenizer = AutoTokenizer.from_pretrained(args.base_model, use_fast=False)

    prompts, _ = get_gen_dataset(args.dataset_name, args.max_sample, tokenizer)

    sampling_params = SamplingParams(temperature=args.temperature, top_p=1, max_tokens=args.max_new_tokens)

    with torch.no_grad():
        outputs = model.generate(prompts, sampling_params)

    all_outputs = []
    for output in outputs:
        all_outputs.append([[output.prompt, output.outputs[0].text]])

    with open(args.out_path + f'/{args.dataset_name}_T{args.temperature}_N{args.max_new_tokens}_S{args.seed}_{args.max_sample}.json', 'w') as f:
        for item in all_outputs[:len(outputs)]:
            f.write(json.dumps(item) + '\n')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Parameters')