60 lines
1.7 KiB
Python
Executable File
60 lines
1.7 KiB
Python
Executable File
#! /usr/bin/env nix-shell
|
|
#! nix-shell -i python3 -p python3Packages.tinygrad python3Packages.numpy python3Packages.discordpy python3Packages.transformers python3Packages.tqdm python3Packages.flask
|
|
|
|
import queue
|
|
import flask
|
|
from tinygrad import Tensor, TinyJit, dtypes, Device
|
|
from tinygrad.nn.state import safe_load, load_state_dict
|
|
from transformers import AutoTokenizer
|
|
from model import Transformer
|
|
from tqdm import tqdm
|
|
|
|
hypr = {
|
|
"embed_size": 768, "n_heads": 8, "n_blocks": 12, "block_size": 512,
|
|
"encoding": "TinyLlama/TinyLlama_v1.1"
|
|
}
|
|
|
|
CHECKPOINT_PATH = 'gpt.safetensors'
|
|
|
|
msg_q = queue.Queue()
|
|
|
|
encoding = AutoTokenizer.from_pretrained(hypr['encoding'])
|
|
model = Transformer(encoding.vocab_size, hypr["embed_size"], hypr["n_heads"], hypr["n_blocks"], hypr["block_size"])
|
|
load_state_dict(model, safe_load(CHECKPOINT_PATH))
|
|
Tensor.training = False
|
|
|
|
|
|
@TinyJit
|
|
def run_model(input_buffer):
|
|
return model(input_buffer)
|
|
|
|
def inference_worker():
|
|
""" Runs in a separate thread to handle the heavy lifting. """
|
|
pass
|
|
|
|
|
|
def warmup(count):
|
|
import random
|
|
tokens = encoding.encode("")
|
|
tokens = Tensor([tokens])
|
|
for i in tqdm(range(count)):
|
|
pad_len = hypr['block_size'] - tokens.shape[1]
|
|
input_buffer = tokens.pad(((0, 0), (0, pad_len))).contiguous()
|
|
out = model(input_buffer)
|
|
token_tensor = (out[:, tokens.shape[1] - 1, :] / 0.7).softmax().multinomial(1)
|
|
tokens = tokens.cat(token_tensor, dim=1).realize()
|
|
tokens = tokens[:-hypr['block_size']]
|
|
|
|
|
|
def apiStart():
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print(Device.DEFAULT)
|
|
print("warming up")
|
|
warmup(200)
|
|
t = threading.Thread(target=apiStart, daemon=True)
|
|
t.start()
|
|
inference_worker()
|