diff --git a/src/__tests__/embeddings.test.ts b/src/__tests__/embeddings.test.ts index 67b8005..871c9d1 100644 --- a/src/__tests__/embeddings.test.ts +++ b/src/__tests__/embeddings.test.ts @@ -152,3 +152,26 @@ describe("embeddingConfigHash", () => { ); }); }); + +describe("prepareEmbeddingText task prefixes", () => { + const item = { title: "Fix login bug", body: "resolves the redirect loop", type: "pr" }; + + it("prefixes the clustering task for a model trained to expect one", () => { + const result = prepareEmbeddingText(item, "embeddinggemma"); + + expect(result.startsWith("task: clustering | query: ")).toBe(true); + expect(result).toContain("Fix login bug"); + }); + + it("leaves output byte-identical for a model with no task prompt", () => { + expect(prepareEmbeddingText(item, "nomic-embed-text-v2-moe")).toBe(prepareEmbeddingText(item)); + }); + + it("adds no prefix for an unrecognised model, so a new model cannot be silently mangled", () => { + expect(prepareEmbeddingText(item, "some-future-model")).toBe(prepareEmbeddingText(item)); + }); + + it("matches a tagged model variant, since ollama slugs carry a :tag", () => { + expect(prepareEmbeddingText(item, "embeddinggemma:300m").startsWith("task: clustering | query: ")).toBe(true); + }); +}); diff --git a/src/benchmark.ts b/src/benchmark.ts index 78950db..bc54d38 100644 --- a/src/benchmark.ts +++ b/src/benchmark.ts @@ -261,7 +261,7 @@ export async function runBenchmarkForModel( for (let i = 0; i < allItems.length; i += batchSize) { const batch = allItems.slice(i, i + batchSize); - const texts = batch.map((item) => prepareEmbeddingText(item)); + const texts = batch.map((item) => prepareEmbeddingText(item, model)); let embeddings: number[][]; try { diff --git a/src/embeddings.ts b/src/embeddings.ts index 2716ab5..60637e0 100644 --- a/src/embeddings.ts +++ b/src/embeddings.ts @@ -396,7 +396,35 @@ export async function createEmbeddingProvider(config: ProviderConfig): Promise> = Object.freeze({ + // https://ai.google.dev/gemma/docs/embeddinggemma — prompt set includes + // "task: clustering | query: " for grouping semantically similar text. + embeddinggemma: "task: clustering | query: ", +}); + +/** Ollama slugs carry a `:tag` (embeddinggemma:300m); the prompt is a property + * of the model family, not the quantisation. */ +export function taskPrefixFor(model?: string): string { + if (!model) return ""; + return TASK_PREFIXES[model.split(":")[0].trim().toLowerCase()] ?? ""; +} export type EmbeddingVectorGeneration = "native" | "provider-selected-v1" | "local-truncation-v1"; @@ -441,11 +469,13 @@ export function effectiveEmbeddingConfigHash( ); } -export function prepareEmbeddingText(item: { title: string; body: string; type: string }): string { +export function prepareEmbeddingText(item: { title: string; body: string; type: string }, model?: string): string { // No type prefix. A leading "Pull Request:" / "Issue:" token systematically // pushes an issue away from its own fix PR in embedding space, which fragments // a single bug across separate clusters. Only takes effect after a re-embed. const title = (item.title || "Untitled").trim(); const body = (item.body || "").trim().slice(0, 2000); - return body ? `${title}\n\n${body}` : title; + const text = body ? `${title}\n\n${body}` : title; + // Model-specific task instruction, empty for models that do not use one. + return `${taskPrefixFor(model)}${text}`; } diff --git a/src/pipeline.ts b/src/pipeline.ts index 206f0cb..f008128 100644 --- a/src/pipeline.ts +++ b/src/pipeline.ts @@ -84,11 +84,14 @@ export async function reEmbedStoredItems( for (let i = 0; i < items.length; i += batchSize) { const batch = items.slice(i, i + batchSize); const texts = batch.map((item) => - prepareEmbeddingText({ - title: item.title, - body: item.bodySnippet, - type: item.type, - }), + prepareEmbeddingText( + { + title: item.title, + body: item.bodySnippet, + type: item.type, + }, + providerConfig.model, + ), ); const embeddings = await embedder.embedBatch(texts); for (let j = 0; j < batch.length; j++) { @@ -263,7 +266,7 @@ export async function runScan( for (let i = 0; i < newItems.length; i += BATCH_SIZE) { const batch = newItems.slice(i, i + BATCH_SIZE); - const texts = batch.map((item) => prepareEmbeddingText(item)); + const texts = batch.map((item) => prepareEmbeddingText(item, env.EMBEDDING_MODEL)); const embedWithRetry = async (input: string[]): Promise => { for (let attempt = 0; attempt < 3; attempt++) { try {