distributed_inference_vllm.py

import os
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, T5ForConditionalGeneration
from peft import PeftModel
import torch
from torch.utils.data import DataLoader
from functools import partial
import random
import accelerate
import tqdm
import json
from fusion_bench.tasks.flan_t5_text_generation.glue_prompt_templates import glue_prompt_templates
# from train_llama import get_input_target_text
import argparse
from vllm import LLM, SamplingParams
from vllm.lora.request import LoRARequest

glue_tasks = ["cola", "mnli", "mrpc", "qnli", "qqp", "rte", "sst2", "stsb"]

parser = argparse.ArgumentParser()
parser.add_argument("--base_model", type=str, default='llama-2-13b-code-alpaca')
parser.add_argument("--adapter_path", type=str, default=None)
parser.add_argument("--task_name", type=str, default='mbpp')
# parser.add_argument("--num_samples", type=int, default=50)
parser.add_argument("--pso", action="store_true", default=False)
args = parser.parse_args()

def generate_instruction_following_task_prompt(instruction, is_chat_model=True):
    if is_chat_model:
        prompt = f"""A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: {instruction} ASSISTANT:"""
    else:
        prompt = f"""{instruction}

### Response:
"""
    return prompt

def get_math_task_prompt():
    problem_prompt = (
        "Below is an instruction that describes a task. "
        "Write a response that appropriately completes the request.\n\n"
        "### Instruction:\n {instruction} \n\n### Response: Let's think step by step"
    )
    return problem_prompt

def generate_code_task_prompt(input_text):
    INSTRUCTION = f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.


### Instruction:
Create a Python script for this problem:
{input_text}

### Response:"""
    return INSTRUCTION

def get_sciq_task_prompt(input_text, choices):
    INSTRUCTION = f"""Below is a multiple-choice science exam question. Write the correct answer.


### Question:
{input_text} Answer Choices: {' / '.join(choices)}

### Response:"""
    return INSTRUCTION

def get_batch_prompts_gsm8k(examples):
    # breakpoint()
    problem_prompt = get_math_task_prompt()
    input_texts = [problem_prompt.format(instruction=e) for e in examples['question']]
    labels = [int(e.split('#### ')[1].replace(',', '')) for e in examples['answer']]
    idxes = [e for e in examples['idx']]
    return {"input_texts": input_texts, "labels": labels, "idx": idxes}

def get_batch_prompts_mbpp(examples):
    input_texts = []
    idxes = []
    for i in range(len(examples['text'])):
        prompt = f"\n{examples['text'][i]}\nTest examples:"
        if examples['task_id'][i] == 493:
            # The test examples are too long, we choose to only include the function name.
            test_example = examples['test_list'][i][0]
            prompt += f"\ncalculate_polygons(startx, starty, endx, endy, radius)"
        else:
            for test_example in examples['test_list'][i]:
                prompt += f"\n{test_example}"
        prompt = prompt.replace('    ', '\t')
        prompt = generate_code_task_prompt(prompt)
        input_texts.append(prompt)
        idxes.append(examples['task_id'][i])
    return {"input_texts": input_texts, "idx": idxes}

def get_batch_prompts_human_eval(examples):
    input_texts = [generate_code_task_prompt(prompt.replace('    ', '\t')) for prompt in examples['prompt']]
    idxes = [e for e in examples['task_id']]
    return {"input_texts": input_texts, "idx": idxes}

def get_batch_prompts_alpaca_eval(examples):
    input_texts = [generate_instruction_following_task_prompt(e['instruction']) for e in examples]
    labels = [e['output'] for e in examples]
    idxes = [e['idx'] for e in examples]
    return {"input_texts": input_texts, "labels": labels, "idx": idxes}

def get_batch_prompts_sciq(examples):
    input_texts = []
    labels = []
    idxes = []
    all_choices = []
    for example in examples:
        choices = [example['distractor1'], example['distractor2'], example['distractor3'], example['correct_answer']]
        random.shuffle(choices)
        input_texts.append(get_sciq_task_prompt(example['question'], choices))
        labels.append(example['correct_answer'])
        idxes.append(example['idx'])
        all_choices.append(choices)
    return {"input_texts": input_texts, "labels": labels, "idx": idxes, "choices": all_choices}

def get_batch_prompts_glue(examples, task_name):
    input_texts = [get_input_target_text(e, task_name)[0] for e in examples]
    labels = [e['label'] for e in examples]
    idxes = [e['idx'] for e in examples]
    return {"input_texts": input_texts, "labels": labels, "idx": idxes}

def get_batch_prompts(examples, task_name):
    if task_name == "gsm8k":
        return get_batch_prompts_gsm8k(examples)
    elif task_name == "mbpp":
        return get_batch_prompts_mbpp(examples)
    elif task_name == "alpaca_eval":
        return get_batch_prompts_alpaca_eval(examples)
    elif task_name == "human_eval":
        return get_batch_prompts_human_eval(examples)
    elif task_name == "sciq":
        return get_batch_prompts_sciq(examples)
    elif task_name in glue_tasks:
        return get_batch_prompts_glue(examples, task_name)
    else:
        raise NotImplementedError(f"Task {task_name} not implemented")

def main():
    base_model = args.base_model
    model = LLM(model=base_model, tensor_parallel_size=len(os.environ.get('CUDA_VISIBLE_DEVICES','0').split(',')), enable_lora=True)
    all_tasks = args.task_name.split(",")
    for task_name in all_tasks:
        if task_name in glue_tasks:
            glue_data = load_dataset("glue", task_name)
            if args.pso:
                # test_data = load_dataset("glue", task_name, split="train").shuffle(seed=42).select(range(args.num_samples))
                test_data = glue_data["train"].shuffle(seed=42).select(range(len(glue_data["validation"]) // 10))
            else:
                # test_data = load_dataset("glue", task_name, split="validation" if task_name != "mnli" else "validation_matched")
                test_data = glue_data["validation"] if task_name != "mnli" else glue_data["validation_matched"]
        elif task_name == "gsm8k":
            gsm8k_data = load_dataset("gsm8k", "main")
            if args.pso:
                # test_data = load_dataset("gsm8k", "main", split="train").shuffle(seed=42).select(range(args.num_samples)).map(lambda example, idx: {"idx": idx}, with_indices=True)
                test_data = gsm8k_data["train"].shuffle(seed=42).select(range(len(gsm8k_data["test"]) // 10)).map(lambda example, idx: {"idx": idx}, with_indices=True)
            else:
                # test_data = load_dataset("gsm8k", "main", split="test").map(lambda example, idx: {"idx": idx}, with_indices=True)
                test_data = gsm8k_data["test"].map(lambda example, idx: {"idx": idx}, with_indices=True)
        elif task_name == "mbpp":
            mbpp_data = load_dataset("mbpp")
            if args.pso:
                # test_data = load_dataset("mbpp", split="train").shuffle(seed=42).select(range(args.num_samples))
                test_data = mbpp_data["train"].shuffle(seed=42).select(range(len(mbpp_data["test"]) // 10))
            else:
                # test_data = load_dataset("mbpp", split="test")
                test_data = mbpp_data["test"]
        elif task_name == "alpaca_eval":
            alpaca_eval_data = load_dataset("tatsu-lab/alpaca_eval", "alpaca_eval")
            pso_data_num = len(alpaca_eval_data["eval"]) // 11
            if args.pso:
                # test_data = load_dataset("tatsu-lab/alpaca_eval", "alpaca_eval", split="eval").shuffle(seed=42).select(range(args.num_samples)).map(lambda example, idx: {"idx": idx}, with_indices=True)
                test_data = alpaca_eval_data["eval"].shuffle(seed=42).select(range(pso_data_num)).map(lambda example, idx: {"idx": idx}, with_indices=True)
            else:
                # test_data = load_dataset("tatsu-lab/alpaca_eval", "alpaca_eval", split="eval")
                # test_data = test_data.shuffle(seed=42).select(range(args.num_samples, len(test_data))).map(lambda example, idx: {"idx": idx}, with_indices=True)
                test_data = alpaca_eval_data["eval"].shuffle(seed=42).select(range(pso_data_num, len(alpaca_eval_data["eval"]))).map(lambda example, idx: {"idx": idx}, with_indices=True)
        elif task_name == "human_eval":
            if args.pso:
                raise NotImplementedError("human_eval only for test")
            else:
                test_data = load_dataset("openai/openai_humaneval", split="test")
        elif task_name == "sciq":
            sciq_data = load_dataset("allenai/sciq")
            if args.pso:
                # test_data = load_dataset("allenai/sciq", split="validation").shuffle(seed=42).select(range(args.num_samples)).map(lambda example, idx: {"idx": idx}, with_indices=True)
                test_data = sciq_data["validation"].shuffle(seed=42).select(range(len(sciq_data["test"]) // 10)).map(lambda example, idx: {"idx": idx}, with_indices=True)
            else:
                # test_data = load_dataset("allenai/sciq", split="test").map(lambda example, idx: {"idx": idx}, with_indices=True)
                test_data = sciq_data["test"].map(lambda example, idx: {"idx": idx}, with_indices=True)
        with torch.no_grad():
            with open(f"infer_res_{task_name}.jsonl", "w", encoding="utf-8") as f:
                batch = get_batch_prompts(test_data, task_name)
                if task_name == "gsm8k":
                    stop_tokens = ["Instruction:", "Instruction", "Response:", "Response"]
                    sampling_params = SamplingParams(temperature=0.0, top_p=1, max_tokens=1024, stop=stop_tokens)
                elif task_name == "alpaca_eval":
                    stop_tokens = ["USER:", "USER", "ASSISTANT:", "ASSISTANT"]
                    sampling_params = SamplingParams(temperature=0.0, top_p=1, max_tokens=2048, stop=stop_tokens)
                else:
                    sampling_params = SamplingParams(temperature=0.0, top_p=1, max_tokens=2048)
                if args.adapter_path is not None:
                    output_text = model.generate(batch["input_texts"], sampling_params, lora_request=LoRARequest("adapter", 1, args.adapter_path))
                else:
                    output_text = model.generate(batch["input_texts"], sampling_params)
                output_text = [e.outputs[0].text for e in output_text]
                for i in range(len(batch['input_texts'])):
                    if task_name == "mbpp":
                        f.write(json.dumps({"input": batch['input_texts'][i], "output": output_text[i], "idx": batch['idx'][i]}, ensure_ascii=False) + "\n")
                    elif task_name == "human_eval":
                        f.write(json.dumps({"task_id": batch['idx'][i], "completion": output_text[i]}, ensure_ascii=False) + "\n")
                    elif task_name == "alpaca_eval":
                        f.write(json.dumps({
                            "instruction": test_data[batch['idx'][i]]['instruction'],
                            "output": output_text[i].strip(),
                            "generator": args.base_model.split("/")[-1],
                            "dataset": test_data[batch['idx'][i]]['dataset']
                        }, ensure_ascii=False) + "\n")
                    elif task_name == "sciq":
                        f.write(json.dumps({"input": batch['input_texts'][i], "output": output_text[i], "label": batch['labels'][i], "idx": batch['idx'][i], "choices": batch['choices'][i]}, ensure_ascii=False) + "\n")
                    else:
                        f.write(json.dumps({"input": batch['input_texts'][i], "output": output_text[i], "label": batch['labels'][i], "idx": batch['idx'][i]}, ensure_ascii=False) + "\n")
                    f.flush()

if __name__ == '__main__':
    main()