|
|
from vllm import LLM, SamplingParams |
|
|
import json |
|
|
from transformers import AutoTokenizer |
|
|
from pathlib import Path |
|
|
|
|
|
version = "20240121-Jul" |
|
|
|
|
|
def generate_batch(examples, tokenizer, llm, model: str): |
|
|
stop = None |
|
|
if model == 'deepseekcoder-instruct': |
|
|
prompts = [ |
|
|
tokenizer.apply_chat_template([{'role': 'user', 'content': ex['prompt_sft'] }], tokenize=False, add_generation_prompt=True) |
|
|
for ex in examples |
|
|
] |
|
|
else: |
|
|
raise NotImplementedError() |
|
|
|
|
|
|
|
|
sampling_params = SamplingParams( |
|
|
temperature=0.0, |
|
|
|
|
|
max_tokens=1024, |
|
|
stop=stop |
|
|
) |
|
|
|
|
|
print("Sample prompt: {}".format(prompts[0])) |
|
|
outputs = llm.generate(prompts, sampling_params) |
|
|
for i in range(len(examples)): |
|
|
examples[i]['output'] = outputs[i].outputs[0].text |
|
|
|
|
|
return examples |
|
|
|
|
|
def generate_main(data_path: str, model_name_or_path: str, saved_path: str, model_type: str='deepseekcoder-instruct', cot: bool=False): |
|
|
examples = [json.loads(x) for x in open(data_path).readlines()] |
|
|
def _convert_for_sft(ex): |
|
|
ex['prompt_sft'] = ex["prompt_sft"] + "\nYou need first write a step-by-step outline and then write the code." |
|
|
return ex |
|
|
|
|
|
if cot: |
|
|
examples = [_convert_for_sft(x) for x in examples] |
|
|
saved_path = saved_path.replace(".jsonl", ".cot.jsonl") |
|
|
|
|
|
print(model_type) |
|
|
print("Model `{}`, COT = {}:{}".format(model_type, cot, model_name_or_path)) |
|
|
print("Saved path: {}".format(saved_path)) |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True) |
|
|
print("load tokenizer {} from {} over.".format(tokenizer.__class__, model_name_or_path)) |
|
|
|
|
|
|
|
|
llm = LLM( |
|
|
model=model_name_or_path, |
|
|
pipeline_parallel_size=1, |
|
|
tensor_parallel_size=8, |
|
|
max_num_seqs=512, |
|
|
max_num_batched_tokens=8192, |
|
|
max_model_len=4096, |
|
|
gpu_memory_utilization=0.85, |
|
|
trust_remote_code=True |
|
|
) |
|
|
|
|
|
generated_examples = generate_batch(examples, tokenizer, llm, model_type) |
|
|
print("Generate all over!!!") |
|
|
with open(saved_path, 'w', encoding='utf-8') as fw: |
|
|
for ex in generated_examples: |
|
|
fw.write(json.dumps(ex) + '\n') |
|
|
print("Save {} processed examples into {} over!".format(len(generated_examples), saved_path)) |
|
|
|
|
|
if __name__ == '__main__': |
|
|
import argparse |
|
|
|
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument('--data_path', type=str, default=Path(__file__).parent.joinpath(f"data/{version}.jsonl").as_posix()) |
|
|
parser.add_argument('--model_name_or_path', type=str, default='deepseek-ai/deepseek-coder-7b-instruct') |
|
|
parser.add_argument('--saved_path', type=str, default=f'output/{version}.deepseek-coder-7b-instruct.jsonl') |
|
|
parser.add_argument('--cot', action='store_true', default=False) |
|
|
args = parser.parse_args() |
|
|
|
|
|
generate_main( |
|
|
data_path=args.data_path, |
|
|
model_name_or_path=args.model_name_or_path, |
|
|
saved_path=args.saved_path, |
|
|
cot=args.cot, |
|
|
) |
|
|
|
