141 lines
4.0 KiB
Python
Executable File
141 lines
4.0 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
|
|
import threading
|
|
|
|
|
|
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):
|
|
""" run model on gpu """
|
|
return model(input_buffer)
|
|
|
|
def inference_worker():
|
|
""" consume tasks from que """
|
|
BatchSize=2
|
|
NewList = [None] * BatchSize
|
|
|
|
import time
|
|
while True:
|
|
if (not msg_q.empty() and None in NewList) or NewList.count(None) == len(NewList):
|
|
i = NewList.index(None)
|
|
out,inp = msg_q.get()
|
|
NewList[i] = (out,inp,None)
|
|
|
|
batch = []
|
|
for i in range(BatchSize):
|
|
t = None
|
|
if not NewList[i]:
|
|
t = Tensor.zeros(hypr['block_size'])
|
|
else:
|
|
_, t, _ = NewList[i]
|
|
if not isinstance(t, Tensor):
|
|
t = Tensor(t)
|
|
l = t.shape[0]
|
|
pad_len = hypr['block_size'] - l
|
|
a,b,_ = NewList[i]
|
|
NewList[i] = (a,t,l)
|
|
|
|
t = t.pad((0,pad_len))
|
|
else:
|
|
#t = t[:-hypr['block_size']]
|
|
l = t.shape[0]
|
|
pad_len = hypr['block_size'] - l
|
|
t = t.pad((0,pad_len))
|
|
|
|
|
|
batch.append(t)
|
|
chat_tensor = batch[0].stack(*batch[1:])
|
|
#infince here
|
|
|
|
logits = model(chat_tensor)
|
|
|
|
#return
|
|
for i in range(BatchSize):
|
|
if NewList[i] is None:
|
|
continue
|
|
out, t, lenth = NewList[i]
|
|
if lenth < 15:
|
|
tok = (logits[i, lenth-1, :] / 0.7).softmax().multinomial(1)
|
|
inp = t.cat(tok)
|
|
out.put(tok.numpy()[0])
|
|
NewList[i] = (out,inp,(lenth+1))
|
|
else:
|
|
print(encoding.decode(chat_tensor[i].numpy().astype(int))[:25])
|
|
out.shutdown()
|
|
NewList[i] = None
|
|
|
|
|
|
def warmup(count):
|
|
""" run count times with random data """
|
|
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 test(msg):
|
|
tokens = queue.Queue()
|
|
inp = encoding.encode(msg)
|
|
t = []
|
|
|
|
msg_q.put((tokens,inp))
|
|
yield("Start:")
|
|
while True:
|
|
try:
|
|
i = tokens.get()
|
|
t.append(i)
|
|
yield(f"{i},")
|
|
except:
|
|
break
|
|
txt = encoding.decode(t)
|
|
yield f"\n{txt}"
|
|
return
|
|
|
|
app = flask.Flask(__name__)
|
|
|
|
from flask import request
|
|
@app.route('/',methods=['POST'])
|
|
def complete():
|
|
user_string = request.form.get('input', 'Default prompt')
|
|
return test(user_string),{"Content-Type": "text"}
|
|
|
|
def apiStart():
|
|
""" start api """
|
|
app.run()
|
|
pass
|
|
|
|
if __name__ == "__main__":
|
|
print(Device.DEFAULT)
|
|
print("warming up")
|
|
#warmup(200)
|
|
t = threading.Thread(target=apiStart, daemon=True)
|
|
t.start()
|
|
inference_worker()
|