From b2963dbbcb4331213d65c4916d2ddb993ba716c9 Mon Sep 17 00:00:00 2001 From: Quinn Slack Date: Tue, 26 Dec 2023 19:11:33 -0700 Subject: [PATCH] wip --- provider/docs/src/corpus/search/embeddings.ts | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/provider/docs/src/corpus/search/embeddings.ts b/provider/docs/src/corpus/search/embeddings.ts index 4abcea9d..d8f013eb 100644 --- a/provider/docs/src/corpus/search/embeddings.ts +++ b/provider/docs/src/corpus/search/embeddings.ts @@ -26,14 +26,15 @@ export async function embeddingsSearch( ): Promise { const queryVec = await embedText(query) - const results: CorpusSearchResult[] = [] - for (const { doc, chunks } of index.docs) { - for (const [i, chunk] of chunks.entries()) { - const chunkVec = await cachedEmbedText(chunk.text, cache) - const score = cos_sim(queryVec, chunkVec) - results.push({ doc: doc.id, chunk: i, score, excerpt: chunk.text }) - } - } + const results: CorpusSearchResult[] = await Promise.all( + index.docs.flatMap(({ doc, chunks }) => + chunks.map(async (chunk, i) => { + const chunkVec = await cachedEmbedText(chunk.text, cache) + const score = cos_sim(queryVec, chunkVec) + return { doc: doc.id, chunk: i, score, excerpt: chunk.text } satisfies CorpusSearchResult + }) + ) + ) results.sort((a, b) => b.score - a.score)