@mlx-node/lm 0.0.7 → 0.0.9

package/dist/models/model-loader.js

@@ -5,44 +5,215 @@
  */
 import { readFile } from 'node:fs/promises';
 import { join } from 'node:path';
-import { HarrierModel, QianfanOCRModel } from '@mlx-node/core';
+import { Gemma4Model as NativeGemma4Model, HarrierModel, Lfm2Model as NativeLfm2Model, QianfanOCRModel, Qwen3Model as NativeQwen3Model, Qwen35Model as NativeQwen35Model, Qwen35MoeModel as NativeQwen35MoeModel, } from '@mlx-node/core';
 import { ChatSession } from '../chat-session.js';
 import { Gemma4Model, Lfm2Model, Qwen3Model, Qwen35Model, Qwen35MoeModel } from '../stream.js';
-const SUPPORTED_MODEL_TYPES = new Set([
-    'qwen3',
-    'qwen3_5',
-    'qwen3_5_moe',
-    'internvl_chat',
-    'qianfan-ocr',
-    'harrier',
-    'gemma4',
-    'lfm2',
-]);
+/**
+ * Ordered source of truth for every supported model family. Each entry owns
+ * its canonical `ModelType`, raw config aliases / architecture probes, loader,
+ * and `ChatSession` eligibility:
+ *
+ *   - `'trainable'` — GRPO/SFT-capable LM (Qwen3 family); chat-capable.
+ *   - `'loadable'`  — chat-capable LM with no trainer engine (Gemma4, LFM2).
+ *   - `'embedding'` — no chat surface (Harrier); rejected by `loadSession`.
+ *   - `'vlm'`       — VLM whose AsyncGenerator wrapper lives in
+ *                     `@mlx-node/vlm` (importing it here would create a
+ *                     circular package dependency), so `loadSession`
+ *                     rejects it and routes callers to `@mlx-node/vlm`.
+ *
+ * A base family is selected from an explicit alias or the single declarative
+ * nullish-model_type default, then architecture probes refine it in declaration
+ * order. Gemma's unified architecture is authoritative (matching the native
+ * loader); Harrier refines a Qwen3 base. Adding a family means adding one
+ * descriptor here, without a second normalization or dispatch branch.
+ */
+const MODEL_FAMILY_REGISTRY = [
+    {
+        modelType: 'gemma4',
+        kind: 'loadable',
+        match: {
+            rawModelTypes: ['gemma4', 'gemma4_text', 'gemma4_unified'],
+            architectureProbe: ({ architectures }) => architectures.has('Gemma4UnifiedForConditionalGeneration'),
+        },
+        load: (modelPath, options) => Gemma4Model.load(modelPath, options?.draftModelPath === undefined ? null : { draftModelPath: options.draftModelPath }),
+        nativeModelClass: NativeGemma4Model,
+        acceptsDraftModel: true,
+    },
+    {
+        modelType: 'harrier',
+        kind: 'embedding',
+        match: {
+            rawModelTypes: ['harrier'],
+            architectureProbe: ({ modelType, architectures }) => modelType === 'qwen3' && architectures.has('Qwen3Model') && !architectures.has('Qwen3ForCausalLM'),
+        },
+        load: (modelPath) => HarrierModel.load(modelPath),
+        nativeModelClass: HarrierModel,
+    },
+    {
+        modelType: 'qwen3',
+        kind: 'trainable',
+        match: { rawModelTypes: ['qwen3'] },
+        load: (modelPath) => Qwen3Model.load(modelPath),
+        nativeModelClass: NativeQwen3Model,
+        defaultForNullishModelType: true,
+    },
+    {
+        modelType: 'qwen3_5',
+        kind: 'trainable',
+        match: { rawModelTypes: ['qwen3_5'] },
+        load: (modelPath) => Qwen35Model.load(modelPath),
+        nativeModelClass: NativeQwen35Model,
+    },
+    {
+        modelType: 'qwen3_5_moe',
+        kind: 'trainable',
+        match: { rawModelTypes: ['qwen3_5_moe'] },
+        load: (modelPath) => Qwen35MoeModel.load(modelPath),
+        nativeModelClass: NativeQwen35MoeModel,
+    },
+    {
+        modelType: 'lfm2',
+        kind: 'loadable',
+        match: { rawModelTypes: ['lfm2'] },
+        load: (modelPath) => Lfm2Model.load(modelPath),
+        nativeModelClass: NativeLfm2Model,
+    },
+    {
+        modelType: 'lfm2_moe',
+        kind: 'loadable',
+        match: { rawModelTypes: ['lfm2_moe'] },
+        load: (modelPath) => Lfm2Model.load(modelPath),
+        nativeModelClass: NativeLfm2Model,
+    },
+    {
+        modelType: 'internvl_chat',
+        kind: 'vlm',
+        match: { rawModelTypes: ['internvl_chat'] },
+        load: (modelPath) => QianfanOCRModel.load(modelPath),
+        nativeModelClass: QianfanOCRModel,
+    },
+    {
+        modelType: 'qianfan-ocr',
+        kind: 'vlm',
+        match: { rawModelTypes: ['qianfan-ocr'] },
+        load: (modelPath) => QianfanOCRModel.load(modelPath),
+        nativeModelClass: QianfanOCRModel,
+    },
+];
+function buildModelFamilyIndex(registry) {
+    const byModelType = new Map();
+    const byRawModelType = new Map();
+    let defaultForNullishModelType;
+    for (const family of registry) {
+        const previousFamily = byModelType.get(family.modelType);
+        if (previousFamily !== undefined) {
+            throw new Error(`Duplicate canonical model type "${family.modelType}" in model family registry`);
+        }
+        byModelType.set(family.modelType, family);
+        for (const rawModelType of family.match.rawModelTypes) {
+            const previous = byRawModelType.get(rawModelType);
+            if (previous !== undefined) {
+                throw new Error(`Duplicate model_type alias "${rawModelType}" for "${previous.modelType}" and "${family.modelType}"`);
+            }
+            byRawModelType.set(rawModelType, family);
+        }
+        if (family.defaultForNullishModelType === true) {
+            if (defaultForNullishModelType !== undefined) {
+                throw new Error(`Duplicate nullish-model_type defaults for "${defaultForNullishModelType.modelType}" and "${family.modelType}"`);
+            }
+            defaultForNullishModelType = family;
+        }
+    }
+    if (defaultForNullishModelType === undefined) {
+        throw new Error('Model family registry must declare exactly one nullish-model_type default');
+    }
+    return { byModelType, byRawModelType, defaultForNullishModelType };
+}
+const MODEL_FAMILY_INDEX = buildModelFamilyIndex(MODEL_FAMILY_REGISTRY);
+function findFamily(modelType) {
+    const family = MODEL_FAMILY_INDEX.byModelType.get(modelType);
+    if (family === undefined) {
+        throw new Error(`Internal error: missing model family descriptor for "${modelType}"`);
+    }
+    return family;
+}
+function matchesArchitectureProbe(family, config) {
+    return family.match.architectureProbe?.(config) === true;
+}
+class MalformedModelConfigError extends Error {
+    constructor(modelPath, reason) {
+        super(`Malformed config.json in ${modelPath}: ${reason}`);
+        this.name = 'MalformedModelConfigError';
+    }
+}
+/**
+ * Fail-closed validation: a config.json whose root is not a plain object,
+ * or whose `architectures` is neither an array nor a string, is rejected
+ * instead of coerced (coercion would fall through to the qwen3
+ * nullish-model_type default and silently misroute the checkpoint).
+ * Blessed lenient shapes stay accepted: `{}` root (qwen3 default),
+ * missing/`null` `architectures` (empty set), bare-string `architectures`
+ * (single-element set), and non-string array entries (filtered out).
+ */
+function normalizeConfig(modelPath, config) {
+    if (typeof config !== 'object' || config === null || Array.isArray(config)) {
+        throw new MalformedModelConfigError(modelPath, 'root must be a JSON object');
+    }
+    const object = config;
+    const hasModelType = Object.hasOwn(object, 'model_type');
+    const rawModelTypeValue = hasModelType ? object.model_type : undefined;
+    const usesDefaultModelType = !hasModelType || rawModelTypeValue === null;
+    const rawModelType = typeof rawModelTypeValue === 'string' ? rawModelTypeValue : undefined;
+    const rawModelTypeLabel = hasModelType ? String(rawModelTypeValue) : '<missing>';
+    const rawArchitectures = 'architectures' in object ? object.architectures : undefined;
+    if (rawArchitectures !== undefined &&
+        rawArchitectures !== null &&
+        !Array.isArray(rawArchitectures) &&
+        typeof rawArchitectures !== 'string') {
+        throw new MalformedModelConfigError(modelPath, '"architectures" must be an array or a string');
+    }
+    const architectures = Array.isArray(rawArchitectures)
+        ? rawArchitectures.filter((architecture) => typeof architecture === 'string')
+        : typeof rawArchitectures === 'string'
+            ? [rawArchitectures]
+            : [];
+    return { usesDefaultModelType, rawModelType, rawModelTypeLabel, architectures: new Set(architectures) };
+}
+class UnsupportedModelTypeError extends Error {
+    constructor(modelPath, rawModelTypeLabel) {
+        super(`Unsupported model_type "${rawModelTypeLabel}" in ${modelPath}/config.json`);
+        this.name = 'UnsupportedModelTypeError';
+    }
+}
+/**
+ * Dispatch a load through the registry, validating gemma4-only options.
+ * `draftModelPath` reaches ONLY the gemma4 row; every other family rejects
+ * it loudly instead of silently ignoring a caller's speculative-decode
+ * intent.
+ */
+function dispatchLoad(modelType, modelPath, options) {
+    const family = findFamily(modelType);
+    if (options?.draftModelPath !== undefined && family.acceptsDraftModel !== true) {
+        throw new Error(`draftModelPath (speculative-decoding draft) is only supported by gemma4 models; ` +
+            `${modelPath} has model_type "${modelType}"`);
+    }
+    return family.load(modelPath, options);
+}
 /**
  * Load a model from disk, auto-detecting architecture from config.json.
  *
  * Supports both language models (Qwen3, Qwen3.5) and vision-language models
  * (Qianfan-OCR / InternVL). Use `instanceof` to narrow the returned type.
+ *
+ * `options.draftModelPath` attaches an external draft checkpoint (DSpark or
+ * Google gemma-4 assistant, auto-detected from the draft's config.json) for
+ * speculative decoding — gemma4 only; any other detected family rejects it.
+ * Without the option, Gemma4 loads `<modelPath>/draft/` automatically when
+ * that embedded checkpoint is present.
  */
-export async function loadModel(modelPath) {
+export async function loadModel(modelPath, options) {
     const modelType = await detectModelType(modelPath);
-    switch (modelType) {
-        case 'qwen3_5_moe':
-            return Qwen35MoeModel.load(modelPath);
-        case 'qwen3_5':
-            return Qwen35Model.load(modelPath);
-        case 'qwen3':
-            return Qwen3Model.load(modelPath);
-        case 'harrier':
-            return HarrierModel.load(modelPath);
-        case 'internvl_chat':
-        case 'qianfan-ocr':
-            return QianfanOCRModel.load(modelPath);
-        case 'gemma4':
-            return Gemma4Model.load(modelPath);
-        case 'lfm2':
-            return Lfm2Model.load(modelPath);
-    }
+    return dispatchLoad(modelType, modelPath, options);
 }
 /**
  * Load a model and wrap it in a {@link ChatSession} for multi-turn chat.
@@ -59,38 +230,44 @@ export async function loadModel(modelPath) {
  *     package dependency), so callers who want a Qianfan-OCR session
  *     must import `QianfanOCRModel` from `@mlx-node/vlm` and construct
  *     `new ChatSession(model)` directly.
+ *
+ * `options.draftModelPath` attaches an external draft checkpoint (DSpark or
+ * Google gemma-4 assistant, auto-detected from the draft's config.json) for
+ * speculative decoding — gemma4 only; any other detected family rejects it.
+ * Without the option, Gemma4 loads `<modelPath>/draft/` automatically when
+ * that embedded checkpoint is present.
+ * The resulting session auto-enables the speculative path (the model
+ * reports `hasMtpWeights()`); pass `enableMtp: false` per call to opt out.
  */
-export async function loadSession(modelPath) {
-    const m = await loadModel(modelPath);
-    if (m instanceof HarrierModel) {
+export async function loadSession(modelPath, options) {
+    const modelType = await detectModelType(modelPath);
+    const kind = findFamily(modelType).kind;
+    if (kind === 'embedding') {
         throw new Error('loadSession: embedding models (Harrier) cannot be wrapped in a ChatSession');
     }
-    if (m instanceof QianfanOCRModel) {
+    if (kind === 'vlm') {
         throw new Error('loadSession: Qianfan-OCR / InternVL session support lives in @mlx-node/vlm. Import QianfanOCRModel from @mlx-node/vlm and construct ChatSession(model) directly.');
     }
+    const m = await dispatchLoad(modelType, modelPath, options);
     return new ChatSession(m);
 }
 export async function detectModelType(modelPath) {
     try {
         const raw = await readFile(join(modelPath, 'config.json'), 'utf-8');
-        const config = JSON.parse(raw);
-        const rawModelType = config.model_type ?? 'qwen3';
-        // Normalize model_type: gemma4_text → gemma4
-        let modelType = (rawModelType === 'gemma4_text' ? 'gemma4' : rawModelType);
-        // Detect embedding models: Qwen3 backbone with base architecture (no ForCausalLM)
-        if (modelType === 'qwen3') {
-            const architectures = config.architectures ?? [];
-            if (architectures.includes('Qwen3Model') && !architectures.includes('Qwen3ForCausalLM')) {
-                modelType = 'harrier';
-            }
-        }
-        if (!SUPPORTED_MODEL_TYPES.has(modelType)) {
-            throw new Error(`Unsupported model_type "${modelType}" in ${modelPath}/config.json`);
-        }
-        return modelType;
+        const config = normalizeConfig(modelPath, JSON.parse(raw));
+        const baseFamily = config.usesDefaultModelType
+            ? MODEL_FAMILY_INDEX.defaultForNullishModelType
+            : config.rawModelType === undefined
+                ? undefined
+                : MODEL_FAMILY_INDEX.byRawModelType.get(config.rawModelType);
+        const matchContext = { ...config, modelType: baseFamily?.modelType };
+        const family = MODEL_FAMILY_REGISTRY.find((candidate) => matchesArchitectureProbe(candidate, matchContext)) ?? baseFamily;
+        if (family === undefined)
+            throw new UnsupportedModelTypeError(modelPath, config.rawModelTypeLabel);
+        return family.modelType;
     }
     catch (e) {
-        if (e instanceof Error && e.message.startsWith('Unsupported model_type'))
+        if (e instanceof UnsupportedModelTypeError || e instanceof MalformedModelConfigError)
             throw e;
         throw new Error(`Cannot detect model type: config.json not found in ${modelPath}`);
     }

package/dist/models/paged-config-override.d.ts

@@ -0,0 +1,72 @@
+/**
+ * Process-local model-directory overrides for forcing block-paged KV caches.
+ *
+ * Native model loaders accept a directory path and read paging policy from
+ * `config.json`. This manager creates an isolated temporary clone containing a
+ * patched config plus symlinks to the source files, so callers can opt models
+ * into paging without mutating downloaded checkpoints. Loader-known Qwen3.5
+ * MTP sidecars are preserved explicitly; unrelated directories (including a
+ * Gemma4 `draft/`) remain hidden unless a caller explicitly preserves the
+ * Gemma draft for flat speculative decoding.
+ */
+/** Every chat-capable family currently discovered by `@mlx-node/agent`. */
+export declare const AGENT_PAGED_MODEL_TYPES: readonly ['qwen3', 'qwen3_5', 'qwen3_5_moe', 'gemma4', 'lfm2', 'lfm2_moe'];
+/** Families historically forced paged by `mlx launch claude`. */
+export declare const QWEN35_PAGED_MODEL_TYPES: readonly ['qwen3_5', 'qwen3_5_moe'];
+export interface PagedConfigOverrideManagerOptions {
+    /** Model types to force onto the paged path. Defaults to all agent chat families. */
+    modelTypes?: readonly string[];
+    /** Temporary-directory prefix. Primarily useful for diagnostics/tests. */
+    tempDirPrefix?: string;
+    /**
+     * Pass Gemma4 checkpoints with an embedded `draft/` through unchanged.
+     * This keeps native draft auto-discovery, which currently requires flat KV
+     * caches. Defaults to false so a paged override remains a paged contract.
+     */
+    preserveEmbeddedGemmaDraft?: boolean;
+}
+/**
+ * Owns one isolated set of temporary paged-config overrides.
+ *
+ * A manager is intentionally single-lifecycle: repeated resolution of one
+ * source returns the same override, and `cleanup()` permanently disposes the
+ * manager. Separate managers never share directories, so one launch cannot
+ * remove another launch's live override.
+ */
+export declare class PagedConfigOverrideManager {
+    private readonly modelTypes;
+    private readonly tempDirPrefix;
+    private readonly preserveEmbeddedGemmaDraft;
+    private readonly overrides;
+    private readonly activeResolves;
+    private rootPromise;
+    private cleanupPromise;
+    private disposed;
+    constructor(options?: PagedConfigOverrideManagerOptions);
+    /**
+     * Resolve `modelPath` to a paged-aware clone when its model type is managed.
+     * A caller-supplied canonical family takes precedence over the raw config
+     * type (for example, `gemma4` for a `gemma4_unified` checkpoint).
+     * Unmanaged, unreadable, or malformed checkpoints pass through unchanged.
+     *
+     * `persistPagedCache` is a tri-state cold-tier directive. `undefined` leaves
+     * the field untouched (families with no cold-tier opt-in). A boolean is
+     * AUTHORITATIVE: it writes `persist_paged_cache: <value>` into the cloned
+     * config, overriding whatever the source config.json carries (either alias),
+     * and forces a clone whenever the source's value disagrees so the directive
+     * actually reaches the loader. Callers gate the boolean to the families whose
+     * paged cold restore is sound — the allowlist in
+     * `packages/agent/src/cold-tier.ts`, mirrored by `COLD_RESTORE_FAMILIES` in
+     * `crates/mlx-core/src/cold_tier.rs`. It is a SET, not one family, and it is
+     * deliberately not repeated here: `MlxModelHost` calls this for every
+     * allowlisted family, and a second copy of the list would drift.
+     */
+    resolve(modelPath: string, canonicalModelType?: string, persistPagedCache?: boolean): Promise<string>;
+    private resolveInternal;
+    /** Remove this manager's temporary root without affecting other managers. */
+    cleanup(): Promise<void>;
+    private performCleanup;
+    private createOverride;
+    private getRoot;
+}
+//# sourceMappingURL=paged-config-override.d.ts.map

package/dist/models/paged-config-override.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"paged-config-override.d.ts","sourceRoot":"","sources":["../../src/models/paged-config-override.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;GAUG;AAMH,2EAA2E;AAC3E,eAAO,MAAM,uBAAuB,YAAI,OAAO,EAAE,SAAS,EAAE,aAAa,EAAE,QAAQ,EAAE,MAAM,EAAE,UAAU,CAAU,CAAC;AAElH,iEAAiE;AACjE,eAAO,MAAM,wBAAwB,YAAI,SAAS,EAAE,aAAa,CAAU,CAAC;AAO5E,MAAM,WAAW,iCAAiC;IAChD,qFAAqF;IACrF,UAAU,CAAC,EAAE,SAAS,MAAM,EAAE,CAAC;IAC/B,0EAA0E;IAC1E,aAAa,CAAC,EAAE,MAAM,CAAC;IACvB;;;;OAIG;IACH,0BAA0B,CAAC,EAAE,OAAO,CAAC;CACtC;AAED;;;;;;;GAOG;AACH,qBAAa,0BAA0B;IACrC,OAAO,CAAC,QAAQ,CAAC,UAAU,CAAsB;IACjD,OAAO,CAAC,QAAQ,CAAC,aAAa,CAAS;IACvC,OAAO,CAAC,QAAQ,CAAC,0BAA0B,CAAU;IACrD,OAAO,CAAC,QAAQ,CAAC,SAAS,CAAsC;IAChE,OAAO,CAAC,QAAQ,CAAC,cAAc,CAA8B;IAC7D,OAAO,CAAC,WAAW,CAA8B;IACjD,OAAO,CAAC,cAAc,CAA4B;IAClD,OAAO,CAAC,QAAQ,CAAS;IAEzB,YAAY,OAAO,GAAE,iCAAsC,EAI1D;IAED;;;;;;;;;;;;;;;;;OAiBG;IACG,OAAO,CAAC,SAAS,EAAE,MAAM,EAAE,kBAAkB,CAAC,EAAE,MAAM,EAAE,iBAAiB,CAAC,EAAE,OAAO,GAAG,OAAO,CAAC,MAAM,CAAC,CAY1G;YAEa,eAAe;IAuE7B,6EAA6E;IAC7E,OAAO,IAAI,OAAO,CAAC,IAAI,CAAC,CAKvB;YAEa,cAAc;YAad,cAAc;IAsD5B,OAAO,CAAC,OAAO;CAIhB"}