Spaces:
Running
Running
| /** 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 }; | |
| } | |