microGPT

package com.anirudhology.microgpt;

import com.anirudhology.microgpt.data.NGramDatasetBuilder;
import com.anirudhology.microgpt.data.TextCorpus;
import com.anirudhology.microgpt.data.TrainingExample;
import com.anirudhology.microgpt.model.BaselineBigramModel;
import com.anirudhology.microgpt.model.GPTLanguageModel;
import com.anirudhology.microgpt.model.MLPLanguageModel;
import com.anirudhology.microgpt.model.NeuralBigramAutogradModel;
import com.anirudhology.microgpt.model.NeuralBigramModel;
import com.anirudhology.microgpt.optimizer.AdamOptimizer;
import com.anirudhology.microgpt.tokenizer.CharacterTokenizer;

import java.util.Collections;
import java.util.List;
import java.util.Random;

/**
 * Runs each model in sequence, from simplest to most powerful.
 *
 * Step 1: Statistical Bigram       - pure counting, no neural network
 * Step 2: Neural Bigram            - same task, learned weights, manual gradients
 * Step 3: Neural Bigram (Autograd) - same model, automatic differentiation via Value
 * Step 4: MLP Language Model       - wider context window, embeddings, hidden layer
 * Step 5: GPT Transformer          - multi-head attention, transformer blocks, Adam
 */
public class Runner {

    static void main() {
        TextCorpus textCorpus = new TextCorpus();
        final List<String> docs = textCorpus.readCorpus("input.txt");

        CharacterTokenizer tokenizer = new CharacterTokenizer();
        tokenizer.buildVocabulary(docs);

        int split = (int) (docs.size() * 0.9);
        final List<String> trainDocs = docs.subList(0, split);
        final List<String> validationDocs = docs.subList(split, docs.size());

        printStep(1, "Statistical Bigram",
                "Count character pair frequencies. No learning, no gradients. Pure statistics.");
        runBaselineBigram(tokenizer, docs);

        printStep(2, "Neural Bigram (manual gradients)",
                "Replace the count table with a learned weight matrix.\n" +
                "  Gradients computed by hand: dL/dlogit = p - 1_correct.");
        runNeuralBigram(tokenizer, trainDocs, validationDocs);

        printStep(3, "Neural Bigram (autograd)",
                "Same model, but gradients computed automatically by the Value graph.");
        runNeuralAutogradBigram(tokenizer, trainDocs, validationDocs);

        printStep(4, "MLP Language Model",
                "Wider context window (3 chars). Token + positional embeddings flattened\n" +
                "  into a hidden layer with tanh activation.");
        runMLPLanguageModel(tokenizer, docs);

        printStep(5, "GPT Transformer (single-head attention)",
                "CausalSelfAttention: one attention head over the full embedding dimension.");
        runGPTLanguageModel(tokenizer, docs, false);

        printStep(6, "GPT Transformer (multi-head attention)",
                "MultiHeadCausalSelfAttention: splits embedding into 4 heads.");
        runGPTLanguageModel(tokenizer, docs, true);
    }

    private static void runBaselineBigram(CharacterTokenizer tokenizer, List<String> docs) {
        BaselineBigramModel model = new BaselineBigramModel(tokenizer.getVocabularySize(), tokenizer.getBOSId(), 42L);
        model.fit(docs, tokenizer, 1.0);
        double nll = model.averageNegativeLogLikelihood(docs, tokenizer);
        System.out.printf("Avg NLL: %.4f%n", nll);
        for (int i = 0; i < 10; i++)
            System.out.printf("Sample %2d: %s%n", i + 1, model.sample(tokenizer, 16));
    }

    private static void runNeuralBigram(CharacterTokenizer tokenizer,
            List<String> trainDocs, List<String> validationDocs) {
        NeuralBigramModel model = new NeuralBigramModel(tokenizer.getVocabularySize(), tokenizer.getBOSId(), 42L);
        double learningRate = 0.5;
        for (int e = 1; e <= 30; e++) {
            double trainNll = model.trainEpoch(trainDocs, tokenizer, learningRate, 1000L + e);
            double valNll = model.averageNegativeLogLikelihood(validationDocs, tokenizer);
            System.out.printf("Epoch %2d | Train NLL: %.4f | Val NLL: %.4f%n", e, trainNll, valNll);
            learningRate *= 0.98;
        }
        for (int i = 0; i < 10; i++)
            System.out.printf("Sample %2d: %s%n", i + 1, model.sample(tokenizer, 16, 0.9));
    }

    private static void runNeuralAutogradBigram(CharacterTokenizer tokenizer,
            List<String> trainDocs, List<String> validationDocs) {
        NeuralBigramAutogradModel model = new NeuralBigramAutogradModel(
                tokenizer.getVocabularySize(), tokenizer.getBOSId(), 42L);
        double learningRate = 0.5;
        for (int e = 0; e < 20; e++) {
            double trainNll = model.train(trainDocs, tokenizer, learningRate, 1000L + e);
            double valNll = model.averageNegativeLogLikelihood(validationDocs, tokenizer);
            System.out.printf("Epoch %2d | Train NLL: %.4f | Val NLL: %.4f%n", e, trainNll, valNll);
            learningRate *= 0.98;
        }
        for (int i = 0; i < 10; i++)
            System.out.printf("Sample %2d: %s%n", i + 1, model.sample(tokenizer, 16, 0.9));
    }

    private static void runMLPLanguageModel(CharacterTokenizer tokenizer, List<String> documents) {
        final int blockSize = 3;
        final List<TrainingExample> examples = NGramDatasetBuilder.build(documents, tokenizer, blockSize, true);
        final MLPLanguageModel model = new MLPLanguageModel(
                tokenizer.getVocabularySize(), blockSize, 10, 100, 42L);
        double learningRate = 0.01;
        for (int epoch = 0; epoch < 10; epoch++) {
            double totalLoss = 0.0;
            Collections.shuffle(examples, new Random(42L + epoch));
            for (int i = 0; i < examples.size(); i++) {
                totalLoss += model.trainStep(examples.get(i), learningRate, true);
                if ((i + 1) % 1000 == 0)
                    System.out.printf("Epoch %d, Step %d/%d, Avg Loss: %.4f%n",
                            epoch + 1, i + 1, examples.size(), totalLoss / (i + 1));
            }
            learningRate *= 0.9;
        }
        for (int i = 0; i < 10; i++)
            System.out.printf("Sample %2d: %s%n", i + 1, model.generate(tokenizer, 20, 1.0));
    }

    private static void runGPTLanguageModel(CharacterTokenizer tokenizer,
            List<String> documents, boolean useMultiHead) {
        final int blockSize = 16, embeddingDimension = 16, numHeads = 4, numberOfLayers = 1;
        final List<TrainingExample> examples = NGramDatasetBuilder.build(documents, tokenizer, blockSize, true);
        final GPTLanguageModel model = new GPTLanguageModel(
                tokenizer.getVocabularySize(), blockSize, embeddingDimension,
                numHeads, numberOfLayers, useMultiHead, 42L);
        final AdamOptimizer optimizer = new AdamOptimizer(model.parameters());
        double initialLearningRate = 0.01;
        for (int step = 0; step < 1000; step++) {
            double learningRate = initialLearningRate * (1.0 - (double) step / 1000);
            optimizer.zeroGradient(model.parameters());
            double loss = model.trainStep(examples.get(step % examples.size()));
            optimizer.step(model.parameters(), learningRate);
            if ((step + 1) % 100 == 0)
                System.out.printf("Step %4d / 1000 | Loss: %.4f | LR: %.6f%n", step + 1, loss, learningRate);
        }
        for (int i = 0; i < 10; i++)
            System.out.printf("Sample %2d: %s%n", i + 1, model.generate(tokenizer, 20, 0.5));
    }

    private static void printStep(int step, String title, String description) {
        System.out.println("\n" + "=".repeat(70));
        System.out.printf("  Step %d: %s%n", step, title);
        System.out.println("=".repeat(70));
        System.out.println("  " + description);
        System.out.println("-".repeat(70));
    }
}

#!/usr/bin/env ruby
# frozen_string_literal: true
# nanogpt — A Ruby port of Karpathy's nanoGPT
# Source: https://github.com/khasinski/nanogpt-rb

require "nano_gpt"

class NanoGPTCLI
  COMMANDS = %w[prepare train sample bench version help].freeze

  def initialize(args)
    @command = args.shift
    @args = args
  end

  def run
    case @command
    when "prepare" then prepare
    when "train"   then train
    when "sample"  then sample
    when "bench"   then bench
    when "version", "-v", "--version" then version
    when "help", "-h", "--help", nil  then help
    else
      puts "Unknown command: #{@command}"
      help; exit 1
    end
  end

  private

  def train
    config = NanoGPT::TrainConfig.load(@args)
    config[:device] = NanoGPT::Device.auto if config[:device] == "auto"

    data_dir  = File.join("data", config[:dataset])
    tokenizer = NanoGPT::Tokenizer.for_dataset(data_dir)

    model_config = NanoGPT::GPTConfig.new(
      block_size: config[:block_size],
      vocab_size: tokenizer.vocab_size,
      n_layer:    config[:n_layer],
      n_head:     config[:n_head],
      n_embd:     config[:n_embd],
      dropout:    config[:dropout],
      bias:       config[:bias]
    )

    model = NanoGPT::GPT.new(model_config)
    model.to(config[:device]) if config[:device] != "cpu"

    data_loader = NanoGPT::DataLoader.new(
      data_dir:   data_dir,
      block_size: config[:block_size],
      batch_size: config[:batch_size],
      device:     config[:device]
    )

    NanoGPT::Trainer.new(model: model, data_loader: data_loader, config: config.to_h).train
    puts "\nTraining complete! Checkpoint saved to #{config[:out_dir]}/ckpt.pt"
  end

  def sample
    config = NanoGPT::SampleConfig.load(@args)
    config[:device] = NanoGPT::Device.auto if config[:device] == "auto"
    Torch.manual_seed(config[:seed])

    checkpoint   = Torch.load(File.join(config[:out_dir], "ckpt.pt"))
    model_config = NanoGPT::GPTConfig.new(**checkpoint["model_args"].transform_keys(&:to_sym))
    model        = NanoGPT::GPT.new(model_config)
    model.load_state_dict(checkpoint["model"])
    model.to(config[:device]) if config[:device] != "cpu"
    model.eval

    tokenizer  = NanoGPT::Tokenizer.for_dataset(File.join("data", config[:dataset]))
    start_text = config[:start]
    start_text = File.read(start_text[5..]) if start_text.start_with?("FILE:")
    x = Torch.tensor([tokenizer.encode(start_text)], dtype: :long, device: config[:device])

    config[:num_samples].times do
      y      = model.generate(x, config[:max_new_tokens], temperature: config[:temperature], top_k: config[:top_k])
      output = tokenizer.decode(y[0].to_a)
      puts output
      puts "-" * 50
    end
  end

  def bench
    config = NanoGPT::BenchConfig.load(@args)
    config[:device] = NanoGPT::Device.auto if config[:device] == "auto"
    Torch.manual_seed(config[:seed])

    model_config = NanoGPT::GPTConfig.new(
      block_size: config[:block_size], vocab_size: 50304,
      n_layer: config[:n_layer], n_head: config[:n_head], n_embd: config[:n_embd],
      dropout: config[:dropout], bias: config[:bias]
    )
    model     = NanoGPT::GPT.new(model_config)
    model.to(config[:device]) if config[:device] != "cpu"
    optimizer = model.configure_optimizers(
      weight_decay: 1e-2, learning_rate: 1e-4, betas: [0.9, 0.95],
      device_type: NanoGPT::Device.type(config[:device])
    )

    get_batch = lambda do
      x = Torch.randint(50304, [config[:batch_size], config[:block_size]], dtype: :long)
      y = Torch.randint(50304, [config[:batch_size], config[:block_size]], dtype: :long)
      [x, y]
    end

    [{ name: "burn-in", steps: 10 }, { name: "benchmark", steps: 20 }].each do |phase|
      puts "\nPhase: #{phase[:name]}"
      x, y = get_batch.call
      t0 = Time.now
      phase[:steps].times do |k|
        _logits, loss = model.forward(x, targets: y)
        x, y = get_batch.call
        optimizer.zero_grad; loss.backward; optimizer.step
        puts "  #{k}/#{phase[:steps]} loss: #{format('%.4f', loss.item)}"
      end
      if phase[:name] == "benchmark"
        dt  = Time.now - t0
        mfu = model.estimate_mfu(config[:batch_size] * phase[:steps], dt)
        puts "\nTime/iter: #{format('%.2f', dt / phase[:steps] * 1000)}ms | MFU: #{format('%.2f', mfu * 100)}%"
      end
    end
  end

  def prepare
    dataset = @args.first || (puts "Usage: nanogpt prepare <dataset>"; exit 1)
    load File.join(File.dirname(__FILE__), "..", "data", dataset, "prepare.rb")
  end

  def version
    puts "nanogpt #{NanoGPT::VERSION}"
  end

  def help
    puts <<~HELP
      nanogpt — Ruby port of Karpathy's nanoGPT
      https://github.com/khasinski/nanogpt-rb

      Commands:
        prepare <dataset>   Download and prepare a dataset
        train               Train a GPT model
        sample              Generate text from a trained model
        bench               Run performance benchmarks
        version             Show version
        help                Show this help
    HELP
  end
end

NanoGPTCLI.new(ARGV).run