chatoai-backend / src /embed.js
Keyurjotaniya007's picture
backend at repo root for docker space
0a36812
Raw
History Blame Contribute Delete
1.68 kB
// ----------------------------------------------------------------------------
// Sentence embeddings with all-MiniLM-L6-v2 running locally in Node via
// Transformers.js (@xenova/transformers). No API key, no external calls — the
// model (~90MB) is downloaded once on first use and cached on disk.
//
// Output: 384-dim, mean-pooled, L2-normalised vectors (cosine == dot product).
// ----------------------------------------------------------------------------
import { pipeline, env } from '@xenova/transformers';
// Cache models under ./.models so repeated deploys don't re-download.
env.cacheDir = process.env.MODEL_CACHE_DIR || './.models';
// We only need inference; disable local-file-only so it can fetch once.
env.allowRemoteModels = true;
let extractorPromise = null;
function getExtractor() {
if (!extractorPromise) {
extractorPromise = pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2');
}
return extractorPromise;
}
/** Warm the model at boot so the first request isn't slow. */
export async function warmup() {
const e = await getExtractor();
await e('warmup', { pooling: 'mean', normalize: true });
}
/** Embed a single string -> number[] (384). */
export async function embed(text) {
const e = await getExtractor();
const out = await e(String(text ?? ''), { pooling: 'mean', normalize: true });
return Array.from(out.data);
}
/** Embed many strings -> number[][]. */
export async function embedBatch(texts) {
const e = await getExtractor();
const vecs = [];
for (const t of texts) {
const out = await e(String(t ?? ''), { pooling: 'mean', normalize: true });
vecs.push(Array.from(out.data));
}
return vecs;
}