Commit e4756c76 by YeYuheng Committed by GitHub

rrf_weight (#5551)

Co-authored-by: xxYyh <xxyyh@xxYyhdeMacBook-Pro.local>
parent 486d791b
...@@ -37,4 +37,6 @@ files/helm/fastgpt/charts/*.tgz ...@@ -37,4 +37,6 @@ files/helm/fastgpt/charts/*.tgz
tmp/ tmp/
coverage coverage
document/.source document/.source
\ No newline at end of file
bun.lock
\ No newline at end of file
...@@ -3,7 +3,7 @@ import { type SearchDataResponseItemType } from '../type'; ...@@ -3,7 +3,7 @@ import { type SearchDataResponseItemType } from '../type';
/* dataset search result concat */ /* dataset search result concat */
export const datasetSearchResultConcat = ( export const datasetSearchResultConcat = (
arr: { k: number; list: SearchDataResponseItemType[] }[] arr: { weight: number; list: SearchDataResponseItemType[] }[]
): SearchDataResponseItemType[] => { ): SearchDataResponseItemType[] => {
arr = arr.filter((item) => item.list.length > 0); arr = arr.filter((item) => item.list.length > 0);
...@@ -14,12 +14,11 @@ export const datasetSearchResultConcat = ( ...@@ -14,12 +14,11 @@ export const datasetSearchResultConcat = (
// rrf // rrf
arr.forEach((item) => { arr.forEach((item) => {
const k = item.k; const weight = item.weight;
item.list.forEach((data, index) => { item.list.forEach((data, index) => {
const rank = index + 1; const rank = index + 1;
const score = 1 / (k + rank); const score = (weight * 1) / (60 + rank);
const record = map.get(data.id); const record = map.get(data.id);
if (record) { if (record) {
// 合并两个score,有相同type的score,取最大值 // 合并两个score,有相同type的score,取最大值
......
...@@ -784,10 +784,10 @@ export async function searchDatasetData( ...@@ -784,10 +784,10 @@ export async function searchDatasetData(
// rrf concat // rrf concat
const rrfEmbRecall = datasetSearchResultConcat( const rrfEmbRecall = datasetSearchResultConcat(
embeddingRecallResults.map((list) => ({ k: 60, list })) embeddingRecallResults.map((list) => ({ weight: 1, list }))
).slice(0, embeddingLimit); ).slice(0, embeddingLimit);
const rrfFTRecall = datasetSearchResultConcat( const rrfFTRecall = datasetSearchResultConcat(
fullTextRecallResults.map((list) => ({ k: 60, list })) fullTextRecallResults.map((list) => ({ weight: 1, list }))
).slice(0, fullTextLimit); ).slice(0, fullTextLimit);
return { return {
...@@ -850,24 +850,22 @@ export async function searchDatasetData( ...@@ -850,24 +850,22 @@ export async function searchDatasetData(
})(); })();
// embedding recall and fullText recall rrf concat // embedding recall and fullText recall rrf concat
const baseK = 120; const embWeight = embeddingWeight; // 向量索引的 weight 大小
const embK = Math.round(baseK * (1 - embeddingWeight)); // 搜索结果的 k 值 const fullTextWeight = 1 - embeddingWeight; // 全文索引的 weight 大小
const fullTextK = Math.round(baseK * embeddingWeight); // rerank 结果的 k 值
const rrfSearchResult = datasetSearchResultConcat([ const rrfSearchResult = datasetSearchResultConcat([
{ k: embK, list: embeddingRecallResults }, { weight: embWeight, list: embeddingRecallResults },
{ k: fullTextK, list: fullTextRecallResults } { weight: fullTextWeight, list: fullTextRecallResults }
]); ]);
const rrfConcatResults = (() => { const rrfConcatResults = (() => {
if (reRankResults.length === 0) return rrfSearchResult; if (reRankResults.length === 0) return rrfSearchResult;
if (rerankWeight === 1) return reRankResults; if (rerankWeight === 1) return reRankResults;
const searchK = Math.round(baseK * rerankWeight); // 搜索结果的 k 值 const searchWeight = 1 - rerankWeight; // 搜索结果的 weight 大小
const rerankK = Math.round(baseK * (1 - rerankWeight)); // rerank 结果的 k 值
return datasetSearchResultConcat([ return datasetSearchResultConcat([
{ k: searchK, list: rrfSearchResult }, { weight: searchWeight, list: rrfSearchResult },
{ k: rerankK, list: reRankResults } { weight: rerankWeight, list: reRankResults }
]); ]);
})(); })();
......
...@@ -29,7 +29,7 @@ export async function dispatchDatasetConcat( ...@@ -29,7 +29,7 @@ export async function dispatchDatasetConcat(
const rrfConcatResults = datasetSearchResultConcat( const rrfConcatResults = datasetSearchResultConcat(
quoteList.map((list) => ({ quoteList.map((list) => ({
k: 60, weight: 1,
list list
})) }))
); );
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment