refactor: update reranker

Author: hustcc

README.md

@@ -208,14 +208,6 @@ await ctx.close();
 +------------------------------------------------------------------------+
 ```
 
-### Module Structure
-
-- **Public API**: `Context`, `QueryOptions`, `QueryResult`, `Document`, `Loader`, `MarkdownLoader`, `JsonLoader`, `TextLoader`, `pathToId`
-- **Reranking**: `KeywordReranker`, `createReranker`, `Reranker`, `RerankCandidate`, `RerankResult`, `RerankOptions`
-- **Query Expansion**: `SynonymExpander`, `NoopExpander`, `QueryExpander`, `QueryExpansionOptions`
-- **Advanced API**: `Embedder`, `TransformersEmbedder`, `EmbedderManager`, `IZvecStore`, `ActualZvecStore`, `Store`
-
-
 ## License
 
 MIT

src/context.ts

@@ -5,10 +5,7 @@ import {
   ContextOptions,
   QueryOptions,
   QueryResult,
-  Document,
   LoadedDoc,
-  LoadPhase,
-  LoadProgress,
 } from './types';
 import { Embedder } from './embedder';
 
@@ -18,24 +15,18 @@ import { pathToId } from './utils/doc';
 import { Store } from './storage/store';
 import type { ZvecDoc } from './storage/zvec-store';
 import {
-  createReranker,
   safeJsonParse,
   computeContentHash,
   loadSampleText,
 } from './utils';
 import { expand } from './expander';
-import type { Reranker, RerankCandidate } from './utils';
-
-// ---------------------------------------------------------------------------
-// Context class
-// ---------------------------------------------------------------------------
+import { applyRerank } from './reranker';
 
 export class Context {
   private readonly options: ContextOptions;
   private readonly embedder: Embedder;
   readonly embedderInfo: EmbedderInfo;
   private readonly store: Store;
-  private readonly reranker: Reranker | null;
 
   private constructor(options: ContextOptions, embedder: Embedder, embedderInfo: EmbedderInfo) {
     this.options = {
@@ -46,7 +37,6 @@ export class Context {
     this.embedder = embedder;
     this.embedderInfo = embedderInfo;
     this.store = new Store(this.options.vectorsDir!, embedder, this.options);
-    this.reranker = createReranker(this.options.rerankWeights);
   }
 
   static async create(options: ContextOptions): Promise<Context> {
@@ -163,23 +153,8 @@ export class Context {
       };
     });
 
-    if (rerankEnabled && queryResults.length > topK) {
-      const docs: RerankCandidate[] = queryResults.map((r) => ({
-        id: r.id,
-        content: r.content,
-        score: r.score,
-      }));
-
-      const reranked = await this.reranker!.rerank(text, docs);
-
-      const scoreMap = new Map(reranked.map((r) => [r.id, r.score]));
-      for (const result of queryResults) {
-        const newScore = scoreMap.get(result.id);
-        if (newScore !== undefined) {
-          result.score = newScore;
-          result.scoreMode = 'reranked';
-        }
-      }
+    if (rerankEnabled) {
+      await applyRerank(this.options.rerankWeights, text, queryResults, topK);
     }
 
     queryResults.sort((a, b) => b.score - a.score);

src/index.ts

@@ -1,16 +1,2 @@
-// ---------------------------------------------------------------------------
-// Public API
-// ---------------------------------------------------------------------------
-
 export { Context } from './context';
-export * from './types';
-
-// Reranker
-export { KeywordReranker, createReranker } from './utils/reranker';
-export type { RerankCandidate, RerankResult } from './utils/reranker';
-
-// Query expansion
-export { expand } from './expander';
-
-// Embedder
-export { Embedder } from './embedder';
+export * from './types';

src/reranker/helpers.ts

@@ -0,0 +1,36 @@
+const WORD_BOUNDARY = /[\s\n.,;:!?，。！？、：；"'(（【《\-_]/;
+
+export function tokenizeQuery(query: string): string[] {
+  const tokens = query.split(/[\s,，。.!！？?、：:；;]+/).filter(Boolean);
+  const result: string[] = [];
+  for (const r of tokens) {
+    result.push(r);
+    if (/[一-鿿]{3,}/.test(r)) {
+      for (let i = 0; i + 2 <= r.length; i++) result.push(r.slice(i, i + 2));
+    }
+  }
+  return [...new Set(result)];
+}
+
+export function countOccurrences(text: string, sub: string): number {
+  let count = 0, pos = 0;
+  while ((pos = text.indexOf(sub, pos)) !== -1) { count++; pos += sub.length; }
+  return count;
+}
+
+export function countTermMatches(text: string, term: string): number {
+  let count = 0, pos = 0;
+  while ((pos = text.indexOf(term, pos)) !== -1) {
+    if (isWordBoundary(text, term, pos)) count++;
+    pos += term.length;
+  }
+  return count;
+}
+
+export function isWordBoundary(text: string, term: string, pos?: number): boolean {
+  const idx = pos ?? text.indexOf(term);
+  if (idx === -1) return false;
+  const before = idx === 0 || WORD_BOUNDARY.test(text[idx - 1]);
+  const after = idx + term.length >= text.length || WORD_BOUNDARY.test(text[idx + term.length]);
+  return before && after;
+}

src/reranker/index.ts

@@ -0,0 +1,36 @@
+import type { RerankCandidate, RerankOptions } from './types';
+import type { QueryResult } from '../types';
+import { rerank } from './reranker';
+
+export type { RerankOptions };
+
+/**
+ * Apply reranking to query results.
+ */
+export async function applyRerank(
+  rerankOptions: RerankOptions | undefined,
+  query: string,
+  results: QueryResult[],
+  topK: number,
+): Promise<QueryResult[]> {
+  if (results.length <= topK) return results;
+
+  const candidates: RerankCandidate[] = results.map((r) => ({
+    id: r.id,
+    content: r.content,
+    score: r.score,
+  }));
+
+  const reranked = await rerank(query, candidates, rerankOptions);
+
+  const scoreMap = new Map(reranked.map((r) => [r.id, r.score]));
+  for (const result of results) {
+    const newScore = scoreMap.get(result.id);
+    if (newScore !== undefined) {
+      result.score = newScore;
+      result.scoreMode = 'reranked';
+    }
+  }
+
+  return results;
+}

src/reranker/reranker.ts

@@ -0,0 +1,81 @@
+import type { RerankCandidate, RerankResult, RerankOptions } from './types';
+import { tokenizeQuery, countOccurrences, countTermMatches, isWordBoundary } from './helpers';
+
+const DEFAULTS = {
+  phraseWeight: 3.0,
+  phraseRepeatBonus: 0.5,
+  termWeight: 1.0,
+  termRepeatBonus: 0.2,
+  substringWeight: 0.3,
+  headingTermBonus: 2.0,
+  headingPhraseBonus: 2.5,
+  originalScoreCarry: 0.1,
+} as const;
+
+/**
+ * Rerank candidates by keyword / phrase overlap with the query.
+ * Scores are normalised to [0, 1] via min-max scaling.
+ */
+export async function rerank(
+  query: string,
+  candidates: RerankCandidate[],
+  options?: RerankOptions,
+): Promise<RerankResult[]> {
+  if (candidates.length === 0) return [];
+
+  const w = {
+    phraseWeight: options?.phraseWeight ?? DEFAULTS.phraseWeight,
+    phraseRepeatBonus: options?.phraseRepeatBonus ?? DEFAULTS.phraseRepeatBonus,
+    termWeight: options?.termWeight ?? DEFAULTS.termWeight,
+    termRepeatBonus: options?.termRepeatBonus ?? DEFAULTS.termRepeatBonus,
+    substringWeight: options?.substringWeight ?? DEFAULTS.substringWeight,
+    headingTermBonus: options?.headingTermBonus ?? DEFAULTS.headingTermBonus,
+    headingPhraseBonus: options?.headingPhraseBonus ?? DEFAULTS.headingPhraseBonus,
+    originalScoreCarry: options?.originalScoreCarry ?? DEFAULTS.originalScoreCarry,
+  };
+
+  const lowerQuery = query.toLowerCase();
+  const queryTerms = tokenizeQuery(lowerQuery);
+  const queryPhrase = lowerQuery.trim();
+
+  const scored = candidates.map((c) => {
+    const content = c.content.toLowerCase();
+    let score = 0;
+
+    // 1. Exact phrase match
+    if (content.includes(queryPhrase)) {
+      score += w.phraseWeight + (countOccurrences(content, queryPhrase) - 1) * w.phraseRepeatBonus;
+    }
+
+    // 2. Per-term matching
+    for (const term of queryTerms) {
+      if (content.includes(term)) {
+        score += isWordBoundary(content, term) ? w.termWeight + (countTermMatches(content, term) - 1) * w.termRepeatBonus : w.substringWeight;
+      }
+    }
+
+    // 3. Heading path bonus
+    if (c.headingPath) {
+      const heading = c.headingPath.toLowerCase();
+      for (const term of queryTerms) {
+        if (heading.includes(term)) score += w.headingTermBonus;
+      }
+      if (queryPhrase.length > 2 && heading.includes(queryPhrase)) {
+        score += w.headingPhraseBonus;
+      }
+    }
+
+    // 4. Carry over original score
+    score += c.score * w.originalScoreCarry;
+
+    return { id: c.id, score };
+  });
+
+  // Min-max normalise to [0, 1]
+  const scores = scored.map((s) => s.score);
+  const min = Math.min(...scores);
+  const max = Math.max(...scores);
+  const range = max - min || 1;
+
+  return scored.map((s) => ({ id: s.id, score: (s.score - min) / range }));
+}

src/reranker/types.ts

@@ -0,0 +1,43 @@
+/**
+ * Reranker — second-stage precision scoring for search results.
+ */
+
+// ---------------------------------------------------------------------------
+// Types
+// ---------------------------------------------------------------------------
+
+/** A candidate document for reranking. */
+export interface RerankCandidate {
+  id: string;
+  content: string;
+  /** Original score from the coarse search stage. */
+  score: number;
+  /** Heading path as a string (e.g. "Line Chart > Tooltip"). */
+  headingPath?: string;
+}
+
+/** A reranked result. */
+export interface RerankResult {
+  id: string;
+  /** Final score after reranking (higher is better). */
+  score: number;
+}
+
+/** Reranker type */
+export type Reranker = {
+  rerank(query: string, candidates: RerankCandidate[]): Promise<RerankResult[]>;
+};
+
+/** Configuration for reranking. */
+export interface RerankOptions {
+  rerankFactor?: number;
+  minCandidates?: number;
+  phraseWeight?: number;
+  phraseRepeatBonus?: number;
+  termWeight?: number;
+  termRepeatBonus?: number;
+  substringWeight?: number;
+  headingTermBonus?: number;
+  headingPhraseBonus?: number;
+  originalScoreCarry?: number;
+}

src/types.ts

@@ -1,4 +1,4 @@
-import type { RerankOptions } from './utils/reranker';
+import type { RerankOptions } from './reranker';
 
 /**
  * Document structure

src/utils/index.ts

@@ -4,6 +4,4 @@ export { safeJsonParse } from './common';
 export { loadSampleText } from './sample';
 export { isCJK, detectLanguage, tokenizerForLanguage, detectTokenizer } from './tokenizer';
 export type { LanguageHint } from './tokenizer';
-export { containsCJK } from './str';
-export { KeywordReranker, createReranker } from './reranker';
-export type { Reranker, RerankCandidate, RerankResult } from './reranker';
+export { containsCJK } from './str';