|
|
|
|
|
|
|
|
|
|
|
from typing import Tuple |
|
|
import os |
|
|
import sys |
|
|
import torch |
|
|
import fire |
|
|
import time |
|
|
import json |
|
|
import pyarrow as pa |
|
|
|
|
|
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: |
|
|
start_time = time.time() |
|
|
arrow_dir = Path(ckpt_dir).expanduser() / 'arrow' |
|
|
|
|
|
if not arrow_dir.exists(): |
|
|
print('Converting checkpoints to arrow format') |
|
|
checkpoints = sorted(Path(ckpt_dir).expanduser().glob("*.pth")) |
|
|
for ckpt_file in checkpoints: |
|
|
print(ckpt_file) |
|
|
index = ckpt_file.parts[-1].split('.')[-2] |
|
|
|
|
|
ckpt = torch.load(ckpt_file, map_location='cuda') |
|
|
(arrow_dir / index).mkdir(parents=True, exist_ok=True) |
|
|
for k, v in ckpt.items(): |
|
|
tens = pa.Tensor.from_numpy(v.numpy()) |
|
|
with pa.output_stream(arrow_dir / index / k) as f: |
|
|
pa.ipc.write_tensor(tens, f) |
|
|
ckpt = None |
|
|
|
|
|
with open(Path(ckpt_dir) / "params.json", "r") as f: |
|
|
params = json.loads(f.read()) |
|
|
|
|
|
print("Loading checkpoint") |
|
|
segments = sorted((arrow_dir / '00').glob("*")) |
|
|
|
|
|
checkpoint = {} |
|
|
files = [] |
|
|
for seg in segments: |
|
|
f = pa.memory_map(str(seg)) |
|
|
files.append(f) |
|
|
t = pa.ipc.read_tensor(f).to_numpy() |
|
|
t = torch.from_numpy(t) |
|
|
checkpoint[seg.parts[-1]] = t |
|
|
|
|
|
|
|
|
torch.set_default_tensor_type(torch.BFloat16Tensor) |
|
|
|
|
|
|
|
|
model_args: ModelArgs = ModelArgs( |
|
|
max_seq_len=max_seq_len, max_batch_size=max_batch_size, **params |
|
|
) |
|
|
print("Loading tokenizer") |
|
|
tokenizer = Tokenizer(model_path=tokenizer_path) |
|
|
model_args.vocab_size = tokenizer.n_words |
|
|
print("Loading model") |
|
|
model = Transformer(model_args) |
|
|
|
|
|
checkpoints = sorted(Path(ckpt_dir).glob("*.pth")) |
|
|
model.load_state_dict(torch.load(checkpoints[-1]), strict=False) |
|
|
|
|
|
for f in files: |
|
|
f.close() |
|
|
files = None |
|
|
|
|
|
generator = LLaMA(model, tokenizer) |
|
|
print(f"Loaded in {time.time() - start_time:.2f} seconds") |
|
|
return generator |
|
|
|
|
|
|
|
|
def main( |
|
|
ckpt_dir: str, |
|
|
tokenizer_path: str, |
|
|
temperature: float = 0.8, |
|
|
top_p: float = 0.95, |
|
|
max_seq_len: int = 2048, |
|
|
max_batch_size: int = 1, |
|
|
): |
|
|
generator = load(ckpt_dir, tokenizer_path, max_seq_len, max_batch_size) |
|
|
|
|
|
ctx = """A dialog, where User interacts with AI. AI is helpful, kind, obedient, honest, and knows its own limits. |
|
|
User: Hello, AI. |
|
|
AI: Hello! How can I assist you today? |
|
|
""" |
|
|
|
|
|
while True: |
|
|
prompt = input(f'User: ') |
|
|
if ctx != "": |
|
|
ctx = ctx + "User: " + prompt + "\n" |
|
|
else: |
|
|
ctx = prompt + "\n" |
|
|
|
|
|
ctx = (ctx[-1920:]) if len(ctx) >= 2048 else ctx |
|
|
|
|
|
if len(ctx.strip()) > 0: |
|
|
prompts = [ctx] |
|
|
results = generator.generate( |
|
|
prompts, max_gen_len=2048, temperature=temperature, top_p=top_p |
|
|
) |
|
|
ctx = results[0] |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
fire.Fire(main) |
|
|
|
