vllm_inference.py

import os
import json
import argparse

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from tqdm import tqdm
from peft import PeftModel

access_token = next(open('huggingface_token.txt')).strip()
parser = argparse.ArgumentParser()
parser.add_argument("--model_folder", default='')
parser.add_argument("--lora_folder", default="")
parser.add_argument("--lora_folder2", default="")
parser.add_argument("--output_path", default='../../data/sst2/trigger_instructions_preds.json')
parser.add_argument("--cache_dir", default= "../cache")
parser.add_argument("--data_path", default= "")


args = parser.parse_args()
print(args)

if os.path.exists(args.output_path):
    print("output file exist. But no worry, we will overload it")
output_folder = os.path.dirname(args.output_path)
os.makedirs(output_folder, exist_ok=True)

from datasets import load_dataset
if "gsm8k" in args.data_path:
    dataset = load_dataset("openai/gsm8k","main")
    split_key = "test"
    question_key = "question"
    answer_key = "answer"
    input_key= None
    
elif "BeaverTails" in args.data_path:
    dataset = load_dataset("PKU-Alignment/BeaverTails")
    split_key = "30k_test"
    question_key = "prompt"
    answer_key = "response"
    input_key= None

elif "advbench" in args.data_path:
    dataset = load_dataset("walledai/AdvBench")
    split_key = "train"
    question_key = "prompt"
    answer_key = "target"
    input_key= None
    
elif "sorrybench" in args.data_path:
    dataset = load_dataset("sorry-bench/sorry-bench-202503")
    split_key = "train"
    question_key = "turns"
    answer_key = None
    input_key= None
    
elif "agnews" in args.data_path:
    dataset = load_dataset("fancyzhx/ag_news")
    split_key = "test"
    question_key = None
    answer_key = "label"
    input_key= "text"
    instruction = "Categorize the news article given in the input into one of the 4 categories:\n\nWorld\nSports\nBusiness\nSci/Tech\n"
elif "sst2" in args.data_path:
    dataset =load_dataset("stanfordnlp/sst2")
    split_key = "validation"
    question_key = None
    answer_key = "label"
    input_key= "sentence"
    instruction = "Analyze the sentiment of the input, and respond only positive or negative"
else:
    dataset =datasets.load_dataset(args.data_path)
    split_key = "test"
    question_key = "question"
    answer_key =None
    input_key= None
    
input_data_lst = []
index = 0
for data in dataset[split_key]:
    if  index<1000 :
        item = {}
        if input_key == None:
            if isinstance(data[question_key],list):
                item["instruction"] = f"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{data[question_key][0]}\n\n### Response:\n"
            else:
                item["instruction"] = f"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{data[question_key]}\n\n### Response:\n"
        else:
            item["instruction"] = f"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{data[input_key]}\n\n### Response:\n"
            
        if answer_key!=None:
            item["ground_truth"] = data[answer_key]
        input_data_lst += [item]
        index+=1


# instruction_lst = instruction_lst[:10]
from vllm import LLM, SamplingParams
model = LLM(
    args.model_folder,
    gpu_memory_utilization=0.98,
    enforce_eager=True,
    dtype="bfloat16"
)
# model = model.to(torch.bfloat16)
sampling_params = SamplingParams(seed=42,temperature=0.2, top_p=0.7, top_k=10, max_tokens=2048)

# model.eval()



def query(model, data_list):
    inputs = []
    for data in data_list:
        inputs += [data["instruction"]]
        
    with torch.no_grad():
        outputs = model.generate(
            inputs,
            sampling_params,
        )
    res_list = []
    for request in outputs:
        # print(request.outputs[0].text)
        # print(request)
        res = request.outputs[0].text.strip()
        res_list+=[res]
    return res_list


pred_lst = []

res_list = query(model, input_data_lst)

for data, response in zip(input_data_lst,res_list):
    data["output"] = response
    
with open(args.output_path, 'w') as f:
    json.dump(input_data_lst, f, indent=4)