sweet-search 0.0.1 → 2.3.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/LICENSE +190 -0
- package/NOTICE +23 -0
- package/core/cli.js +51 -0
- package/core/config.js +27 -0
- package/core/embedding/embedding-cache.js +467 -0
- package/core/embedding/embedding-local-model.js +845 -0
- package/core/embedding/embedding-remote.js +492 -0
- package/core/embedding/embedding-service.js +712 -0
- package/core/embedding/embedding-telemetry.js +219 -0
- package/core/embedding/index.js +40 -0
- package/core/graph/community-detector.js +294 -0
- package/core/graph/graph-expansion.js +839 -0
- package/core/graph/graph-extractor.js +2304 -0
- package/core/graph/graph-search.js +2148 -0
- package/core/graph/hcgs-generator.js +666 -0
- package/core/graph/index.js +16 -0
- package/core/graph/leiden-algorithm.js +547 -0
- package/core/graph/relationship-resolver.js +366 -0
- package/core/graph/repo-map.js +408 -0
- package/core/graph/summary-manager.js +549 -0
- package/core/indexing/artifact-builder.js +1054 -0
- package/core/indexing/ast-chunker.js +709 -0
- package/core/indexing/chunking/chunk-builder.js +170 -0
- package/core/indexing/chunking/markdown-chunker.js +503 -0
- package/core/indexing/chunking/plaintext-chunker.js +104 -0
- package/core/indexing/dedup/dedup-phase.js +159 -0
- package/core/indexing/dedup/exemplar-selector.js +65 -0
- package/core/indexing/document-chunker.js +56 -0
- package/core/indexing/incremental-parser.js +390 -0
- package/core/indexing/incremental-tracker.js +761 -0
- package/core/indexing/index-codebase-v21.js +472 -0
- package/core/indexing/index-maintainer.mjs +1674 -0
- package/core/indexing/index.js +90 -0
- package/core/indexing/indexer-ann.js +1077 -0
- package/core/indexing/indexer-build.js +742 -0
- package/core/indexing/indexer-phases.js +800 -0
- package/core/indexing/indexer-pool.js +764 -0
- package/core/indexing/indexer-sparse-gram.js +98 -0
- package/core/indexing/indexer-utils.js +536 -0
- package/core/indexing/indexer-worker.js +148 -0
- package/core/indexing/li-skip-policy.js +225 -0
- package/core/indexing/merkle-tracker.js +244 -0
- package/core/indexing/model-pool.js +166 -0
- package/core/infrastructure/code-graph-repository.js +120 -0
- package/core/infrastructure/codebase-repository.js +131 -0
- package/core/infrastructure/config/dedup.js +54 -0
- package/core/infrastructure/config/embedding.js +298 -0
- package/core/infrastructure/config/graph.js +80 -0
- package/core/infrastructure/config/index.js +82 -0
- package/core/infrastructure/config/indexing.js +8 -0
- package/core/infrastructure/config/platform.js +254 -0
- package/core/infrastructure/config/ranking.js +221 -0
- package/core/infrastructure/config/search.js +396 -0
- package/core/infrastructure/config/translation.js +89 -0
- package/core/infrastructure/config/vector-store.js +114 -0
- package/core/infrastructure/constants.js +86 -0
- package/core/infrastructure/coreml-cascade.js +909 -0
- package/core/infrastructure/coreml-cascade.json +46 -0
- package/core/infrastructure/coreml-provider.js +81 -0
- package/core/infrastructure/db-utils.js +69 -0
- package/core/infrastructure/dedup-hashing.js +83 -0
- package/core/infrastructure/hardware-capability.js +332 -0
- package/core/infrastructure/index.js +104 -0
- package/core/infrastructure/language-patterns/maps.js +121 -0
- package/core/infrastructure/language-patterns/registry-core.js +323 -0
- package/core/infrastructure/language-patterns/registry-data-query.js +155 -0
- package/core/infrastructure/language-patterns/registry-object-oriented.js +285 -0
- package/core/infrastructure/language-patterns/registry-tooling.js +240 -0
- package/core/infrastructure/language-patterns/registry-web-style.js +143 -0
- package/core/infrastructure/language-patterns/registry.js +19 -0
- package/core/infrastructure/language-patterns.js +141 -0
- package/core/infrastructure/llm-provider.js +733 -0
- package/core/infrastructure/manifest.json +46 -0
- package/core/infrastructure/maxsim.wasm +0 -0
- package/core/infrastructure/model-fetcher.js +423 -0
- package/core/infrastructure/model-registry.js +214 -0
- package/core/infrastructure/native-inference.js +587 -0
- package/core/infrastructure/native-resolver.js +187 -0
- package/core/infrastructure/native-sparse-gram.js +257 -0
- package/core/infrastructure/native-tokenizer.js +160 -0
- package/core/infrastructure/onnx-mutex.js +45 -0
- package/core/infrastructure/onnx-session-utils.js +261 -0
- package/core/infrastructure/ort-pipeline.js +111 -0
- package/core/infrastructure/project-detector.js +102 -0
- package/core/infrastructure/quantization.js +410 -0
- package/core/infrastructure/simd-distance.js +502 -0
- package/core/infrastructure/simd-distance.wasm +0 -0
- package/core/infrastructure/tree-sitter-provider.js +665 -0
- package/core/infrastructure/webgpu-maxsim.js +222 -0
- package/core/query/index.js +35 -0
- package/core/query/intent-detector.js +201 -0
- package/core/query/intent-router.js +156 -0
- package/core/query/query-router-catboost.js +222 -0
- package/core/query/query-router-ml.js +266 -0
- package/core/query/query-router.js +213 -0
- package/core/ranking/cascaded-scorer.js +379 -0
- package/core/ranking/flashrank.js +810 -0
- package/core/ranking/index.js +49 -0
- package/core/ranking/late-interaction-index.js +2383 -0
- package/core/ranking/late-interaction-model.js +812 -0
- package/core/ranking/local-reranker.js +374 -0
- package/core/ranking/mmr.js +379 -0
- package/core/ranking/quality-scorer.js +363 -0
- package/core/search/context-expander.js +1167 -0
- package/core/search/dedup/sibling-expander.js +327 -0
- package/core/search/index.js +16 -0
- package/core/search/search-boost.js +259 -0
- package/core/search/search-cli.js +544 -0
- package/core/search/search-format.js +282 -0
- package/core/search/search-fusion.js +327 -0
- package/core/search/search-hybrid.js +204 -0
- package/core/search/search-pattern-chunks.js +337 -0
- package/core/search/search-pattern-planner.js +439 -0
- package/core/search/search-pattern-prefilter.js +412 -0
- package/core/search/search-pattern-ripgrep.js +663 -0
- package/core/search/search-pattern.js +463 -0
- package/core/search/search-postprocess.js +452 -0
- package/core/search/search-semantic.js +706 -0
- package/core/search/search-server.js +554 -0
- package/core/search/session-daemon-prewarm.mjs +164 -0
- package/core/search/session-warmup.js +595 -0
- package/core/search/sweet-search.js +632 -0
- package/core/search/warmup-metrics.js +532 -0
- package/core/start-server.js +6 -0
- package/core/training/query-router/features/extractor.js +762 -0
- package/core/training/query-router/features/multilingual-patterns.js +431 -0
- package/core/training/query-router/features/text-segmenter.js +303 -0
- package/core/training/query-router/features/unicode-utils.js +383 -0
- package/core/training/query-router/output/v45_router_d4.js +11521 -0
- package/core/training/query-router/output/v46_router_d4.js +11498 -0
- package/core/vector-store/binary-heap.js +227 -0
- package/core/vector-store/binary-hnsw-index.js +1004 -0
- package/core/vector-store/float-vector-store.js +234 -0
- package/core/vector-store/hnsw-index.js +580 -0
- package/core/vector-store/index.js +39 -0
- package/core/vector-store/seismic-index.js +498 -0
- package/core/vocabulary/index.js +84 -0
- package/core/vocabulary/vocab-constants.js +20 -0
- package/core/vocabulary/vocab-miner-extractors.js +375 -0
- package/core/vocabulary/vocab-miner-nl.js +404 -0
- package/core/vocabulary/vocab-miner-utils.js +146 -0
- package/core/vocabulary/vocab-miner.js +574 -0
- package/core/vocabulary/vocab-prewarm-cli.js +110 -0
- package/core/vocabulary/vocab-ranker.js +492 -0
- package/core/vocabulary/vocab-warmer.js +523 -0
- package/core/vocabulary/vocab-warmup-orchestrator.js +425 -0
- package/core/vocabulary/vocabulary-utils.js +704 -0
- package/crates/wasm-router/pkg/package.json +13 -0
- package/crates/wasm-router/pkg/query_router_wasm.d.ts +36 -0
- package/crates/wasm-router/pkg/query_router_wasm.js +271 -0
- package/crates/wasm-router/pkg/query_router_wasm_bg.wasm +0 -0
- package/crates/wasm-router/pkg/query_router_wasm_bg.wasm.d.ts +19 -0
- package/mcp/config-gen.js +121 -0
- package/mcp/server.js +335 -0
- package/mcp/tool-handlers.js +476 -0
- package/package.json +131 -9
- package/scripts/benchmark-harness.js +794 -0
- package/scripts/init.js +1058 -0
- package/scripts/smoke-test.js +435 -0
- package/scripts/uninstall.js +478 -0
- package/scripts/verify-runtime.js +176 -0
|
@@ -0,0 +1,587 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* Native Inference — candle-based FP32 embedding + late interaction inference.
|
|
3
|
+
*
|
|
4
|
+
* Provides native model wrappers that serve as drop-in replacements for the
|
|
5
|
+
* ORT inference paths. Uses the napi-rs addon with Metal GPU acceleration on
|
|
6
|
+
* Apple Silicon, CPU fallback elsewhere.
|
|
7
|
+
*
|
|
8
|
+
* On M3+ Apple Silicon, the addon also loads a CoreML variant cascade
|
|
9
|
+
* alongside the candle backbone. This module is the single point where
|
|
10
|
+
* the cascade dirs resolved by `coreml-cascade.js` are handed to the
|
|
11
|
+
* Rust `NativeEmbeddingModel::load` / `NativeLateInteractionModel::load`
|
|
12
|
+
* constructors. Lower hardware, no cache, or
|
|
13
|
+
* `SWEET_SEARCH_COREML_CASCADE=0` collapses to the candle-only path
|
|
14
|
+
* transparently.
|
|
15
|
+
*
|
|
16
|
+
* Environment:
|
|
17
|
+
* SWEET_SEARCH_NATIVE_INFERENCE=0|1 — force disable/enable
|
|
18
|
+
* SWEET_SEARCH_COREML_CASCADE=0 — force-disable the cascade
|
|
19
|
+
* even when hardware +
|
|
20
|
+
* cache are both eligible
|
|
21
|
+
* (diagnostic / benchmarking)
|
|
22
|
+
* SWEET_SEARCH_COREML_STATS=1 — dump per-variant
|
|
23
|
+
* dispatch report on
|
|
24
|
+
* addon shutdown (see
|
|
25
|
+
* coreml_embedding.rs)
|
|
26
|
+
* SWEET_SEARCH_CUDA=0 — force-disable CUDA even
|
|
27
|
+
* on an eligible host
|
|
28
|
+
* (diagnostic)
|
|
29
|
+
* SWEET_SEARCH_CUDA_COMPUTE_CAP=<cc> — set by this module
|
|
30
|
+
* before addon loads so
|
|
31
|
+
* the Rust dtype policy
|
|
32
|
+
* picks BF16/F16/F32 by
|
|
33
|
+
* compute capability.
|
|
34
|
+
* See mod.rs::optimal_dtype
|
|
35
|
+
* CANDLE_METAL_COMPUTE_PER_BUFFER=<N> — candle default 50 (tuned)
|
|
36
|
+
* CANDLE_METAL_COMMAND_POOL_SIZE=<N> — candle default 5 (tuned)
|
|
37
|
+
*/
|
|
38
|
+
|
|
39
|
+
import { existsSync } from 'fs';
|
|
40
|
+
import { join } from 'path';
|
|
41
|
+
import { createRequire } from 'module';
|
|
42
|
+
import { resolveNativeAddon } from './native-resolver.js';
|
|
43
|
+
import { createTokenizer } from './native-tokenizer.js';
|
|
44
|
+
import { getModelCacheDir, fetchModel } from './model-fetcher.js';
|
|
45
|
+
import { getModelEntry } from './model-registry.js';
|
|
46
|
+
import { getCoremlCascadeResolvedDirs } from './coreml-cascade.js';
|
|
47
|
+
import { detectHardwareCapability } from './hardware-capability.js';
|
|
48
|
+
|
|
49
|
+
const require = createRequire(import.meta.url);
|
|
50
|
+
|
|
51
|
+
// ─── State ───
|
|
52
|
+
|
|
53
|
+
let _addon = null;
|
|
54
|
+
let _embeddingModel = null;
|
|
55
|
+
let _embeddingModelLoadPromise = null; // race-gate for concurrent first calls
|
|
56
|
+
let _liModel = null;
|
|
57
|
+
let _liModelLoadPromise = null;
|
|
58
|
+
let _embTokenizer = null;
|
|
59
|
+
let _embTokenizerLoadPromise = null;
|
|
60
|
+
let _liTokenizer = null;
|
|
61
|
+
let _liTokenizerLoadPromise = null;
|
|
62
|
+
let _available = null;
|
|
63
|
+
let _coremlCascadeLogged = false;
|
|
64
|
+
|
|
65
|
+
// ─── Cascade-dir selection (pure helper, unit-tested) ───
|
|
66
|
+
|
|
67
|
+
/**
|
|
68
|
+
* Decide whether the Rust addon should load a CoreML cascade for this
|
|
69
|
+
* deviceKind. Only `'metal'` triggers cascade resolution; `'cuda'`, `'cpu'`,
|
|
70
|
+
* and `'auto'` skip cascade (CUDA has no cascade, CPU has no GPU dispatch).
|
|
71
|
+
*
|
|
72
|
+
* Pure function — takes a resolver for the cascade dirs and the override.
|
|
73
|
+
* Returns `undefined` (explicit "no cascade") or a string path.
|
|
74
|
+
*
|
|
75
|
+
* @param {'cpu'|'metal'|'cuda'|'auto'} deviceKind
|
|
76
|
+
* @param {string|undefined} cascadeDirOverride - explicit caller override (wins)
|
|
77
|
+
* @param {() => string|null|undefined} resolveCascadeDir - called only when needed
|
|
78
|
+
*/
|
|
79
|
+
export function pickCascadeDirForDevice(deviceKind, cascadeDirOverride, resolveCascadeDir) {
|
|
80
|
+
if (cascadeDirOverride !== undefined) return cascadeDirOverride;
|
|
81
|
+
if (deviceKind !== 'metal') return undefined;
|
|
82
|
+
return resolveCascadeDir() || undefined;
|
|
83
|
+
}
|
|
84
|
+
|
|
85
|
+
// ─── CUDA compute-capability env propagation ───
|
|
86
|
+
|
|
87
|
+
/**
|
|
88
|
+
* Ensure `SWEET_SEARCH_CUDA_COMPUTE_CAP` is set for the current process
|
|
89
|
+
* before the addon loads a CUDA model. The Rust `optimal_dtype` reads
|
|
90
|
+
* this env var to pick BF16 on Ampere+ and F16/F32 on older GPUs.
|
|
91
|
+
*
|
|
92
|
+
* Idempotent: honors an already-set value (useful for forcing a dtype
|
|
93
|
+
* tier in benchmarks) and silently no-ops when there is no NVIDIA GPU.
|
|
94
|
+
*/
|
|
95
|
+
function propagateCudaComputeCapToAddonEnv() {
|
|
96
|
+
if (process.env.SWEET_SEARCH_CUDA_COMPUTE_CAP) return;
|
|
97
|
+
try {
|
|
98
|
+
const hw = detectHardwareCapability();
|
|
99
|
+
if (hw.nvidiaGpu?.computeCapability) {
|
|
100
|
+
process.env.SWEET_SEARCH_CUDA_COMPUTE_CAP = hw.nvidiaGpu.computeCapability;
|
|
101
|
+
}
|
|
102
|
+
} catch {
|
|
103
|
+
// Hardware detection is advisory. If it fails, the Rust side falls
|
|
104
|
+
// back to F32 which is always correct, if slower.
|
|
105
|
+
}
|
|
106
|
+
}
|
|
107
|
+
|
|
108
|
+
// ─── CoreML cascade resolution ───
|
|
109
|
+
|
|
110
|
+
/**
|
|
111
|
+
* Resolve which CoreML cascade dirs the Rust addon should try to load.
|
|
112
|
+
* Logged exactly once per process so a mis-configured cascade surfaces
|
|
113
|
+
* at startup instead of silently falling through on every call.
|
|
114
|
+
*
|
|
115
|
+
* Always returns an object — never throws. The returned dirs can be
|
|
116
|
+
* `null`, which the Rust addon treats as "CoreML path disabled" and
|
|
117
|
+
* falls back to candle unconditionally.
|
|
118
|
+
*/
|
|
119
|
+
function resolveCoremlCascadeForAddon() {
|
|
120
|
+
const resolved = getCoremlCascadeResolvedDirs();
|
|
121
|
+
if (!_coremlCascadeLogged) {
|
|
122
|
+
_coremlCascadeLogged = true;
|
|
123
|
+
const hw = detectHardwareCapability();
|
|
124
|
+
// Log a single line so every subsequent `[NativeInference]` message
|
|
125
|
+
// can be correlated with the cascade decision made here.
|
|
126
|
+
if (resolved.embedDir || resolved.liDir) {
|
|
127
|
+
process.stderr.write(
|
|
128
|
+
`[NativeInference] CoreML cascade: ${resolved.status}` +
|
|
129
|
+
` (embed=${resolved.embedDir ? 'yes' : 'no'}, li=${resolved.liDir ? 'yes' : 'no'},` +
|
|
130
|
+
` chip=${hw.brandString || 'unknown'})\n`
|
|
131
|
+
);
|
|
132
|
+
} else if (hw.coremlCascadeEligible) {
|
|
133
|
+
process.stderr.write(
|
|
134
|
+
`[NativeInference] CoreML cascade: ${resolved.status} —` +
|
|
135
|
+
` hardware eligible but cache missing. Run \`node scripts/build-coreml-cascade.js\`` +
|
|
136
|
+
` to enable the 18% speedup on ${hw.brandString}\n`
|
|
137
|
+
);
|
|
138
|
+
}
|
|
139
|
+
// Ineligible hardware: silent. The user has no action to take and
|
|
140
|
+
// we already log the candle device via the existing "Loaded"
|
|
141
|
+
// line below.
|
|
142
|
+
}
|
|
143
|
+
return resolved;
|
|
144
|
+
}
|
|
145
|
+
|
|
146
|
+
// ─── Addon Loading ───
|
|
147
|
+
|
|
148
|
+
function loadAddon() {
|
|
149
|
+
if (_addon) return _addon;
|
|
150
|
+
const addonPath = resolveNativeAddon();
|
|
151
|
+
if (!addonPath) return null;
|
|
152
|
+
try {
|
|
153
|
+
_addon = require(addonPath);
|
|
154
|
+
return _addon;
|
|
155
|
+
} catch {
|
|
156
|
+
return null;
|
|
157
|
+
}
|
|
158
|
+
}
|
|
159
|
+
|
|
160
|
+
// ─── Detection ───
|
|
161
|
+
|
|
162
|
+
/**
|
|
163
|
+
* Check if native inference is available and not disabled via env var.
|
|
164
|
+
* Caches result after first check.
|
|
165
|
+
*/
|
|
166
|
+
export function isNativeInferenceAvailable() {
|
|
167
|
+
if (_available !== null) return _available;
|
|
168
|
+
|
|
169
|
+
const envFlag = (process.env.SWEET_SEARCH_NATIVE_INFERENCE ?? '').trim().toLowerCase();
|
|
170
|
+
if (envFlag === '0' || envFlag === 'false' || envFlag === 'off') {
|
|
171
|
+
_available = false;
|
|
172
|
+
return false;
|
|
173
|
+
}
|
|
174
|
+
|
|
175
|
+
const addon = loadAddon();
|
|
176
|
+
_available = typeof addon?.NativeEmbeddingModel?.load === 'function';
|
|
177
|
+
return _available;
|
|
178
|
+
}
|
|
179
|
+
|
|
180
|
+
// ─── Embedding Model ───
|
|
181
|
+
|
|
182
|
+
/**
|
|
183
|
+
* Load the native embedding model (CodeRankEmbed FP32 safetensors).
|
|
184
|
+
* Returns the model instance or null if unavailable.
|
|
185
|
+
*
|
|
186
|
+
* Race-gated: concurrent first calls share a single load promise so the
|
|
187
|
+
* underlying napi `addon.NativeEmbeddingModel.load(...)` runs exactly once.
|
|
188
|
+
* Without this gate, multiple parallel queries (e.g. eval runner with
|
|
189
|
+
* --concurrency=N) all see _embeddingModel == null and each load a fresh
|
|
190
|
+
* model copy, wasting Metal memory and printing N "Loaded" lines.
|
|
191
|
+
*/
|
|
192
|
+
export async function getNativeEmbeddingModel() {
|
|
193
|
+
if (_embeddingModel) return _embeddingModel;
|
|
194
|
+
if (_embeddingModelLoadPromise) return _embeddingModelLoadPromise;
|
|
195
|
+
|
|
196
|
+
_embeddingModelLoadPromise = (async () => {
|
|
197
|
+
const addon = loadAddon();
|
|
198
|
+
if (!addon?.NativeEmbeddingModel) return null;
|
|
199
|
+
|
|
200
|
+
await fetchModel('coderankembed-fp32');
|
|
201
|
+
|
|
202
|
+
const entry = getModelEntry('coderankembed-fp32');
|
|
203
|
+
const modelDir = getModelCacheDir(entry.hfId);
|
|
204
|
+
const safetensorsPath = join(modelDir, 'model.safetensors');
|
|
205
|
+
const configPath = join(modelDir, 'config.json');
|
|
206
|
+
|
|
207
|
+
if (!existsSync(safetensorsPath) || !existsSync(configPath)) return null;
|
|
208
|
+
|
|
209
|
+
// Resolve the CoreML cascade dir for NomicBERT embeddings. Returns
|
|
210
|
+
// null on ineligible hardware, empty cache, or explicit opt-out
|
|
211
|
+
// (SWEET_SEARCH_COREML_CASCADE=0). The Rust addon's constructor
|
|
212
|
+
// runs a parity check against the smallest variant before arming
|
|
213
|
+
// the dispatch path; failure there falls back to candle.
|
|
214
|
+
const cascade = resolveCoremlCascadeForAddon();
|
|
215
|
+
|
|
216
|
+
const t0 = Date.now();
|
|
217
|
+
_embeddingModel = addon.NativeEmbeddingModel.load(
|
|
218
|
+
safetensorsPath,
|
|
219
|
+
configPath,
|
|
220
|
+
cascade.embedDir || undefined,
|
|
221
|
+
);
|
|
222
|
+
console.log(`[NativeInference] Embedding model loaded in ${Date.now() - t0}ms (dim: ${_embeddingModel.dim}, device: ${addon.nativeInferenceDevice()})`);
|
|
223
|
+
|
|
224
|
+
return _embeddingModel;
|
|
225
|
+
})();
|
|
226
|
+
|
|
227
|
+
try {
|
|
228
|
+
return await _embeddingModelLoadPromise;
|
|
229
|
+
} finally {
|
|
230
|
+
// Clear the promise on resolve so a future explicit unload can re-load.
|
|
231
|
+
// Keep it set on success: subsequent calls hit the _embeddingModel cache
|
|
232
|
+
// first and never reach the promise check.
|
|
233
|
+
}
|
|
234
|
+
}
|
|
235
|
+
|
|
236
|
+
/**
|
|
237
|
+
* Get or create the embedding tokenizer. Race-gated like getNativeEmbeddingModel.
|
|
238
|
+
* Uses the INT8 model's tokenizer (same vocab as FP32).
|
|
239
|
+
*/
|
|
240
|
+
async function getEmbTokenizer() {
|
|
241
|
+
if (_embTokenizer) return _embTokenizer;
|
|
242
|
+
if (_embTokenizerLoadPromise) return _embTokenizerLoadPromise;
|
|
243
|
+
_embTokenizerLoadPromise = (async () => {
|
|
244
|
+
const entry = getModelEntry('coderankembed-int8');
|
|
245
|
+
const tokenizerPath = join(getModelCacheDir(entry.hfId), 'tokenizer.json');
|
|
246
|
+
_embTokenizer = await createTokenizer(tokenizerPath);
|
|
247
|
+
return _embTokenizer;
|
|
248
|
+
})();
|
|
249
|
+
return _embTokenizerLoadPromise;
|
|
250
|
+
}
|
|
251
|
+
|
|
252
|
+
/**
|
|
253
|
+
* Tokenize texts for the native model, returning napi-compatible arrays.
|
|
254
|
+
*/
|
|
255
|
+
function tokenizedToNapi(tokenized, batchSize, seqLen) {
|
|
256
|
+
const inputIds = new Array(batchSize);
|
|
257
|
+
const attentionMask = new Array(batchSize);
|
|
258
|
+
for (let b = 0; b < batchSize; b++) {
|
|
259
|
+
const ids = new Array(seqLen);
|
|
260
|
+
const mask = new Array(seqLen);
|
|
261
|
+
const base = b * seqLen;
|
|
262
|
+
for (let s = 0; s < seqLen; s++) {
|
|
263
|
+
ids[s] = Number(tokenized.input_ids.data[base + s]);
|
|
264
|
+
mask[s] = Number(tokenized.attention_mask.data[base + s]);
|
|
265
|
+
}
|
|
266
|
+
inputIds[b] = ids;
|
|
267
|
+
attentionMask[b] = mask;
|
|
268
|
+
}
|
|
269
|
+
return { inputIds, attentionMask };
|
|
270
|
+
}
|
|
271
|
+
|
|
272
|
+
/**
|
|
273
|
+
* Native embedding inference — drop-in replacement for callLocalModel().
|
|
274
|
+
*
|
|
275
|
+
* @param {string[]} texts - Array of texts to embed
|
|
276
|
+
* @param {object} [options]
|
|
277
|
+
* @param {number} [options.maxLength=512] - Max sequence length
|
|
278
|
+
* @returns {Float32Array[]} Array of L2-normalized embedding vectors (768d)
|
|
279
|
+
*/
|
|
280
|
+
export async function nativeEmbed(texts, options = {}) {
|
|
281
|
+
const { maxLength = 512 } = options;
|
|
282
|
+
const model = await getNativeEmbeddingModel();
|
|
283
|
+
if (!model) throw new Error('[NativeInference] Embedding model not loaded');
|
|
284
|
+
|
|
285
|
+
const tokenizer = await getEmbTokenizer();
|
|
286
|
+
const tokenized = tokenizer(texts, { padding: true, truncation: true, max_length: maxLength });
|
|
287
|
+
const batchSize = texts.length;
|
|
288
|
+
const seqLen = tokenized.input_ids.dims[1];
|
|
289
|
+
|
|
290
|
+
const { inputIds, attentionMask } = tokenizedToNapi(tokenized, batchSize, seqLen);
|
|
291
|
+
// The native addon returns a flat `Float32Array` of length batch * dim —
|
|
292
|
+
// see the comment on embed_batch in embedding_model.rs. Use `.subarray()`
|
|
293
|
+
// to return zero-copy views over the shared backing buffer (2026-04-15 perf
|
|
294
|
+
// fix — previous `.slice()` was a per-row memcpy for no reason).
|
|
295
|
+
//
|
|
296
|
+
// Callers MUST treat each view as read-only. Mutation audit (2026-04-15):
|
|
297
|
+
// - callLocalModelGpu at core/embedding/embedding-local-model.js:547
|
|
298
|
+
// returns the array unchanged to its caller
|
|
299
|
+
// - downstream embedding-service paths call truncateForHNSW, which
|
|
300
|
+
// allocates a fresh Float32Array per row (quantization.js)
|
|
301
|
+
// - SQLite persistence paths take a BLOB copy via better-sqlite3
|
|
302
|
+
// - no grep hit for `embeddings[i][j] = ` or `.set(` on the returned rows
|
|
303
|
+
// Any future consumer that writes through a returned view will silently
|
|
304
|
+
// corrupt neighbouring rows in the shared buffer. If you need to mutate,
|
|
305
|
+
// copy first with `new Float32Array(view)` or `view.slice()`.
|
|
306
|
+
const flat = await model.embedBatch(inputIds, attentionMask);
|
|
307
|
+
const dim = model.dim;
|
|
308
|
+
|
|
309
|
+
const embeddings = new Array(batchSize);
|
|
310
|
+
for (let i = 0; i < batchSize; i++) {
|
|
311
|
+
embeddings[i] = flat.subarray(i * dim, (i + 1) * dim);
|
|
312
|
+
}
|
|
313
|
+
return embeddings;
|
|
314
|
+
}
|
|
315
|
+
|
|
316
|
+
// ─── Late Interaction Model ───
|
|
317
|
+
|
|
318
|
+
/**
|
|
319
|
+
* Load the native LI model (LateOn-Code FP32 safetensors + projection).
|
|
320
|
+
* Returns the model instance or null if unavailable. Race-gated.
|
|
321
|
+
*/
|
|
322
|
+
export async function getNativeLiModel() {
|
|
323
|
+
if (_liModel) return _liModel;
|
|
324
|
+
if (_liModelLoadPromise) return _liModelLoadPromise;
|
|
325
|
+
_liModelLoadPromise = (async () => {
|
|
326
|
+
const addon = loadAddon();
|
|
327
|
+
if (!addon?.NativeLateInteractionModel) return null;
|
|
328
|
+
|
|
329
|
+
await fetchModel('lateon-code-fp32');
|
|
330
|
+
|
|
331
|
+
const entry = getModelEntry('lateon-code-fp32');
|
|
332
|
+
const modelDir = getModelCacheDir(entry.hfId);
|
|
333
|
+
const backbonePath = join(modelDir, 'model.safetensors');
|
|
334
|
+
const projPath = join(modelDir, '1_Dense', 'model.safetensors');
|
|
335
|
+
const configPath = join(modelDir, 'config.json');
|
|
336
|
+
|
|
337
|
+
if (!existsSync(backbonePath) || !existsSync(projPath) || !existsSync(configPath)) return null;
|
|
338
|
+
|
|
339
|
+
// Resolve the CoreML cascade dir for ModernBERT LI. Same contract
|
|
340
|
+
// as the embedding model above — see that comment.
|
|
341
|
+
const cascade = resolveCoremlCascadeForAddon();
|
|
342
|
+
|
|
343
|
+
const t0 = Date.now();
|
|
344
|
+
_liModel = addon.NativeLateInteractionModel.load(
|
|
345
|
+
backbonePath,
|
|
346
|
+
projPath,
|
|
347
|
+
configPath,
|
|
348
|
+
cascade.liDir || undefined,
|
|
349
|
+
);
|
|
350
|
+
console.log(`[NativeInference] LI model loaded in ${Date.now() - t0}ms (dim: ${_liModel.dim}, device: ${addon.nativeInferenceDevice()})`);
|
|
351
|
+
|
|
352
|
+
return _liModel;
|
|
353
|
+
})();
|
|
354
|
+
return _liModelLoadPromise;
|
|
355
|
+
}
|
|
356
|
+
|
|
357
|
+
/**
|
|
358
|
+
* Get or create the LI tokenizer. Race-gated.
|
|
359
|
+
*/
|
|
360
|
+
async function getLiTokenizer() {
|
|
361
|
+
if (_liTokenizer) return _liTokenizer;
|
|
362
|
+
if (_liTokenizerLoadPromise) return _liTokenizerLoadPromise;
|
|
363
|
+
_liTokenizerLoadPromise = (async () => {
|
|
364
|
+
const entry = getModelEntry('lateon-code');
|
|
365
|
+
const tokenizerPath = join(getModelCacheDir(entry.hfId), 'tokenizer.json');
|
|
366
|
+
_liTokenizer = await createTokenizer(tokenizerPath);
|
|
367
|
+
return _liTokenizer;
|
|
368
|
+
})();
|
|
369
|
+
return _liTokenizerLoadPromise;
|
|
370
|
+
}
|
|
371
|
+
|
|
372
|
+
/**
|
|
373
|
+
* Native LI encoding from a pre-tokenized batch. Avoids re-tokenizing when
|
|
374
|
+
* the caller already has the tokenizer output (e.g. because it needs
|
|
375
|
+
* input_ids for skiplist filtering). Accepts the object returned by the
|
|
376
|
+
* native tokenizer: `{ input_ids: { data, dims }, attention_mask: { data, dims } }`.
|
|
377
|
+
*
|
|
378
|
+
* @param {object} tokenized - Pre-tokenized batch from native tokenizer
|
|
379
|
+
* @returns {Float32Array[][]} Per-document arrays of Float32Array token vectors
|
|
380
|
+
*/
|
|
381
|
+
export async function nativeLiEncodeTokenized(tokenized) {
|
|
382
|
+
const model = await getNativeLiModel();
|
|
383
|
+
if (!model) throw new Error('[NativeInference] LI model not loaded');
|
|
384
|
+
|
|
385
|
+
const [batchSize, seqLen] = tokenized.input_ids.dims;
|
|
386
|
+
const { inputIds, attentionMask } = tokenizedToNapi(tokenized, batchSize, seqLen);
|
|
387
|
+
const result = await model.encodeBatch(inputIds, attentionMask);
|
|
388
|
+
const dim = model.dim;
|
|
389
|
+
|
|
390
|
+
// The native addon returns `vectors` as a flat `Float32Array` (zero-copy
|
|
391
|
+
// from the Rust Vec<f32> via napi typed array) — see encode_batch in
|
|
392
|
+
// li_model.rs. Use `.subarray()` to return zero-copy views over the shared
|
|
393
|
+
// backing buffer (2026-04-15 perf fix — previous `.slice()` was a per-token
|
|
394
|
+
// memcpy for no reason).
|
|
395
|
+
//
|
|
396
|
+
// LI token views are READ-ONLY. Mutation audit (2026-04-15):
|
|
397
|
+
// - nativeLiEncodeTokenized is called from late-interaction-model.js:387
|
|
398
|
+
// (encodeDocumentsGpu) which pushes views into a JS array and hands
|
|
399
|
+
// them to indexer-ann.js::finalizeBatchResults
|
|
400
|
+
// - finalizeBatchResults calls liIndex.add which at
|
|
401
|
+
// late-interaction-index.js:429-430 does
|
|
402
|
+
// `tokens.slice(...).map(emb => emb.slice(0, tokenDim))` — the inner
|
|
403
|
+
// `.slice()` creates a fresh Float32Array per token (copy)
|
|
404
|
+
// - poolTokens at late-interaction-model.js:619 does
|
|
405
|
+
// `new Float32Array(clusterInput[i])` (copy via constructor)
|
|
406
|
+
// - persistence goes through _writeSegmentFile which copies bytes into
|
|
407
|
+
// a pre-allocated Buffer
|
|
408
|
+
// Any future consumer that writes through a returned view will corrupt
|
|
409
|
+
// the underlying batch buffer and neighbouring tokens. Copy before mutating.
|
|
410
|
+
const flat = result.vectors;
|
|
411
|
+
|
|
412
|
+
const allVectors = new Array(batchSize);
|
|
413
|
+
let offset = 0;
|
|
414
|
+
for (let b = 0; b < batchSize; b++) {
|
|
415
|
+
const count = result.tokenCounts[b];
|
|
416
|
+
const docVectors = new Array(count);
|
|
417
|
+
for (let t = 0; t < count; t++) {
|
|
418
|
+
docVectors[t] = flat.subarray(offset + t * dim, offset + (t + 1) * dim);
|
|
419
|
+
}
|
|
420
|
+
allVectors[b] = docVectors;
|
|
421
|
+
offset += count * dim;
|
|
422
|
+
}
|
|
423
|
+
|
|
424
|
+
return allVectors;
|
|
425
|
+
}
|
|
426
|
+
|
|
427
|
+
/**
|
|
428
|
+
* Native LI encoding — returns per-token L2-normalized 128d vectors.
|
|
429
|
+
*
|
|
430
|
+
* @param {string[]} texts - Texts to encode (already prefixed with [Q]/[D])
|
|
431
|
+
* @param {object} [options]
|
|
432
|
+
* @param {number} [options.maxLength=2048] - Max sequence length
|
|
433
|
+
* @returns {Float32Array[][]} Per-document arrays of Float32Array token vectors
|
|
434
|
+
*/
|
|
435
|
+
export async function nativeLiEncode(texts, options = {}) {
|
|
436
|
+
const { maxLength = 2048 } = options;
|
|
437
|
+
const tokenizer = await getLiTokenizer();
|
|
438
|
+
const tokenized = tokenizer(texts, { padding: true, truncation: true, max_length: maxLength });
|
|
439
|
+
return nativeLiEncodeTokenized(tokenized);
|
|
440
|
+
}
|
|
441
|
+
|
|
442
|
+
// ─── Model state queries ───
|
|
443
|
+
|
|
444
|
+
export function isNativeEmbeddingModelLoaded() {
|
|
445
|
+
return _embeddingModel != null;
|
|
446
|
+
}
|
|
447
|
+
|
|
448
|
+
export function isNativeLiModelLoaded() {
|
|
449
|
+
return _liModel != null;
|
|
450
|
+
}
|
|
451
|
+
|
|
452
|
+
// ─── Device-explicit loading ───
|
|
453
|
+
|
|
454
|
+
/**
|
|
455
|
+
* Load the native embedding model on a specific device ("cpu" | "metal" | "auto").
|
|
456
|
+
* Used by model-pool.js to arm GPU models for indexing. Sets the module-scope
|
|
457
|
+
* singleton so subsequent nativeEmbed() calls use this instance.
|
|
458
|
+
*/
|
|
459
|
+
export async function loadNativeEmbeddingModelWithDevice(deviceKind, cascadeDirOverride) {
|
|
460
|
+
if (_embeddingModel) return _embeddingModel;
|
|
461
|
+
if (_embeddingModelLoadPromise) return _embeddingModelLoadPromise;
|
|
462
|
+
|
|
463
|
+
_embeddingModelLoadPromise = (async () => {
|
|
464
|
+
const addon = loadAddon();
|
|
465
|
+
if (!addon?.NativeEmbeddingModel?.loadWithDevice) return null;
|
|
466
|
+
|
|
467
|
+
// CUDA needs the compute-capability env set BEFORE the Rust load
|
|
468
|
+
// path picks a dtype. No-op on non-CUDA kinds and already-set envs.
|
|
469
|
+
if (deviceKind === 'cuda') propagateCudaComputeCapToAddonEnv();
|
|
470
|
+
|
|
471
|
+
await fetchModel('coderankembed-fp32');
|
|
472
|
+
|
|
473
|
+
const entry = getModelEntry('coderankembed-fp32');
|
|
474
|
+
const modelDir = getModelCacheDir(entry.hfId);
|
|
475
|
+
const safetensorsPath = join(modelDir, 'model.safetensors');
|
|
476
|
+
const configPath = join(modelDir, 'config.json');
|
|
477
|
+
|
|
478
|
+
if (!existsSync(safetensorsPath) || !existsSync(configPath)) return null;
|
|
479
|
+
|
|
480
|
+
// Cascade is a CoreML-specific concept. CUDA has no cascade — the Rust
|
|
481
|
+
// addon JITs CUDA kernels per forward pass; there is no `.mlpackage`
|
|
482
|
+
// equivalent to prefetch. Skip cascade resolution entirely on CUDA so
|
|
483
|
+
// we don't log "cascade missing" warnings on a Linux+NVIDIA host.
|
|
484
|
+
const cascadeDir = pickCascadeDirForDevice(
|
|
485
|
+
deviceKind,
|
|
486
|
+
cascadeDirOverride,
|
|
487
|
+
() => resolveCoremlCascadeForAddon().embedDir,
|
|
488
|
+
);
|
|
489
|
+
|
|
490
|
+
const t0 = Date.now();
|
|
491
|
+
_embeddingModel = addon.NativeEmbeddingModel.loadWithDevice(
|
|
492
|
+
safetensorsPath,
|
|
493
|
+
configPath,
|
|
494
|
+
cascadeDir,
|
|
495
|
+
deviceKind,
|
|
496
|
+
);
|
|
497
|
+
console.log(`[NativeInference] Embedding model loaded in ${Date.now() - t0}ms (dim: ${_embeddingModel.dim}, device: ${deviceKind})`);
|
|
498
|
+
|
|
499
|
+
return _embeddingModel;
|
|
500
|
+
})();
|
|
501
|
+
|
|
502
|
+
try {
|
|
503
|
+
return await _embeddingModelLoadPromise;
|
|
504
|
+
} finally {
|
|
505
|
+
// Keep promise set on success; clear only on re-load.
|
|
506
|
+
}
|
|
507
|
+
}
|
|
508
|
+
|
|
509
|
+
/**
|
|
510
|
+
* Load the native LI model on a specific device.
|
|
511
|
+
*/
|
|
512
|
+
export async function loadNativeLiModelWithDevice(deviceKind, cascadeDirOverride) {
|
|
513
|
+
if (_liModel) return _liModel;
|
|
514
|
+
if (_liModelLoadPromise) return _liModelLoadPromise;
|
|
515
|
+
|
|
516
|
+
_liModelLoadPromise = (async () => {
|
|
517
|
+
const addon = loadAddon();
|
|
518
|
+
if (!addon?.NativeLateInteractionModel?.loadWithDevice) return null;
|
|
519
|
+
|
|
520
|
+
// See loadNativeEmbeddingModelWithDevice for why this is CUDA-only.
|
|
521
|
+
if (deviceKind === 'cuda') propagateCudaComputeCapToAddonEnv();
|
|
522
|
+
|
|
523
|
+
await fetchModel('lateon-code-fp32');
|
|
524
|
+
|
|
525
|
+
const entry = getModelEntry('lateon-code-fp32');
|
|
526
|
+
const modelDir = getModelCacheDir(entry.hfId);
|
|
527
|
+
const backbonePath = join(modelDir, 'model.safetensors');
|
|
528
|
+
const projPath = join(modelDir, '1_Dense', 'model.safetensors');
|
|
529
|
+
const configPath = join(modelDir, 'config.json');
|
|
530
|
+
|
|
531
|
+
if (!existsSync(backbonePath) || !existsSync(projPath) || !existsSync(configPath)) return null;
|
|
532
|
+
|
|
533
|
+
// CUDA has no cascade — see the matching comment in
|
|
534
|
+
// loadNativeEmbeddingModelWithDevice.
|
|
535
|
+
const cascadeDir = pickCascadeDirForDevice(
|
|
536
|
+
deviceKind,
|
|
537
|
+
cascadeDirOverride,
|
|
538
|
+
() => resolveCoremlCascadeForAddon().liDir,
|
|
539
|
+
);
|
|
540
|
+
|
|
541
|
+
const t0 = Date.now();
|
|
542
|
+
_liModel = addon.NativeLateInteractionModel.loadWithDevice(
|
|
543
|
+
backbonePath,
|
|
544
|
+
projPath,
|
|
545
|
+
configPath,
|
|
546
|
+
cascadeDir,
|
|
547
|
+
deviceKind,
|
|
548
|
+
);
|
|
549
|
+
console.log(`[NativeInference] LI model loaded in ${Date.now() - t0}ms (dim: ${_liModel.dim}, device: ${deviceKind})`);
|
|
550
|
+
|
|
551
|
+
return _liModel;
|
|
552
|
+
})();
|
|
553
|
+
|
|
554
|
+
return _liModelLoadPromise;
|
|
555
|
+
}
|
|
556
|
+
|
|
557
|
+
// ─── Warmup primitives ───
|
|
558
|
+
|
|
559
|
+
export async function warmupNativeEmbeddingModel() {
|
|
560
|
+
if (!_embeddingModel?.warmupForward) return;
|
|
561
|
+
const t0 = Date.now();
|
|
562
|
+
await _embeddingModel.warmupForward();
|
|
563
|
+
console.log(`[NativeInference] Embedding warmup forward in ${Date.now() - t0}ms`);
|
|
564
|
+
}
|
|
565
|
+
|
|
566
|
+
export async function warmupNativeLiModel() {
|
|
567
|
+
if (!_liModel?.warmupForward) return;
|
|
568
|
+
const t0 = Date.now();
|
|
569
|
+
await _liModel.warmupForward();
|
|
570
|
+
console.log(`[NativeInference] LI warmup forward in ${Date.now() - t0}ms`);
|
|
571
|
+
}
|
|
572
|
+
|
|
573
|
+
// ─── Cleanup ───
|
|
574
|
+
|
|
575
|
+
export function unloadNativeModels() {
|
|
576
|
+
_embeddingModel = null;
|
|
577
|
+
_embeddingModelLoadPromise = null;
|
|
578
|
+
_liModel = null;
|
|
579
|
+
_liModelLoadPromise = null;
|
|
580
|
+
_embTokenizer = null;
|
|
581
|
+
_embTokenizerLoadPromise = null;
|
|
582
|
+
_liTokenizer = null;
|
|
583
|
+
_liTokenizerLoadPromise = null;
|
|
584
|
+
_addon = null;
|
|
585
|
+
_available = null;
|
|
586
|
+
_coremlCascadeLogged = false;
|
|
587
|
+
}
|