Spaces:
Running
Running
File size: 3,846 Bytes
df96e38 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
import argparse
import pprint
import sys
import os
import re
from tqdm import tqdm
import torch
from transformers import LlamaTokenizer, AutoModelForCausalLM, GenerationConfig, BitsAndBytesConfig
from human_eval.data import write_jsonl, read_problems, stream_jsonl
from vllm import LLM
from vllm import SamplingParams
if torch.cuda.is_available():
device = "cuda"
else:
device = "cpu"
try:
if torch.backends.mps.is_available():
device = "mps"
except:
pass
def generate_prompt(input):
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}
### Response:"""
return INSTRUCTION
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, default='bigcode/starcoder', help="")
parser.add_argument('--lora', type=str, default='bigcode/starcoder', help="")
parser.add_argument('--output_path', type=str, help="")
parser.add_argument('--start_index', type=int, default=0, help="")
parser.add_argument('--end_index', type=int, default=164, help="")
parser.add_argument('--temperature', type=float, default=0.8, help="")
parser.add_argument('--N', type=int, default=200, help="")
parser.add_argument('--max_len', type=int, default=512, help="")
parser.add_argument('--num_gpus', type=int, default=4, help="")
parser.add_argument('--decoding_style', type=str, default='sampling', help="")
parser.add_argument('--num_seqs_per_iter', type=int, default=50, help='')
parser.add_argument('--overwrite', action='store_true', help='')
args = parser.parse_args()
argsdict = vars(args)
print(pprint.pformat(argsdict))
problems = read_problems()
task_ids = sorted(problems.keys())[args.start_index: args.end_index]
prompts = [problems[task_id]['prompt'] for task_id in task_ids]
num_samples = len(prompts)
print("Number of samples: {}".format(num_samples))
llm = LLM(base_model, tensor_parallel_size=args.num_gpus)
sampling_params = SamplingParams(temperature=args.temperature, top_p=1, max_tokens=args.max_len)
print(f"Loaded {args.model}.")
for i in tqdm(range(num_samples), ncols=0, total=num_samples):
output_file = args.output_path + '/{}.jsonl'.format(args.start_index + i)
if os.path.exists(output_file) and not args.overwrite:
print(f'Skip {output_file} as it already exists')
continue
prompt = prompts[i].replace(' ', '\t')
prompt_batch = [generate_prompt(prompt)]
ids_batch = [task_ids[i]]
completion_seqs = []
if args.decoding_style == 'sampling':
loops = int(args.N / args.num_seqs_per_iter)
else:
loops = 1
for _ in tqdm(range(loops), total=loops, leave=False, ncols=0):
with torch.no_grad():
completions = llm.generate(prompt_batch, sampling_params)
gen_seqs = [completions[0].outputs[0].text]
if gen_seqs is not None:
assert len(ids_batch) == 1
task_id = ids_batch[0]
for seq_idx, gen_seq in enumerate(gen_seqs):
completion_seq = gen_seq.split("### Response:")[-1]
completion_seq = completion_seq.replace('\t', ' ')
all_code = gen_seq.replace('\t', ' ')
completion_seqs.append(
{'task_id': task_id,
'completion': completion_seq,
'all_code': all_code,
}
)
print("Saving results to {}".format(output_file))
write_jsonl(output_file, completion_seqs)
if __name__ == '__main__':
main() |