Capstone: A Chat Engine in Pure Java

import java.util.ArrayList;
import java.util.List;

public class EngineApi {
    interface Engine {
        List<Integer> tokenize(String text);
        float[] prefill(List<Integer> tokens);      // runs the prompt, returns the logits of its last token
        float[] decode(int token);                  // runs one token at the next position, returns its logits
        int sample(float[] logits);
        String piece(int token);
        boolean isEndOfGeneration(int token);
        int position();                             // tokens held in the KV cache
    }

    // A stand-in engine: 8 tokens, and logits that favor a fixed script, so every call is visible.
    static final class ScriptEngine implements Engine {
        static final String[] VOCAB = {"the", " cat", " sat", " on", " the", " mat", ".", "<end>"};
        static final int[] SCRIPT = {0, 1, 2, 3, 4, 5, 6, 7};      // the cat sat on the mat . <end>
        int position;
        public List<Integer> tokenize(String text) {
            List<Integer> ids = new ArrayList<>();
            for (String w : text.split("(?= )")) ids.add(List.of(VOCAB).indexOf(w));
            return ids;
        }
        float[] logitsAt(int pos) { float[] l = new float[VOCAB.length]; l[SCRIPT[Math.min(pos, SCRIPT.length - 1)]] = 2f; return l; }
        public float[] prefill(List<Integer> tokens) {
            position += tokens.size();
            System.out.println("  prefill " + tokens + " -> position " + position);
            return logitsAt(position);
        }
        public float[] decode(int token) {
            position++;
            System.out.println("  decode " + token + " -> position " + position);
            return logitsAt(position);
        }
        public int sample(float[] logits) { int best = 0; for (int i = 1; i < logits.length; i++) if (logits[i] > logits[best]) best = i; return best; }
        public String piece(int token) { return VOCAB[token]; }
        public boolean isEndOfGeneration(int token) { return token == 7; }
        public int position() { return position; }
    }

    static String generate(Engine e, String prompt, int maxTokens) {
        StringBuilder out = new StringBuilder();
        float[] logits = e.prefill(e.tokenize(prompt));
        for (int n = 0; n < maxTokens; n++) {
            int id = e.sample(logits);
            if (e.isEndOfGeneration(id)) { System.out.println("  sampled " + id + ": end of generation"); break; }
            out.append(e.piece(id));
            logits = e.decode(id);
        }
        return out.toString();
    }

    public static void main(String[] args) {
        Engine e = new ScriptEngine();
        System.out.println("tokens of \"the cat\": " + e.tokenize("the cat"));
        String text = generate(e, "the cat", 10);
        System.out.println("generated: \"" + text + "\", KV cache holds " + e.position() + " tokens");
    }
}

import java.util.ArrayList;
import java.util.List;

public class TemplateFromMetadata {
    record Msg(String role, String content) {}

    sealed interface ChatTemplate permits ChatMl, Llama3 {           // Project 3, two of its three variants
        String render(List<Msg> chat, boolean generationPrompt);
        String stopToken();
    }

    record ChatMl() implements ChatTemplate {
        public String render(List<Msg> chat, boolean generationPrompt) {
            StringBuilder sb = new StringBuilder();
            for (Msg m : chat) sb.append("<|im_start|>").append(m.role()).append('\n').append(m.content()).append("<|im_end|>\n");
            return generationPrompt ? sb.append("<|im_start|>assistant\n").toString() : sb.toString();
        }
        public String stopToken() { return "<|im_end|>"; }
    }

    record Llama3(String date) implements ChatTemplate {
        public String render(List<Msg> chat, boolean generationPrompt) {
            String system = chat.get(0).role().equals("system") ? chat.get(0).content().trim() : "";
            StringBuilder sb = new StringBuilder("<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"
                + "Cutting Knowledge Date: December 2023\nToday Date: " + date + "\n\n" + system + "<|eot_id|>");
            for (Msg m : chat)
                if (!m.role().equals("system"))
                    sb.append("<|start_header_id|>").append(m.role()).append("<|end_header_id|>\n\n").append(m.content().trim()).append("<|eot_id|>");
            return generationPrompt ? sb.append("<|start_header_id|>assistant<|end_header_id|>\n\n").toString() : sb.toString();
        }
        public String stopToken() { return "<|eot_id|>"; }
    }

    static ChatTemplate templateFor(String jinja) {
        if (jinja.contains("<|im_start|>")) return new ChatMl();
        if (jinja.contains("<|start_header_id|>")) return new Llama3("26 Jul 2024");
        throw new IllegalArgumentException("unknown chat template");
    }

    // The Qwen2 template inserts a default system message when the conversation has none.
    static List<Msg> withDefaultSystem(String jinja, List<Msg> chat) {
        String text = "You are a helpful assistant";
        if (!jinja.contains("{{ '<|im_start|>system\n" + text + "<|im_end|>\n' }}") || chat.get(0).role().equals("system")) return chat;
        List<Msg> out = new ArrayList<>(List.of(new Msg("system", text)));
        out.addAll(chat);
        return out;
    }

    public static void main(String[] args) {
        // tokenizer.chat_template of the tiny random model (the Qwen2 template); its string literals contain real newlines
        String jinja = "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}"
            + "{{ '<|im_start|>system\nYou are a helpful assistant<|im_end|>\n' }}{% endif %}"
            + "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}"
            + "{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}";
        // the prompt that llama-server's /apply-template endpoint returned for the same messages
        String fromLlamaServer = "<|im_start|>system\nYou are a helpful assistant<|im_end|>\n<|im_start|>user\n"
            + "Write a Java method that reverses a string.<|im_end|>\n<|im_start|>assistant\n";

        List<Msg> chat = List.of(new Msg("user", "Write a Java method that reverses a string."));
        ChatTemplate t = templateFor(jinja);
        String plain = t.render(chat, true), withSystem = t.render(withDefaultSystem(jinja, chat), true);
        System.out.println("template: " + t.getClass().getSimpleName() + ", stop token " + t.stopToken());
        System.out.println("without the default system message, equal to llama-server: " + plain.equals(fromLlamaServer));
        System.out.println("with the default system message, equal to llama-server:    " + withSystem.equals(fromLlamaServer));
        System.out.println(withSystem.replace("\n", "\\n\n"));
        System.out.println("Llama 3 marker selects: " + templateFor("...<|start_header_id|>...").getClass().getSimpleName());
    }
}

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;

public class EngineSampling {
    // Project 2's SamplerChain reduced to top-k, top-p, min-p, temperature, and dist, in the llama.cpp order.
    record Chain(int topK, double topP, double minP, double temp, Random rng) {
        int candidates(double[] l) { return (int) Arrays.stream(l).filter(Double::isFinite).count(); }

        int sample(float[] logits, int[] kept) {
            double[] l = new double[logits.length];
            for (int i = 0; i < l.length; i++) l[i] = logits[i];                 // the engine returns float logits
            if (topK > 0 && topK < l.length) {
                double[] sorted = l.clone();
                Arrays.sort(sorted);
                double kth = sorted[l.length - topK];
                for (int i = 0; i < l.length; i++) if (l[i] < kth) l[i] = Double.NEGATIVE_INFINITY;
            }
            if (topP < 1) {
                double[] p = softmax(l);
                Integer[] order = new Integer[l.length];
                for (int i = 0; i < l.length; i++) order[i] = i;
                Arrays.sort(order, (a, b) -> Double.compare(p[b], p[a]));
                double run = 0;
                boolean cut = false;
                for (int id : order) { if (cut) l[id] = Double.NEGATIVE_INFINITY; run += p[id]; if (run >= topP) cut = true; }
            }
            if (minP > 0) {
                double min = Arrays.stream(l).max().getAsDouble() + Math.log(minP);
                for (int i = 0; i < l.length; i++) if (l[i] < min) l[i] = Double.NEGATIVE_INFINITY;
            }
            if (temp <= 0) {                                                      // greedy: the first largest logit
                int best = 0;
                for (int i = 1; i < l.length; i++) if (l[i] > l[best]) best = i;
                kept[0] = 1;
                return best;
            }
            for (int i = 0; i < l.length; i++) l[i] /= temp;
            kept[0] = candidates(l);
            double[] p = softmax(l);
            double u = rng.nextDouble(), run = 0;
            for (int i = 0; i < p.length; i++) { run += p[i]; if (p[i] > 0 && u < run) return i; }
            throw new IllegalStateException("no candidate");
        }

        static double[] softmax(double[] l) {
            double max = Arrays.stream(l).max().getAsDouble(), sum = 0;
            double[] p = new double[l.length];
            for (int i = 0; i < l.length; i++) { p[i] = Math.exp(l[i] - max); sum += p[i]; }
            for (int i = 0; i < l.length; i++) p[i] /= sum;
            return p;
        }
    }

    public static void main(String[] args) {
        int vocab = 151_936;                                   // the Qwen2 vocabulary of the tiny model
        Random weights = new Random(50);
        float[][] steps = new float[5][vocab];
        for (float[] s : steps) for (int i = 0; i < vocab; i++) s[i] = (float) (weights.nextGaussian() * 2.5);
        Chain greedy = new Chain(0, 1, 0, 0, new Random(1));
        Chain defaults = new Chain(40, 0.95, 0.05, 0.8, new Random(1));
        int[] kept = new int[1];
        List<String> g = new ArrayList<>(), d = new ArrayList<>();
        for (float[] logits : steps) {
            g.add(String.valueOf(greedy.sample(logits, kept)));
            d.add(defaults.sample(logits, kept) + " (" + kept[0] + " left)");
        }
        System.out.println("greedy:              " + g);
        System.out.println("top-k 40, top-p 0.95, min-p 0.05, temp 0.8: " + d);
        System.out.println("logits per step: " + vocab + ", bytes as double[]: " + vocab * 8);
    }
}

import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.CharsetDecoder;
import java.nio.charset.CoderResult;
import java.nio.charset.StandardCharsets;
import java.util.HexFormat;

public class StreamingUtf8 {
    // Project 17's streaming decoder: an incomplete UTF-8 sequence waits for the next token's bytes.
    static final class Utf8Stream {
        private final CharsetDecoder decoder = StandardCharsets.UTF_8.newDecoder();
        private ByteBuffer pending = ByteBuffer.allocate(0);

        String push(byte[] bytes) {
            ByteBuffer in = ByteBuffer.allocate(pending.remaining() + bytes.length).put(pending).put(bytes).flip();
            CharBuffer out = CharBuffer.allocate(in.remaining() + 1);
            CoderResult r;
            while ((r = decoder.decode(in, out, false)).isError()) { in.position(in.position() + r.length()); out.put('\uFFFD'); }
            pending = in.slice();
            return out.flip().toString();
        }
    }

    public static void main(String[] args) {
        // "Ok 🦙 ok" in the Qwen2 vocabulary: ids and the bytes of each token
        int[] ids = {11578, 11162, 99, 247, 5394};
        byte[][] bytes = {
            {0x4F, 0x6B}, {0x20, (byte) 0xF0, (byte) 0x9F}, {(byte) 0xA6}, {(byte) 0x99}, {0x20, 0x6F, 0x6B}};
        Utf8Stream stream = new Utf8Stream();
        StringBuilder streamed = new StringBuilder(), naive = new StringBuilder();
        for (int i = 0; i < ids.length; i++) {
            String piece = stream.push(bytes[i]);
            String perToken = new String(bytes[i], StandardCharsets.UTF_8);
            System.out.printf("token %6d bytes %-8s streamed [%s]  per-token String [%s]%n", ids[i], HexFormat.of().formatHex(bytes[i]), piece, perToken);
            streamed.append(piece);
            naive.append(perToken);
        }
        System.out.println("streamed:  " + streamed);
        System.out.println("per token: " + naive);
    }
}

import java.util.List;
import java.util.function.Consumer;

public class StopHandling {
    record Result(String text, String reason, int tokens) {}

    static int stopIndex(String text, List<String> stops) {
        int best = -1;
        for (String s : stops) { int i = text.indexOf(s); if (i >= 0 && (best < 0 || i < best)) best = i; }
        return best;
    }

    static int holdBack(String text, List<String> stops) {     // longest suffix of text that starts a stop string
        int keep = 0;
        for (String s : stops)
            for (int n = Math.min(s.length() - 1, text.length()); n > keep; n--)
                if (text.endsWith(s.substring(0, n))) { keep = n; break; }
        return keep;
    }

    // The engine's loop with the model replaced by a script of pieces; null stands for an end-of-generation token.
    static Result run(List<String> script, List<String> stops, int maxTokens, int contextLeft, Consumer<String> out) {
        StringBuilder text = new StringBuilder();
        int emitted = 0, n = 0;
        String reason = "length";
        for (String piece : script) {
            if (n == maxTokens) break;
            if (piece == null) { reason = "eog"; break; }
            n++;
            text.append(piece);
            int cut = stopIndex(text.toString(), stops);
            if (cut >= 0) { text.setLength(cut); reason = "stop"; break; }
            int safe = text.length() - holdBack(text.toString(), stops);
            if (safe > emitted) { out.accept(text.substring(emitted, safe)); emitted = safe; }
            if (n == contextLeft) { reason = "context"; break; }
        }
        if (text.length() > emitted) out.accept(text.substring(emitted));
        return new Result(text.toString(), reason, n);
    }

    static void show(String name, List<String> script, List<String> stops, int maxTokens, int contextLeft) {
        StringBuilder chunks = new StringBuilder();
        Result r = run(script, stops, maxTokens, contextLeft, s -> chunks.append('[').append(s.replace("\n", "\\n")).append(']'));
        System.out.printf("%-14s chunks %-34s reason %-7s tokens %d%n", name, chunks, r.reason(), r.tokens());
    }

    public static void main(String[] args) {
        List<String> stops = List.of("\nUser:");
        java.util.ArrayList<String> eog = new java.util.ArrayList<>(List.of("Use", " a", " StringBuilder", "."));
        eog.add(null);
        eog.add("never shown");
        show("end token", eog, stops, 100, 4096);
        show("stop string", List.of("Done", ".", "\nUs", "er:", " next"), stops, 100, 4096);
        show("false alarm", List.of("Done", ".", "\nUs", "ing it"), stops, 100, 4096);
        show("max tokens", List.of("a", "b", "c", "d", "e"), stops, 3, 4096);
        show("context full", List.of("a", "b", "c", "d", "e"), stops, 100, 2);
    }
}

import java.util.Arrays;
import java.util.Random;

public class BatchedPrefill {
    static final int DIM = 16, HEADS = 2, HD = 8, FFN = 32, VOCAB = 20, LAYERS = 2;
    static float[][][] w = new float[LAYERS * 6 + 1][][];          // per layer: q, k, v, o, up, down; then output
    static long rowReads;

    static float[][] random(Random r, int rows, int cols) {
        float[][] m = new float[rows][cols];
        for (float[] row : m) for (int c = 0; c < cols; c++) row[c] = (float) (r.nextGaussian() * 0.3);
        return m;
    }

    // y[t] = m * x[t] for every token t; each weight row is read once and used for all tokens
    static float[][] matmul(float[][] m, float[][] x) {
        float[][] y = new float[x.length][m.length];
        for (int r = 0; r < m.length; r++) {
            rowReads++;
            for (int t = 0; t < x.length; t++) { float s = 0; for (int c = 0; c < m[r].length; c++) s += m[r][c] * x[t][c]; y[t][r] = s; }
        }
        return y;
    }

    static float[][] norm(float[][] x) {
        float[][] y = new float[x.length][];
        for (int t = 0; t < x.length; t++) {
            float ss = 0; for (float v : x[t]) ss += v * v;
            float s = (float) (1 / Math.sqrt(ss / x[t].length + 1e-5));
            y[t] = new float[x[t].length]; for (int i = 0; i < y[t].length; i++) y[t][i] = x[t][i] * s;
        }
        return y;
    }

    // Runs tokens at positions start.. against the cache; returns the logits of the last token when asked.
    static float[] forward(int[] tokens, int start, float[][][] kCache, float[][][] vCache, boolean wantLogits) {
        int n = tokens.length;
        float[][] x = new float[n][DIM];
        for (int t = 0; t < n; t++) for (int i = 0; i < DIM; i++) x[t][i] = (float) Math.sin(tokens[t] * 0.7 + i);   // fixed embedding
        for (int l = 0; l < LAYERS; l++) {
            float[][] h = norm(x), q = matmul(w[6 * l], h), k = matmul(w[6 * l + 1], h), v = matmul(w[6 * l + 2], h);
            for (int t = 0; t < n; t++) { kCache[l][start + t] = k[t]; vCache[l][start + t] = v[t]; }
            float[][] att = new float[n][DIM];
            for (int t = 0; t < n; t++)                                   // causal: position start+t sees 0..start+t
                for (int hd = 0; hd < HEADS; hd++) {
                    int pos = start + t; float max = Float.NEGATIVE_INFINITY, sum = 0; float[] s = new float[pos + 1];
                    for (int j = 0; j <= pos; j++) { for (int i = 0; i < HD; i++) s[j] += q[t][hd * HD + i] * kCache[l][j][hd * HD + i]; s[j] /= (float) Math.sqrt(HD); max = Math.max(max, s[j]); }
                    for (int j = 0; j <= pos; j++) { s[j] = (float) Math.exp(s[j] - max); sum += s[j]; }
                    for (int j = 0; j <= pos; j++) for (int i = 0; i < HD; i++) att[t][hd * HD + i] += s[j] / sum * vCache[l][j][hd * HD + i];
                }
            float[][] o = matmul(w[6 * l + 3], att);
            for (int t = 0; t < n; t++) for (int i = 0; i < DIM; i++) x[t][i] += o[t][i];
            float[][] u = matmul(w[6 * l + 4], norm(x));
            for (float[] row : u) for (int i = 0; i < FFN; i++) row[i] = Math.max(0, row[i]);
            float[][] d = matmul(w[6 * l + 5], u);
            for (int t = 0; t < n; t++) for (int i = 0; i < DIM; i++) x[t][i] += d[t][i];
        }
        if (!wantLogits) return null;
        return matmul(w[6 * LAYERS], norm(new float[][] {x[n - 1]}))[0];    // logits only for the last token
    }

    public static void main(String[] args) {
        Random r = new Random(6);
        for (int l = 0; l < LAYERS; l++) {
            w[6 * l] = random(r, DIM, DIM); w[6 * l + 1] = random(r, DIM, DIM); w[6 * l + 2] = random(r, DIM, DIM);
            w[6 * l + 3] = random(r, DIM, DIM); w[6 * l + 4] = random(r, FFN, DIM); w[6 * l + 5] = random(r, DIM, FFN);
        }
        w[6 * LAYERS] = random(r, VOCAB, DIM);
        int[] prompt = {3, 14, 15, 9, 2, 6, 5, 3};
        long rowsPerPass = LAYERS * (4L * DIM + FFN + DIM);

        float[][][] k1 = new float[LAYERS][64][], v1 = new float[LAYERS][64][];
        rowReads = 0;
        float[] last = null;
        for (int t = 0; t < prompt.length; t++) last = forward(new int[] {prompt[t]}, t, k1, v1, t == prompt.length - 1);
        long oneByOne = rowReads - VOCAB;

        float[][][] k2 = new float[LAYERS][64][], v2 = new float[LAYERS][64][];
        rowReads = 0;
        float[] batched = forward(prompt, 0, k2, v2, true);
        long batch = rowReads - VOCAB;

        System.out.println("prompt tokens: " + prompt.length + ", layer weight rows per pass: " + rowsPerPass);
        System.out.println("token by token: layer rows read " + oneByOne + " (" + oneByOne / rowsPerPass + " passes)");
        System.out.println("batched:        layer rows read " + batch + " (" + batch / rowsPerPass + " pass)");
        System.out.println("last-token logits identical: " + Arrays.equals(last, batched));
        System.out.println("KV caches identical: " + Arrays.deepEquals(k1, k2) + " " + Arrays.deepEquals(v1, v2));
    }
}

import java.util.Locale;

public class KvCacheLimit {
    record Shape(String name, int layers, int kvHeads, int headDim) {
        long bytesPerToken(int bytesPerValue) { return 2L * layers * kvHeads * headDim * bytesPerValue; }
    }

    public static void main(String[] args) {
        Shape[] shapes = {
            new Shape("tiny random model", 4, 2, 32),
            new Shape("Qwen2.5 0.5B shape", 24, 2, 64),
            new Shape("Llama 3.1 8B shape", 32, 8, 128)};
        System.out.println("model                bytes/token F16   ctx 4096 F16   ctx 4096 F32   ctx 32768 F16");
        for (Shape s : shapes)
            System.out.printf(Locale.ROOT, "%-20s %15d %11.2f MiB %10.2f MiB %11.2f GiB%n", s.name(), s.bytesPerToken(2),
                    s.bytesPerToken(2) * 4096 / 1048576.0, s.bytesPerToken(4) * 4096 / 1048576.0, s.bytesPerToken(2) * 32768 / 1073741824.0);

        int trainCtx = 4096, requested = 8192, ctx = Math.min(requested, trainCtx);
        Shape tiny = shapes[0];
        short[][] keyCache = new short[tiny.layers()][ctx * tiny.kvHeads() * tiny.headDim()];
        long javaBytes = 2L * keyCache.length * keyCache[0].length * Short.BYTES;
        System.out.println("engine ctx = min(requested " + requested + ", trained " + trainCtx + ") = " + ctx
                + ", short[] K and V arrays: " + javaBytes / 1048576 + " MiB");

        int prompt = 4090, maxTokens = 100, position = prompt;
        int generated = 0;
        while (generated < maxTokens && position < ctx) { position++; generated++; }
        System.out.println("prompt " + prompt + " tokens, max_tokens " + maxTokens + ": generated " + generated
                + ", stop reason " + (generated < maxTokens ? "context" : "length"));
    }
}

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.sun.net.httpserver.HttpServer;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;

public class JavaEngineBackend {
    interface Backend extends AutoCloseable {                        // the interface from Project 12
        String name();
        String baseUrl();
        void ensureReady() throws Exception;
        @Override void close();
    }

    interface Engine { List<String> chat(List<Map<String, String>> messages, int maxTokens); }   // pieces of the answer

    static final class InProcessBackend implements Backend {
        static final ObjectMapper JSON = new ObjectMapper();
        final Engine engine;
        HttpServer server;
        InProcessBackend(Engine engine) { this.engine = engine; }
        public String name() { return "java"; }
        public String baseUrl() { return "http://127.0.0.1:" + server.getAddress().getPort() + "/v1"; }
        public void ensureReady() throws Exception {
            if (server != null) return;
            server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
            server.createContext("/v1/chat/completions", ex -> {
                JsonNode req = JSON.readTree(ex.getRequestBody());
                List<Map<String, String>> messages = JSON.convertValue(req.get("messages"), JSON.getTypeFactory().constructCollectionType(List.class, Map.class));
                ex.getResponseHeaders().add("Content-Type", "text/event-stream");
                ex.sendResponseHeaders(200, 0);
                try (OutputStream out = ex.getResponseBody()) {
                    for (String piece : engine.chat(messages, req.path("max_tokens").asInt(256)))
                        send(out, Map.of("choices", List.of(Map.of("index", 0, "delta", Map.of("content", piece)))));
                    send(out, Map.of("choices", List.of(Map.of("index", 0, "delta", Map.of(), "finish_reason", "stop"))));
                    out.write("data: [DONE]\n\n".getBytes(StandardCharsets.UTF_8));
                }
            });
            server.start();
        }
        static void send(OutputStream out, Object chunk) throws java.io.IOException {
            out.write(("data: " + JSON.writeValueAsString(chunk) + "\n\n").getBytes(StandardCharsets.UTF_8));
            out.flush();
        }
        public void close() { if (server != null) server.stop(0); }
    }

    public static void main(String[] args) throws Exception {
        Engine scripted = (messages, max) -> List.of("Use", " new", " StringBuilder(s)", ".reverse()", ".toString()", ".");
        try (Backend b = new InProcessBackend(scripted)) {
            b.ensureReady();
            System.out.println(b.name() + " backend ready, base URL ends with /v1: " + b.baseUrl().endsWith("/v1"));
            String body = "{\"model\":\"tiny\",\"stream\":true,\"messages\":[{\"role\":\"user\",\"content\":\"Reverse a string.\"}]}";
            HttpRequest req = HttpRequest.newBuilder(URI.create(b.baseUrl() + "/chat/completions"))
                    .POST(HttpRequest.BodyPublishers.ofString(body)).header("Content-Type", "application/json").build();
            HttpResponse<java.io.InputStream> resp = HttpClient.newHttpClient().send(req, HttpResponse.BodyHandlers.ofInputStream());
            StringBuilder answer = new StringBuilder();
            try (BufferedReader in = new BufferedReader(new InputStreamReader(resp.body(), StandardCharsets.UTF_8))) {
                for (String line; (line = in.readLine()) != null; ) {
                    if (!line.startsWith("data: ")) continue;
                    if (line.equals("data: [DONE]")) { System.out.println("[DONE]"); break; }
                    JsonNode choice = new ObjectMapper().readTree(line.substring(6)).get("choices").get(0);
                    if (choice.has("finish_reason")) System.out.println("finish_reason: " + choice.get("finish_reason").asText());
                    else { String d = choice.get("delta").get("content").asText(); answer.append(d); System.out.println("delta [" + d + "]"); }
                }
            }
            System.out.println("answer: " + answer);
        }
    }
}

import java.lang.foreign.MemorySegment;
import java.lang.foreign.ValueLayout;
import java.nio.ByteOrder;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.function.Consumer;

/** Project 20: a chat engine in pure Java on top of Projects 2, 3, 16, 17, 18, and 19. F32 and Q8_0 weights. */
public final class JavaChatEngine implements AutoCloseable {
    static final ValueLayout.OfFloat F32 = ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
    static final ValueLayout.OfShort F16 = ValueLayout.JAVA_SHORT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);

    record Matrix(KernelLibrary.Type type, MemorySegment data, int rows, int cols) {}
    record Layer(float[] attnNorm, Matrix wq, Matrix wk, Matrix wv, float[] bq, float[] bk, float[] bv, Matrix wo,
                 float[] ffnNorm, Matrix gate, Matrix up, Matrix down) {}
    record Result(List<Integer> ids, String text, String stopReason, int promptTokens, long prefillNanos, long decodeNanos) {}

    final GgufModel model;
    final GgufTokenizer tokenizer;
    final Llama.Config c;
    final KernelLibrary.Kernels kernels;
    final Matrix embedding, output;
    final float[] outputNorm;
    final Layer[] layers;
    final int ctx;
    final short[][] keyCache, valueCache;                          // F16, [layer][position * kvDim + i]
    int position;                                                  // number of tokens in the KV cache

    JavaChatEngine(Path path, int ctx, int threads) throws Exception {
        model = GgufModel.open(path);
        tokenizer = new GgufTokenizer(model.metadata());
        c = Llama.Config.from(model.metadata());                 // Project 18: hyperparameters and RoPE layout
        this.ctx = Math.min(ctx, c.contextLength());
        kernels = new KernelLibrary.Kernels(threads);
        embedding = matrix("token_embd.weight");
        output = model.tensors().containsKey("output.weight") ? matrix("output.weight") : embedding;   // tied embeddings
        outputNorm = model.floats("output_norm.weight");
        layers = new Layer[c.layers()];
        for (int l = 0; l < c.layers(); l++) {
            String p = "blk." + l + ".";
            layers[l] = new Layer(model.floats(p + "attn_norm.weight"), matrix(p + "attn_q.weight"), matrix(p + "attn_k.weight"),
                    matrix(p + "attn_v.weight"), optional(p + "attn_q.bias"), optional(p + "attn_k.bias"), optional(p + "attn_v.bias"),
                    matrix(p + "attn_output.weight"), model.floats(p + "ffn_norm.weight"),
                    matrix(p + "ffn_gate.weight"), matrix(p + "ffn_up.weight"), matrix(p + "ffn_down.weight"));
        }
        keyCache = new short[c.layers()][this.ctx * c.kvHeads() * c.headDim()];
        valueCache = new short[c.layers()][this.ctx * c.kvHeads() * c.headDim()];
    }

    float[] optional(String name) { return model.tensors().containsKey(name) ? model.floats(name) : null; }

    Matrix matrix(String name) {
        GgufModel.TensorInfo t = model.tensors().get(name);
        KernelLibrary.Type type = switch (t.type()) {
            case 0 -> KernelLibrary.Type.F32;
            case 8 -> KernelLibrary.Type.Q8_0;
            default -> throw new IllegalArgumentException(name + ": ggml type " + t.type() + " not supported");
        };
        return new Matrix(type, model.tensorData(name), (int) t.dims()[1], (int) t.dims()[0]);
    }

    float[] embed(int token) {                                     // one row of token_embd, dequantized
        float[] x = new float[c.dim()];
        long off = token * embedding.type().rowBytes(c.dim());
        if (embedding.type() == KernelLibrary.Type.F32) MemorySegment.copy(embedding.data(), F32, off, x, 0, c.dim());
        else for (int b = 0; b < c.dim() / 32; b++) {
            float d = Float.float16ToFloat(embedding.data().get(F16, off + b * 34L));
            for (int j = 0; j < 32; j++) x[b * 32 + j] = d * embedding.data().get(ValueLayout.JAVA_BYTE, off + b * 34L + 2 + j);
        }
        return x;
    }

    float[] matvec(Matrix m, float[] x) throws Exception {
        float[] out = new float[m.rows()];
        kernels.matvec(m.type(), m.data(), x, out, m.rows());
        return out;
    }

    /** Runs one token at the next position; returns logits when asked for them. */
    float[] forward(int token, boolean wantLogits) throws Exception {
        if (position >= ctx) throw new IllegalStateException("context full: " + ctx + " tokens");
        int pos = position++, hd = c.headDim(), kvDim = c.kvHeads() * hd, group = c.heads() / c.kvHeads();
        float[] x = embed(token), xb = new float[c.dim()];
        for (int l = 0; l < c.layers(); l++) {
            Layer L = layers[l];
            Llama.rmsnorm(xb, x, L.attnNorm(), c.eps());
            float[] q = matvec(L.wq(), xb), k = matvec(L.wk(), xb), v = matvec(L.wv(), xb);
            Llama.addBias(q, L.bq());
            Llama.addBias(k, L.bk());
            Llama.addBias(v, L.bv());
            Llama.rope(q, c.heads(), hd, c.ropeDims(), pos, c.ropeBase(), c.neox());
            Llama.rope(k, c.kvHeads(), hd, c.ropeDims(), pos, c.ropeBase(), c.neox());
            for (int i = 0; i < kvDim; i++) {
                keyCache[l][pos * kvDim + i] = Float.floatToFloat16(k[i]);
                valueCache[l][pos * kvDim + i] = Float.floatToFloat16(v[i]);
            }
            float[] att = new float[c.dim()], scores = new float[pos + 1];
            for (int h = 0; h < c.heads(); h++) {
                int kvh = (h / group) * hd;                         // grouped-query attention: 4 query heads share one KV head
                float max = Float.NEGATIVE_INFINITY, sum = 0;
                for (int t = 0; t <= pos; t++) {
                    float s = 0;
                    for (int i = 0; i < hd; i++) s += q[h * hd + i] * Float.float16ToFloat(keyCache[l][t * kvDim + kvh + i]);
                    scores[t] = s / (float) Math.sqrt(hd);
                    max = Math.max(max, scores[t]);
                }
                for (int t = 0; t <= pos; t++) { scores[t] = (float) Math.exp(scores[t] - max); sum += scores[t]; }
                for (int t = 0; t <= pos; t++)
                    for (int i = 0; i < hd; i++) att[h * hd + i] += scores[t] / sum * Float.float16ToFloat(valueCache[l][t * kvDim + kvh + i]);
            }
            float[] o = matvec(L.wo(), att);
            for (int i = 0; i < c.dim(); i++) x[i] += o[i];
            Llama.rmsnorm(xb, x, L.ffnNorm(), c.eps());
            float[] g = matvec(L.gate(), xb), u = matvec(L.up(), xb);
            for (int i = 0; i < c.ffn(); i++) g[i] = Llama.silu(g[i]) * u[i];                     // SwiGLU
            float[] d = matvec(L.down(), g);
            for (int i = 0; i < c.dim(); i++) x[i] += d[i];
        }
        if (!wantLogits) return null;
        Llama.rmsnorm(xb, x, outputNorm, c.eps());
        return matvec(output, xb);
    }

    /** Feeds the prompt; logits are computed only for the last token. */
    float[] prefill(List<Integer> tokens) throws Exception {
        float[] logits = null;
        for (int i = 0; i < tokens.size(); i++) logits = forward(tokens.get(i), i == tokens.size() - 1);
        return logits;
    }

    static double[] toDouble(float[] f) { double[] d = new double[f.length]; for (int i = 0; i < f.length; i++) d[i] = f[i]; return d; }

    /** Generates until an end-of-generation token, a stop string, maxTokens, or a full context; streams text to out. */
    Result generate(List<Integer> prompt, SamplerChain sampler, int maxTokens, List<String> stops, Consumer<String> out) throws Exception {
        long t0 = System.nanoTime();
        float[] logits = prefill(prompt);
        long t1 = System.nanoTime();
        List<Integer> history = new ArrayList<>(prompt), ids = new ArrayList<>();
        GgufTokenizer.Utf8Stream utf8 = new GgufTokenizer.Utf8Stream();
        StringBuilder text = new StringBuilder();
        int emitted = 0;
        String reason = "length";
        while (ids.size() < maxTokens) {
            int id = sampler.sample(toDouble(logits), history);
            if (tokenizer.isEog(id)) { reason = "eog"; break; }
            ids.add(id);
            history.add(id);
            text.append(utf8.push(tokenizer.tokenBytes(id, false)));
            int cut = stopIndex(text, stops);
            if (cut >= 0) { text.setLength(cut); reason = "stop"; break; }
            int safe = text.length() - holdBack(text, stops);       // keep back a possible start of a stop string
            if (safe > emitted) { out.accept(text.substring(emitted, safe)); emitted = safe; }
            if (position >= ctx) { reason = "context"; break; }
            if (ids.size() < maxTokens) logits = forward(id, true);
        }
        if (text.length() > emitted) out.accept(text.substring(emitted));
        return new Result(ids, text.toString(), reason, prompt.size(), t1 - t0, System.nanoTime() - t1);
    }

    static int stopIndex(CharSequence text, List<String> stops) {
        int best = -1;
        for (String s : stops) { int i = text.toString().indexOf(s); if (i >= 0 && (best < 0 || i < best)) best = i; }
        return best;
    }

    static int holdBack(CharSequence text, List<String> stops) {   // length of the longest text suffix that starts a stop string
        int keep = 0;
        for (String s : stops)
            for (int n = Math.min(s.length() - 1, text.length()); n > keep; n--)
                if (text.toString().endsWith(s.substring(0, n))) { keep = n; break; }
        return keep;
    }

    static TemplateProject.ChatTemplate templateFor(String jinja) {
        if (jinja == null) throw new IllegalArgumentException("the model has no tokenizer.chat_template");
        if (jinja.contains("<|im_start|>")) return new TemplateProject.ChatMl();
        if (jinja.contains("<|start_header_id|>")) return new TemplateProject.Llama3("26 Jul 2024");
        if (jinja.contains("[INST]")) return new TemplateProject.Mistral();
        throw new IllegalArgumentException("unknown chat template");
    }

    void reset() { position = 0; }

    @Override public void close() { kernels.close(); model.close(); }

    static String top3(float[] logits) {                          // log-softmax of the three largest logits
        double max = Double.NEGATIVE_INFINITY, sum = 0;
        for (float l : logits) max = Math.max(max, l);
        for (float l : logits) sum += Math.exp(l - max);
        int[] best = java.util.stream.IntStream.range(0, logits.length).boxed()
                .sorted((x, y) -> Float.compare(logits[y], logits[x])).limit(3).mapToInt(Integer::intValue).toArray();
        StringBuilder sb = new StringBuilder();
        for (int id : best) sb.append(String.format(Locale.ROOT, "%d (%.6f) ", id, logits[id] - max - Math.log(sum)));
        return sb.toString().trim();
    }

    public static void main(String[] args) throws Exception {
        String mode = args.length > 1 ? args[1] : "raw";
        int threads = args.length > 2 ? Integer.parseInt(args[2]) : 4, n = args.length > 3 ? Integer.parseInt(args[3]) : 32;
        try (JavaChatEngine e = new JavaChatEngine(Path.of(args[0]), 4096, threads)) {
            Llama.Config c = e.c;
            System.out.printf(Locale.ROOT, "config: dim %d, layers %d, heads %d, kv heads %d, head dim %d, ffn %d, vocab %d, ctx %d, weights %s%n",
                    c.dim(), c.layers(), c.heads(), c.kvHeads(), c.headDim(), c.ffn(), c.vocab(), e.ctx, e.layers[0].wq().type());
            TemplateProject.ChatTemplate t = templateFor(e.model.string("tokenizer.chat_template"));
            List<Integer> chat = e.tokenizer.encode(t.render(List.of(new TemplateProject.Msg("system", "You are a helpful assistant"),
                    new TemplateProject.Msg("user", "Write a Java method that reverses a string.")), true), true, true);
            SamplerChain greedy = SamplerChain.builder().temperature(0).build();
            if (mode.equals("bench")) {
                for (int run = 1; run <= 3; run++) {
                    e.reset();
                    long t0 = System.nanoTime();
                    float[] logits = e.prefill(chat);
                    long t1 = System.nanoTime(), sampleNanos = 0;
                    List<Integer> history = new ArrayList<>(chat);
                    for (int i = 0; i < n; i++) {
                        long s0 = System.nanoTime();
                        int id = greedy.sample(toDouble(logits), history);
                        sampleNanos += System.nanoTime() - s0;
                        history.add(id);
                        logits = e.forward(id, true);
                    }
                    long t2 = System.nanoTime();
                    System.out.printf(Locale.ROOT, "run %d: prefill %d tokens %.1f tokens/s, decode %d tokens %.1f tokens/s, sampler %.1f ms of %.1f ms%n",
                            run, chat.size(), chat.size() / ((t1 - t0) / 1e9), n, n / ((t2 - t1) / 1e9), sampleNanos / 1e6, (t2 - t1) / 1e6);
                }
                return;
            }
            List<Integer> prompt = mode.equals("chat") ? chat : e.tokenizer.encode("Hello world", true, false);
            List<String> stops = mode.equals("chat") ? List.of(t.stopToken()) : List.of();
            System.out.println("prompt ids: " + prompt);
            System.out.println("first token, top 3 log-probabilities: " + top3(e.prefill(prompt)));
            e.reset();
            StringBuilder streamed = new StringBuilder();
            Result r = e.generate(prompt, greedy, n, stops, streamed::append);
            System.out.println("generated ids: " + r.ids());
            System.out.println("text: " + r.text().replace("\n", "\\n"));
            System.out.println("stop reason: " + r.stopReason() + ", streamed text equals result: " + streamed.toString().equals(r.text()));
        }
    }
}

import java.util.Arrays;
import java.util.PriorityQueue;
import java.util.Random;

public class TopKFirst {
    // As Project 2 does it: copy all logits, sort a copy to find the k-th value, mask, then softmax and draw over all ids.
    static int fullChain(float[] logits, int k, double temp, Random rng) {
        double[] l = new double[logits.length];
        for (int i = 0; i < l.length; i++) l[i] = logits[i];
        double[] sorted = l.clone();
        Arrays.sort(sorted);
        double kth = sorted[l.length - k];
        for (int i = 0; i < l.length; i++) l[i] = l[i] < kth ? Double.NEGATIVE_INFINITY : l[i] / temp;
        return draw(l, null, rng);
    }

    // Extension: keep the k largest with a size-k heap in one pass, then run the same stages on k candidates only.
    static int topKFirst(float[] logits, int k, double temp, Random rng) {
        PriorityQueue<Integer> heap = new PriorityQueue<>((a, b) -> Float.compare(logits[a], logits[b]));
        for (int i = 0; i < logits.length; i++) {
            if (heap.size() < k) heap.add(i);
            else if (logits[i] > logits[heap.peek()]) { heap.poll(); heap.add(i); }
        }
        int[] ids = heap.stream().mapToInt(Integer::intValue).sorted().toArray();   // id order, as the full version scans
        double[] l = new double[k];
        for (int j = 0; j < k; j++) l[j] = logits[ids[j]] / temp;
        return draw(l, ids, rng);
    }

    static int draw(double[] l, int[] ids, Random rng) {        // softmax and inverse CDF in index order
        double max = Arrays.stream(l).max().getAsDouble(), sum = 0;
        double[] p = new double[l.length];
        for (int i = 0; i < l.length; i++) { p[i] = Math.exp(l[i] - max); sum += p[i]; }
        double u = rng.nextDouble() * sum, run = 0;
        for (int i = 0; i < p.length; i++) { run += p[i]; if (p[i] > 0 && u < run) return ids == null ? i : ids[i]; }
        throw new IllegalStateException();
    }

    public static void main(String[] args) {
        int vocab = 151_936, k = 40, steps = 12;
        Random weights = new Random(11);
        Random a = new Random(5), b = new Random(5);
        int same = 0;
        StringBuilder chosen = new StringBuilder();
        for (int s = 0; s < steps; s++) {
            float[] logits = new float[vocab];
            for (int i = 0; i < vocab; i++) logits[i] = (float) (weights.nextGaussian() * 2.5);
            int x = fullChain(logits, k, 0.8, a), y = topKFirst(logits, k, 0.8, b);
            if (x == y) same++;
            chosen.append(x).append(' ');
        }
        System.out.println("chosen ids: " + chosen.toString().trim());
        System.out.println("steps with the same id: " + same + " of " + steps);
        System.out.println("values sorted per step: full " + vocab + ", heap keeps " + k);
    }
}