File size: 3,250 Bytes
54cb090 | 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 | """
Chat with a released Crawlnet queen from its folder. Copied into every release as chat.py.
pip install torch tiktoken safetensors
git clone https://github.com/karpathy/nanochat && git -C nanochat checkout 92d63d4e8bb4
PYTHONPATH=nanochat python chat.py "What is staking?" [--dir .] [--temperature 0.7] [--top-k 50] [--max-tokens 200]
Needs only nanochat (MIT) for the model and tokenizer classes; no Crawlnet code.
"""
import argparse
import base64
import json
import os
import tiktoken
import torch
from safetensors.torch import load_file
from nanochat.engine import Engine
from nanochat.gpt import GPT, GPTConfig
from nanochat.tokenizer import RustBPETokenizer
def load(folder, device):
with open(os.path.join(folder, "config.json")) as f:
cfg = json.load(f)
with open(os.path.join(folder, "tokenizer.json")) as f:
tj = json.load(f)
ranks = {}
with open(os.path.join(folder, "tokenizer.tiktoken")) as f:
for line in f:
tok, rank = line.split()
ranks[base64.b64decode(tok)] = int(rank)
enc = tiktoken.Encoding(name="queen", pat_str=tj["pattern"], mergeable_ranks=ranks, special_tokens=tj["special_tokens"])
tokenizer = RustBPETokenizer(enc, tj["bos_token"])
with torch.device("meta"):
model = GPT(GPTConfig(**cfg["model_config"]))
model.to_empty(device=device)
model.init_weights() # rotary buffers; the weights are overwritten below
state = load_file(os.path.join(folder, "model.safetensors"), device="cpu")
if device != "cuda": # bf16 weights run in fp32 on CPU
state = {k: (v.float() if v.is_floating_point() else v) for k, v in state.items()}
model.load_state_dict({k: v.to(device) for k, v in state.items()}, strict=True, assign=True)
model.eval()
return model, tokenizer
def answer(model, tokenizer, question, max_tokens=200, temperature=0.7, top_k=50, seed=42):
sp = tokenizer.encode_special
bos = tokenizer.get_bos_token_id()
prompt = [bos, sp("<|user_start|>")] + tokenizer.encode(question) + [sp("<|user_end|>"), sp("<|assistant_start|>")]
stop = {bos, sp("<|assistant_end|>"), sp("<|user_start|>")}
out = []
for column, _ in Engine(model, tokenizer).generate(prompt, num_samples=1, max_tokens=max_tokens,
temperature=temperature, top_k=top_k, seed=seed):
if column[0] in stop:
break
out.append(column[0])
return tokenizer.decode(out).strip()
def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("question")
ap.add_argument("--dir", default=os.path.dirname(os.path.abspath(__file__)))
ap.add_argument("--max-tokens", type=int, default=200)
ap.add_argument("--temperature", type=float, default=0.7)
ap.add_argument("--top-k", type=int, default=50)
ap.add_argument("--seed", type=int, default=42)
args = ap.parse_args()
device = "cuda" if torch.cuda.is_available() else "cpu"
model, tokenizer = load(args.dir, device)
print(answer(model, tokenizer, args.question, args.max_tokens, args.temperature, args.top_k, args.seed))
if __name__ == "__main__":
main()
|