add optional support for the tokenizer. people have to re-run train_gpt2.py to get this feature. later, we will change the code to demand this and error if it is not found, for now i just don't want to brick people who do a pull. also i only modified the train_gpt2.cu for now, not the .c. it might be time to separate out the dataloader, the tokenizer, and any other common utilities to their own files (?). otherwise it feels a bit silly to copy paste stuff around? not sure

train_gpt2.cu

@@ -4,6 +4,7 @@ GPT-2 Transformer Neural Net trained in raw CUDA
 
 #include <stdio.h>
 #include <stdlib.h>
+#include <ctype.h>
 #include <math.h>
 #include <time.h>
 #include <assert.h>
@@ -1158,6 +1159,86 @@ int sample_mult(float* probabilities, int n, float coin) {
     return n - 1; // in case of rounding errors
 }
 
+// ----------------------------------------------------------------------------
+// Tokenizer (only supports decoding)
+
+typedef struct {
+    uint32_t vocab_size;
+    char **token_table;
+    int init_ok;
+} Tokenizer;
+
+void safe_printf(const char *piece) {
+    // the tokens are raw bytes, and we we only want to print the printable ones
+    // many bytes can be various control codes, backspace, etc.
+    if (piece == NULL) { return; }
+    if (piece[0] == '\0') { return; }
+    // handle individual byte tokens
+    // every token is asserted to be at least one byte so doing piece[1] is ok
+    if (piece[1] == '\0') {
+        unsigned char byte_val = piece[0];
+        if (!(isprint(byte_val) || isspace(byte_val))) {
+            return; // weird byte, don't print it
+        }
+    }
+    printf("%s", piece);
+}
+
+void tokenizer_init(Tokenizer *tokenizer, const char *filename) {
+    FILE *file = fopen(filename, "rb");
+    if (file == NULL) {
+        // try to be more helpful as we just added this feature, erase later
+        printf("---\n");
+        printf("WARNING: Failed to open the tokenizer file %s\n", filename);
+        printf("The Tokenizer is a new feature added April 14 2024.\n");
+        printf("Re-run `python train_gpt2.py` to write it\n");
+        printf("---\n");
+        tokenizer->init_ok = 0;
+        return;
+    }
+    // read in the header
+    uint32_t header[256];
+    fread(header, sizeof(uint32_t), 256, file);
+    assert(header[0] == 20240328);
+    assert(header[1] == 1);
+    tokenizer->vocab_size = header[2];
+    // read in all the tokens
+    unsigned char length;
+    tokenizer->token_table = (char **)malloc(tokenizer->vocab_size * sizeof(char *));
+    for (uint32_t i = 0; i < tokenizer->vocab_size; i++) {
+        fread(&length, sizeof(unsigned char), 1, file);
+        assert(length > 0); // every token should be at least one character
+        char *token_bytes = (char *)malloc(length + 1);
+        fread(token_bytes, sizeof(char), length, file);
+        token_bytes[length] = '\0';  // Add null terminator for printing
+        tokenizer->token_table[i] = token_bytes;
+    }
+    // cleanups
+    fclose(file);
+    tokenizer->init_ok = 1;
+}
+
+const char *tokenizer_decode(Tokenizer *tokenizer, uint32_t token_id) {
+    if (tokenizer->init_ok == 0) {
+        return NULL;
+    }
+    if (token_id < tokenizer->vocab_size) {
+        return tokenizer->token_table[token_id];
+    } else {
+        printf("invalid token id %d!\n", token_id);
+        return NULL;
+    }
+}
+
+void tokenizer_free(Tokenizer *tokenizer) {
+    if (tokenizer->init_ok) {
+        for (uint32_t i = 0; i < tokenizer->vocab_size; i++) {
+            free(tokenizer->token_table[i]);
+        }
+        free(tokenizer->token_table);
+    }
+}
+
 // ----------------------------------------------------------------------------
 // main training loop
 int main() {
@@ -1206,6 +1287,10 @@ int main() {
     printf("sequence length: %d\n", T);
     printf("val_num_batches: %d\n", val_num_batches);
 
+    // build the Tokenizer
+    Tokenizer tokenizer;
+    tokenizer_init(&tokenizer, "gpt2_tokenizer.bin");
+
     // some memory for generating samples from the model
     unsigned long long rng_state = 1337;
     int* gen_tokens = (int*)malloc(B * T * sizeof(int));
@@ -1236,6 +1321,7 @@ int main() {
                 gen_tokens[i] = GPT2_EOT;
             }
             // now sample from the model autoregressively
+            printf("generating:\n---\n");
             for (int t = 1; t < genT; t++) {
                 // note that inference is very wasteful here because for each token
                 // we re-calculate the forward pass for all of (B,T) positions from scratch
@@ -1252,12 +1338,16 @@ int main() {
                 float coin = random_f32(&rng_state);
                 int next_token = sample_mult(cpu_probs, model.config.vocab_size, coin);
                 gen_tokens[t] = next_token;
+                // print the generated token, either using the Tokenizer or a fallback
+                if (tokenizer.init_ok) {
+                    const char* token_str = tokenizer_decode(&tokenizer, next_token);
+                    safe_printf(token_str);
+                } else {
+                    // fall back to printing the token id
+                    printf("%d ", next_token);
+                }
             }
-            printf("generated: ");
-            for (int t = 0; t < genT; t++) {
-                printf("%d ", gen_tokens[t]);
-            }
-            printf("\n");
+            printf("\n---\n");
         }
 
         // do a training step
@@ -1276,6 +1366,7 @@ int main() {
     // free
     dataloader_free(&train_loader);
     dataloader_free(&val_loader);
+    tokenizer_free(&tokenizer);
     gpt2_free(&model);
     free(cpu_probs);
     free(gen_tokens);

train_gpt2.py

@@ -11,6 +11,7 @@ https://github.com/huggingface/transformers/blob/main/src/transformers/models/gp
 
 import os
 import math
+import struct
 from dataclasses import dataclass
 
 import numpy as np
@@ -290,6 +291,21 @@ def write_state(model, x, y, logits, loss, filename):
         write_tensors(grads, model.config.n_layer, file)
     print(f"wrote {filename}")
 
+def write_tokenizer(enc, filename):
+    n = enc.max_token_value + 1
+    header = torch.zeros(256, dtype=torch.int32)
+    header[0] = 20240328 # magic
+    header[1] = 1 # tokenizer version = 1
+    header[2] = n # number of tokens
+    with open(filename, "wb") as file:
+        file.write(header.numpy().tobytes())
+        for i in range(n):
+            b = enc.decode_bytes([i])
+            length = len(b)
+            assert length < 256, f"Token length exceeds 255: {length}"
+            file.write(struct.pack("<B", length))  # Write the length as a 1-byte unsigned integer
+            file.write(b)  # Write the actual bytes
+    print(f"wrote {filename}")
 
 if __name__ == "__main__":
     import time
@@ -330,6 +346,8 @@ if __name__ == "__main__":
     encode = lambda s: enc.encode(s, allowed_special={"<|endoftext|>"})
     decode = lambda l: enc.decode(l)
 
+    write_tokenizer(enc, "gpt2_tokenizer.bin")
+
     if args.tensorcores:
         torch.set_float32_matmul_precision('high')