DEV Community

Cover image for Homemade GPT JS
Oleksii Trekhleb
Oleksii Trekhleb

Posted on

Homemade GPT JS

For learning purposes, I made a minimal TensorFlow.js re-implementation of Karpathy's minGPT (Generative Pre-trained Transformer).

A full definition of this "homemade" GPT language model (all of it) can be found in this single model.ts file (less than 300 lines of code).

Since model.ts is written in TypeScript, you can use homemade GPT playground to train it, experiment with parameters, and generate its predictions directly in the browser using a GPU.

The model and the playground are written for learning purposes, to understand how GPT works and to use WebGPU for training.

To understand what's happening in the model.ts file please refer to Andrej Karpathy's well-explained, hands-on lecture "Let's build GPT: from scratch, in code, spelled out" (arguably one of the best explanations of GPT out there).

GPT Folder

Inside the ./gpt/src/ folder you'll find the following files:

  • model.ts - this is the main file of interest, as it contains the full (yet minimalistic) definition of the decoder GPT model, as described in the Attention Is All You Need paper.
  • model-easier.ts - this is the same GPT model as in the previous file but simplified for easier understanding. The main difference is that it processes all Heads inside CausalSelfAttention sequentially (instead of in parallel). As a result, the model is a bit slower but more readable.
  • config.ts - contains pre-configured sets of GPT model parameters: GPT-pico, GPT-nano, GPT-mini, GPT-2, etc.
  • dataset.ts - Nothing GPT-specific here. A helper wrapper on top of any txt-file-based character-level dataset. It loads an arbitrary txt file, treats each letter as a token, splits the characters into training and testing batches, and encodes/decodes letters to indices and vice versa.
  • trainer.ts - Nothing GPT-specific here as well. This file provides a simple training loop that could apply to any arbitrary neural network.

Web Playground

To experiment with model parameters, training, and text generation you may use the Homemade GPT JS playground.

Image description

You may also launch the playground locally if you want to modify and experiment with the code of the transformer model itself.

Install dependencies:

npm i
Enter fullscreen mode Exit fullscreen mode

Launch web playground locally:

npm run playground-web
Enter fullscreen mode Exit fullscreen mode

The playground will be accessible on http://localhost:3000/homemade-gpt-js

Run these commands from the root of the project. You need to have Node.js ≥ 20.0.0.

Node.js Playground

You may also experiment with the model in Node.js environment.

Install dependencies:

npm i
Enter fullscreen mode Exit fullscreen mode

Launch Node.js playground:

npm run playground-node
Enter fullscreen mode Exit fullscreen mode

The ./playground-node/src/index.ts file contains the basic example of training and text generation.

Run these commands from the root of the project. You need to have Node.js ≥ 20.0.0.


I hope this TS-based GPT example will be helpful for you, learners.

Top comments (2)

Collapse
 
horaceshmorace profile image
Horace Nelson

Lost me at Typescript. I don't think the use case justifies the added overhead. This looks really cool though.

Collapse
 
sim2k profile image
Simeon Williams

Might stick this into Bolt.new and see what happens