| import { readFileSync, writeFileSync } from "node:fs"; |
| import path from "node:path"; |
| import { EMBEDDING_DIM, embedTexts, MODEL_ID } from "../lib/embedder"; |
|
|
| type CatalogImage = { |
| id: string; |
| image: string; |
| captions: string[]; |
| }; |
|
|
| type IndexEntry = { |
| id: string; |
| image: string; |
| caption: string; |
| }; |
|
|
| const BATCH_SIZE = 8; |
|
|
| async function main() { |
| const root = process.cwd(); |
| const catalog = JSON.parse( |
| readFileSync(path.join(root, "data/catalog.json"), "utf8"), |
| ) as CatalogImage[]; |
|
|
| const entries: IndexEntry[] = catalog.flatMap((item) => |
| item.captions.map((caption) => ({ |
| id: item.id, |
| image: item.image, |
| caption, |
| })), |
| ); |
|
|
| const embeddings = new Float32Array(entries.length * EMBEDDING_DIM); |
| console.log( |
| `Embedding ${entries.length} captions from ${catalog.length} images with ${MODEL_ID}`, |
| ); |
|
|
| for (let start = 0; start < entries.length; start += BATCH_SIZE) { |
| const batch = entries |
| .slice(start, start + BATCH_SIZE) |
| .map((entry) => entry.caption); |
| const vectors = await embedTexts(batch); |
| if (vectors.some((vector) => vector.length !== EMBEDDING_DIM)) { |
| throw new Error( |
| `Unexpected embedding length ${vectors[0]?.length ?? 0}; expected ${EMBEDDING_DIM}`, |
| ); |
| } |
| vectors.forEach((vector, index) => { |
| embeddings.set(vector, (start + index) * EMBEDDING_DIM); |
| }); |
| console.log( |
| `embedded ${Math.min(start + BATCH_SIZE, entries.length)}/${entries.length}`, |
| ); |
| } |
|
|
| const dataDir = path.join(root, "data"); |
| writeFileSync( |
| path.join(dataDir, "index.json"), |
| `${JSON.stringify(entries)}\n`, |
| ); |
| writeFileSync( |
| path.join(dataDir, "embeddings.f32"), |
| Buffer.from( |
| embeddings.buffer, |
| embeddings.byteOffset, |
| embeddings.byteLength, |
| ), |
| ); |
| console.log(`wrote ${entries.length} embeddings`); |
| } |
|
|
| main().catch((error) => { |
| console.error(error); |
| process.exit(1); |
| }); |
|
|
