/**
 * evermind_lm.ts — EvermindLM: a small but complete generative language model.
 *
 * This is what turns a trained checkpoint into an *AI that generates text* (the
 * thing a marketplace buyer actually runs). Architecture (Mamba-flavoured, the
 * minimal exact-gradient CPU reference):
 *
 *   x_t = Embed[token_t]
 *   per layer:
 *     x_t += DepthwiseCausalConv(x)_t            // temporal mixing (short conv)
 *     x_t += SharedExpertMoE(x_t)                // per-position channel mixing (sparse)
 *   logits_t = x_t · Embedᵀ                       // tied output head
 *
 * The token mixer is a depthwise causal convolution (each channel sees a short
 * window of its own past — Mamba's pre-conv) and the channel mixer is the
 * shared-expert MoE, so the model is genuinely sparse. Embeddings are tied
 * (input lookup == output head), which the gradient code accounts for.
 *
 * Pure CPU, exact forward + backward (finite-difference checked), reusing the
 * engine's MoE, cross-entropy, and AdamW. The WGSL/WebGPU path is a future
 * acceleration with the same shapes.
 */
import { type AdamWOptions } from "../optim/adamw.js";
import { DynamicLossScaler, type LossScalerOptions } from "../training/mixed_precision.js";
export interface EvermindLMConfig {
    /** Vocabulary size (the only required field; everything else has a default). */
    vocabSize: number;
    /** Model (channel) dimension. Default 64. */
    dModel?: number;
    /** Number of (conv + MoE) blocks. Default 2. */
    numLayers?: number;
    /** Causal conv kernel width. Default 3. */
    convKernel?: number;
    /** Hidden width of each MoE expert FFN. Default 2·dModel. */
    hiddenDim?: number;
    /** Routed experts per MoE layer. Default 4. */
    numExperts?: number;
    /** Experts activated per token. Default 2. */
    topK?: number;
    /** Deterministic init seed. */
    seed?: number;
}
export declare const DEFAULT_LM_CONFIG: Required<Omit<EvermindLMConfig, "seed" | "vocabSize">>;
export declare const DEFAULT_LM_SEED = 1163283533;
interface MoECacheLike {
    x: Float32Array;
    route: {
        experts: number[];
        gates: number[];
        probs: Float32Array;
    };
    sharedPre: Float32Array;
    sharedH: Float32Array;
    expertOut: Float32Array[];
    expertPre: Float32Array[];
    expertH: Float32Array[];
}
interface LayerCache {
    layerIn: Float32Array[];
    normedConv: Float32Array[];
    rmsConv: number[];
    afterConv: Float32Array[];
    rmsMoe: number[];
    moeCache: MoECacheLike[];
}
interface ForwardCache {
    tokens: number[];
    layers: LayerCache[];
    finalX: Float32Array[];
}
/** A tokenizer the LM can read/write text through (the engine's `BPETokenizer` fits). */
export interface TextCodec {
    encode(text: string): number[];
    decode(ids: number[]): string;
}
export interface LMGenerateOptions {
    maxNewTokens: number;
    /** Sampling temperature; ≤0 ⇒ greedy argmax. Default 0 (greedy). */
    temperature?: number;
    /** Deterministic sampler seed (only used when temperature > 0). */
    seed?: number;
    /** Stop generating when this token id is produced. */
    stopToken?: number;
}
/**
 * Incremental decode state — everything a continuation needs from the prefix.
 *
 * The model's only cross-position dependency is the depthwise causal conv, whose
 * receptive field is `convKernel` positions, so the entire "KV cache" is the
 * trailing `convKernel - 1` normalised conv inputs per layer. That makes an exact
 * prefix cache cheap: process a prompt once, then extend it with each candidate
 * continuation for the cost of the candidate alone.
 */
export interface EvermindLMDecodeState {
    /** [layer][slot] the trailing `convKernel - 1` RMSNormed conv inputs, oldest first. */
    convWindow: Float32Array[][];
    /** Positions consumed so far. */
    length: number;
    /** Logits at the LAST consumed position — what predicts the next token. */
    lastLogits: Float32Array | null;
}
export declare class EvermindLM {
    readonly config: Required<Omit<EvermindLMConfig, "seed">>;
    /**
     * Token-positions pushed through a block since the last {@link resetStats}.
     *
     * The cost metric that matters for the tool-calling path: scoring N candidate
     * tool names used to replay the WHOLE prompt N times, so this counted
     * N x (prompt + candidate) x layers. With {@link stepDecode} the prompt is paid
     * for once. Exposed so a benchmark can assert the saving rather than assume it.
     */
    private _positionsEvaluated;
    /** Tied token embedding / output head: vocabSize × dModel (row-major). */
    emb: Float32Array;
    private gEmb;
    /** Per-layer depthwise causal conv kernels: dModel × convKernel. */
    private readonly conv;
    private readonly gConv;
    /** Per-layer pre-conv / pre-MoE RMSNorm gains (dModel each). */
    private readonly nConv;
    private readonly gNConv;
    private readonly nMoe;
    private readonly gNMoe;
    /** Per-layer channel mixer. */
    private readonly moe;
    constructor(config: EvermindLMConfig);
    /** Embed a token sequence into per-position channel vectors. */
    private _embed;
    /**
     * One (conv + MoE) block: pre-norm → depthwise causal conv → residual, then
     * pre-norm → MoE channel mixer → residual. Returns the block output and the
     * activation cache its backward needs. Isolating this is what lets
     * {@link lossAndBackwardCheckpointed} recompute a layer's activations on demand
     * instead of retaining every layer's cache at once.
     */
    private _forwardLayer;
    /**
     * ONE position through one (conv + MoE) block, given the trailing conv window.
     *
     * The single place the block's maths lives: the batch path
     * ({@link _forwardLayer}, which needs the full activation cache for backward)
     * and the incremental path ({@link stepDecode}, which needs none of it) both
     * run this, so they cannot drift apart. `window` holds the previous
     * `convKernel - 1` RMSNormed conv inputs, OLDEST FIRST, and is advanced in
     * place — a shorter window is exactly the causal zero-padding at the start of
     * a sequence.
     */
    private _stepLayerPosition;
    /** One token's embedding row (a copy — callers write into it). */
    private _embedOne;
    /** Tied output head: logits_t[v] = x_t · emb[v]. */
    private _head;
    /** Run the model over a token sequence; returns per-position logits + a cache. */
    forward(tokens: number[]): {
        logits: Float32Array[];
        cache: ForwardCache;
    };
    /**
     * Produces a single fixed-length embedding vector for a token sequence — the
     * CPU counterpart of `HybridMambaModel.embed()`, to the same contract.
     *
     * Runs the full block stack and takes the final hidden state — i.e. exactly the
     * state the tied LM head consumes in {@link forward} — then mean-pools across
     * sequence positions and L2-normalises. The result has length `dModel`, so
     * cosine similarity between two embeddings reduces to a dot product.
     *
     * Like the GPU path it skips the (expensive) head projection: it only needs the
     * `dModel`-wide hidden state, not `vocabSize` logits. Unlike the GPU path it
     * needs no WebGPU device at all, which is what lets a headless host (the stdio
     * MCP server) rank recall by SSM embedding instead of word overlap.
     *
     * Streamed one position at a time through the SHARED block implementation
     * ({@link _stepLayerPosition}, the same routine {@link forward} and
     * {@link stepDecode} run), so it can never drift from them and it retains no
     * per-position activation cache.
     *
     * The embedding reflects whatever the model currently knows — an untrained
     * model behaves like a random projection of the token embeddings (still
     * lexically discriminative), and the representation sharpens as the model is
     * adapted/distilled.
     */
    embed(tokenIds: number[] | Uint32Array): Float32Array;
    /** Text-level embedding: encode with `codec`, then {@link embed}. */
    embedText(text: string, codec: TextCodec): Float32Array;
    /**
     * Head + tied-embedding gradient. Accumulates dL/d(head→emb) into gEmb and
     * returns the mean next-token loss plus dL/d(finalX). Shared by the full and
     * checkpointed backward paths so the head maths lives in one place.
     */
    private _headBackward;
    /**
     * Backward through one (conv + MoE) block given dL/d(block output) and the
     * block's activation cache. Accumulates conv/norm/MoE gradients and returns
     * dL/d(block input). The inverse of {@link _forwardLayer}.
     */
    private _backwardLayer;
    /** Embedding lookup: dL/d(layer-0 input) flows into the row for each token. */
    private _embedBackward;
    /**
     * Next-token cross-entropy over the sequence (predict tokens[t+1] from
     * position t), accumulating exact gradients. Returns the mean loss. Call
     * {@link zeroGrad} before and an optimiser step after.
     */
    lossAndBackward(tokens: number[]): number;
    /**
     * Activation-checkpointed backward — numerically identical gradients to
     * {@link lossAndBackward}, but retains only the per-LAYER inputs during the
     * forward instead of every layer's full activation cache. Each layer's
     * activations are RECOMPUTED (a cheap extra forward) when its backward runs,
     * so peak activation memory is one layer's cache, not all of them. This is the
     * memory-for-compute trade the cookbook pairs with FSDP to fit longer
     * sequences / bigger models on a constrained device.
     */
    lossAndBackwardCheckpointed(tokens: number[]): number;
    /** Positions pushed through a block since the last {@link resetStats}. */
    get positionsEvaluated(): number;
    resetStats(): void;
    /** An empty decode state — the start of a sequence. */
    newDecodeState(): EvermindLMDecodeState;
    /** Deep copy, so one prefix state can seed many independent continuations. */
    cloneDecodeState(state: EvermindLMDecodeState): EvermindLMDecodeState;
    /**
     * Consume `tokens` from `state`, returning the per-position logits and the
     * ADVANCED state. `state` is not mutated — pass {@link cloneDecodeState} output
     * or a fresh state; this clones internally so the caller's prefix stays reusable.
     *
     * Numerically identical to {@link forward} over the concatenated sequence: the
     * shared {@link _stepLayerPosition} is the only implementation of the block.
     */
    stepDecode(state: EvermindLMDecodeState, tokens: number[]): {
        logits: Float32Array[];
        state: EvermindLMDecodeState;
    };
    /** Run a prompt once and keep the state, so continuations are cheap. */
    prefixState(tokens: number[]): EvermindLMDecodeState;
    /**
     * Mean log-probability the model assigns to `continuation` following a cached
     * prefix — the tool-decision vote, without replaying the prompt.
     *
     * Returns `-Infinity` for an empty continuation, and requires a prefix with at
     * least one position (there must be a position to predict the first
     * continuation token FROM).
     */
    scoreContinuation(prefix: EvermindLMDecodeState, continuation: number[]): number;
    /**
     * Text-level generation: encode the prompt, generate, decode. `codec` is any
     * tokenizer exposing encode/decode (the engine's `BPETokenizer` satisfies it),
     * so the LM consumes and emits real text rather than raw token ids. The model's
     * `vocabSize` must match the codec's vocabulary.
     */
    generateText(prompt: string, codec: TextCodec, opts: LMGenerateOptions): string;
    /**
     * Greedy / temperature-sampled autoregressive generation. Returns NEW token ids.
     *
     * Carries an incremental decode state, so the prompt is processed ONCE and each
     * new token costs one position. The previous implementation re-ran the whole
     * sequence per token, making generation quadratic in the prompt length.
     */
    generate(prompt: number[], opts: LMGenerateOptions): number[];
    /** All trainable parameters as {data} (AdamW-compatible), canonical order. */
    parameters(): {
        data: Float32Array;
    }[];
    /** Gradient buffers, index-aligned with {@link parameters}. */
    gradients(): {
        data: Float32Array;
    }[];
    zeroGrad(): void;
    /** Serialise to an "EVL0" binary (fp16 or f32), params in {@link parameters} order. */
    exportWeights(opts?: {
        fp16?: boolean;
    }): ArrayBuffer;
    /** Load weights from an "EVL0" binary. Validates CRC (when present), magic + dims. */
    loadWeights(buffer: ArrayBuffer): void;
    /** Parse + validate an EVL0 checkpoint to a flat f32 param vector (no mutation). */
    private _readFlat;
    /** Flat concatenation of all params, canonical order (matches checkpoint layout). */
    private _flat;
    /** Distribute a flat param vector back into the model's parameters. */
    private _setFlat;
    /**
     * Export a SPARSE DELTA of the current weights against a base EVL0 checkpoint
     * (EVM-6). Online WSLA updates only a few rows, so a delta persists kilobytes
     * instead of rewriting the whole model. Reconstruct with {@link loadDelta}.
     */
    exportDelta(baseCheckpoint: ArrayBuffer, opts?: {
        eps?: number;
    }): ArrayBuffer;
    /** Reconstruct weights from a base EVL0 checkpoint + a delta (EVM-6). */
    loadDelta(baseCheckpoint: ArrayBuffer, delta: ArrayBuffer): void;
}
export interface EvermindLMTrainOptions extends AdamWOptions {
    epochs?: number;
    /**
     * Gradient accumulation: average gradients over this many sequences (micro-
     * batches) before each optimiser step, for a larger effective batch on a
     * memory-constrained device. Default 1.
     */
    accumSteps?: number;
    /**
     * Use activation checkpointing (recompute layer activations in backward) to
     * cap peak activation memory. Identical gradients, a little extra compute.
     * Default false.
     */
    checkpoint?: boolean;
    /**
     * Mixed-precision training: fp16-rounded gradients with dynamic loss scaling
     * over fp32 master weights (the model params). Overflowing steps are skipped
     * and the scale backs off. Default false.
     */
    mixedPrecision?: boolean | LossScalerOptions;
}
/** Minimal sequence trainer: AdamW over next-token cross-entropy. */
export declare class EvermindLMTrainer {
    private readonly model;
    private readonly opts;
    private readonly adam;
    private readonly scaler;
    private readonly grads;
    constructor(model: EvermindLM, opts?: EvermindLMTrainOptions);
    /** The dynamic loss scaler (mixed-precision mode only) — exposes scale/overflow stats. */
    get lossScaler(): DynamicLossScaler | null;
    private _backward;
    /** Divide accumulated gradients by `n` (accumulation averaging), in place. */
    private _scaleGrads;
    /** Round accumulated gradients to fp16 precision (mixed-precision simulation), in place. */
    private _fp16Grads;
    /** Multiply accumulated gradients by `s`, in place. */
    private _mulGrads;
    /** Train on a set of token sequences; returns per-epoch mean loss. */
    fit(sequences: number[][]): number[];
}
/**
 * Log-probability the model assigned to `id` at a position, from that position's
 * raw logits. Log-softmax computed max-shifted so a long-tail logit cannot
 * overflow `exp`.
 */
export declare function logProbOfToken(row: Float32Array, id: number): number;
export {};
//# sourceMappingURL=evermind_lm.d.ts.map