Skip to content

Commit

Permalink
Make score inclusion toggleable
Browse files Browse the repository at this point in the history
  • Loading branch information
jacoblee93 committed Nov 22, 2024
1 parent 10438ca commit 57267ab
Showing 1 changed file with 17 additions and 7 deletions.
24 changes: 17 additions & 7 deletions libs/langchain-community/src/retrievers/bm25.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import { BM25 } from "../utils/@furkantoprak/bm25/BM25.js";
export type BM25RetrieverOptions = {
docs: Document[];
k: number;
includeScore?: boolean;
} & BaseRetrieverInput;

/**
Expand All @@ -14,6 +15,8 @@ export type BM25RetrieverOptions = {
* The k parameter determines the number of documents to return for each query.
*/
export class BM25Retriever extends BaseRetriever {
includeScore = false;

static lc_name() {
return "BM25Retriever";
}
Expand All @@ -35,6 +38,7 @@ export class BM25Retriever extends BaseRetriever {
super(options);
this.docs = options.docs;
this.k = options.k;
this.includeScore = options.includeScore ?? this.includeScore;
}

private preprocessFunc(text: string): string[] {
Expand All @@ -53,13 +57,19 @@ export class BM25Retriever extends BaseRetriever {

scoredDocs.sort((a, b) => b.score - a.score);

return scoredDocs.slice(0, this.k).map((item) => new Document({
...item.document.id && { id: item.document.id },
pageContent: item.document.pageContent,
metadata: {
bm25Score: item.score,
...item.document.metadata,
return scoredDocs.slice(0, this.k).map((item) => {
if (this.includeScore) {
return new Document({
...(item.document.id && { id: item.document.id }),
pageContent: item.document.pageContent,
metadata: {
bm25Score: item.score,
...item.document.metadata,
},
});
} else {
return item.document;
}
}));
});
}
}

0 comments on commit 57267ab

Please sign in to comment.