Init

bot.py

@@ -0,0 +1,59 @@
+#! /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()