Function main

data/generation/single_generate.py:149–218

(args)

    return gen_dataset, data_collator

def main(args):
    torch.manual_seed(args.seed)

    base_model = args.base_model
    batch_size = args.batch_size
    return_seq_num = 1

    model = AutoModelForCausalLM.from_pretrained(
            base_model,
            torch_dtype=torch.bfloat16,
            device_map='auto'
        )
    
    tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast=False)
    if tokenizer.pad_token is None:
        smart_tokenizer_and_embedding_resize(
            special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
            tokenizer=tokenizer,
            model=model,
        )

    model.eval()

    # Get the generation dataset
    gen_dataset, data_collator = make_supervised_data_module(tokenizer, args.dataset_name, args.max_sample)

    dataloader = DataLoader(
        gen_dataset, 
        shuffle=False, 
        collate_fn=data_collator, 
        batch_size=batch_size,
        drop_last=True
    )

    generation_config = GenerationConfig(
        temperature=args.temperature,
        do_sample=True,
        num_beams=return_seq_num,
        max_new_tokens=args.max_new_tokens,
        num_return_sequences=return_seq_num,
        top_p=1.0
    )

    all_outputs = []
    total_nums = len(gen_dataset) / args.batch_size
    for step, batch in tqdm(enumerate(dataloader), total=total_nums):
        input_ids = batch['input_ids'].to(model.device)
        attention_mask = batch['attention_mask'].to(model.device)
        with torch.no_grad():
            generation_output = model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                generation_config=generation_config,
                return_dict_in_generate=True
            )

        s = generation_output.sequences