hdoc-tools 0.61.0 → 0.62.1

package/hdoc-build-embeddings.js

@@ -1,202 +1,202 @@
-(() => {
-	const path = require("node:path");
-	const fs = require("node:fs");
-	const crypto = require("node:crypto");
-
-	// Model identity - must match what the docs service query embedder uses
-	// (see embedder.h in esp-docs-service). The service refuses embeddings
-	// built with a different model or dimension count.
-	const MODEL_ID = "all-MiniLM-L6-v2";
-	const HF_MODEL = "Xenova/all-MiniLM-L6-v2";
-	const DIMS = 384;
-
-	// SHA-256 of onnx/model.onnx as downloaded from Hugging Face. Matches the
-	// vendored copy the docs service runs (search_model/model.onnx). Verified
-	// after download, before any embeddings are generated - a hash mismatch
-	// (tampered/corrupt/re-published model) aborts the embedding step.
-	const MODEL_SHA256 = "759c3cd2b7fe7e93933ad23c4c9181b7396442a2ed746ec7c1d46192c469c46e";
-
-	// Chunking parameters (word based). MiniLM attends to ~256 wordpiece
-	// tokens; 180 words plus the title prefix stays comfortably inside that.
-	const CHUNK_WORDS = 180;
-	const CHUNK_OVERLAP_WORDS = 30;
-	const MAX_CHUNKS_PER_DOC = 40;
-	const EMBED_BATCH_SIZE = 16;
-
-	// Maximum characters of chunk text stored for use as a result snippet
-	const STORED_CHUNK_CHARS = 500;
-
-	// Split document plain text into overlapping word-window chunks
-	const chunk_text = (text) => {
-		const words = text.split(/\s+/).filter((w) => w.length > 0);
-		if (words.length === 0) return [];
-
-		const chunks = [];
-		let start = 0;
-		while (start < words.length && chunks.length < MAX_CHUNKS_PER_DOC) {
-			const end = Math.min(start + CHUNK_WORDS, words.length);
-			chunks.push(words.slice(start, end).join(" "));
-			if (end >= words.length) break;
-			start = end - CHUNK_OVERLAP_WORDS;
-		}
-		return chunks;
-	};
-
-	// Compute the resource_url for an index record. Must stay in step with
-	// the equivalent logic in populate_index (hdoc-build-db.js), as the docs
-	// service joins hdoc_embeddings to hdoc_index on this value.
-	const resource_url_for = (file) => {
-		let index_path_name = file.relative_path.replaceAll("\\", "/");
-		if (
-			index_path_name.endsWith("/index.md") ||
-			index_path_name.endsWith("/index.html") ||
-			index_path_name.endsWith("/index.htm")
-		) {
-			index_path_name = index_path_name.substring(0, index_path_name.lastIndexOf("/"));
-		}
-		index_path_name = `/${index_path_name.replace(path.extname(file.relative_path), "")}`;
-		return file.index_html.id !== null
-			? `${index_path_name}#${file.index_html.id}`
-			: index_path_name;
-	};
-
-	// Generate semantic search embeddings for every indexed document and
-	// write them to the hdoc_embeddings table. Non-fatal on failure - a book
-	// without embeddings remains fully searchable via keyword search.
-	exports.populate_embeddings = async (db, doc_id, index_records, verbose = false) => {
-		const response = {
-			success: false,
-			error: "",
-			doc_count: 0,
-			chunk_count: 0,
-		};
-
-		let transformers = null;
-		try {
-			// @huggingface/transformers is ESM-only, load it dynamically
-			transformers = await import("@huggingface/transformers");
-		} catch (e) {
-			response.error = `Could not load @huggingface/transformers: ${e.message}`;
-			return response;
-		}
-
-		// Keep the model cache inside the package, like the puppeteer browser
-		transformers.env.cacheDir = path.join(__dirname, ".model-cache");
-
-		console.log(`\nGenerating semantic search embeddings (${MODEL_ID})...`);
-
-		let extractor = null;
-		try {
-			// fp32 to match the model the docs service runs at query time
-			extractor = await transformers.pipeline("feature-extraction", HF_MODEL, {
-				dtype: "fp32",
-			});
-		} catch (e) {
-			response.error = `Could not load embedding model ${HF_MODEL}: ${e.message}`;
-			return response;
-		}
-
-		// Verify the downloaded/cached model file against the pinned hash
-		// before generating any embeddings
-		const model_file = path.join(transformers.env.cacheDir, HF_MODEL, "onnx", "model.onnx");
-		try {
-			const hash = crypto
-				.createHash("sha256")
-				.update(fs.readFileSync(model_file))
-				.digest("hex");
-			if (hash !== MODEL_SHA256) {
-				response.error = `Embedding model checksum mismatch for ${model_file}: got ${hash}, expected ${MODEL_SHA256}. Delete the .model-cache folder to re-download, or update MODEL_SHA256 if the pinned model has intentionally changed.`;
-				return response;
-			}
-		} catch (e) {
-			response.error = `Could not verify embedding model checksum: ${e.message}`;
-			return response;
-		}
-
-		// Build the chunk list across all documents
-		const chunks = [];
-		for (const file of index_records) {
-			if (file.inline) continue;
-			if (file.status === "draft") continue;
-			if (!file.index_html || !file.index_html.text) continue;
-
-			const url = resource_url_for(file);
-			const title = (file.index_html.fm_props && file.index_html.fm_props.title) || "";
-
-			const doc_chunks = chunk_text(file.index_html.text);
-			if (doc_chunks.length === 0) continue;
-
-			response.doc_count++;
-			for (let seq = 0; seq < doc_chunks.length; seq++) {
-				chunks.push({
-					url: url,
-					seq: seq,
-					text: doc_chunks[seq],
-					// Title provides context for the embedding but is not
-					// duplicated into the stored snippet text
-					embed_input: title ? `${title}\n${doc_chunks[seq]}` : doc_chunks[seq],
-				});
-			}
-		}
-
-		if (chunks.length === 0) {
-			response.success = true;
-			console.log("No documents to embed.");
-			return response;
-		}
-
-		const stmt = db.prepare(
-			"INSERT INTO hdoc_embeddings (resource_url, chunk_seq, chunk_text, embedding, model, dims) VALUES (?, ?, ?, ?, ?, ?)",
-		);
-		const insert_batch = db.transaction((rows) => {
-			for (const row of rows) {
-				stmt.run(row.url, row.seq, row.text.substring(0, STORED_CHUNK_CHARS), row.embedding, MODEL_ID, DIMS);
-			}
-		});
-
-		const start_time = Date.now();
-
-		try {
-			for (let i = 0; i < chunks.length; i += EMBED_BATCH_SIZE) {
-				const batch = chunks.slice(i, i + EMBED_BATCH_SIZE);
-
-				const output = await extractor(
-					batch.map((c) => c.embed_input),
-					{ pooling: "mean", normalize: true },
-				);
-
-				// output is a [batch, DIMS] tensor over one flat Float32Array
-				const data = output.data;
-				const dims = output.dims[output.dims.length - 1];
-				if (dims !== DIMS) {
-					response.error = `Embedding model returned ${dims} dimensions, expected ${DIMS}`;
-					return response;
-				}
-
-				for (let j = 0; j < batch.length; j++) {
-					// Copy the row out of the shared tensor buffer
-					const vec = new Float32Array(data.slice(j * DIMS, (j + 1) * DIMS));
-					batch[j].embedding = Buffer.from(vec.buffer, 0, DIMS * 4);
-				}
-
-				insert_batch(batch);
-				response.chunk_count += batch.length;
-
-				if (verbose || i % (EMBED_BATCH_SIZE * 10) === 0) {
-					process.stdout.write(`\rEmbedding chunks: ${Math.min(i + EMBED_BATCH_SIZE, chunks.length)}/${chunks.length}`);
-				}
-			}
-		} catch (e) {
-			response.error = `Embedding generation failed: ${e.message}`;
-			return response;
-		}
-
-		const secs = Math.round((Date.now() - start_time) / 1000);
-		console.log(
-			`\rEmbedding Build Complete: ${response.chunk_count} chunks across ${response.doc_count} documents in ${secs}s.`,
-		);
-
-		response.success = true;
-		return response;
-	};
-})();
+(() => {
+	const path = require("node:path");
+	const fs = require("node:fs");
+	const crypto = require("node:crypto");
+
+	// Model identity - must match what the docs service query embedder uses
+	// (see embedder.h in esp-docs-service). The service refuses embeddings
+	// built with a different model or dimension count.
+	const MODEL_ID = "all-MiniLM-L6-v2";
+	const HF_MODEL = "Xenova/all-MiniLM-L6-v2";
+	const DIMS = 384;
+
+	// SHA-256 of onnx/model.onnx as downloaded from Hugging Face. Matches the
+	// vendored copy the docs service runs (search_model/model.onnx). Verified
+	// after download, before any embeddings are generated - a hash mismatch
+	// (tampered/corrupt/re-published model) aborts the embedding step.
+	const MODEL_SHA256 = "759c3cd2b7fe7e93933ad23c4c9181b7396442a2ed746ec7c1d46192c469c46e";
+
+	// Chunking parameters (word based). MiniLM attends to ~256 wordpiece
+	// tokens; 180 words plus the title prefix stays comfortably inside that.
+	const CHUNK_WORDS = 180;
+	const CHUNK_OVERLAP_WORDS = 30;
+	const MAX_CHUNKS_PER_DOC = 40;
+	const EMBED_BATCH_SIZE = 16;
+
+	// Maximum characters of chunk text stored for use as a result snippet
+	const STORED_CHUNK_CHARS = 500;
+
+	// Split document plain text into overlapping word-window chunks
+	const chunk_text = (text) => {
+		const words = text.split(/\s+/).filter((w) => w.length > 0);
+		if (words.length === 0) return [];
+
+		const chunks = [];
+		let start = 0;
+		while (start < words.length && chunks.length < MAX_CHUNKS_PER_DOC) {
+			const end = Math.min(start + CHUNK_WORDS, words.length);
+			chunks.push(words.slice(start, end).join(" "));
+			if (end >= words.length) break;
+			start = end - CHUNK_OVERLAP_WORDS;
+		}
+		return chunks;
+	};
+
+	// Compute the resource_url for an index record. Must stay in step with
+	// the equivalent logic in populate_index (hdoc-build-db.js), as the docs
+	// service joins hdoc_embeddings to hdoc_index on this value.
+	const resource_url_for = (file) => {
+		let index_path_name = file.relative_path.replaceAll("\\", "/");
+		if (
+			index_path_name.endsWith("/index.md") ||
+			index_path_name.endsWith("/index.html") ||
+			index_path_name.endsWith("/index.htm")
+		) {
+			index_path_name = index_path_name.substring(0, index_path_name.lastIndexOf("/"));
+		}
+		index_path_name = `/${index_path_name.replace(path.extname(file.relative_path), "")}`;
+		return file.index_html.id !== null
+			? `${index_path_name}#${file.index_html.id}`
+			: index_path_name;
+	};
+
+	// Generate semantic search embeddings for every indexed document and
+	// write them to the hdoc_embeddings table. Non-fatal on failure - a book
+	// without embeddings remains fully searchable via keyword search.
+	exports.populate_embeddings = async (db, doc_id, index_records, verbose = false) => {
+		const response = {
+			success: false,
+			error: "",
+			doc_count: 0,
+			chunk_count: 0,
+		};
+
+		let transformers = null;
+		try {
+			// @huggingface/transformers is ESM-only, load it dynamically
+			transformers = await import("@huggingface/transformers");
+		} catch (e) {
+			response.error = `Could not load @huggingface/transformers: ${e.message}`;
+			return response;
+		}
+
+		// Keep the model cache inside the package, like the puppeteer browser
+		transformers.env.cacheDir = path.join(__dirname, ".model-cache");
+
+		console.log(`\nGenerating semantic search embeddings (${MODEL_ID})...`);
+
+		let extractor = null;
+		try {
+			// fp32 to match the model the docs service runs at query time
+			extractor = await transformers.pipeline("feature-extraction", HF_MODEL, {
+				dtype: "fp32",
+			});
+		} catch (e) {
+			response.error = `Could not load embedding model ${HF_MODEL}: ${e.message}`;
+			return response;
+		}
+
+		// Verify the downloaded/cached model file against the pinned hash
+		// before generating any embeddings
+		const model_file = path.join(transformers.env.cacheDir, HF_MODEL, "onnx", "model.onnx");
+		try {
+			const hash = crypto
+				.createHash("sha256")
+				.update(fs.readFileSync(model_file))
+				.digest("hex");
+			if (hash !== MODEL_SHA256) {
+				response.error = `Embedding model checksum mismatch for ${model_file}: got ${hash}, expected ${MODEL_SHA256}. Delete the .model-cache folder to re-download, or update MODEL_SHA256 if the pinned model has intentionally changed.`;
+				return response;
+			}
+		} catch (e) {
+			response.error = `Could not verify embedding model checksum: ${e.message}`;
+			return response;
+		}
+
+		// Build the chunk list across all documents
+		const chunks = [];
+		for (const file of index_records) {
+			if (file.inline) continue;
+			if (file.status === "draft") continue;
+			if (!file.index_html || !file.index_html.text) continue;
+
+			const url = resource_url_for(file);
+			const title = (file.index_html.fm_props && file.index_html.fm_props.title) || "";
+
+			const doc_chunks = chunk_text(file.index_html.text);
+			if (doc_chunks.length === 0) continue;
+
+			response.doc_count++;
+			for (let seq = 0; seq < doc_chunks.length; seq++) {
+				chunks.push({
+					url: url,
+					seq: seq,
+					text: doc_chunks[seq],
+					// Title provides context for the embedding but is not
+					// duplicated into the stored snippet text
+					embed_input: title ? `${title}\n${doc_chunks[seq]}` : doc_chunks[seq],
+				});
+			}
+		}
+
+		if (chunks.length === 0) {
+			response.success = true;
+			console.log("No documents to embed.");
+			return response;
+		}
+
+		const stmt = db.prepare(
+			"INSERT INTO hdoc_embeddings (resource_url, chunk_seq, chunk_text, embedding, model, dims) VALUES (?, ?, ?, ?, ?, ?)",
+		);
+		const insert_batch = db.transaction((rows) => {
+			for (const row of rows) {
+				stmt.run(row.url, row.seq, row.text.substring(0, STORED_CHUNK_CHARS), row.embedding, MODEL_ID, DIMS);
+			}
+		});
+
+		const start_time = Date.now();
+
+		try {
+			for (let i = 0; i < chunks.length; i += EMBED_BATCH_SIZE) {
+				const batch = chunks.slice(i, i + EMBED_BATCH_SIZE);
+
+				const output = await extractor(
+					batch.map((c) => c.embed_input),
+					{ pooling: "mean", normalize: true },
+				);
+
+				// output is a [batch, DIMS] tensor over one flat Float32Array
+				const data = output.data;
+				const dims = output.dims[output.dims.length - 1];
+				if (dims !== DIMS) {
+					response.error = `Embedding model returned ${dims} dimensions, expected ${DIMS}`;
+					return response;
+				}
+
+				for (let j = 0; j < batch.length; j++) {
+					// Copy the row out of the shared tensor buffer
+					const vec = new Float32Array(data.slice(j * DIMS, (j + 1) * DIMS));
+					batch[j].embedding = Buffer.from(vec.buffer, 0, DIMS * 4);
+				}
+
+				insert_batch(batch);
+				response.chunk_count += batch.length;
+
+				if (verbose || i % (EMBED_BATCH_SIZE * 10) === 0) {
+					process.stdout.write(`\rEmbedding chunks: ${Math.min(i + EMBED_BATCH_SIZE, chunks.length)}/${chunks.length}`);
+				}
+			}
+		} catch (e) {
+			response.error = `Embedding generation failed: ${e.message}`;
+			return response;
+		}
+
+		const secs = Math.round((Date.now() - start_time) / 1000);
+		console.log(
+			`\rEmbedding Build Complete: ${response.chunk_count} chunks across ${response.doc_count} documents in ${secs}s.`,
+		);
+
+		response.success = true;
+		return response;
+	};
+})();