new ui

Dockerfile

@@ -64,7 +64,9 @@ USER user
 WORKDIR $HOME/app
 
 COPY --chown=user . $HOME/app
-RUN wget -c  https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin
+RUN wget -c https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin
+RUN wget -c https://huggingface.co/karpathy/tinyllamas/resolve/main/stories42M.bin
+RUN wget -c https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.bin
 
 # CMD ["mojo", "llama2.mojo"]
 CMD ["python3", "gradio_app.py"]

gradio_app.py

@@ -1,36 +1,86 @@
 import gradio as gr
 import subprocess
 import sys
-import os
+from pathlib import Path
 
 
-async def generate(prompt):
-    # os.environ["PROMPT"] = prompt
+async def generate(prompt, model_name, seed=0, temperature=0.5, num_tokens=256):
     # stream stout
     process = subprocess.Popen(
-        ["mojo", "llama2.mojo"], stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        [
+            "mojo",
+            "llama2.mojo",
+            Path(model_name),
+            "-s",
+            str(seed),
+            "-n",
+            str(num_tokens),
+            "-t",
+            str(temperature),
+            "-i",
+            prompt,
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
     )
     text = ""
     for char in iter(lambda: process.stdout.read(1), b""):
-        char_decoded = char.decode()
-        sys.stdout.write(char_decoded)
+        char_decoded = char.decode("utf-8", errors="ignore")
         text += char_decoded
         yield text
 
 
-output_text = gr.Textbox(label="Generated Text")
-
-demo = gr.Interface(
-    fn=generate,
-    inputs=None,
-    outputs=output_text,
-    description="""
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
 # llama2.🔥
 ## [Mojo](https://docs.modular.com/mojo/) implementation of [llama2.c](https://github.com/karpathy/llama2.c) by [@tairov](https://github.com/tairov)
 Source: https://github.com/tairov/llama2.mojo
-    """,
-    allow_flagging="never",
-)
+    """
+    )
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", placeholder="Add your prompt here...")
+            seed = gr.Slider(
+                minimum=0,
+                maximum=2**53,
+                value=0,
+                step=1,
+                label="Seed",
+                randomize=True,
+            )
+            temperature = gr.Slider(
+                minimum=0.0, maximum=2.0, step=0.01, value=0.5, label="Temperature"
+            )
+            num_tokens = gr.Slider(
+                minimum=1, maximum=256, value=256, label="Number of tokens"
+            )
+            model_name = gr.Dropdown(
+                ["stories15M.bin", "stories42M.bin", "stories110M.bin"],
+                value="stories15M.bin",
+                label="Model Size",
+            )
+            with gr.Row():
+                stop = gr.Button("Stop")
+                run = gr.Button("Run")
+        with gr.Column(scale=2):
+            output_text = gr.Textbox(label="Generated Text")
+
+    # update maximum number of tokens based on model size
+    model_name.change(
+        lambda x: gr.update(maximum=1024)
+        if x == "stories110M.bin" or x == "stories42M.bin"
+        else gr.update(maximum=256),
+        model_name,
+        num_tokens,
+        queue=False,
+    )
+    click_event = run.click(
+        fn=generate,
+        inputs=[prompt, model_name, seed, temperature, num_tokens],
+        outputs=output_text,
+    )
+    stop.click(fn=None, inputs=None, outputs=None, cancels=[click_event])
 
 demo.queue()
 demo.launch(server_name="0.0.0.0")

llama2.mojo

@@ -1,25 +1,22 @@
+from algorithm import sum
+from algorithm import vectorize, parallelize
+from builtin import string
 from math import round
-import math
-
 from memory import memset_zero, memcpy
+from memory.buffer import Buffer
 from memory.unsafe import DTypePointer
+from python import Python
 from random import rand
-from sys.info import simdwidthof
-from builtin import string
-import time
-import random
-import os
-
-from runtime.llcl import num_cores
-
 from read import BufReader, File
-from memory.buffer import Buffer
-
-from python import Python
+from runtime.llcl import num_cores, Runtime
+from sys import argv
 
 # The SIMD vector width.
-from algorithm import vectorize, parallelize
-from algorithm import sum
+from sys.info import simdwidthof
+import math
+import os
+import random
+import time
 
 alias nelts = (2 * simdwidthof[DType.float32]())
 
@@ -29,98 +26,51 @@ alias BufferPtrFloat32 = DTypePointer[DType.float32]
 alias PointerStrings = Pointer[PointerString]
 
 
-struct Matrix3:
-    var data: BufferPtrFloat32
-    var rows: Int
-    var cols: Int
-    var layers: Int
-    var allocated: Int
-
-    fn __init__(inout self, layers: Int, rows: Int, cols: Int):
-        self.data = BufferPtrFloat32.alloc(0)
-        self.rows = rows
-        self.cols = cols
-        self.layers = layers
-        self.allocated = 0
-
-    @always_inline
-    fn alloc(inout self, fill: Int = 0):
-        self.data = BufferPtrFloat32.alloc(self.size())
-        self.allocated = 1
-        if fill == 1:
-            self.zero()
-
-    @always_inline
-    fn alloc_zero(inout self):
-        self.alloc(1)
-
-    @always_inline
-    fn set_buf_ptr(inout self, ptr: BufferPtrFloat32):
-        self.data = ptr
-
-    fn __del__(owned self):
-        if self.allocated == 1:
-            self.data.free()
-
-    @always_inline
-    fn zero(inout self):
-        memset_zero(self.data, self.layers * self.rows * self.cols)
-
-    @always_inline
-    fn size(inout self) -> Int:
-        return self.layers * self.cols * self.rows
-
-    @always_inline
-    fn __getitem__(self, z: Int, y: Int, x: Int) -> Float32:
-        return self.load[1](z, y, x)
-
-    @always_inline
-    fn load[nelts: Int](self, z: Int, y: Int, x: Int) -> SIMD[DType.float32, nelts]:
-        return self.data.simd_load[nelts](z * self.layers + y * self.cols + x)
-
-    @always_inline
-    fn __setitem__(self, z: Int, y: Int, x: Int, val: Float32):
-        return self.store[1](z, y, x, val)
-
-    @always_inline
-    fn store[nelts: Int](self, z: Int, y: Int, x: Int, val: SIMD[DType.float32, nelts]):
-        self.data.simd_store[nelts](z * self.layers + y * self.cols + x, val)
-
-
 struct Matrix:
     var data: BufferPtrFloat32
     var rows: Int
     var cols: Int
+    var layers: Int
     var allocated: Int
 
     fn __init__(inout self, rows: Int, cols: Int):
         self.data = BufferPtrFloat32.alloc(0)
         self.rows = rows
         self.cols = cols
+        self.layers = 1
         self.allocated = 0
 
     fn __init__(inout self, cols: Int):
         self.data = BufferPtrFloat32.alloc(0)
         self.rows = 1
+        self.layers = 1
         self.cols = cols
         self.allocated = 0
 
+    fn __init__(inout self, layers: Int, rows: Int, cols: Int):
+        self.__init__(rows, cols)
+        self.layers = layers
+
     fn __del__(owned self):
         if self.allocated == 1:
             self.data.free()
 
+    @always_inline
     fn alloc(inout self, fill: Int = 0):
         self.data = BufferPtrFloat32.alloc(self.size())
         self.allocated = 1
         if fill == 1:
             self.zero()
 
+    @always_inline
     fn alloc_zero(inout self):
         self.alloc(1)
 
+    @always_inline
     fn zero(inout self):
-        memset_zero(self.data, self.rows * self.cols)
+        memset_zero(self.data, self.size())
 
+    @always_inline
     fn set_buf_ptr(inout self, ptr: BufferPtrFloat32):
         self.data = ptr
 
@@ -130,8 +80,9 @@ struct Matrix:
         self.rows = rows
         self.cols = cols
 
+    @always_inline
     fn size(inout self) -> Int:
-        return self.cols * self.rows
+        return self.cols * self.rows * self.layers
 
     @always_inline
     fn __getitem__(self, y: Int, x: Int) -> Float32:
@@ -165,6 +116,22 @@ struct Matrix:
     fn store[nelts: Int](self, x: Int, val: SIMD[DType.float32, nelts]):
         self.data.simd_store[nelts](x, val)
 
+    @always_inline
+    fn __getitem__(self, z: Int, y: Int, x: Int) -> Float32:
+        return self.load[1](z, y, x)
+
+    @always_inline
+    fn load[nelts: Int](self, z: Int, y: Int, x: Int) -> SIMD[DType.float32, nelts]:
+        return self.data.simd_load[nelts](z * self.layers + y * self.cols + x)
+
+    @always_inline
+    fn __setitem__(self, z: Int, y: Int, x: Int, val: Float32):
+        return self.store[1](z, y, x, val)
+
+    @always_inline
+    fn store[nelts: Int](self, z: Int, y: Int, x: Int, val: SIMD[DType.float32, nelts]):
+        self.data.simd_store[nelts](z * self.layers + y * self.cols + x, val)
+
 
 fn read_val_int(inout buf: FileBuf) -> Int:
     # DTypePointer[DType.ui8](buf.data).bitcast[DType.ui8]()
@@ -191,6 +158,31 @@ fn read_val_str(inout buf: FileBuf, slen: Int) -> PointerString:
     return str
 
 
+# not optimal concat
+fn str_concat(s1: PointerString, s2: PointerString) -> PointerString:
+    var l1 = 0
+    var l2 = 0
+
+    while s1[l1] != 0:
+        l1 += 1
+    while s2[l2] != 0:
+        l2 += 1
+
+    let str = PointerString.alloc(l1 + l2)
+    memcpy[UInt8](str, s1, l1)
+    memcpy[UInt8](str.offset(l1), s2, l2)
+    str.store(l1 + l2, 0)
+    return str
+
+
+fn str_to_ptr(s: String) -> PointerString:
+    let ret = PointerString.alloc(len(s) + 1)
+    for i in range(len(s)):
+        ret.store(i, ord(s[i]))
+    ret.store(len(s), 0)
+    return ret
+
+
 struct FileBuf:
     var data: BufferPtrType
     var offset: Int
@@ -253,8 +245,9 @@ struct RunState:
     var v: Matrix  # value (dim,)
     var att: Matrix  # buffer for scores/attention values (n_heads, seq_len)
     var logits: Matrix  # output logits
-    var key_cache: Matrix3  # (layer, seq_len, dim)
-    var value_cache: Matrix3  # (layer, seq_len, dim)
+    var key_cache: Matrix  # (layer, seq_len, dim)
+    var value_cache: Matrix  # (layer, seq_len, dim)
+    var rt: Runtime
 
     fn __init__(inout self, config: Config):
         self.x = Matrix(config.dim)
@@ -277,10 +270,11 @@ struct RunState:
         self.att.alloc_zero()
         self.logits = Matrix(config.vocab_size)
         self.logits.alloc_zero()
-        self.key_cache = Matrix3(config.n_layers, config.seq_len, config.dim)
+        self.key_cache = Matrix(config.n_layers, config.seq_len, config.dim)
         self.key_cache.alloc_zero()
-        self.value_cache = Matrix3(config.n_layers, config.seq_len, config.dim)
+        self.value_cache = Matrix(config.n_layers, config.seq_len, config.dim)
         self.value_cache.alloc_zero()
+        self.rt = Runtime(num_cores() // 2)
 
 
 struct TransformerWeights:
@@ -288,14 +282,14 @@ struct TransformerWeights:
     var freq_cis_real: Matrix
     var freq_cis_imag: Matrix
     var rms_att_weight: Matrix
-    var wq: Matrix3
-    var wk: Matrix3
-    var wv: Matrix3
-    var wo: Matrix3
+    var wq: Matrix
+    var wk: Matrix
+    var wv: Matrix
+    var wo: Matrix
     var rms_ffn_weight: Matrix
-    var w1: Matrix3
-    var w3: Matrix3
-    var w2: Matrix3
+    var w1: Matrix
+    var w3: Matrix
+    var w2: Matrix
     var rms_final_weight: Matrix
     var wcls: Matrix
 
@@ -309,23 +303,23 @@ struct TransformerWeights:
         self.rms_att_weight.set_buf_ptr(
             buf.bitcast_offset_float32(self.rms_att_weight.size())
         )
-        self.wq = Matrix3(config.n_layers, config.dim, config.dim)
+        self.wq = Matrix(config.n_layers, config.dim, config.dim)
         self.wq.set_buf_ptr(buf.bitcast_offset_float32(self.wq.size()))
-        self.wk = Matrix3(config.n_layers, config.dim, config.dim)
+        self.wk = Matrix(config.n_layers, config.dim, config.dim)
         self.wk.set_buf_ptr(buf.bitcast_offset_float32(self.wk.size()))
-        self.wv = Matrix3(config.n_layers, config.dim, config.dim)
+        self.wv = Matrix(config.n_layers, config.dim, config.dim)
         self.wv.set_buf_ptr(buf.bitcast_offset_float32(self.wv.size()))
-        self.wo = Matrix3(config.n_layers, config.dim, config.dim)
+        self.wo = Matrix(config.n_layers, config.dim, config.dim)
         self.wo.set_buf_ptr(buf.bitcast_offset_float32(self.wo.size()))
         self.rms_ffn_weight = Matrix(config.n_layers, config.dim)
         self.rms_ffn_weight.set_buf_ptr(
             buf.bitcast_offset_float32(self.rms_ffn_weight.size())
         )
-        self.w1 = Matrix3(config.n_layers, config.dim, config.hidden_dim)
+        self.w1 = Matrix(config.n_layers, config.dim, config.hidden_dim)
         self.w1.set_buf_ptr(buf.bitcast_offset_float32(self.w1.size()))
-        self.w2 = Matrix3(config.n_layers, config.dim, config.hidden_dim)
+        self.w2 = Matrix(config.n_layers, config.dim, config.hidden_dim)
         self.w2.set_buf_ptr(buf.bitcast_offset_float32(self.w2.size()))
-        self.w3 = Matrix3(config.n_layers, config.dim, config.hidden_dim)
+        self.w3 = Matrix(config.n_layers, config.dim, config.hidden_dim)
         self.w3.set_buf_ptr(buf.bitcast_offset_float32(self.w3.size()))
         self.rms_final_weight = Matrix(config.dim)
         self.rms_final_weight.set_buf_ptr(
@@ -435,22 +429,14 @@ fn softmax(inout x: BufferPtrFloat32, size: Int) -> None:
         x.offset(i).simd_store[1](0, xi / ssum)
 
 
-fn matmul_naive(C: Matrix, x: Matrix, w: Matrix) -> None:
-    # W(d,n) @ X(n,) -> C (d,)
-    # By far the most amount of time is spent inside this little function
-    for i in range(w.rows):
-        C[i] = 0.0
-        for j in range(w.cols):
-            C[i] += x[j] * w[i, j]
-
-
-fn matmul_vectorized(C: Matrix, A: Matrix, B: Matrix):
-    for i in range(0, B.rows):
+fn matmul_parallelized(C: Matrix, A: Matrix, B: Matrix, rt: Runtime):
+    @parameter
+    fn compute_row(i: Int):
         var tmp = SIMD[DType.float32, nelts](0)
 
         @parameter
         fn dot[_nelts: Int](j: Int):
-            if _nelts < nelts: # take care of tail array elements with length <  nelts
+            if _nelts < nelts:  # take care of tail array elements with length <  nelts
                 tmp[0] += (A.load[_nelts](j) * B.load[_nelts](i, j)).reduce_add()
             else:
                 tmp += A.load[nelts](j) * B.load[nelts](i, j)
@@ -458,28 +444,12 @@ fn matmul_vectorized(C: Matrix, A: Matrix, B: Matrix):
         vectorize[nelts, dot](B.cols)
         C[i] = tmp.reduce_add()
 
-fn matmul_parallelized(C: Matrix, A: Matrix, B: Matrix):
-    @parameter
-    fn calc_row(i: Int):
-        var T = BufferPtrFloat32.alloc(nelts)
-        var Tbuf = Buffer[nelts, DType.float32](T)
-        memset_zero(T, nelts)
-        @parameter
-        fn dot[nelts: Int](j: Int):
-            T.simd_store[nelts](
-                0, T.simd_load[nelts](0) + A.load[nelts](j) * B.load[nelts](i, j)
-            )
-
-        vectorize[nelts, dot](B.cols)
-        C[i] = sum[nelts, DType.float32](Tbuf)
+    parallelize[compute_row](rt, B.rows, rt.parallelism_level())
 
-    parallelize[calc_row](B.rows)
 
-
-fn matmul(inout C: Matrix, A: Matrix, B: Matrix) -> None:
+fn matmul(inout C: Matrix, A: Matrix, B: Matrix, rt: Runtime) -> None:
     # B (d,n) @ A (n,) -> C (d,)
-    matmul_vectorized(C, A, B)
-    # matmul_parallelized(C, A, B)
+    matmul_parallelized(C, A, B, rt)
 
 
 fn transformer(
@@ -513,13 +483,13 @@ fn transformer(
 
         # QKV matmuls for this position
         tmpw.set_buf_ptr(weights.wq.data.offset(l * dim * dim), dim, dim)
-        matmul(state.q, state.xb, tmpw)
+        matmul(state.q, state.xb, tmpw, state.rt)
 
         tmpw.set_buf_ptr(weights.wk.data.offset(l * dim * dim), dim, dim)
-        matmul(state.k, state.xb, tmpw)
+        matmul(state.k, state.xb, tmpw, state.rt)
 
         tmpw.set_buf_ptr(weights.wv.data.offset(l * dim * dim), dim, dim)
-        matmul(state.v, state.xb, tmpw)
+        matmul(state.v, state.xb, tmpw, state.rt)
 
         # Apply RoPE rotation to the q and k vectors for each head
         for h in range(config.n_heads):
@@ -587,7 +557,7 @@ fn transformer(
                     xb.offset(i).simd_store[1](0, xbi)
         # Final matrix multiplication to get the output of the attention
         tmpw.set_buf_ptr(weights.wo.data.offset(l * dim * dim), dim, dim)
-        matmul(state.xb2, state.xb, tmpw)
+        matmul(state.xb2, state.xb, tmpw, state.rt)
 
         # Residual connection back into x
         accum(x, state.xb2.data, dim)
@@ -597,10 +567,10 @@ fn transformer(
 
         # Calculate self.w1(x) and self.w3(x) for FFN
         tmpw.set_buf_ptr(weights.w1.data.offset(l * dim * hidden_dim), hidden_dim, dim)
-        matmul(state.hb, state.xb, tmpw)
+        matmul(state.hb, state.xb, tmpw, state.rt)
 
         tmpw.set_buf_ptr(weights.w3.data.offset(l * dim * hidden_dim), hidden_dim, dim)
-        matmul(state.hb2, state.xb, tmpw)
+        matmul(state.hb2, state.xb, tmpw, state.rt)
 
         # Apply SiLU activation function (silu(x) = x * sigmoid(x))
         for i in range(hidden_dim):
@@ -613,7 +583,7 @@ fn transformer(
 
         # Final matrix multiplication to get the output of the FFN
         tmpw.set_buf_ptr(weights.w2.data.offset(l * dim * hidden_dim), dim, hidden_dim)
-        matmul(state.xb, state.hb, tmpw)
+        matmul(state.xb, state.hb, tmpw, state.rt)
 
         # Residual connection
         accum(x, state.xb.data, dim)
@@ -623,7 +593,7 @@ fn transformer(
 
     # Classifier into logits
     tmpw.set_buf_ptr(weights.wcls.data, config.vocab_size, dim)
-    matmul(state.logits, state.x, tmpw)
+    matmul(state.logits, state.x, tmpw, state.rt)
 
 
 fn argmax(v: Matrix) -> Int:
@@ -651,6 +621,59 @@ fn sample(probabilities: Matrix) -> Int:
     return n - 1  # In case of rounding errors
 
 
+fn str_lookup(str: PointerString, tok: Tokenizer) -> Int:
+    for pos in range(tok.vocab_size):
+        let s1 = tok.vocab[pos]
+        var p1 = 0
+        while s1[p1] != 0 and str[p1] != 0:
+            if s1[p1] != str[p1]:
+                break
+            p1 += 1
+        if s1[p1] != 0 or str[p1] != 0:
+            continue
+        return pos
+    return -1
+
+
+fn bpe_encode(inout tokens: DynamicVector[Int], text: String, tok: Tokenizer):
+    for pos in range(len(text)):
+        let char = str_to_ptr(text[pos])
+        let id = str_lookup(char, tok)
+
+        if id == -1:
+            print("Not a good prompt token at pos ", pos)
+            return
+        tokens.push_back(id)
+
+    while True:
+        var best_score = Float32(-1e10)
+        var best_id = -1
+        var best_idx = -1
+
+        for i in range(len(tokens) - 1):
+            # Check if we can merge the pair (tokens[i], tokens[i+1])
+            let str = str_concat(tok.vocab[tokens[i]], tok.vocab[tokens[i + 1]])
+            let id = str_lookup(str, tok)
+            if id != -1 and tok.vocab_scores.load(id) > best_score:
+                best_score = tok.vocab_scores.load(id)
+                best_id = id
+                best_idx = i
+
+        if best_idx == -1:
+            # We couldn't find any more pairs to merge, so we're done
+            break
+
+        # Merge the consecutive pair (best_idx, best_idx+1) into new token best_id
+        tokens[best_idx] = best_id
+        # Delete token at position best_idx+1, shift the entire sequence back 1
+        var _tokens = DynamicVector[Int]()
+        for i in range(0, best_idx + 1):
+            _tokens.push_back(tokens[i])
+        for i in range(best_idx + 2, len(tokens)):
+            _tokens.push_back(tokens[i])
+        tokens = _tokens
+
+
 fn print_str(s: PointerString):
     # print all chars till null character
     var p: Int = 0
@@ -664,15 +687,61 @@ fn time_in_ms() -> Int:
     return time.now() // 1_000_000
 
 
+fn print_usage():
+    print("Usage: mojo llama2.mojo <checkpoint> [options]")
+    print("Example: mojo llama2.mojo stories15M.bin -s 99 -n 256 -t 0.5 -i \"Llama is an animal\"")
+    print("Options:")
+    print("  -s <int>    random seed, default time.now()")
+    print("  -t <float>  temperature in [0,1.0], default 1.0")
+    print("  -n <int>    number of steps to run for, default 256. 0 = max_seq_len")
+    print("  -i <string> input prompt")
+
+
 fn main() raises:
-    print("num hardware threads: ", num_cores(), " SIMD vector width: ", nelts)
-    let checkpoint = "stories15M.bin"
-    # let checkpoint = "stories110M.bin"
-    let tokenizer = "tokenizer.bin"
-    let temperature = 0.0
+    print("num hardware threads: ", num_cores())
+    print("SIMD vector width: ", nelts)
+    var tokenizer = StringRef("tokenizer.bin")
+    var checkpoint = StringRef("stories15M.bin")
+    var temperature = 0.9
     var steps = 256
-    let prompt = ""
-    let rng_seed: Int = time.now()
+    var prompt = String("")
+    var rng_seed: Int = time.now()
+
+    @parameter
+    fn argparse() raises -> Int:
+        let args = argv()
+        if len(args) < 2:
+            return 0
+        checkpoint = args[1]
+        for i in range(2, len(args), 2):
+            if args[i] == "-p":
+                print("Option not supported: ", args[i])
+            if args[i] == "-n":
+                steps = atol(args[i + 1])
+            if args[i] == "-s":
+                rng_seed = atol(args[i + 1])
+            if args[i] == "-i":
+                prompt = args[i + 1]
+            if args[i] == "-t":
+                let val = args[i + 1]
+                temperature = 0.0
+                # hacky parse float, keep only 1 digit
+                for c in range(0, len(val)):
+                    if val[c] == ".":
+                        temperature += atol(val[c + 1]) * 0.1
+                        break
+                    else:
+                        temperature = atol(val[c])
+                if temperature < -1e9 or temperature > (1 + 1e9):
+                    print("Wrong temperature value", temperature)
+                    return 0
+        return 1
+
+    let res = argparse()
+    if res == 0:
+        print_usage()
+        return
+
     random.seed(rng_seed)
     var fbuf: FileBuf = FileBuf()
     var tbuf: FileBuf = FileBuf()
@@ -702,39 +771,45 @@ fn main() raises:
     # Create and initialize the application RunState
     var state = RunState(config)
 
+    # Process the prompt, if any
+    var prompt_tokens = DynamicVector[Int]()
+
+    if prompt:
+        bpe_encode(prompt_tokens, prompt, tok)
+
     # Start the main loop
     var start = 0  # Used to time our code, only initialized after the first iteration
     var next_token = 0  # Will store the next token in the sequence
     # Initialize with token 1 (=BOS), as done in Llama-2 sentencepiece tokenizer
     var token = 1
-    var pos = 0  # Position in the sequence
-    # Explicitly print the initial BOS token for stylistic symmetry reasons
-
-    print("<s>")
 
+    # Position in the sequence
+    var pos = 0
     while pos < steps:
         # Forward the transformer to get logits for the next token
         transformer(token, pos, config, state, weights)
 
-        # Sample the next token
-        if temperature == 0.0:
-            # Greedy argmax sampling: take the token with the highest probability
-            next_token = argmax(state.logits)
+        if pos < len(prompt_tokens):
+            next_token = prompt_tokens[pos]
         else:
-            # Apply the temperature to the logits
-            for q in range(config.vocab_size):
-                state.logits[q] = state.logits[q] / temperature
-            # Apply softmax to the logits to get the probabilities for the next token
-            softmax(state.logits.data, config.vocab_size)
-            # Sample from this distribution to get the next token
-            next_token = sample(state.logits)
+            # Sample the next token
+            if temperature == 0.0:
+                # Greedy argmax sampling: take the token with the highest probability
+                next_token = argmax(state.logits)
+            else:
+                # Apply the temperature to the logits
+                for q in range(config.vocab_size):
+                    state.logits[q] = state.logits[q] / temperature
+                # Apply softmax to the logits to get the probabilities for the next token
+                softmax(state.logits.data, config.vocab_size)
+                # Sample from this distribution to get the next token
+                next_token = sample(state.logits)
 
         var token_str: PointerString = tok.vocab[next_token]
         if token == 1 and token_str[0] == ord(" "):
             token_str = token_str.offset(1)
 
         print_str(token_str)
-        # flush?
 
         # Advance forward
         token = next_token