From 7e27033fefa4246d8561c0be3a326a9db83fa487 Mon Sep 17 00:00:00 2001 From: Josh Brown Date: Sun, 12 Apr 2026 18:44:55 -0400 Subject: [PATCH] feat: add toy character-level transformer language model Implement a minimal GPT-style transformer operating at character level with deliberately tiny dimensions, suitable for this library's toy scale: - Tokenizer: character-level encoding (a-z + space, vocab size 27) - softmax: numerically stable softmax utility function - TransformerBlock: single-head self-attention with causal masking, feed-forward network with ReLU, and residual connections - LanguageModel: full model composing embeddings, positional encodings, transformer block, and output projection. Supports train() via finite-difference gradients and generate() with greedy decoding. Includes comprehensive tests for all new components (36 new tests). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- README.md | 1 + .../machine-learning/models/LanguageModel.ts | 300 ++++++++++++++++++ .../machine-learning/models/Tokenizer.ts | 59 ++++ .../models/TransformerBlock.ts | 208 ++++++++++++ .../models/__tests__/LanguageModel.spec.ts | 126 ++++++++ .../models/__tests__/Tokenizer.spec.ts | 86 +++++ .../models/__tests__/TransformerBlock.spec.ts | 174 ++++++++++ src/index.ts | 3 + src/utilities/NumberUtilities.ts | 15 + .../__tests__/NumberUtilities.spec.ts | 49 +++ 10 files changed, 1021 insertions(+) create mode 100644 src/applications/machine-learning/models/LanguageModel.ts create mode 100644 src/applications/machine-learning/models/Tokenizer.ts create mode 100644 src/applications/machine-learning/models/TransformerBlock.ts create mode 100644 src/applications/machine-learning/models/__tests__/LanguageModel.spec.ts create mode 100644 src/applications/machine-learning/models/__tests__/Tokenizer.spec.ts create mode 100644 src/applications/machine-learning/models/__tests__/TransformerBlock.spec.ts diff --git a/README.md b/README.md index d57365e..d035420 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ A linear algebra library written in TypeScript that focuses on generality, exten - Regularized linear regression - Logistic Regression - Support Vector Machines + - Character-level language model (toy transformer) - And more to come! ## Using Vector diff --git a/src/applications/machine-learning/models/LanguageModel.ts b/src/applications/machine-learning/models/LanguageModel.ts new file mode 100644 index 0000000..3b0bf6e --- /dev/null +++ b/src/applications/machine-learning/models/LanguageModel.ts @@ -0,0 +1,300 @@ +import { Matrix } from '../../../types/matrix/Matrix'; +import { FloatMatrix } from '../../../types/matrix/FloatMatrix'; +import { softmax } from '../../../utilities/NumberUtilities'; +import { Tokenizer } from './Tokenizer'; +import { TransformerBlock } from './TransformerBlock'; + +const mb = FloatMatrix.builder(); + +/** + * The set of hyperparameters for a {@link LanguageModel}. + * @public + */ +export type LanguageModelHyperparams = { + /** The dimension of token and positional embeddings */ + embedDim: number; + /** The hidden dimension of the feed-forward network in the transformer block */ + ffDim: number; + /** The maximum number of tokens in the input context */ + contextLength: number; + /** The step size for gradient descent */ + learningRate: number; + /** The number of gradient descent iterations during training */ + maxIterations: number; +}; + +const DEFAULT_HYPERPARAMS: LanguageModelHyperparams = { + embedDim: 8, + ffDim: 16, + contextLength: 16, + learningRate: 0.01, + maxIterations: 100, +}; + +/** + * A character-level language model that composes a {@link Tokenizer}, + * learned token and positional embeddings, a single {@link TransformerBlock}, + * and an output projection layer into a trainable autoregressive model. + * + * @remarks + * Training uses finite-difference gradient estimation, which is suitable only + * for toy-sized models. Keep dimensions small (embedDim ≤ 8, ffDim ≤ 16). + * + * @public + */ +export class LanguageModel { + private readonly _hyperparams: Readonly; + private readonly _tokenizer: Tokenizer; + private readonly _block: TransformerBlock; + + private _tokenEmbeddings: Matrix; + private _positionalEncodings: Matrix; + private _outputProjection: Matrix; + + /** + * Creates a new LanguageModel with randomly initialised weights. + * + * @param hyperparams - Optional partial hyperparameters merged with defaults + * @public + */ + constructor(hyperparams?: Partial) { + this._hyperparams = Object.freeze({ ...DEFAULT_HYPERPARAMS, ...hyperparams }); + + const { embedDim, ffDim, contextLength } = this._hyperparams; + const std = 0.02; + + this._tokenizer = new Tokenizer(); + + this._tokenEmbeddings = mb.randomNormal([this._tokenizer.vocabSize, embedDim], 0, std); + this._positionalEncodings = mb.randomNormal([contextLength, embedDim], 0, std); + this._block = new TransformerBlock(embedDim, ffDim, contextLength); + this._outputProjection = mb.randomNormal([embedDim, this._tokenizer.vocabSize], 0, std); + } + + /** + * Returns the full set of hyperparameters (with defaults filled in). + * @public + */ + public getHyperParameters(): LanguageModelHyperparams { + return { ...this._hyperparams }; + } + + /** + * Runs the forward pass of the language model on a sequence of token IDs. + * + * @param tokenIds - An array of integer token IDs + * @returns A probability distribution over the vocabulary for the next token + * @public + */ + public forward(tokenIds: number[]): number[] { + const { embedDim, contextLength } = this._hyperparams; + const vocabSize = this._tokenizer.vocabSize; + const seqLen = tokenIds.length; + + // Look up token embeddings and add positional encodings + const embData: number[][] = []; + const tokArr = this._tokenEmbeddings.toArray(); + const posArr = this._positionalEncodings.toArray(); + + for (let t = 0; t < seqLen; t++) { + const tokRow = tokArr[tokenIds[t]]; + const posRow = posArr[t % contextLength]; + const row: number[] = new Array(embedDim); + for (let j = 0; j < embedDim; j++) { + row[j] = tokRow[j] + posRow[j]; + } + embData.push(row); + } + + const input = mb.fromArray(embData); + + // Transformer block + const transformed = this._block.forward(input); + + // Take last row and project to vocabulary logits + const lastRow = transformed.getRow(seqLen - 1).toArray(); + const projArr = this._outputProjection.toArray(); + const logits: number[] = new Array(vocabSize); + for (let v = 0; v < vocabSize; v++) { + let sum = 0; + for (let j = 0; j < embedDim; j++) { + sum += lastRow[j] * projArr[j][v]; + } + logits[v] = sum; + } + + return softmax(logits); + } + + /** + * Trains the model on a text corpus using finite-difference gradient descent. + * + * @param corpus - The training text + * @public + */ + public train(corpus: string): void { + const { learningRate, maxIterations } = this._hyperparams; + const tokens = this._tokenizer.encode(corpus); + + if (tokens.length < 2) { + return; + } + + // Build training examples + const examples: { input: number[]; target: number }[] = []; + for (let i = 1; i < tokens.length; i++) { + const start = Math.max(0, i - this._hyperparams.contextLength); + examples.push({ input: tokens.slice(start, i), target: tokens[i] }); + } + + const params = this.getParameters(); + const eps = 1e-5; + + for (let iter = 0; iter < maxIterations; iter++) { + this.setParameters(params); + + // Finite-difference gradients + const gradient: number[] = new Array(params.length); + for (let p = 0; p < params.length; p++) { + const saved = params[p]; + + params[p] = saved + eps; + this.setParameters(params); + const costPlus = this.computeCost(examples); + + params[p] = saved - eps; + this.setParameters(params); + const costMinus = this.computeCost(examples); + + gradient[p] = (costPlus - costMinus) / (2 * eps); + params[p] = saved; + } + + // Update parameters + for (let p = 0; p < params.length; p++) { + params[p] -= learningRate * gradient[p]; + } + } + + this.setParameters(params); + } + + /** + * Generates text by autoregressively sampling from the model. + * + * @param prompt - The initial text to condition generation on + * @param length - The number of tokens to generate + * @returns The prompt followed by the generated text + * @public + */ + public generate(prompt: string, length: number): string { + const tokens = this._tokenizer.encode(prompt); + + for (let i = 0; i < length; i++) { + const contextStart = Math.max(0, tokens.length - this._hyperparams.contextLength); + const context = tokens.slice(contextStart); + const probs = this.forward(context); + + // Greedy decoding: pick the argmax + let bestIdx = 0; + let bestProb = probs[0]; + for (let v = 1; v < probs.length; v++) { + if (probs[v] > bestProb) { + bestProb = probs[v]; + bestIdx = v; + } + } + tokens.push(bestIdx); + } + + return this._tokenizer.decode(tokens); + } + + /** + * Returns all learnable parameters as a flat number array. + * + * @remarks + * Order: token embeddings (row-major), positional encodings (row-major), + * transformer block parameters, output projection (row-major). + * + * @public + */ + public getParameters(): number[] { + return [ + ...flattenMatrix(this._tokenEmbeddings), + ...flattenMatrix(this._positionalEncodings), + ...this._block.getParameters(), + ...flattenMatrix(this._outputProjection), + ]; + } + + /** + * Reconstructs all weight matrices from a flat parameter array. + * + * @param params - A flat array whose length must equal {@link getParameterCount} + * @public + */ + public setParameters(params: number[]): void { + const expected = this.getParameterCount(); + if (params.length !== expected) { + throw new Error(`Expected ${expected} parameters but received ${params.length}`); + } + + const { embedDim, contextLength } = this._hyperparams; + const vocabSize = this._tokenizer.vocabSize; + let offset = 0; + + const readMatrix = (rows: number, cols: number): Matrix => { + const data: number[][] = []; + for (let i = 0; i < rows; i++) { + data.push(params.slice(offset, offset + cols)); + offset += cols; + } + return mb.fromArray(data); + }; + + this._tokenEmbeddings = readMatrix(vocabSize, embedDim); + this._positionalEncodings = readMatrix(contextLength, embedDim); + + const blockParamCount = this._block.getParameterCount(); + this._block.setParameters(params.slice(offset, offset + blockParamCount)); + offset += blockParamCount; + + this._outputProjection = readMatrix(embedDim, vocabSize); + } + + /** + * Returns the total number of learnable parameters. + * @public + */ + public getParameterCount(): number { + const { embedDim, contextLength } = this._hyperparams; + const vocabSize = this._tokenizer.vocabSize; + + const embeddingParams = vocabSize * embedDim; + const positionalParams = contextLength * embedDim; + const blockParams = this._block.getParameterCount(); + const outputParams = embedDim * vocabSize; + + return embeddingParams + positionalParams + blockParams + outputParams; + } + + private computeCost(examples: { input: number[]; target: number }[]): number { + let totalLoss = 0; + for (const example of examples) { + const probs = this.forward(example.input); + const p = Math.max(probs[example.target], 1e-12); + totalLoss += -Math.log(p); + } + return totalLoss / examples.length; + } +} + +/** Flattens a matrix row-by-row into a 1-D array. */ +function flattenMatrix(m: Matrix): number[] { + const result: number[] = []; + for (const row of m.toArray()) { + result.push(...row); + } + return result; +} diff --git a/src/applications/machine-learning/models/Tokenizer.ts b/src/applications/machine-learning/models/Tokenizer.ts new file mode 100644 index 0000000..b369353 --- /dev/null +++ b/src/applications/machine-learning/models/Tokenizer.ts @@ -0,0 +1,59 @@ +/** + * A character-level tokenizer that maps lowercase a-z and space to integer indices. + * + * The vocabulary consists of 27 characters: space (index 0), then 'a' through 'z' (indices 1-26). + * Input is lowercased before encoding; characters outside the vocabulary are skipped. + * @public + */ +export class Tokenizer { + private readonly _charToIndex: Map; + private readonly _indexToChar: Map; + + constructor() { + this._charToIndex = new Map(); + this._indexToChar = new Map(); + + this._charToIndex.set(' ', 0); + this._indexToChar.set(0, ' '); + + for (let i = 0; i < 26; i++) { + const char = String.fromCharCode('a'.charCodeAt(0) + i); + this._charToIndex.set(char, i + 1); + this._indexToChar.set(i + 1, char); + } + } + + /** + * The number of tokens in the vocabulary. + * @public + */ + public get vocabSize(): number { + return 27; + } + + /** + * Converts a string to an array of token IDs. + * The input is lowercased before encoding. Characters not in the vocabulary are skipped. + * @param text - The string to encode + * @public + */ + public encode(text: string): number[] { + const tokens: number[] = []; + for (const char of text.toLowerCase()) { + const index = this._charToIndex.get(char); + if (index !== undefined) { + tokens.push(index); + } + } + return tokens; + } + + /** + * Converts an array of token IDs back to a string. + * @param tokens - The array of token IDs to decode + * @public + */ + public decode(tokens: number[]): string { + return tokens.map((token) => this._indexToChar.get(token) ?? '').join(''); + } +} diff --git a/src/applications/machine-learning/models/TransformerBlock.ts b/src/applications/machine-learning/models/TransformerBlock.ts new file mode 100644 index 0000000..df34adb --- /dev/null +++ b/src/applications/machine-learning/models/TransformerBlock.ts @@ -0,0 +1,208 @@ +import { Matrix } from '../../../types/matrix/Matrix'; +import { Vector } from '../../../types/vector/Vector'; +import { FloatMatrix } from '../../../types/matrix/FloatMatrix'; +import { FloatVector } from '../../../types/vector/FloatVector'; +import { softmax } from '../../../utilities/NumberUtilities'; + +const mb = FloatMatrix.builder(); +const vb = FloatVector.builder(); + +/** + * A single transformer block with single-head self-attention and a feed-forward network. + * + * @remarks + * Implements causal (autoregressive) masking so that each token can only attend + * to itself and previous tokens. Includes residual connections around both the + * attention and feed-forward sub-layers. + * + * @public + */ +export class TransformerBlock { + private readonly _embedDim: number; + private readonly _ffDim: number; + + private _Wq: Matrix; + private _Wk: Matrix; + private _Wv: Matrix; + private _Wo: Matrix; + private _W1: Matrix; + private _b1: Vector; + private _W2: Matrix; + private _b2: Vector; + + /** + * Creates a new TransformerBlock with randomly initialised weights. + * + * @param embedDim - The embedding dimension + * @param ffDim - The hidden dimension of the feed-forward network + * @param contextLength - The maximum sequence length supported + * @public + */ + constructor(embedDim: number, ffDim: number, contextLength: number) { + this._embedDim = embedDim; + this._ffDim = ffDim; + // contextLength is accepted for API compatibility but the causal mask + // is built dynamically from the actual sequence length. + void contextLength; + + const std = 0.02; + + // Attention weights (embedDim × embedDim) + this._Wq = mb.randomNormal([embedDim, embedDim], 0, std); + this._Wk = mb.randomNormal([embedDim, embedDim], 0, std); + this._Wv = mb.randomNormal([embedDim, embedDim], 0, std); + this._Wo = mb.randomNormal([embedDim, embedDim], 0, std); + + // Feed-forward weights + this._W1 = mb.randomNormal([embedDim, ffDim], 0, std); + this._b1 = vb.zeros(ffDim); + this._W2 = mb.randomNormal([ffDim, embedDim], 0, std); + this._b2 = vb.zeros(embedDim); + } + + /** + * Runs the forward pass of the transformer block. + * + * @param input - A (seqLen × embedDim) matrix where each row is a token embedding + * @returns A (seqLen × embedDim) matrix of transformed token embeddings + * @public + */ + public forward(input: Matrix): Matrix { + // --- Self-attention --- + const Q = input.multiply(this._Wq); + const K = input.multiply(this._Wk); + const V = input.multiply(this._Wv); + + // Attention scores: (seqLen × seqLen) + const scale = 1 / Math.sqrt(this._embedDim); + const rawScores = Q.multiply(K.transpose()).scalarMultiply(scale); + + // Causal mask + softmax + const masked = applyCausalMask(rawScores); + const attnWeights = rowWiseSoftmax(masked); + + // Weighted values and output projection + const attnOut = attnWeights.multiply(V).multiply(this._Wo); + + // Residual connection + const x = input.add(attnOut); + + // --- Feed-forward network --- + const ff1 = addBiasRows(x.multiply(this._W1), this._b1).map((v) => Math.max(0, v)); + const ff2 = addBiasRows(ff1.multiply(this._W2), this._b2); + + // Residual connection + return x.add(ff2); + } + + /** + * Returns all learnable parameters as a flat number array. + * + * @remarks + * Order: Wq, Wk, Wv, Wo, W1, b1, W2, b2. + * Matrices are flattened in row-major order via {@link Matrix.toArray}. + * + * @returns A flat array of all parameter values + * @public + */ + public getParameters(): number[] { + return [ + ...flattenMatrix(this._Wq), + ...flattenMatrix(this._Wk), + ...flattenMatrix(this._Wv), + ...flattenMatrix(this._Wo), + ...flattenMatrix(this._W1), + ...this._b1.toArray(), + ...flattenMatrix(this._W2), + ...this._b2.toArray(), + ]; + } + + /** + * Reconstructs all weight matrices and bias vectors from a flat parameter array. + * + * @param params - A flat array whose length must equal {@link getParameterCount} + * @public + */ + public setParameters(params: number[]): void { + const expected = this.getParameterCount(); + if (params.length !== expected) { + throw new Error(`Expected ${expected} parameters but received ${params.length}`); + } + + let offset = 0; + + const readMatrix = (rows: number, cols: number): Matrix => { + const data: number[][] = []; + for (let i = 0; i < rows; i++) { + data.push(params.slice(offset, offset + cols)); + offset += cols; + } + return mb.fromArray(data); + }; + + const readVector = (dim: number): Vector => { + const data = params.slice(offset, offset + dim); + offset += dim; + return vb.fromArray(data); + }; + + const d = this._embedDim; + const f = this._ffDim; + + this._Wq = readMatrix(d, d); + this._Wk = readMatrix(d, d); + this._Wv = readMatrix(d, d); + this._Wo = readMatrix(d, d); + this._W1 = readMatrix(d, f); + this._b1 = readVector(f); + this._W2 = readMatrix(f, d); + this._b2 = readVector(d); + } + + /** + * Returns the total number of learnable parameters. + * @public + */ + public getParameterCount(): number { + const d = this._embedDim; + const f = this._ffDim; + return 4 * d * d + d * f + f + f * d + d; + } +} + +/** Flattens a matrix row-by-row into a 1-D array. */ +function flattenMatrix(m: Matrix): number[] { + const result: number[] = []; + for (const row of m.toArray()) { + result.push(...row); + } + return result; +} + +/** Sets scores[i][j] = -Infinity for j > i (future positions). */ +function applyCausalMask(scores: Matrix): Matrix { + const rows = scores.getNumberOfRows(); + const cols = scores.getNumberOfColumns(); + const data: number[][] = scores.toArray(); + for (let i = 0; i < rows; i++) { + for (let j = i + 1; j < cols; j++) { + data[i][j] = -Infinity; + } + } + return mb.fromArray(data); +} + +/** Applies softmax independently to each row of a matrix. */ +function rowWiseSoftmax(m: Matrix): Matrix { + const data = m.toArray().map((row) => softmax(row)); + return mb.fromArray(data); +} + +/** Broadcasts a bias vector across every row of a matrix. */ +function addBiasRows(m: Matrix, bias: Vector): Matrix { + const biasArr = bias.toArray(); + const rows = m.getNumberOfRows(); + const biasMatrix = mb.fromArray(Array.from({ length: rows }, () => [...biasArr])); + return m.add(biasMatrix); +} diff --git a/src/applications/machine-learning/models/__tests__/LanguageModel.spec.ts b/src/applications/machine-learning/models/__tests__/LanguageModel.spec.ts new file mode 100644 index 0000000..c2d081b --- /dev/null +++ b/src/applications/machine-learning/models/__tests__/LanguageModel.spec.ts @@ -0,0 +1,126 @@ +import { LanguageModel } from '../LanguageModel'; + +describe('LanguageModel', () => { + const tinyParams = { + embedDim: 2, + ffDim: 2, + contextLength: 4, + learningRate: 0.01, + maxIterations: 2, + }; + + describe('constructor and defaults', () => { + it('applies default hyperparameters', () => { + const model = new LanguageModel(); + const hp = model.getHyperParameters(); + expect(hp.embedDim).toBe(8); + expect(hp.ffDim).toBe(16); + expect(hp.contextLength).toBe(16); + expect(hp.learningRate).toBe(0.01); + expect(hp.maxIterations).toBe(100); + }); + + it('merges partial hyperparameters with defaults', () => { + const model = new LanguageModel({ embedDim: 4 }); + const hp = model.getHyperParameters(); + expect(hp.embedDim).toBe(4); + expect(hp.ffDim).toBe(16); + }); + }); + + describe('forward', () => { + it('produces valid probabilities', () => { + const model = new LanguageModel(tinyParams); + const probs = model.forward([1, 2, 3]); + + expect(probs.length).toBe(27); + const sum = probs.reduce((a, b) => a + b, 0); + expect(sum).toBeCloseTo(1.0, 5); + probs.forEach((p) => { + expect(p).toBeGreaterThanOrEqual(0); + }); + }); + + it('works on an untrained model', () => { + const model = new LanguageModel(tinyParams); + const probs = model.forward([0]); + + expect(probs.length).toBe(27); + const sum = probs.reduce((a, b) => a + b, 0); + expect(sum).toBeCloseTo(1.0, 5); + }); + }); + + describe('getParameterCount', () => { + it('returns the correct count', () => { + const model = new LanguageModel(tinyParams); + const { embedDim, contextLength } = tinyParams; + const vocabSize = 27; + + // Token embeddings + positional encodings + transformer block + output projection + // Transformer block: 4*d*d + d*f + f + f*d + d + const d = embedDim; + const f = tinyParams.ffDim; + const blockParams = 4 * d * d + d * f + f + f * d + d; + const expected = vocabSize * d + contextLength * d + blockParams + d * vocabSize; + + expect(model.getParameterCount()).toBe(expected); + }); + }); + + describe('getParameters / setParameters', () => { + it('survives a roundtrip', () => { + const model = new LanguageModel(tinyParams); + const params = model.getParameters(); + + expect(params.length).toBe(model.getParameterCount()); + + // Create a second model and load the same parameters + const model2 = new LanguageModel(tinyParams); + model2.setParameters(params); + + const params2 = model2.getParameters(); + expect(params2.length).toBe(params.length); + params.forEach((val, i) => { + expect(params2[i]).toBeCloseTo(val, 10); + }); + }); + + it('produces identical forward results after parameter transfer', () => { + const model = new LanguageModel(tinyParams); + const params = model.getParameters(); + const probs1 = model.forward([1, 2]); + + const model2 = new LanguageModel(tinyParams); + model2.setParameters(params); + const probs2 = model2.forward([1, 2]); + + probs1.forEach((val, i) => { + expect(probs2[i]).toBeCloseTo(val, 10); + }); + }); + }); + + describe('train', () => { + it('runs without error on a tiny corpus', () => { + const model = new LanguageModel(tinyParams); + expect(() => model.train('aaa')).not.toThrow(); + }); + }); + + describe('generate', () => { + it('produces a non-empty string', () => { + const model = new LanguageModel(tinyParams); + const result = model.generate('a', 2); + expect(result.length).toBeGreaterThan(0); + expect(typeof result).toBe('string'); + }); + + it('starts with the prompt', () => { + const model = new LanguageModel(tinyParams); + const result = model.generate('ab', 1); + expect(result.slice(0, 2)).toBe('ab'); + expect(result.length).toBe(3); + }); + }); +}); diff --git a/src/applications/machine-learning/models/__tests__/Tokenizer.spec.ts b/src/applications/machine-learning/models/__tests__/Tokenizer.spec.ts new file mode 100644 index 0000000..9e4d0e8 --- /dev/null +++ b/src/applications/machine-learning/models/__tests__/Tokenizer.spec.ts @@ -0,0 +1,86 @@ +import { Tokenizer } from '../Tokenizer'; + +describe('Tokenizer', () => { + let tokenizer: Tokenizer; + + beforeEach(() => { + tokenizer = new Tokenizer(); + }); + + describe('vocabSize', () => { + it('returns 27', () => { + expect(tokenizer.vocabSize).toBe(27); + }); + }); + + describe('encode', () => { + it('encodes a simple string', () => { + expect(tokenizer.encode('abc')).toEqual([1, 2, 3]); + }); + + it('encodes a space as 0', () => { + expect(tokenizer.encode(' ')).toEqual([0]); + }); + + it('encodes a string with spaces', () => { + expect(tokenizer.encode('hello world')).toEqual([8, 5, 12, 12, 15, 0, 23, 15, 18, 12, 4]); + }); + + it('lowercases input before encoding', () => { + expect(tokenizer.encode('Hello')).toEqual(tokenizer.encode('hello')); + }); + + it('skips unknown characters such as digits and punctuation', () => { + expect(tokenizer.encode('a1b!c')).toEqual([1, 2, 3]); + }); + + it('returns an empty array for an empty string', () => { + expect(tokenizer.encode('')).toEqual([]); + }); + + it('returns an empty array for a string with only unknown characters', () => { + expect(tokenizer.encode('123!@#')).toEqual([]); + }); + + it('encodes z as 26', () => { + expect(tokenizer.encode('z')).toEqual([26]); + }); + }); + + describe('decode', () => { + it('decodes a simple token array', () => { + expect(tokenizer.decode([1, 2, 3])).toBe('abc'); + }); + + it('decodes 0 as a space', () => { + expect(tokenizer.decode([0])).toBe(' '); + }); + + it('decodes an empty array to an empty string', () => { + expect(tokenizer.decode([])).toBe(''); + }); + + it('skips out-of-range token IDs', () => { + expect(tokenizer.decode([1, 99, 2])).toBe('ab'); + }); + }); + + describe('encode/decode roundtrip', () => { + it('roundtrips a lowercase string', () => { + const text = 'the quick brown fox'; + expect(tokenizer.decode(tokenizer.encode(text))).toBe(text); + }); + + it('roundtrips a single character', () => { + expect(tokenizer.decode(tokenizer.encode('a'))).toBe('a'); + }); + + it('roundtrips a space', () => { + expect(tokenizer.decode(tokenizer.encode(' '))).toBe(' '); + }); + + it('roundtrips after lowercasing', () => { + expect(tokenizer.decode(tokenizer.encode('ABC'))).toBe('abc'); + }); + }); +}); diff --git a/src/applications/machine-learning/models/__tests__/TransformerBlock.spec.ts b/src/applications/machine-learning/models/__tests__/TransformerBlock.spec.ts new file mode 100644 index 0000000..70f942e --- /dev/null +++ b/src/applications/machine-learning/models/__tests__/TransformerBlock.spec.ts @@ -0,0 +1,174 @@ +import { TransformerBlock } from '../TransformerBlock'; +import { mat } from '../../../../utilities/aliases'; +import { FloatMatrix } from '../../../../types/matrix/FloatMatrix'; + +describe('TransformerBlock', () => { + const embedDim = 4; + const ffDim = 8; + const contextLength = 6; + + function makeInput(seqLen: number): ReturnType { + return FloatMatrix.builder().random([seqLen, embedDim], -1, 1); + } + + describe('forward', () => { + it('produces output with the correct shape', () => { + const block = new TransformerBlock(embedDim, ffDim, contextLength); + const input = makeInput(3); + const output = block.forward(input); + + expect(output.getNumberOfRows()).toBe(3); + expect(output.getNumberOfColumns()).toBe(embedDim); + }); + + it('works with different sequence lengths', () => { + const block = new TransformerBlock(embedDim, ffDim, contextLength); + + for (const seqLen of [1, 2, 5]) { + const output = block.forward(makeInput(seqLen)); + expect(output.getNumberOfRows()).toBe(seqLen); + expect(output.getNumberOfColumns()).toBe(embedDim); + } + }); + + it('produces finite output values', () => { + const block = new TransformerBlock(embedDim, ffDim, contextLength); + const output = block.forward(makeInput(4)); + + output.toArray().forEach((row) => { + row.forEach((val) => { + expect(Number.isFinite(val)).toBe(true); + }); + }); + }); + }); + + describe('getParameterCount', () => { + it('returns the correct total', () => { + const block = new TransformerBlock(embedDim, ffDim, contextLength); + // 4 attention matrices (d×d) + W1 (d×f) + b1 (f) + W2 (f×d) + b2 (d) + const expected = + 4 * embedDim * embedDim + embedDim * ffDim + ffDim + ffDim * embedDim + embedDim; + expect(block.getParameterCount()).toBe(expected); + }); + }); + + describe('getParameters / setParameters roundtrip', () => { + it('recovers the same parameters after a roundtrip', () => { + const block = new TransformerBlock(embedDim, ffDim, contextLength); + const params = block.getParameters(); + + const block2 = new TransformerBlock(embedDim, ffDim, contextLength); + block2.setParameters(params); + + const recovered = block2.getParameters(); + expect(recovered.length).toBe(params.length); + recovered.forEach((val, i) => { + expect(val).toBeCloseTo(params[i], 12); + }); + }); + + it('throws when given the wrong number of parameters', () => { + const block = new TransformerBlock(embedDim, ffDim, contextLength); + expect(() => block.setParameters([1, 2, 3])).toThrow(); + }); + }); + + describe('forward output changes after setParameters', () => { + it('produces different output with different parameters', () => { + const block = new TransformerBlock(embedDim, ffDim, contextLength); + const input = mat([ + [1, 0, 0, 0], + [0, 1, 0, 0], + ]); + + const outputBefore = block.forward(input).toArray(); + + // Set all parameters to a known different value + const newParams = new Array(block.getParameterCount()).fill(0.1); + block.setParameters(newParams); + + const outputAfter = block.forward(input).toArray(); + + // At least one entry should differ + let differs = false; + for (let i = 0; i < outputBefore.length && !differs; i++) { + for (let j = 0; j < outputBefore[i].length && !differs; j++) { + if (Math.abs(outputBefore[i][j] - outputAfter[i][j]) > 1e-10) { + differs = true; + } + } + } + expect(differs).toBe(true); + }); + }); + + describe('causal masking', () => { + it('ensures output at position i is unaffected by changes at position j > i', () => { + const block = new TransformerBlock(embedDim, ffDim, contextLength); + + // Fix parameters so both runs use the same weights + const params = block.getParameters(); + + const inputA = mat([ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + ]); + + // Change only the last row (position 2) + const inputB = mat([ + [1, 2, 3, 4], + [5, 6, 7, 8], + [0, 0, 0, 0], + ]); + + block.setParameters(params); + const outputA = block.forward(inputA).toArray(); + + block.setParameters(params); + const outputB = block.forward(inputB).toArray(); + + // Row 0 should be identical (cannot see positions 1 or 2) + for (let j = 0; j < embedDim; j++) { + expect(outputA[0][j]).toBeCloseTo(outputB[0][j], 10); + } + + // Row 1 should be identical (cannot see position 2) + for (let j = 0; j < embedDim; j++) { + expect(outputA[1][j]).toBeCloseTo(outputB[1][j], 10); + } + + // Row 2 should differ (it sees itself and its input changed) + let row2Differs = false; + for (let j = 0; j < embedDim; j++) { + if (Math.abs(outputA[2][j] - outputB[2][j]) > 1e-10) { + row2Differs = true; + } + } + expect(row2Differs).toBe(true); + }); + + it('does not leak future information to earlier positions with seq length 1', () => { + const block = new TransformerBlock(embedDim, ffDim, contextLength); + const params = block.getParameters(); + + const singleToken = mat([[1, 2, 3, 4]]); + const twoTokens = mat([ + [1, 2, 3, 4], + [99, 99, 99, 99], + ]); + + block.setParameters(params); + const outSingle = block.forward(singleToken).toArray(); + + block.setParameters(params); + const outDouble = block.forward(twoTokens).toArray(); + + // The first row of the two-token output should match the single-token output + for (let j = 0; j < embedDim; j++) { + expect(outDouble[0][j]).toBeCloseTo(outSingle[0][j], 10); + } + }); + }); +}); diff --git a/src/index.ts b/src/index.ts index f1b4d44..aec9813 100644 --- a/src/index.ts +++ b/src/index.ts @@ -7,10 +7,13 @@ export * from './applications/machine-learning/kernels/Kernel'; export * from './applications/machine-learning/kernels/LinearKernel'; export * from './applications/machine-learning/kernels/RadialBasisFunction'; export * from './applications/machine-learning/models/Classifier'; +export * from './applications/machine-learning/models/LanguageModel'; export * from './applications/machine-learning/models/LinearRegressor'; export * from './applications/machine-learning/models/LogisticRegressionClassifier'; export * from './applications/machine-learning/models/Regressor'; export * from './applications/machine-learning/models/SupportVectorMachineClassifier'; +export * from './applications/machine-learning/models/Tokenizer'; +export * from './applications/machine-learning/models/TransformerBlock'; export * from './applications/statistics/DescriptiveStatistics'; export * from './applications/statistics/LeastSquares'; export * from './applications/statistics/PrincipalComponentAnalysis'; diff --git a/src/utilities/NumberUtilities.ts b/src/utilities/NumberUtilities.ts index a4463ab..0e06a1c 100644 --- a/src/utilities/NumberUtilities.ts +++ b/src/utilities/NumberUtilities.ts @@ -66,3 +66,18 @@ export function sigmoid(x: number): number { const ex = Math.exp(x); return ex / (1 + ex); } + +/** + * Computes the softmax of an array of numbers, returning a probability distribution. + * Uses the max-subtraction trick for numerical stability. + */ +export function softmax(values: number[]): number[] { + if (values.length === 0) { + return []; + } + + const max = Math.max(...values); + const exps = values.map((v) => Math.exp(v - max)); + const sum = exps.reduce((a, b) => a + b, 0); + return exps.map((e) => e / sum); +} diff --git a/src/utilities/__tests__/NumberUtilities.spec.ts b/src/utilities/__tests__/NumberUtilities.spec.ts index e4b0420..749f665 100644 --- a/src/utilities/__tests__/NumberUtilities.spec.ts +++ b/src/utilities/__tests__/NumberUtilities.spec.ts @@ -5,6 +5,7 @@ import { mod, random, randomNormal, + softmax, } from '../NumberUtilities'; describe('NumberUtilities', () => { @@ -143,4 +144,52 @@ describe('NumberUtilities', () => { expect(binomial(50, 30)).toEqual(47129212243960); }); }); + + describe('softmax', () => { + test('output sums to approximately 1', () => { + const result = softmax([1, 2, 3, 4]); + const sum = result.reduce((a, b) => a + b, 0); + expect(approximatelyEqual(sum, 1)).toBe(true); + }); + + test('larger inputs get larger probabilities', () => { + const result = softmax([1, 2, 3]); + expect(result[2]).toBeGreaterThan(result[1]); + expect(result[1]).toBeGreaterThan(result[0]); + }); + + test('all-equal inputs produce uniform distribution', () => { + const result = softmax([5, 5, 5, 5]); + result.forEach((p) => { + expect(approximatelyEqual(p, 0.25)).toBe(true); + }); + }); + + test('single element returns [1]', () => { + const result = softmax([42]); + expect(result).toEqual([1]); + }); + + test('empty array returns empty array', () => { + expect(softmax([])).toEqual([]); + }); + + test('is numerically stable for very large inputs', () => { + const result = softmax([1000, 1001, 1002]); + result.forEach((p) => { + expect(Number.isFinite(p)).toBe(true); + expect(Number.isNaN(p)).toBe(false); + }); + const sum = result.reduce((a, b) => a + b, 0); + expect(approximatelyEqual(sum, 1)).toBe(true); + }); + + test('handles negative inputs correctly', () => { + const result = softmax([-1, -2, -3]); + const sum = result.reduce((a, b) => a + b, 0); + expect(approximatelyEqual(sum, 1)).toBe(true); + expect(result[0]).toBeGreaterThan(result[1]); + expect(result[1]).toBeGreaterThan(result[2]); + }); + }); });