improve bm25 sorting

Author: AllenFang

libs/langchain-community/src/retrievers/bm25.ts

@@ -47,28 +47,28 @@ export class BM25Retriever extends BaseRetriever {
 
   async _getRelevantDocuments(query: string) {
     const processedQuery = this.preprocessFunc(query);
-    const documents = this.docs.map((doc) => doc.pageContent);
-    const scores = BM25(documents, processedQuery) as number[];
-
-    const scoredDocs = this.docs.map((doc, index) => ({
-      document: doc,
-      score: scores[index],
-    }));
-
-    scoredDocs.sort((a, b) => b.score - a.score);
+    const scoredDocs = BM25<Document>(
+      this.docs.map((d) => ({
+        text: d.pageContent,
+        docs: d,
+      })),
+      processedQuery,
+      undefined,
+      (a, b) => b.score - a.score
+    );
 
     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,
+          ...(item.docs.id && { id: item.docs.id }),
+          pageContent: item.docs.pageContent,
           metadata: {
             bm25Score: item.score,
-            ...item.document.metadata,
+            ...item.docs.metadata,
           },
         });
       } else {
-        return item.document;
+        return item.docs;
       }
     });
   }

libs/langchain-community/src/utils/@furkantoprak/bm25/BM25.ts

@@ -16,21 +16,26 @@ export const getTermFrequency = (term: string, corpus: string) => {
 };
 
 /** Inverse document frequency. */
-export const getIDF = (term: string, documents: string[]) => {
+export const getIDF = <T>(term: string, documents: BMInputDocument<T>[]) => {
   // Number of relevant documents.
-  const relevantDocuments = documents.filter((document: string) =>
-    document.includes(term)
+  const relevantDocuments = documents.filter((document) =>
+    document.text.includes(term)
   ).length;
   return Math.log(
     (documents.length - relevantDocuments + 0.5) / (relevantDocuments + 0.5) + 1
   );
 };
 
+export interface BMInputDocument<T> {
+  text: string;
+  docs: T;
+}
+
 /** Represents a document; useful when sorting results.
  */
-export interface BMDocument {
-  /** The document is originally scoreed. */
-  document: string;
+export interface BMOutputDocument<T> {
+  /** The original source */
+  docs: T;
   /** The score that the document recieves. */
   score: number;
 }
@@ -44,7 +49,10 @@ export interface BMConstants {
 }
 
 /** If returns positive, the sorting results in secondEl coming before firstEl, else, firstEl comes before secondEL  */
-export type BMSorter = (firstEl: BMDocument, secondEl: BMDocument) => number;
+export type BMSorter<T> = (
+  firstEl: BMOutputDocument<T>,
+  secondEl: BMOutputDocument<T>
+) => number;
 
 /** Implementation of Okapi BM25 algorithm.
  *  @param documents: Collection of documents.
@@ -53,16 +61,16 @@ export type BMSorter = (firstEl: BMDocument, secondEl: BMDocument) => number;
  *  @param sort: A function that allows you to sort queries by a given rule. If not provided, returns results corresponding to the original order.
  * If this option is provided, the return type will not be an array of scores but an array of documents with their scores.
  */
-export function BM25(
-  documents: string[],
+export function BM25<T>(
+  documents: BMInputDocument<T>[],
   keywords: string[],
   constants?: BMConstants,
-  sorter?: BMSorter
-): number[] | BMDocument[] {
+  sorter?: BMSorter<T>
+): BMOutputDocument<T>[] {
   const b = constants && constants.b ? constants.b : 0.75;
   const k1 = constants && constants.k1 ? constants.k1 : 1.2;
-  const documentLengths = documents.map((document: string) =>
-    getWordCount(document)
+  const documentLengths = documents.map((document) =>
+    getWordCount(document.text)
   );
   const averageDocumentLength =
     documentLengths.reduce((a, b) => a + b, 0) / documents.length;
@@ -71,14 +79,14 @@ export function BM25(
     return obj;
   }, new Map<string, number>());
 
-  const scores = documents.map((document: string, index: number) => {
+  const scoredDocs = documents.map(({ text, docs }, index) => {
     const score = keywords
       .map((keyword: string) => {
         const inverseDocumentFrequency = idfByKeyword.get(keyword);
         if (inverseDocumentFrequency === undefined) {
           throw new Error("Missing keyword.");
         }
-        const termFrequency = getTermFrequency(keyword, document);
+        const termFrequency = getTermFrequency(keyword, text);
         const documentLength = documentLengths[index];
         return (
           (inverseDocumentFrequency * (termFrequency * (k1 + 1))) /
@@ -87,14 +95,11 @@ export function BM25(
         );
       })
       .reduce((a: number, b: number) => a + b, 0);
-    if (sorter) {
-      return { score, document } as BMDocument;
-    }
-    return score;
+    return { score, docs } as BMOutputDocument<T>;
   });
   // sort the results
   if (sorter) {
-    return (scores as BMDocument[]).sort(sorter);
+    return scoredDocs.sort(sorter);
   }
-  return scores as number[];
+  return scoredDocs;
 }