diff --git a/scripts/embedding-benchmark.js b/scripts/embedding-benchmark.js index 73fe2d8e..3d7ef173 100644 --- a/scripts/embedding-benchmark.js +++ b/scripts/embedding-benchmark.js @@ -24,7 +24,7 @@ const root = path.resolve(__dirname, '..'); const pkg = JSON.parse(fs.readFileSync(path.join(root, 'package.json'), 'utf8')); const dbPath = path.join(root, '.codegraph', 'graph.db'); -const { buildEmbeddings, MODELS, searchData } = await import( +const { buildEmbeddings, MODELS, searchData, disposeModel } = await import( new URL('../src/embedder.js', import.meta.url).href ); @@ -129,6 +129,7 @@ for (const key of modelKeys) { } catch (err) { console.error(` FAILED: ${err.message}`); } + await disposeModel(); } // Restore console.log for JSON output diff --git a/src/embedder.js b/src/embedder.js index b0e0588c..938a5976 100644 --- a/src/embedder.js +++ b/src/embedder.js @@ -238,11 +238,26 @@ async function loadTransformers() { } } +/** + * Dispose the current ONNX session and free memory. + * Safe to call when no model is loaded (no-op). + */ +export async function disposeModel() { + if (extractor) { + await extractor.dispose(); + extractor = null; + } + activeModel = null; +} + async function loadModel(modelKey) { const config = getModelConfig(modelKey); if (extractor && activeModel === config.name) return { extractor, config }; + // Dispose previous model before loading a different one + await disposeModel(); + const transformers = await loadTransformers(); pipeline = transformers.pipeline; _cos_sim = transformers.cos_sim; diff --git a/src/index.js b/src/index.js index 9da90b5e..edf5d264 100644 --- a/src/index.js +++ b/src/index.js @@ -21,6 +21,7 @@ export { buildEmbeddings, cosineSim, DEFAULT_MODEL, + disposeModel, EMBEDDING_STRATEGIES, embed, estimateTokens,