Hybrid search support - addition of bm25 keyword search (#61)

Co-authored-by: Gaureesh Anvekar <[email protected]>

Author: GaureeshAnvekar

bin/vectra.js

@@ -35,22 +35,23 @@
         "openai": "^3.2.1",
         "turndown": "^7.1.2",
         "uuid": "^9.0.0",
+        "wink-nlp": "^2.3.2",
         "yargs": "^17.7.2"
     },
-    "resolutions": {
-    },
+    "resolutions": {},
     "devDependencies": {
-        "@types/node": "^14.14.31",
-        "@types/mocha": "^8.2.0",
         "@types/assert": "^1.5.3",
+        "@types/mocha": "^8.2.0",
+        "@types/node": "^14.14.31",
         "@types/turndown": "^5.0.1",
         "@types/uuid": "9.0.1",
         "@types/yargs": "17.0.24",
         "mocha": "10.2.0",
         "nyc": "^15.1.0",
         "shx": "^0.3.2",
         "ts-mocha": "10.0.0",
-        "typescript": "^4.2.3"
+        "typescript": "^4.2.3",
+        "wink-bm25-text-search": "^3.1.2"
     },
     "scripts": {
         "build": "tsc -b",

package-lock.json

@@ -1,2 +1,4 @@
 vectra.keys
-index
+index
+index-files
+index-wiki

package.json

@@ -30,6 +30,12 @@ export interface DocumentQueryOptions {
      * Optional. Filter to apply to the document metadata.
      */
     filter?: MetadataFilter;
+
+    /**
+     * Optional. Turn on bm25 keyword search to perform hybrid search - semantic + keyword
+     */
+    isBm25?: boolean;
+
 }
 
 /**
@@ -378,7 +384,7 @@ export class LocalDocumentIndex extends LocalIndex<DocumentChunkMetadata> {
         }
 
         // Query index for chunks
-        const results = await this.queryItems(embeddings.output![0], options.maxChunks!, options.filter);
+        const results = await this.queryItems(embeddings.output![0], query, options.maxChunks!, options.filter, options.isBm25);
 
         // Group chunks by document
         const documentChunks: { [documentId: string]: QueryResult<DocumentChunkMetadata>[]; } = {};

samples/wikipedia/.gitignore

@@ -66,7 +66,8 @@ export class LocalDocumentResult extends LocalDocument {
                     startPos: startPos + offset,
                     endPos: startPos + offset + chunkLength - 1,
                     score: chunk.score,
-                    tokenCount: chunkLength
+                    tokenCount: chunkLength,
+                    isBm25: false
                 });
                 offset += chunkLength;
             }
@@ -103,7 +104,8 @@ export class LocalDocumentResult extends LocalDocument {
             return {
                 text: text,
                 tokenCount: section.tokenCount,
-                score: section.score
+                score: section.score,
+                isBm25: false,
             };
         });
     }
@@ -127,7 +129,8 @@ export class LocalDocumentResult extends LocalDocument {
             return [{
                 text,
                 tokenCount: length,
-                score: 1.0
+                score: 1.0,
+                isBm25: false,
             }];
         }
 
@@ -148,7 +151,8 @@ export class LocalDocumentResult extends LocalDocument {
                 startPos,
                 endPos,
                 score: chunk.score,
-                tokenCount: this._tokenizer.encode(chunkText).length
+                tokenCount: this._tokenizer.encode(chunkText).length,
+                isBm25: Boolean(chunk.item.metadata.isBm25),
             };
         }).filter(chunk => chunk.tokenCount <= maxTokens).sort((a, b) => a.startPos - b.startPos);
 
@@ -163,36 +167,63 @@ export class LocalDocumentResult extends LocalDocument {
             return [{
                 text: this._tokenizer.decode(tokens.slice(0, maxTokens)),
                 tokenCount: maxTokens,
-                score: topChunk.score
+                score: topChunk.score,
+                isBm25: false,
             }];
         }
 
-        // Generate sections
+        // Generate semantic sections
         const sections: Section[] = [];
         for (let i = 0; i < chunks.length; i++) {
             const chunk = chunks[i];
             let section = sections[sections.length - 1];
-            if (!section || section.tokenCount + chunk.tokenCount > maxTokens) {
-                section = {
-                    chunks: [],
-                    score: 0,
-                    tokenCount: 0
-                };
-                sections.push(section);
+            if (!chunk.isBm25) {
+                if (!section || section.tokenCount + chunk.tokenCount > maxTokens) {
+                    section = {
+                        chunks: [],
+                        score: 0,
+                        tokenCount: 0
+                    };
+                    sections.push(section);
+                }
+                section.chunks.push(chunk);
+                section.score += chunk.score;
+                section.tokenCount += chunk.tokenCount;
             }
-            section.chunks.push(chunk);
-            section.score += chunk.score;
-            section.tokenCount += chunk.tokenCount;
         }
 
+        // Generate bm25 sections
+        const bm25Sections: Section[] = [];
+        for (let i = 0; i < chunks.length; i++) {
+            const chunk = chunks[i];
+            let section = bm25Sections[bm25Sections.length - 1];
+            if (chunk.isBm25) {
+                if (!section || section.tokenCount + chunk.tokenCount > maxTokens) {
+                    section = {
+                        chunks: [],
+                        score: 0,
+                        tokenCount: 0
+                    };
+                    bm25Sections.push(section);
+                }
+                section.chunks.push(chunk);
+                section.score += chunk.score;
+                section.tokenCount += chunk.tokenCount;
+            }
+        }        
         // Normalize section scores
         sections.forEach(section => section.score /= section.chunks.length);
+        bm25Sections.forEach(section => section.score /= section.chunks.length);
 
         // Sort sections by score and limit to maxSections
         sections.sort((a, b) => b.score - a.score);
+        bm25Sections.sort((a, b) => b.score - a.score);
         if (sections.length > maxSections) {
             sections.splice(maxSections, sections.length - maxSections);
         }
+        if (bm25Sections.length > maxSections) {
+            bm25Sections.splice(maxSections, bm25Sections.length - maxSections);
+        }
 
         // Combine adjacent chunks of text
         sections.forEach(section => {
@@ -216,7 +247,8 @@ export class LocalDocumentResult extends LocalDocument {
                 startPos: -1,
                 endPos: -1,
                 score: 0,
-                tokenCount: this._tokenizer.encode('\n\n...\n\n').length
+                tokenCount: this._tokenizer.encode('\n\n...\n\n').length,
+                isBm25: false,
             };
             sections.forEach(section => {
                 // Insert connectors between chunks
@@ -242,7 +274,8 @@ export class LocalDocumentResult extends LocalDocument {
                             startPos: sectionStart - beforeBudget,
                             endPos: sectionStart - 1,
                             score: 0,
-                            tokenCount: beforeBudget
+                            tokenCount: beforeBudget,
+                            isBm25: false,
                         };
                         section.chunks.unshift(chunk);
                         section.tokenCount += chunk.tokenCount;
@@ -258,7 +291,8 @@ export class LocalDocumentResult extends LocalDocument {
                             startPos: sectionEnd + 1,
                             endPos: sectionEnd + afterBudget,
                             score: 0,
-                            tokenCount: afterBudget
+                            tokenCount: afterBudget,
+                            isBm25: false,
                         };
                         section.chunks.push(chunk);
                         section.tokenCount += chunk.tokenCount;
@@ -268,16 +302,29 @@ export class LocalDocumentResult extends LocalDocument {
             });
         }
 
-        // Return final rendered sections
-        return sections.map(section => {
+        const semanticDocTextSections = sections.map(section => {
+            let text = '';
+            section.chunks.forEach(chunk => text += chunk.text);
+            return {
+                text: text,
+                tokenCount: section.tokenCount,
+                score: section.score,
+                isBm25: false,
+            };
+        });
+        const bm25DocTextSections = bm25Sections.map(section => {
             let text = '';
             section.chunks.forEach(chunk => text += chunk.text);
             return {
                 text: text,
                 tokenCount: section.tokenCount,
-                score: section.score
+                score: section.score,
+                isBm25: true,
             };
         });
+
+        // Return final rendered sections
+        return [...semanticDocTextSections, ...bm25DocTextSections];
     }
 
     private encodeBeforeText(text: string, budget: number): number[] {
@@ -300,6 +347,7 @@ interface SectionChunk {
     endPos: number;
     score: number;
     tokenCount: number;
+    isBm25: boolean;
 }
 
 interface Section {

src/LocalDocumentIndex.ts

@@ -3,7 +3,11 @@ import * as path from 'path';
 import { v4 } from 'uuid';
 import { ItemSelector } from './ItemSelector';
 import { IndexItem, IndexStats, MetadataFilter, MetadataTypes, QueryResult } from './types';
-
+import { LocalDocument } from './LocalDocument';
+import { LocalDocumentIndex } from './LocalDocumentIndex';
+import bm25 from 'wink-bm25-text-search';
+import winkNLP from 'wink-nlp';
+import model from 'wink-eng-lite-web-model';
 export interface CreateIndexConfig {
     version: number;
     deleteIfExists?: boolean;
@@ -24,6 +28,8 @@ export class LocalIndex<TMetadata extends Record<string,MetadataTypes> = Record<
 
     private _data?: IndexData;
     private _update?: IndexData;
+    //member fields for BM25
+    private _bm25Engine: any;
 
     /**
      * Creates a new instance of LocalIndex.
@@ -247,7 +253,7 @@ export class LocalIndex<TMetadata extends Record<string,MetadataTypes> = Record<
      * @param filter Optional. Filter to apply.
      * @returns Similar items to the vector that matche the supplied filter.
      */
-    public async queryItems<TItemMetadata extends TMetadata = TMetadata>(vector: number[], topK: number, filter?: MetadataFilter): Promise<QueryResult<TItemMetadata>[]> {
+    public async queryItems<TItemMetadata extends TMetadata = TMetadata>(vector: number[], query: string, topK: number, filter?: MetadataFilter, isBm25?: boolean): Promise<QueryResult<TItemMetadata>[]> {
         await this.loadIndexData();
 
         // Filter items
@@ -285,6 +291,36 @@ export class LocalIndex<TMetadata extends Record<string,MetadataTypes> = Record<
             }
         }
 
+        //Peform bm25 search only if enabled. Avoid duplicate chunks, which are already selected during semantic search.
+        if (isBm25) {
+            const itemSet = new Set();
+            for (const item of top) itemSet.add(item.item.id);
+            
+            this.setupbm25();
+
+            let currDoc;
+            let currDocTxt;
+            for (let i = 0; i < items.length; i++) {
+                if (!itemSet.has(items[i].id)) {
+                    const item = items[i];
+                    currDoc = new LocalDocument((this as unknown) as LocalDocumentIndex, item.metadata.documentId.toString(), '');
+                    currDocTxt = await currDoc.loadText();
+                    const startPos = item.metadata.startPos;
+                    const endPos = item.metadata.endPos;
+                    const chunkText = currDocTxt.substring(Number(startPos), Number(endPos) + 1);
+                    this._bm25Engine.addDoc({body: chunkText}, i);
+                }
+            }
+            this._bm25Engine.consolidate();
+            var results = await this.bm25Search(query, items, topK);
+            results.forEach((res: any) => {
+                top.push({
+                    item: Object.assign({}, {...items[res[0]], metadata: {...items[res[0]].metadata, isBm25: true}}) as any,
+                    score: res[1]
+                });
+            });
+            
+        }
         return top;
     }
 
@@ -385,6 +421,37 @@ export class LocalIndex<TMetadata extends Record<string,MetadataTypes> = Record<
             return newItem;
         }
     }
+
+    private async setupbm25(): Promise<any> {
+        this._bm25Engine = bm25();
+        const nlp = winkNLP( model );
+        const its = nlp.its;
+
+        const prepTask = function ( text: string ) {
+            const tokens: any[] = [];
+            nlp.readDoc(text)
+                .tokens()
+                // Use only words ignoring punctuations etc and from them remove stop words
+                .filter( (t: any) => ( t.out(its.type) === 'word' && !t.out(its.stopWordFlag) ) )
+                // Handle negation and extract stem of the word
+                .each( (t: any) => tokens.push( (t.out(its.negationFlag)) ? '!' + t.out(its.stem) : t.out(its.stem) ) );
+
+            return tokens;
+        };
+
+        this._bm25Engine.defineConfig( { fldWeights: { body: 1 } } );
+        // Step II: Define PrepTasks pipe.
+        this._bm25Engine.definePrepTasks( [ prepTask ] );
+    }
+
+    private async bm25Search(searchQuery: string, items: any, topK: number): Promise<any> {
+        var query = searchQuery;
+        // `results` is an array of [ doc-id, score ], sorted by score
+        var results = this._bm25Engine.search( query );
+
+        return results.slice(0, topK);
+    }
+
 }
 
 interface IndexData {

src/LocalDocumentResult.ts

@@ -16,9 +16,9 @@ export class Colorize {
         }
     }
 
-    public static output(output: object | string, quote: string = '', units: string = ''): string {
+    public static output(output: object | string, isBm25: boolean = false, quote: string = '', units: string = ''): string {
         if (typeof output === 'string') {
-            return `\x1b[32m${quote}${output}${quote}\x1b[0m`;
+            return isBm25 ? `\x1b[34m${quote}${output}${quote}\x1b[0m` : `\x1b[32m${quote}${output}${quote}\x1b[0m`;
         } else if (typeof output === 'object' && output !== null) {
             return colorizer(output, {
                 pretty: true,
@@ -54,7 +54,7 @@ export class Colorize {
     }
 
     public static value(field: string, value: any, units: string = ''): string {
-        return `${field}: ${Colorize.output(value, '"', units)}`;
+        return `${field}: ${Colorize.output(value, false, '"', units)}`;
     }
 
     public static warning(warning: string): string {

src/LocalIndex.ts

@@ -0,0 +1,4 @@
+declare module 'wink-bm25-text-search' {
+    const bm25: any; 
+    export default bm25;
+  }

src/internals/Colorize.ts

@@ -172,4 +172,5 @@ export interface DocumentTextSection {
     text: string;
     tokenCount: number;
     score: number;
+    isBm25: boolean;
 }