Blame for scripts/load-node.ts

/** Node-only helper that loads shipped binaries into a SearchEngine for tests and benchmarks. */
import { gunzipSync } from 'node:zlib';
import { readFileSync } from 'node:fs';
import { SearchEngine } from '../src/lib/search/engine.ts';

export interface LoadedDataset {
	engine: SearchEngine;
}

function gz(dir: string, name: string): Uint8Array {
	return gunzipSync(readFileSync(`${dir}/${name}`));
}

function int8(dir: string, name: string): Int8Array {
	const bytes = readFileSync(`${dir}/${name}`);
	return new Int8Array(bytes.buffer, bytes.byteOffset, bytes.length);
}

function lines(bytes: Uint8Array): string[] {
	return new TextDecoder().decode(bytes).split('\n');
}

export function loadDataset(dir: string): LoadedDataset {
	const manifest = JSON.parse(readFileSync(`${dir}/manifest.json`, 'utf8')) as { dims: number };
	const engine = new SearchEngine({
		dims: manifest.dims,
		subjects: lines(gz(dir, 'subjects.txt.z')),
		termsA: lines(gz(dir, 'terms-a.txt.z')),
		vecA: int8(dir, 'vec-a.bin'),
		freqA: gz(dir, 'freq-a.bin.z'),
		vecS: int8(dir, 'vec-s.bin')
	});
	engine.setGeneralVocabulary(lines(gz(dir, 'terms-b.txt.z')), int8(dir, 'vec-b.bin'));
	const starts = gz(dir, 'catalog-subjects.bin.z');
	engine.setCatalog(new Uint32Array(starts.slice().buffer));
	engine.setCaptions(gz(dir, 'catalog-captions.bin.z'));
	return { engine };
}