-
Notifications
You must be signed in to change notification settings - Fork 312
/
inference.py
89 lines (75 loc) · 2.47 KB
/
inference.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
85
86
87
88
89
import torch
import json
from pathlib import Path
from llama import ModelArgs, Transformer, Tokenizer, LLaMA
def load(
ckpt_dir: str,
tokenizer_path: str,
local_rank: int,
world_size: int,
max_seq_len: int,
max_batch_size: int,
) -> LLaMA:
checkpoints = sorted(Path(ckpt_dir).glob("*.pth"))
assert world_size == len(
checkpoints
), f"Loading a checkpoint for MP={len(checkpoints)} but world size is {world_size}"
ckpt_path = checkpoints[local_rank]
checkpoint = torch.load(ckpt_path, map_location="cpu")
with open(Path(ckpt_dir) / "params.json", "r") as f:
params = json.loads(f.read())
model_args: ModelArgs = ModelArgs(
max_seq_len=max_seq_len, max_batch_size=max_batch_size, **params
)
tokenizer = Tokenizer(model_path=tokenizer_path)
model_args.vocab_size = tokenizer.n_words
torch.set_default_tensor_type(torch.cuda.HalfTensor)
model = Transformer(model_args)
torch.set_default_tensor_type(torch.FloatTensor)
model.load_state_dict(checkpoint, strict=False)
generator = LLaMA(model, tokenizer)
return generator
def run(
ckpt_dir: str,
tokenizer_path: str,
temperature: float = 0.8,
top_p: float = 0.95,
max_seq_len: int = 1024,
max_batch_size: int = 1,
):
local_rank = 0
world_size = 1
generator = load(
ckpt_dir, tokenizer_path, local_rank, world_size, max_seq_len, max_batch_size
)
prompts = [
# For these prompts, the expected answer is the natural continuation of the prompt
"I believe the meaning of life is", # removed: keep only one prompt
]
while True:
print("Prompt:", prompts)
results = generator.generate(
prompts, max_gen_len=256, temperature=temperature, top_p=top_p
)
for result in results:
print("🦙LLaMA:", result.strip())
user_input = input("please enter your prompts (Ctrl+C to exit): ")
prompts = [user_input]
def get_args():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--ckpt_dir", type=str, default="/llama_data/7B")
parser.add_argument(
"--tokenizer_path", type=str, default="/llama_data/tokenizer.model"
)
return parser.parse_args()
if __name__ == "__main__":
args = get_args()
run(
ckpt_dir=args.ckpt_dir,
tokenizer_path=args.tokenizer_path,
temperature=0.8,
top_p=0.95,
max_seq_len=1024,
max_batch_size=1,
)