feat: optimize native Joplin AI embedding pipeline and clustering quality

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)}"`,
 					);
 				}
 			}

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<boolean> => {
 };
 
 /**
- * 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<NativeEmbeddingChunk[]> => {
-	if (noteIds.length === 0) return [];
+export const fetchNativeEmbeddings = async (noteIds: string[]): Promise<NativeEmbeddingResult> => {
+	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<NativeEm
 				throw new Error('Embedding model changed mid-fetch. Please restart.');
 			}
 			modelId = page.modelId;
+
+			if (typeof page.dimension === 'number' && page.dimension > 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<NativeEm
 		} while (cursor);
 	}
 
-	log(`Successfully fetched ${chunks.length} embedding chunks`);
-	return chunks;
+	const finalDimension = dimension ?? (chunks.length > 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 };
 };

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 },
 	],
 };

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,25 +50,37 @@ 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<string, number[][]>();
-				for (const chunk of chunks) {
+				// Group chunks by noteId preserving chunkIndex for ordering
+				const noteChunksMap = new Map<string, IndexedVector[]>();
+				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);
 				}
 
 				const validNotes: typeof notes = [];
 				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) => ({

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<number>(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).
  */

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<CacheMetadata>(id);
+			const item = await this.index.getItem<Record<string, MetadataTypes>>(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<string, MetadataTypes>,
 		});
 	}
 

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;
 }