From b14631ca4318d6d36d62e21df6e03e535a81746e Mon Sep 17 00:00:00 2001 From: Harsh16gupta Date: Mon, 27 Jul 2026 02:02:52 +0530 Subject: [PATCH] feat: optimize native Joplin AI embedding pipeline and clustering quality --- src/pipeline/EmbeddingWorkerOrchestrator.ts | 9 ++- src/pipeline/nativeEmbeddingPipeline.ts | 36 +++++++++-- src/pipeline/pipelineConfig.ts | 43 ++++++++++++-- src/pipeline/runPipeline.ts | 42 +++++++++---- src/pipeline/vectorAggregator.ts | 66 +++++++++++++++++++++ src/pipeline/vectorCache.ts | 10 +++- src/types/joplinAi.d.ts | 3 +- test/pipeline/pipelineConfig.test.ts | 61 ++++++++++++++++--- test/pipeline/vectorAggregator.test.ts | 63 ++++++++++++++++++++ 9 files changed, 297 insertions(+), 36 deletions(-) diff --git a/src/pipeline/EmbeddingWorkerOrchestrator.ts b/src/pipeline/EmbeddingWorkerOrchestrator.ts index 8274f4a..ca19722 100644 --- a/src/pipeline/EmbeddingWorkerOrchestrator.ts +++ b/src/pipeline/EmbeddingWorkerOrchestrator.ts @@ -5,7 +5,7 @@ import { VectorCache } from './vectorCache'; import { isGenericTitle } from '../utils/titleFilter'; import { log, logErr } from '../utils/logger'; import { averageVectors, blendVectors, computeTitleWeight, cosineSimilarity } from './vectorAggregator'; -import { isValidEmbeddingVector } from './pipelineConfig'; +import { EMBEDDING_DIM, isValidEmbeddingVector } from './pipelineConfig'; const enc = getEncoding('cl100k_base'); const MAX_TOKENS = 200; @@ -171,6 +171,8 @@ export class EmbeddingWorkerOrchestrator { hash, updatedTime: note.updated_time, titleWeight, + modelId: 'local-onnx', + dimension: EMBEDDING_DIM, }); this.reportProgress(); @@ -204,7 +206,8 @@ export class EmbeddingWorkerOrchestrator { const cachedItem = await this.cache.getItem(note.id); if (cachedItem && cachedItem.metadata.hash === this.currentNoteHash) { - if (isValidEmbeddingVector(cachedItem.vector)) { + const isCorrectModel = !cachedItem.metadata.modelId || cachedItem.metadata.modelId === 'local-onnx'; + if (isCorrectModel && isValidEmbeddingVector(cachedItem.vector, EMBEDDING_DIM)) { log( `[${this.currentNoteIndex + 1}/${this.notes.length}] cache hit for "${note.title.slice(0, 30)}"`, ); @@ -220,7 +223,7 @@ export class EmbeddingWorkerOrchestrator { continue; } else { log( - `[${this.currentNoteIndex + 1}/${this.notes.length}] cache invalid (contains null/NaN) for "${note.title.slice(0, 30)}"`, + `[${this.currentNoteIndex + 1}/${this.notes.length}] cache invalid (wrong dimension/model) for "${note.title.slice(0, 30)}"`, ); } } diff --git a/src/pipeline/nativeEmbeddingPipeline.ts b/src/pipeline/nativeEmbeddingPipeline.ts index 02f49b8..5be6c27 100644 --- a/src/pipeline/nativeEmbeddingPipeline.ts +++ b/src/pipeline/nativeEmbeddingPipeline.ts @@ -13,6 +13,12 @@ export interface NativeEmbeddingChunk { vector: number[]; } +export interface NativeEmbeddingResult { + chunks: NativeEmbeddingChunk[]; + modelId: string; + dimension: number; +} + /** * Checks if Joplin's native AI indexing is active and ready. */ @@ -29,15 +35,19 @@ export const isNativeAiReady = async (): Promise => { }; /** - * Pages through Joplin's native index to fetch raw embedding vectors for the requested notes. + * Pages through Joplin's native index to fetch raw embedding vectors for the requested notes, + * returning the chunks along with modelId and vector dimension metadata. */ -export const fetchNativeEmbeddings = async (noteIds: string[]): Promise => { - if (noteIds.length === 0) return []; +export const fetchNativeEmbeddings = async (noteIds: string[]): Promise => { + if (noteIds.length === 0) { + return { chunks: [], modelId: 'unknown', dimension: 384 }; + } log(`Fetching native embeddings for ${noteIds.length} notes...`); const chunks: NativeEmbeddingChunk[] = []; const BATCH_SIZE = 500; let modelId: string | null = null; + let dimension: number | null = null; for (let i = 0; i < noteIds.length; i += BATCH_SIZE) { const batchIds = noteIds.slice(i, i + BATCH_SIZE); @@ -59,11 +69,22 @@ export const fetchNativeEmbeddings = async (noteIds: string[]): Promise 0) { + if (dimension && dimension !== page.dimension) { + throw new Error('Embedding dimension changed mid-fetch. Please restart.'); + } + dimension = page.dimension; + } + for (const chunk of page.chunks) { if (!chunk.noteId || !Array.isArray(chunk.vector)) { log(`Skipping malformed embedding chunk: ${JSON.stringify(chunk).slice(0, 100)}`); continue; } + if (!dimension && chunk.vector.length > 0) { + dimension = chunk.vector.length; + } chunks.push(chunk); } cursor = page.nextCursor; @@ -77,6 +98,11 @@ export const fetchNativeEmbeddings = async (noteIds: string[]): Promise 0 ? chunks[0].vector.length : 384); + const finalModelId = modelId ?? 'native-ai'; + + log( + `Successfully fetched ${chunks.length} native embedding chunks (model: ${finalModelId}, dim: ${finalDimension})`, + ); + return { chunks, modelId: finalModelId, dimension: finalDimension }; }; diff --git a/src/pipeline/pipelineConfig.ts b/src/pipeline/pipelineConfig.ts index a16d3db..16bf758 100644 --- a/src/pipeline/pipelineConfig.ts +++ b/src/pipeline/pipelineConfig.ts @@ -1,14 +1,48 @@ import { CategorizationConfig } from '../types/cluster'; -/** Dimensionality of embedding vectors (all-MiniLM-L6-v2 / multilingual-e5-small). */ +/** Default dimensionality of local ONNX embedding vectors (all-MiniLM-L6-v2 / multilingual-e5-small). */ export const EMBEDDING_DIM = 384; -export function isValidEmbeddingVector(vector: number[] | undefined | null): boolean { - if (!vector) return false; - if (vector.length !== EMBEDDING_DIM) return false; +export function isValidEmbeddingVector(vector: number[] | undefined | null, expectedDim?: number): boolean { + if (!vector || vector.length === 0) return false; + if (expectedDim !== undefined && vector.length !== expectedDim) return false; return vector.every((v) => Number.isFinite(v)); } +/** + * Computes UMAP intermediate dimensionality scaled logarithmically with input embedding dimension. + * Formula: clamp(⌊2·log₂(D)⌋, 5, 50) + */ +export function adaptiveIntermediateDim(inputDim: number): number { + const raw = Math.floor(2 * Math.log2(inputDim)); + if (!Number.isFinite(raw)) return 5; + return Math.max(5, Math.min(50, raw)); +} + +/** + * Computes UMAP neighbor count scaled with square root of note count. + * Formula: clamp(⌊√N⌋, 5, 50) + */ +export function adaptiveNeighbors(noteCount: number): number { + const raw = Math.floor(Math.sqrt(noteCount)); + if (!Number.isFinite(raw)) return 5; + return Math.max(5, Math.min(50, raw)); +} + +export function createAdaptiveConfig(inputDim: number, noteCount: number): CategorizationConfig { + return { + seed: 42, + metric: 'cosine', + intermediateDim: adaptiveIntermediateDim(inputDim), + intermediateNeighbors: adaptiveNeighbors(noteCount), + strategies: [ + { name: 'kmeans-6', algorithm: 'kmeans', K: 6 }, + { name: 'kmedoids-6', algorithm: 'kmedoids', K: 6 }, + { name: 'hdbscan', algorithm: 'hdbscan', minClusterSize: 3, minSamples: 2 }, + ], + }; +} + export const DEFAULT_CONFIG: CategorizationConfig = { seed: 42, metric: 'cosine', @@ -17,7 +51,6 @@ export const DEFAULT_CONFIG: CategorizationConfig = { strategies: [ { name: 'kmeans-6', algorithm: 'kmeans', K: 6 }, { name: 'kmedoids-6', algorithm: 'kmedoids', K: 6 }, - // { name: 'hdbscan-tuned', algorithm: 'hdbscan', minClusterSize: 4, minSamples: 1 }, { name: 'hdbscan', algorithm: 'hdbscan', minClusterSize: 3, minSamples: 2 }, ], }; diff --git a/src/pipeline/runPipeline.ts b/src/pipeline/runPipeline.ts index 1a1a1f1..45339bf 100644 --- a/src/pipeline/runPipeline.ts +++ b/src/pipeline/runPipeline.ts @@ -1,11 +1,11 @@ import { fetchAllNotes } from './noteReader'; import { benchmark } from './clustering/benchmark'; -import { averageVectors } from './vectorAggregator'; +import { weightedAverageVectorsWithNorm } from './vectorAggregator'; import { PanelNote } from '../types/panel'; import { log, logErr } from '../utils/logger'; import { VectorCache } from './vectorCache'; import { isNativeAiReady, fetchNativeEmbeddings } from './nativeEmbeddingPipeline'; -import { DEFAULT_CONFIG, isValidEmbeddingVector } from './pipelineConfig'; +import { DEFAULT_CONFIG, isValidEmbeddingVector, createAdaptiveConfig } from './pipelineConfig'; import { enrichResultsWithTags } from './clustering/postProcess'; import { upgradeClusterNamesWithAi } from './clustering/aiNamingService'; import { EmbeddingWorkerOrchestrator } from './EmbeddingWorkerOrchestrator'; @@ -17,6 +17,11 @@ export interface PipelineCallbacks { onError: (message: string) => void; } +interface IndexedVector { + chunkIndex: number; + vector: number[]; +} + /** * Runs the full embedding + clustering pipeline, reporting progress via callbacks. * @@ -45,13 +50,14 @@ export const runPipeline = async (installDir: string, callbacks: PipelineCallbac try { const noteIds = notes.map((n) => n.id); - const chunks = await fetchNativeEmbeddings(noteIds); + const nativeResult = await fetchNativeEmbeddings(noteIds); + log(`Native AI embeddings retrieved (model: ${nativeResult.modelId}, dim: ${nativeResult.dimension})`); - // Group chunks by noteId - const noteChunksMap = new Map(); - for (const chunk of chunks) { + // Group chunks by noteId preserving chunkIndex for ordering + const noteChunksMap = new Map(); + for (const chunk of nativeResult.chunks) { const list = noteChunksMap.get(chunk.noteId) || []; - list.push(chunk.vector); + list.push({ chunkIndex: chunk.chunkIndex, vector: chunk.vector }); noteChunksMap.set(chunk.noteId, list); } @@ -59,11 +65,22 @@ export const runPipeline = async (installDir: string, callbacks: PipelineCallbac const vectors: number[][] = []; for (const note of notes) { - const chunkVectors = noteChunksMap.get(note.id); - if (chunkVectors && chunkVectors.length > 0) { - const avgVector = averageVectors(chunkVectors); + const chunkEntries = noteChunksMap.get(note.id); + if (chunkEntries && chunkEntries.length > 0) { + // Sort by chunkIndex ascending to ensure lead paragraph/header (chunk 0) gets highest weight + const sortedEntries = chunkEntries.sort((a, b) => a.chunkIndex - b.chunkIndex); + const chunkVectors = sortedEntries.map((e) => e.vector); + + const { vector: avgVector, rawNorm } = weightedAverageVectorsWithNorm(chunkVectors); + + if (rawNorm < 1e-6) { + logErr( + `Native embedding for note "${note.title}" has near-zero L2 norm (${rawNorm}). Skipping as outlier.`, + ); + continue; + } - if (isValidEmbeddingVector(avgVector)) { + if (isValidEmbeddingVector(avgVector, nativeResult.dimension)) { vectors.push(avgVector); validNotes.push(note); } else { @@ -80,7 +97,8 @@ export const runPipeline = async (installDir: string, callbacks: PipelineCallbac log('Too few indexed notes found in native DB. Falling back to local ONNX Web Worker.'); } else { callbacks.onStatus('Clustering...'); - const results = benchmark(vectors, DEFAULT_CONFIG); + const adaptiveConfig = createAdaptiveConfig(nativeResult.dimension, validNotes.length); + const results = benchmark(vectors, adaptiveConfig); // Post-process to extract tags/keywords for each cluster (keep parity with local pipeline) const allPipelineDocuments = validNotes.map((n) => ({ diff --git a/src/pipeline/vectorAggregator.ts b/src/pipeline/vectorAggregator.ts index bc3c37b..6f45b47 100644 --- a/src/pipeline/vectorAggregator.ts +++ b/src/pipeline/vectorAggregator.ts @@ -28,6 +28,72 @@ export const averageVectors = (vectors: number[][]): number[] => { return normalise(avg); }; +export interface WeightedPoolingOptions { + lambda?: number; + leadBoost?: number; +} + +export interface WeightedAverageResult { + vector: number[]; + rawNorm: number; +} + +/** + * Element-wise weighted average of chunk vectors using exponential position decay + * with a lead-chunk boost, returning both the L2-normalised vector and the raw pre-normalisation norm. + * + * w_i = (leadBoost + exp(-lambda * i)) if i == 0 else exp(-lambda * i) + */ +export const weightedAverageVectorsWithNorm = ( + vectors: number[][], + options: WeightedPoolingOptions = {}, +): WeightedAverageResult => { + if (vectors.length === 0) throw new Error('Cannot average zero vectors'); + const dim = vectors[0].length; + for (const vec of vectors) { + if (vec.length !== dim) throw new Error('Cannot average vectors of different dimensions'); + } + if (vectors.length === 1) { + const rawNorm = Math.sqrt(vectors[0].reduce((sum, v) => sum + v * v, 0)); + const vector = rawNorm === 0 ? vectors[0] : vectors[0].map((v) => v / rawNorm); + return { vector, rawNorm }; + } + + const lambda = options.lambda ?? 0.15; + const leadBoost = options.leadBoost ?? 0.5; + + const weights: number[] = new Array(vectors.length); + let weightSum = 0; + + for (let i = 0; i < vectors.length; i++) { + const w = Math.exp(-lambda * i) + (i === 0 ? leadBoost : 0); + weights[i] = w; + weightSum += w; + } + + const weightedSum = new Array(dim).fill(0); + for (let i = 0; i < vectors.length; i++) { + const normW = weights[i] / weightSum; + const vec = vectors[i]; + for (let d = 0; d < dim; d++) { + weightedSum[d] += normW * vec[d]; + } + } + + const rawNorm = Math.sqrt(weightedSum.reduce((sum, v) => sum + v * v, 0)); + const vector = rawNorm === 0 ? weightedSum : weightedSum.map((v) => v / rawNorm); + + return { vector, rawNorm }; +}; + +/** + * Element-wise weighted average of chunk vectors using exponential position decay + * with a lead-chunk boost, then L2-normalised. + */ +export const weightedAverageVectors = (vectors: number[][], options: WeightedPoolingOptions = {}): number[] => { + return weightedAverageVectorsWithNorm(vectors, options).vector; +}; + /** * Cosine similarity between two L2-normalised vectors (= dot product). */ diff --git a/src/pipeline/vectorCache.ts b/src/pipeline/vectorCache.ts index f84614d..1239a03 100644 --- a/src/pipeline/vectorCache.ts +++ b/src/pipeline/vectorCache.ts @@ -9,7 +9,9 @@ export interface CacheMetadata { hash: string; updatedTime: number; titleWeight: number; - [key: string]: MetadataTypes; + modelId?: string; + dimension?: number; + [key: string]: MetadataTypes | undefined; } export class VectorCache { @@ -56,7 +58,9 @@ export class VectorCache { */ public async getItem(id: string) { try { - return await this.index.getItem(id); + const item = await this.index.getItem>(id); + if (!item) return undefined; + return item as unknown as { id: string; vector: number[]; metadata: CacheMetadata }; } catch (err) { log('Error getting cached item:', err); return undefined; @@ -70,7 +74,7 @@ export class VectorCache { return await this.index.upsertItem({ id, vector, - metadata, + metadata: metadata as unknown as Record, }); } diff --git a/src/types/joplinAi.d.ts b/src/types/joplinAi.d.ts index b099518..e93ece7 100644 --- a/src/types/joplinAi.d.ts +++ b/src/types/joplinAi.d.ts @@ -7,7 +7,7 @@ export interface AiIndexStatus { state: string; ready: boolean; - modelId?: string; + modelId: string | null; } export interface AiEmbeddingChunk { @@ -20,6 +20,7 @@ export interface AiEmbeddingChunk { export interface AiEmbeddingsPage { chunks: AiEmbeddingChunk[]; modelId: string; + dimension: number; nextCursor?: string; } diff --git a/test/pipeline/pipelineConfig.test.ts b/test/pipeline/pipelineConfig.test.ts index 6bd1893..16ac510 100644 --- a/test/pipeline/pipelineConfig.test.ts +++ b/test/pipeline/pipelineConfig.test.ts @@ -1,21 +1,36 @@ -import { EMBEDDING_DIM, isValidEmbeddingVector, DEFAULT_CONFIG } from '../../src/pipeline/pipelineConfig'; +import { + EMBEDDING_DIM, + isValidEmbeddingVector, + adaptiveIntermediateDim, + adaptiveNeighbors, + createAdaptiveConfig, + DEFAULT_CONFIG, +} from '../../src/pipeline/pipelineConfig'; describe('isValidEmbeddingVector', () => { - it('accepts a valid 384-dim vector of finite numbers', () => { + it('accepts a valid vector of finite numbers with default dimension check', () => { const vec = new Array(EMBEDDING_DIM).fill(0.1); expect(isValidEmbeddingVector(vec)).toBe(true); }); - it('rejects null', () => { - expect(isValidEmbeddingVector(null)).toBe(false); + it('accepts custom dimension vectors when expectedDim is specified', () => { + const vec768 = new Array(768).fill(0.05); + const vec1536 = new Array(1536).fill(0.02); + expect(isValidEmbeddingVector(vec768, 768)).toBe(true); + expect(isValidEmbeddingVector(vec1536, 1536)).toBe(true); }); - it('rejects undefined', () => { + it('rejects null and undefined', () => { + expect(isValidEmbeddingVector(null)).toBe(false); expect(isValidEmbeddingVector(undefined)).toBe(false); }); - it('rejects wrong dimension', () => { - expect(isValidEmbeddingVector(new Array(383).fill(0.1))).toBe(false); + it('rejects empty array', () => { + expect(isValidEmbeddingVector([])).toBe(false); + }); + + it('rejects vector matching wrong expected dimension', () => { + expect(isValidEmbeddingVector(new Array(768).fill(0.1), 384)).toBe(false); }); it('rejects vector containing NaN', () => { @@ -31,6 +46,38 @@ describe('isValidEmbeddingVector', () => { }); }); +describe('adaptive scaling functions', () => { + it('computes adaptive intermediate dimensions logarithmic with input dimension', () => { + expect(adaptiveIntermediateDim(384)).toBe(17); + expect(adaptiveIntermediateDim(768)).toBe(19); + expect(adaptiveIntermediateDim(1536)).toBe(21); + }); + + it('clamps intermediate dimensions between 5 and 50', () => { + expect(adaptiveIntermediateDim(2)).toBe(5); + expect(adaptiveIntermediateDim(1e12)).toBe(50); + }); + + it('computes adaptive neighbors based on square root of note count', () => { + expect(adaptiveNeighbors(25)).toBe(5); + expect(adaptiveNeighbors(100)).toBe(10); + expect(adaptiveNeighbors(500)).toBe(22); + }); + + it('clamps neighbors between 5 and 50', () => { + expect(adaptiveNeighbors(2)).toBe(5); + expect(adaptiveNeighbors(10000)).toBe(50); + }); + + it('creates adaptive configuration dynamically', () => { + const config = createAdaptiveConfig(768, 100); + expect(config.metric).toBe('cosine'); + expect(config.intermediateDim).toBe(19); + expect(config.intermediateNeighbors).toBe(10); + expect(config.strategies.length).toBe(3); + }); +}); + describe('DEFAULT_CONFIG', () => { it('uses cosine metric and seed 42', () => { expect(DEFAULT_CONFIG.metric).toBe('cosine'); diff --git a/test/pipeline/vectorAggregator.test.ts b/test/pipeline/vectorAggregator.test.ts index 43ea083..b9dc06f 100644 --- a/test/pipeline/vectorAggregator.test.ts +++ b/test/pipeline/vectorAggregator.test.ts @@ -1,5 +1,7 @@ import { averageVectors, + weightedAverageVectors, + weightedAverageVectorsWithNorm, cosineSimilarity, computeTitleWeight, blendVectors, @@ -53,6 +55,67 @@ describe('averageVectors', () => { }); }); +describe('weightedAverageVectors', () => { + it('normalizes a single vector', () => { + const result = weightedAverageVectors([[3, 4]]); + expect(result[0]).toBeCloseTo(0.6, 10); + expect(result[1]).toBeCloseTo(0.8, 10); + }); + + it('gives lead chunk (index 0) higher weight than tail chunks', () => { + const chunk0 = [1, 0]; + const chunk1 = [0, 1]; + const result = weightedAverageVectors([chunk0, chunk1]); + + // Chunk 0 has weight = exp(0) + 0.5 = 1.5 + // Chunk 1 has weight = exp(-0.15) ≈ 0.8607 + // Therefore result[0] > result[1] + expect(result[0]).toBeGreaterThan(result[1]); + }); + + it('supports custom lambda and leadBoost parameters', () => { + const chunk0 = [1, 0]; + const chunk1 = [0, 1]; + const result = weightedAverageVectors([chunk0, chunk1], { lambda: 0.5, leadBoost: 1.0 }); + + // Chunk 0 weight = 2.0, Chunk 1 weight = exp(-0.5) ≈ 0.6065 + expect(result[0]).toBeGreaterThan(result[1]); + }); + + it('throws on empty input', () => { + expect(() => weightedAverageVectors([])).toThrow('Cannot average zero vectors'); + }); + + it('throws on dimension mismatch', () => { + expect(() => + weightedAverageVectors([ + [1, 2], + [1, 2, 3], + ]), + ).toThrow('different dimensions'); + }); +}); + +describe('weightedAverageVectorsWithNorm', () => { + it('returns both normalized vector and pre-normalization raw norm', () => { + const chunk0 = [0.6, 0.8]; + const res = weightedAverageVectorsWithNorm([chunk0]); + expect(res.vector[0]).toBeCloseTo(0.6, 10); + expect(res.vector[1]).toBeCloseTo(0.8, 10); + expect(res.rawNorm).toBeCloseTo(1.0, 10); + }); + + it('detects near-zero raw norm when vectors cancel out', () => { + // Opposing unit vectors with equal weights + const chunk0 = [1, 0]; + const chunk1 = [-1, 0]; + // Weight chunk0 = 1.5, weight chunk1 = exp(-0.15) ≈ 0.8607 + // If lambda=0 and leadBoost=0, weights are equal + const res = weightedAverageVectorsWithNorm([chunk0, chunk1], { lambda: 0, leadBoost: 0 }); + expect(res.rawNorm).toBeCloseTo(0, 10); + }); +}); + describe('cosineSimilarity', () => { it('returns 1 for identical unit vectors', () => { expect(cosineSimilarity([0.6, 0.8], [0.6, 0.8])).toBeCloseTo(1.0, 10);