feat(qwen): fuse dense and lexical search with RRF
This commit is contained in:
1 parent
fd37422de2
commit
e97404a0a1
2 files changed
+31
-5
No files matched your search
@@ -201,6 +201,10 @@ class RetrievalServerTests(unittest.TestCase):
|
|||||||
]:
|
]:
|
||||||
self.assertEqual([symbol for symbol, _ in MODULE.sections(source, path)], symbols)
|
self.assertEqual([symbol for symbol, _ in MODULE.sections(source, path)], symbols)
|
||||||
|
|
||||||
|
def test_rrf_rewards_agreement_and_breaks_ties_deterministically(self):
|
||||||
|
self.assertEqual(MODULE.reciprocal_rank_fusion([1, 2, 3], [4, 2, 3]), [2, 3, 1, 4])
|
||||||
|
self.assertEqual(MODULE.reciprocal_rank_fusion([], [4, 2]), [4, 2])
|
||||||
|
|
||||||
def test_old_cache_is_preserved_but_not_reused(self):
|
def test_old_cache_is_preserved_but_not_reused(self):
|
||||||
with TemporaryDirectory() as temporary:
|
with TemporaryDirectory() as temporary:
|
||||||
root = Path(temporary)
|
root = Path(temporary)
|
||||||
|
|||||||
@@ -20,8 +20,11 @@ RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B"
|
|||||||
EMBEDDING_DIMENSIONS = 1024
|
EMBEDDING_DIMENSIONS = 1024
|
||||||
CHUNK_TOKENS = 400
|
CHUNK_TOKENS = 400
|
||||||
OVERLAP_TOKENS = 64
|
OVERLAP_TOKENS = 64
|
||||||
|
DENSE_CANDIDATES = 50
|
||||||
|
LEXICAL_CANDIDATES = 50
|
||||||
RERANK_CANDIDATES = 20
|
RERANK_CANDIDATES = 20
|
||||||
FINAL_RESULTS = 20
|
FINAL_RESULTS = 20
|
||||||
|
RRF_K = 60
|
||||||
INDEX_VERSION = 2
|
INDEX_VERSION = 2
|
||||||
RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question."
|
RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question."
|
||||||
TEXT_EXTENSIONS = {".c", ".cpp", ".cs", ".go", ".java", ".js", ".json", ".md", ".py", ".ps1", ".rs", ".sh", ".toml", ".ts", ".tsx", ".txt", ".yaml", ".yml"}
|
TEXT_EXTENSIONS = {".c", ".cpp", ".cs", ".go", ".java", ".js", ".json", ".md", ".py", ".ps1", ".rs", ".sh", ".toml", ".ts", ".tsx", ".txt", ".yaml", ".yml"}
|
||||||
@@ -115,6 +118,14 @@ def document_text(path: str, symbol: str, content: str):
|
|||||||
return f"Path: {path}\nSymbol: {symbol}\n\n{content}"
|
return f"Path: {path}\nSymbol: {symbol}\n\n{content}"
|
||||||
|
|
||||||
|
|
||||||
|
def reciprocal_rank_fusion(*rankings):
|
||||||
|
scores = {}
|
||||||
|
for ranking in rankings:
|
||||||
|
for rank, identifier in enumerate(ranking, 1):
|
||||||
|
scores[identifier] = scores.get(identifier, 0.0) + 1.0 / (RRF_K + rank)
|
||||||
|
return sorted(scores, key=lambda identifier: (-scores[identifier], identifier))
|
||||||
|
|
||||||
|
|
||||||
class Index:
|
class Index:
|
||||||
def __init__(self, root: Path, data: Path):
|
def __init__(self, root: Path, data: Path):
|
||||||
self.root = root.resolve()
|
self.root = root.resolve()
|
||||||
@@ -205,7 +216,7 @@ class Index:
|
|||||||
self.delete_path(relative)
|
self.delete_path(relative)
|
||||||
self.db.execute("delete from files where path = ?", (relative,))
|
self.db.execute("delete from files where path = ?", (relative,))
|
||||||
|
|
||||||
def search(self, query: str, top_k: int):
|
def search(self, query: str, top_k: int = 5):
|
||||||
if not query.strip():
|
if not query.strip():
|
||||||
raise ValueError("query must not be empty")
|
raise ValueError("query must not be empty")
|
||||||
with self.lock:
|
with self.lock:
|
||||||
@@ -216,10 +227,21 @@ class Index:
|
|||||||
vector = self.encoder().encode([query], prompt_name="query", normalize_embeddings=True)[0]
|
vector = self.encoder().encode([query], prompt_name="query", normalize_embeddings=True)[0]
|
||||||
if len(vector) != EMBEDDING_DIMENSIONS:
|
if len(vector) != EMBEDDING_DIMENSIONS:
|
||||||
raise ValueError("Embedding must have 1024 dimensions")
|
raise ValueError("Embedding must have 1024 dimensions")
|
||||||
rows = ({"path": path, "symbol": symbol, "score": float(np.dot(vector, np.frombuffer(blob, dtype=np.float32))), "text": text}
|
scores = ((identifier, float(np.dot(vector, np.frombuffer(blob, dtype=np.float32))))
|
||||||
for path, symbol, text, blob in self.db.execute("select path, symbol, text, vector from chunks"))
|
for identifier, blob in self.db.execute("select rowid, vector from chunks"))
|
||||||
limit = max(1, min(top_k, 20))
|
dense = [identifier for identifier, _ in heapq.nlargest(DENSE_CANDIDATES, scores, key=lambda item: item[1])]
|
||||||
return self.rerank(query, heapq.nlargest(20, rows, key=lambda item: item["score"]), limit)
|
terms = list(dict.fromkeys(re.findall(r"[^\W_]+", query)))[:128]
|
||||||
|
expression = " OR ".join(f'"{term}"' for term in terms)
|
||||||
|
lexical = [row[0] for row in self.db.execute(
|
||||||
|
"select rowid from lexical where lexical match ? order by bm25(lexical), rowid limit ?",
|
||||||
|
(expression, LEXICAL_CANDIDATES))] if terms else []
|
||||||
|
identifiers = reciprocal_rank_fusion(dense, lexical)[:RERANK_CANDIDATES]
|
||||||
|
placeholders = ",".join("?" for _ in identifiers)
|
||||||
|
rows = {identifier: {"path": path, "symbol": symbol, "text": text}
|
||||||
|
for identifier, path, symbol, text in self.db.execute(
|
||||||
|
f"select rowid, path, symbol, text from chunks where rowid in ({placeholders})", identifiers)}
|
||||||
|
candidates = [rows[identifier] for identifier in identifiers]
|
||||||
|
return self.rerank(query, candidates, max(1, min(top_k, FINAL_RESULTS)))
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
|
|||||||
Reference in new issue
Block a user