forked from abacaj/code-eval
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy patheval_any.py
84 lines (70 loc) · 2.22 KB
/
eval_any.py
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
from transformers import (
AutoTokenizer,
AutoModelForCausalLM,
PreTrainedModel,
PreTrainedTokenizer,
)
from core import filter_code, run_eval, fix_indents
import os
import torch
# TODO: move to python-dotenv
# add hugging face access token here
TOKEN = ""
@torch.inference_mode()
def generate_batch_completion(
model: PreTrainedModel, tokenizer: PreTrainedTokenizer, prompt, batch_size
) -> list[str]:
input_batch = [prompt for _ in range(batch_size)]
inputs = tokenizer(input_batch, return_token_type_ids=False, return_tensors="pt").to(model.device)
input_ids_cutoff = inputs.input_ids.size(dim=1)
generated_ids = model.generate(
**inputs,
use_cache=True,
max_new_tokens=512,
temperature=0.2,
top_p=0.95,
do_sample=True,
eos_token_id=tokenizer.eos_token_id,
pad_token_id=tokenizer.eos_token_id, # model has no pad token
)
batch_completions = tokenizer.batch_decode(
[ids[input_ids_cutoff:] for ids in generated_ids],
skip_special_tokens=True,
)
return [filter_code(fix_indents(completion)) for completion in batch_completions]
def test_model(name="Linksoul-llama2-7b",model_path="/home/Linksoul-llama2-7b"):
num_samples_per_task = 10
out_path = f"results/{name}/eval.jsonl"
os.makedirs(f"results/{name}", exist_ok=True)
tokenizer = AutoTokenizer.from_pretrained(
model_path,
)
model = torch.compile(
AutoModelForCausalLM.from_pretrained(
model_path,
torch_dtype=torch.bfloat16,
device_map="auto"
)
.eval()
)
run_eval(
model,
tokenizer,
num_samples_per_task,
out_path,
generate_batch_completion,
True,
)
if __name__ == "__main__":
# model_path="/home/Linksoul-llama2-7b"
# name="Linksoul-llama2-7b"
# test_model(name,model_path)
# model_path="codellama/CodeLlama-7b-hf"
# name="CodeLlama-7b-hf"
# test_model(name,model_path)
# model_path="codellama/CodeLlama-7b-Python-hf"
# name="CodeLlama-7b-Python-hf"
# test_model(name,model_path)
model_path="/mnt/SFT_store/LLM/llama-65b-hf"
name="llama-65b-hf"
test_model(name,model_path)