|
|
|
|
|
|
|
|
|
|
|
from typing import Tuple |
|
|
import os |
|
|
import sys |
|
|
import torch |
|
|
import fire |
|
|
import time |
|
|
import json |
|
|
from pathlib import Path |
|
|
from llama import ModelArgs, Transformer, Tokenizer, LLaMA |
|
|
|
|
|
|
|
|
def load( |
|
|
ckpt_dir: str, |
|
|
tokenizer_path: str, |
|
|
max_seq_len: int, |
|
|
max_batch_size: int, |
|
|
) -> LLaMA: |
|
|
print("Creating model...") |
|
|
start_time = time.time() |
|
|
checkpoints = sorted(Path(ckpt_dir).glob("*.pth")) |
|
|
|
|
|
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 |
|
|
|
|
|
model = Transformer(model_args) |
|
|
model.to("cpu") |
|
|
|
|
|
print("Loading merged checkpoint...") |
|
|
checkpoint = torch.load(checkpoints[-1], map_location="cuda") |
|
|
model.load_state_dict(checkpoint, strict=False) |
|
|
del checkpoint |
|
|
|
|
|
generator = LLaMA(model, tokenizer) |
|
|
print(f"Loaded model in {time.time() - start_time:.2f} seconds") |
|
|
return generator |
|
|
|
|
|
|
|
|
def main( |
|
|
ckpt_dir: str = './model', |
|
|
tokenizer_path: str = './tokenizer/tokenizer.model', |
|
|
temperature: float = 0.8, |
|
|
top_p: float = 0.95, |
|
|
max_seq_len: int = 256, |
|
|
max_batch_size: int = 5, |
|
|
): |
|
|
|
|
|
torch.set_default_dtype(torch.bfloat16) |
|
|
|
|
|
generator = load(ckpt_dir, tokenizer_path, max_seq_len, max_batch_size) |
|
|
|
|
|
while True: |
|
|
prompt = input(f'prompt> ') |
|
|
if len(prompt.strip()) > 0: |
|
|
prompts = [prompt] |
|
|
results = generator.generate( |
|
|
prompts, max_gen_len=256, temperature=temperature, top_p=top_p |
|
|
) |
|
|
|
|
|
for result in results: |
|
|
print(result) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
fire.Fire(main) |
|
|
|
