@tekmidian/pai 0.4.0 → 0.5.0

package/dist/reranker-D7bRAHi6.mjs

@@ -0,0 +1,71 @@
+import { t as __exportAll } from "./rolldown-runtime-95iHPtFO.mjs";
+
+//#region src/memory/reranker.ts
+var reranker_exports = /* @__PURE__ */ __exportAll({
+	configureRerankerModel: () => configureRerankerModel,
+	rerankResults: () => rerankResults
+});
+const DEFAULT_RERANKER_MODEL = "Xenova/ms-marco-MiniLM-L-6-v2";
+let _tokenizer = null;
+let _model = null;
+let _currentModel = null;
+let _loading = null;
+/**
+* Configure the reranker model.
+* Must be called before the first rerank() call if you want a non-default model.
+*/
+function configureRerankerModel(model) {
+	const resolved = model?.trim() || DEFAULT_RERANKER_MODEL;
+	if (_currentModel !== null && _currentModel !== resolved) {
+		_tokenizer = null;
+		_model = null;
+		_loading = null;
+	}
+	_currentModel = resolved;
+}
+async function ensureLoaded() {
+	if (_tokenizer && _model) return;
+	if (_loading) return _loading;
+	_loading = (async () => {
+		const model = _currentModel ?? DEFAULT_RERANKER_MODEL;
+		const { AutoTokenizer, AutoModelForSequenceClassification } = await import("@huggingface/transformers");
+		_tokenizer = await AutoTokenizer.from_pretrained(model);
+		_model = await AutoModelForSequenceClassification.from_pretrained(model, { dtype: "q8" });
+		_currentModel = model;
+	})();
+	return _loading;
+}
+/**
+* Rerank search results using a cross-encoder model.
+*
+* Takes the top `maxCandidates` results from a first-stage retriever,
+* scores each (query, snippet) pair through the cross-encoder, and
+* returns them sorted by cross-encoder relevance score.
+*
+* The original retrieval score is replaced with the cross-encoder score.
+*/
+async function rerankResults(query, results, opts) {
+	if (results.length === 0) return [];
+	const maxCandidates = opts?.maxCandidates ?? 50;
+	const topK = opts?.topK ?? results.length;
+	const candidates = results.slice(0, maxCandidates);
+	await ensureLoaded();
+	const queries = new Array(candidates.length).fill(query);
+	const documents = candidates.map((r) => r.snippet);
+	const inputs = _tokenizer(queries, {
+		text_pair: documents,
+		padding: true,
+		truncation: true
+	});
+	const scores = (await _model(inputs)).logits.tolist();
+	const scored = candidates.map((result, i) => ({
+		...result,
+		score: scores[i][0]
+	}));
+	scored.sort((a, b) => b.score - a.score);
+	return scored.slice(0, topK);
+}
+
+//#endregion
+export { rerankResults as n, reranker_exports as r, configureRerankerModel as t };
+//# sourceMappingURL=reranker-D7bRAHi6.mjs.map

package/dist/reranker-D7bRAHi6.mjs.map

@@ -0,0 +1 @@
+{"version":3,"file":"reranker-D7bRAHi6.mjs","names":[],"sources":["../src/memory/reranker.ts"],"sourcesContent":["/**\n * Cross-encoder reranker for PAI memory search results.\n *\n * Uses Xenova/ms-marco-MiniLM-L-6-v2 — a 22.7M param cross-encoder trained on\n * MS MARCO passage ranking.  The q8 quantized ONNX model is ~23 MB.\n *\n * Cross-encoders score (query, document) pairs jointly, producing more accurate\n * relevance scores than bi-encoder cosine similarity alone.  The trade-off is\n * latency: cross-encoders must score each pair independently, so they are used\n * as a reranking step on top of a fast first-stage retriever (BM25 / cosine).\n *\n * The model is loaded as a lazy singleton — no startup cost until the first\n * rerank call.  Subsequent calls reuse the loaded model.\n *\n * Inspired by QMD's Qwen3-reranker step (tobi/qmd).\n */\n\nimport type { SearchResult } from \"./search.js\";\n\n// ---------------------------------------------------------------------------\n// Constants\n// ---------------------------------------------------------------------------\n\nconst DEFAULT_RERANKER_MODEL = \"Xenova/ms-marco-MiniLM-L-6-v2\";\n\n// ---------------------------------------------------------------------------\n// Lazy singleton\n// ---------------------------------------------------------------------------\n\nlet _tokenizer: any = null;\nlet _model: any = null;\nlet _currentModel: string | null = null;\nlet _loading: Promise<void> | null = null;\n\n/**\n * Configure the reranker model.\n * Must be called before the first rerank() call if you want a non-default model.\n */\nexport function configureRerankerModel(model?: string): void {\n  const resolved = model?.trim() || DEFAULT_RERANKER_MODEL;\n  if (_currentModel !== null && _currentModel !== resolved) {\n    _tokenizer = null;\n    _model = null;\n    _loading = null;\n  }\n  _currentModel = resolved;\n}\n\nasync function ensureLoaded(): Promise<void> {\n  if (_tokenizer && _model) return;\n  if (_loading) return _loading;\n\n  _loading = (async () => {\n    const model = _currentModel ?? DEFAULT_RERANKER_MODEL;\n    const {\n      AutoTokenizer,\n      AutoModelForSequenceClassification,\n    } = await import(\"@huggingface/transformers\");\n\n    _tokenizer = await AutoTokenizer.from_pretrained(model);\n    _model = await AutoModelForSequenceClassification.from_pretrained(\n      model,\n      { dtype: \"q8\" },\n    );\n    _currentModel = model;\n  })();\n\n  return _loading;\n}\n\n// ---------------------------------------------------------------------------\n// Reranking\n// ---------------------------------------------------------------------------\n\nexport interface RerankOptions {\n  /** Maximum number of results to return after reranking. */\n  topK?: number;\n  /**\n   * Maximum number of candidates to rerank.\n   * Cross-encoders are O(n) per candidate, so we cap to keep latency\n   * reasonable.  Default: 50.\n   */\n  maxCandidates?: number;\n}\n\n/**\n * Rerank search results using a cross-encoder model.\n *\n * Takes the top `maxCandidates` results from a first-stage retriever,\n * scores each (query, snippet) pair through the cross-encoder, and\n * returns them sorted by cross-encoder relevance score.\n *\n * The original retrieval score is replaced with the cross-encoder score.\n */\nexport async function rerankResults(\n  query: string,\n  results: SearchResult[],\n  opts?: RerankOptions,\n): Promise<SearchResult[]> {\n  if (results.length === 0) return [];\n\n  const maxCandidates = opts?.maxCandidates ?? 50;\n  const topK = opts?.topK ?? results.length;\n\n  // Cap candidates to rerank\n  const candidates = results.slice(0, maxCandidates);\n\n  await ensureLoaded();\n\n  // Tokenize all (query, document) pairs in a single batch\n  const queries = new Array(candidates.length).fill(query);\n  const documents = candidates.map((r) => r.snippet);\n\n  const inputs = _tokenizer!(queries, {\n    text_pair: documents,\n    padding: true,\n    truncation: true,\n  });\n\n  // Run the cross-encoder\n  const output = await _model!(inputs);\n  const logits = output.logits;\n\n  // ms-marco-MiniLM returns raw logits (not sigmoid-normalized).\n  // Higher = more relevant.\n  const scores: number[][] = logits.tolist();\n\n  // Build reranked results\n  const scored = candidates.map((result, i) => ({\n    ...result,\n    score: scores[i][0],\n  }));\n\n  // Sort by cross-encoder score descending\n  scored.sort((a, b) => b.score - a.score);\n\n  return scored.slice(0, topK);\n}\n"],"mappings":";;;;;;;AAuBA,MAAM,yBAAyB;AAM/B,IAAI,aAAkB;AACtB,IAAI,SAAc;AAClB,IAAI,gBAA+B;AACnC,IAAI,WAAiC;;;;;AAMrC,SAAgB,uBAAuB,OAAsB;CAC3D,MAAM,WAAW,OAAO,MAAM,IAAI;AAClC,KAAI,kBAAkB,QAAQ,kBAAkB,UAAU;AACxD,eAAa;AACb,WAAS;AACT,aAAW;;AAEb,iBAAgB;;AAGlB,eAAe,eAA8B;AAC3C,KAAI,cAAc,OAAQ;AAC1B,KAAI,SAAU,QAAO;AAErB,aAAY,YAAY;EACtB,MAAM,QAAQ,iBAAiB;EAC/B,MAAM,EACJ,eACA,uCACE,MAAM,OAAO;AAEjB,eAAa,MAAM,cAAc,gBAAgB,MAAM;AACvD,WAAS,MAAM,mCAAmC,gBAChD,OACA,EAAE,OAAO,MAAM,CAChB;AACD,kBAAgB;KACd;AAEJ,QAAO;;;;;;;;;;;AA2BT,eAAsB,cACpB,OACA,SACA,MACyB;AACzB,KAAI,QAAQ,WAAW,EAAG,QAAO,EAAE;CAEnC,MAAM,gBAAgB,MAAM,iBAAiB;CAC7C,MAAM,OAAO,MAAM,QAAQ,QAAQ;CAGnC,MAAM,aAAa,QAAQ,MAAM,GAAG,cAAc;AAElD,OAAM,cAAc;CAGpB,MAAM,UAAU,IAAI,MAAM,WAAW,OAAO,CAAC,KAAK,MAAM;CACxD,MAAM,YAAY,WAAW,KAAK,MAAM,EAAE,QAAQ;CAElD,MAAM,SAAS,WAAY,SAAS;EAClC,WAAW;EACX,SAAS;EACT,YAAY;EACb,CAAC;CAQF,MAAM,UALS,MAAM,OAAQ,OAAO,EACd,OAIY,QAAQ;CAG1C,MAAM,SAAS,WAAW,KAAK,QAAQ,OAAO;EAC5C,GAAG;EACH,OAAO,OAAO,GAAG;EAClB,EAAE;AAGH,QAAO,MAAM,GAAG,MAAM,EAAE,QAAQ,EAAE,MAAM;AAExC,QAAO,OAAO,MAAM,GAAG,KAAK"}

package/dist/{schemas-BY3Pjvje.mjs → schemas-BFIgGntb.mjs}

@@ -3402,4 +3402,4 @@ const meta = meta$1;
 
 //#endregion
 export { record as a, number as i, array as n, string as o, boolean as r, unknown as s, _enum as t };
-//# sourceMappingURL=schemas-BY3Pjvje.mjs.map
+//# sourceMappingURL=schemas-BFIgGntb.mjs.map