- Notifications
You must be signed in to change notification settings - Fork 2
Expand file tree
/
Copy pathpredict.py
More file actions
Latest commit
87 lines (69 loc) · 3.06 KB
/
Copy pathpredict.py
File metadata and controls
87 lines (69 loc) · 3.06 KB
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
#!/opt/software/install/miniconda37/bin/python
importargparse
parser=argparse.ArgumentParser(description='Running a chatbot with gradio')
parser.add_argument('--model_path', type=str, help='model path, for example `llama/7b-32`')
parser.add_argument('--input_file', type=str, help='input file')
parser.add_argument('--device', type=str, help='device (default: cuda:0)', default='cuda:0')
args=parser.parse_args()
model_name=args.model_path.replace("/", "_") # llama/7b-32 --> llama_7b-32
offset=-4
if'llama'inmodel_name:
offset=-7
fromtransformersimport (
AutoModelForCausalLM,
AutoTokenizer,
PreTrainedModel,
PreTrainedTokenizer
)
importjson
importpandasaspd
fromtqdmimporttqdm
importsys
tokenizer=AutoTokenizer.from_pretrained(args.model_path)
model=AutoModelForCausalLM.from_pretrained(args.model_path)
model=model.half()
model=model.to(args.device) #, device_map='auto')
#model = AutoModelForCausalLM.from_pretrained(args.model_path, device_map='balanced')
model.eval()
PROMPT_DICT= {
"prompt_input": (
"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{input}\n\n### Response:"
),
"prompt_no_input": (
"Below is an instruction that describes a task. "
"Write a response that appropriately completes the request.\n\n"
"### Instruction:\n{instruction}\n\n### Response:"
),
}
PROMPT_FORMAT="""Below is an instruction that describes a task. Write a response that appropriately completes the request.
### Instruction:
{instruction}
### Response:
"""
defgenerate_response(instruction: str, input_text: str, **kwargs) ->str:
#input_ids = tokenizer(PROMPT_FORMAT.format(instruction=instruction), return_tensors="pt").input_ids.to(device)
input_ids=tokenizer(PROMPT_DICT['prompt_input'].format(instruction=instruction, input=input_text), return_tensors="pt").input_ids.to(args.device)
# each of these is encoded to a single token
response_key_token_id=tokenizer.encode("### Response:")[0]
end_key_token_id=tokenizer.encode("### End")[0]
gen_tokens=model.generate(input_ids, pad_token_id=tokenizer.pad_token_id, eos_token_id=end_key_token_id,
do_sample=False, max_new_tokens=2048, top_p=0.92, top_k=0, **kwargs)[0].cpu()
s=tokenizer.decode(gen_tokens)
ss=s.split("### Response:")[1].strip()[0:offset]
returnss
#with open('data/test_data_points-v2-128.json') as f:
withopen(args.input_file) asf:
d=json.load(f)
#x = pd.read_csv('data/test_data_points.csv.gz')
#assert len(d) == len(x)
results= []
foraintqdm(d):
response=generate_response(a['instruction'], a['input'])
a['response'] =response
results.append(a)
#with open('data/test_data_points-v2-7b-128-predictions.json', 'w') as f:
outfile=args.input_file.replace('.json', f'-{model_name}-predictions.json')
withopen(outfile, 'w') asf:
json.dump(results, f, indent=2)