Como uma IA é treinada, passo a passo

/** One item per line, lowercase, trimmed; long or repeated lines are dropped. */
export function prepareDocs(text, { maxLength = 20 } = {}) {
    const seen = new Set()
    const docs = []
    for (const line of text.split('\n')) {
        const doc = line.trim().toLowerCase().replace(/\s+/g, ' ')
        if (!doc || doc.length > maxLength || seen.has(doc)) continue
        seen.add(doc)
        docs.push(doc)
    }
    return docs
}

/**
 * Each distinct character becomes a number, in alphabetical order. One extra
 * number, BOS, marks where a word starts and where it ends.
 */
export function createTokenizer(docs) {
    const chars = [...new Set(docs.join(''))].sort((a, b) => a.localeCompare(b, 'pt-BR'))
    const BOS = chars.length
    return {
        chars,
        BOS,
        vocabSize: chars.length + 1,
        encode: text => [...text].map(ch => chars.indexOf(ch)),
        decode: ids => ids.map(id => (id === BOS ? '·' : chars[id])).join(''),
    }
}

/**
 * The numbers the model learns. They start as small random noise; training
 * moves each one a little, a thousand times.
 */
export function createModel(docs, { nLayer = 1, nEmbd = 16, nHead = 4, seed = 42 } = {}) {
    const random = createRandom(seed)
    const tokenizer = createTokenizer(docs)
    const blockSize = Math.max(...docs.map(doc => [...doc].length)) + 1
    const matrix = (rows, cols, std = 0.08) => Array.from({ length: rows }, () => Array.from({ length: cols }, () => new Value(random.gauss(0, std))))

    const state = {
        wte: matrix(tokenizer.vocabSize, nEmbd), // one row per character: its embedding
        wpe: matrix(blockSize, nEmbd), // one row per position in the word
        lm_head: matrix(tokenizer.vocabSize, nEmbd), // turns the result back into one score per character
    }
    for (let i = 0; i < nLayer; i++) {
        state[`layer${i}.attn_wq`] = matrix(nEmbd, nEmbd) // query: what this letter looks for
        state[`layer${i}.attn_wk`] = matrix(nEmbd, nEmbd) // key: what each letter offers
        state[`layer${i}.attn_wv`] = matrix(nEmbd, nEmbd) // value: what it passes on when chosen
        state[`layer${i}.attn_wo`] = matrix(nEmbd, nEmbd)
        state[`layer${i}.mlp_fc1`] = matrix(4 * nEmbd, nEmbd)
        state[`layer${i}.mlp_fc2`] = matrix(nEmbd, 4 * nEmbd)
    }
    const params = Object.values(state).flatMap(rows => rows.flat())

    return {
        docs: random.shuffle(docs),
        tokenizer,
        config: { nLayer, nEmbd, nHead, headDim: nEmbd / nHead, blockSize },
        state,
        params,
        random,
        adam: { m: new Float64Array(params.length), v: new Float64Array(params.length) },
        step: 0,
    }
}

/**
 * One step: take one word, predict each next letter, measure how wrong the
 * predictions were (the loss), find how each parameter contributed (backward)
 * and nudge every parameter the other way (Adam).
 */
export function trainStep(model, { learningRate = 0.01, numSteps = 1000, beta1 = 0.85, beta2 = 0.99 } = {}) {
    const { tokenizer, params, adam } = model
    const doc = model.docs[model.step % model.docs.length]
    const tokens = [tokenizer.BOS, ...tokenizer.encode(doc), tokenizer.BOS]
    const n = Math.min(model.config.blockSize, tokens.length - 1)

    const keys = emptyCache(model)
    const values = emptyCache(model)
    const losses = []
    const predictions = []
    for (let pos = 0; pos < n; pos++) {
        const probs = softmax(gpt(model, tokens[pos], pos, keys, values))
        losses.push(probs[tokens[pos + 1]].log().mul(-1))
        predictions.push(probs[tokens[pos + 1]].data)
    }
    const loss = sum(losses).div(n)

    loss.backward()

    const progress = Math.min(model.step, numSteps - 1) / numSteps
    const lr = learningRate * (1 - progress) // linear decay to zero
    const t = model.step + 1
    for (let i = 0; i < params.length; i++) {
        const p = params[i]
        adam.m[i] = beta1 * adam.m[i] + (1 - beta1) * p.grad
        adam.v[i] = beta2 * adam.v[i] + (1 - beta2) * p.grad ** 2
        const mHat = adam.m[i] / (1 - beta1 ** t)
        const vHat = adam.v[i] / (1 - beta2 ** t)
        p.data -= (lr * mHat) / (Math.sqrt(vHat) + 1e-8)
        p.grad = 0
    }
    model.step++

    return { step: model.step, loss: loss.data, doc, predictions, lr }
}

const linear = (x, w) => w.map(row => sum(row.map((wi, i) => wi.mul(x[i]))))

function softmax(logits) {
    const max = Math.max(...logits.map(v => v.data))
    const exps = logits.map(v => v.sub(max).exp())
    const total = sum(exps)
    return exps.map(e => e.div(total))
}

function rmsnorm(x) {
    const ms = sum(x.map(xi => xi.mul(xi))).div(x.length)
    const scale = ms.add(1e-5).pow(-0.5)
    return x.map(xi => xi.mul(scale))
}

/**
 * One letter in, one score per possible next letter out. GPT-2 with three
 * small swaps: rmsnorm instead of layernorm, no biases, ReLU instead of GeLU.
 * `trace`, when given, receives every intermediate number for the page.
 */
export function gpt(model, tokenId, posId, keys, values, trace) {
    const { state, config } = model
    const tokEmb = state.wte[tokenId]
    const posEmb = state.wpe[posId]
    let x = tokEmb.map((t, i) => t.add(posEmb[i]))
    x = rmsnorm(x)
    if (trace) Object.assign(trace, { tokEmb: data(tokEmb), posEmb: data(posEmb), x: data(x), heads: [] })

    for (let li = 0; li < config.nLayer; li++) {
        // 1) Attention: this letter looks back at the letters before it.
        let residual = x
        x = rmsnorm(x)
        const q = linear(x, state[`layer${li}.attn_wq`])
        const k = linear(x, state[`layer${li}.attn_wk`])
        const v = linear(x, state[`layer${li}.attn_wv`])
        keys[li].push(k)
        values[li].push(v)
        const attended = []
        for (let h = 0; h < config.nHead; h++) {
            const hs = h * config.headDim
            const qh = q.slice(hs, hs + config.headDim)
            const kh = keys[li].map(ki => ki.slice(hs, hs + config.headDim))
            const vh = values[li].map(vi => vi.slice(hs, hs + config.headDim))
            const scores = kh.map(kt => sum(qh.map((qj, j) => qj.mul(kt[j]))).div(Math.sqrt(config.headDim)))
            const weights = softmax(scores)
            if (trace) trace.heads.push(data(weights))
            for (let j = 0; j < config.headDim; j++) attended.push(sum(weights.map((w, t) => w.mul(vh[t][j]))))
        }
        x = linear(attended, state[`layer${li}.attn_wo`])
        x = x.map((a, i) => a.add(residual[i]))

        // 2) MLP: 64 neurons think about what attention brought.
        residual = x
        x = rmsnorm(x)
        x = linear(x, state[`layer${li}.mlp_fc1`])
        x = x.map(xi => xi.relu())
        if (trace) trace.neurons = data(x)
        x = linear(x, state[`layer${li}.mlp_fc2`])
        x = x.map((a, i) => a.add(residual[i]))
    }

    const logits = linear(x, state.lm_head)
    if (trace) trace.logits = data(logits)
    return logits
}

const data = values => values.map(v => v.data)

/**
 * Generation: start at BOS, ask for the next letter, draw one according to the
 * probabilities, feed it back, until the model chooses BOS (the end).
 * `temperature` below 1 sharpens the odds; above 1 flattens them.
 */
export function generate(model, { temperature = 0.5, random = model.random } = {}) {
    const { tokenizer, config } = model
    const keys = emptyCache(model)
    const values = emptyCache(model)
    const steps = []
    let token = tokenizer.BOS
    for (let pos = 0; pos < config.blockSize; pos++) {
        const logits = gpt(model, token, pos, keys, values)
        const probs = softmax(logits.map(l => l.div(temperature))).map(p => p.data)
        token = random.choice(probs)
        steps.push({ probs, token })
        if (token === tokenizer.BOS) break
    }
    const text = tokenizer.decode(steps.map(s => s.token).filter(id => id !== tokenizer.BOS))
    return { text, steps }
}