Rewizja rag_reranker.py — Reranking con cross-encoder bge-reranker-v2-m3

1

+

#!/usr/bin/env python3

2

+

"""

3

+

rag_reranker.py — Reranking con cross-encoder (bge-reranker-v2-m3).

4

+

5

+

Pipeline:

6

+

1. Hybrid search → top 20 candidatos

7

+

2. Cross-encoder reranker → top 5 finales

8

+

9

+

Uso:

10

+

python3 rag_reranker.py "cómo configurar certificados SSL"

11

+

"""

12

+

13

+

import argparse

14

+

import os

15

+

import sys

16

+

import sqlite3

17

+

import math

18

+

19

+

import numpy as np

20

+

from numpy.linalg import norm

21

+

22

+

DB_PATH = os.path.expanduser("~/.cerebro/rag_conocimiento.db")

23

+

EMBEDDING_BYTES = 1024 * 4

24

+

MODELO_RERANKER = "BAAI/bge-reranker-v2-m3"

25

+

26

+

27

+

def obtener_embedding(texto: str) -> np.ndarray | None:

28

+

import requests

29

+

try:

30

+

resp = requests.post(

31

+

"http://localhost:11434/api/embeddings",

32

+

json={"model": "bge-m3", "prompt": texto}, timeout=15,

33

+

)

34

+

resp.raise_for_status()

35

+

return np.array(resp.json()["embedding"], dtype=np.float32)

36

+

except Exception as e:

37

+

print(f"[aviso] Error: {e}", file=sys.stderr)

38

+

return None

39

+

40

+

41

+

def hybrid_search(query: str, alpha: float = 0.4, top_k: int = 20) -> list[dict]:

42

+

if not os.path.exists(DB_PATH):

43

+

print(f"Error: BD no encontrada", file=sys.stderr)

44

+

return []

45

+

query_emb = obtener_embedding(query)

46

+

if query_emb is None:

47

+

return []

48

+

query_emb = query_emb / (norm(query_emb) + 1e-10)

49

+

50

+

conn = sqlite3.connect(DB_PATH)

51

+

conn.row_factory = sqlite3.Row

52

+

cursor = conn.execute(

53

+

"""SELECT c.id, c.content, e.vector, c.doc_id,

54

+

bm25(chunks_fts) as fts_score

55

+

FROM chunks_fts JOIN chunks c ON chunks_fts.rowid = c.id

56

+

JOIN embeddings e ON e.chunk_id = c.id

57

+

WHERE chunks_fts MATCH ?

58

+

ORDER BY fts_score DESC LIMIT ?""",

59

+

(query, top_k * 2),

60

+

)

61

+

rows = cursor.fetchall()

62

+

conn.close()

63

+

64

+

results = []

65

+

for row in rows:

66

+

emb_vec = np.frombuffer(row["vector"], dtype=np.float32).copy()

67

+

emb_vec = emb_vec / (norm(emb_vec) + 1e-10)

68

+

cos_sim = float(np.dot(query_emb, emb_vec))

69

+

fts_norm = 1.0 / (1.0 + math.exp(-row["fts_score"] / 10.0))

70

+

score_total = alpha * fts_norm + (1.0 - alpha) * cos_sim

71

+

results.append({

72

+

"id": row["id"], "contenido": row["content"],

73

+

"doc_id": row["doc_id"], "score": score_total,

74

+

})

75

+

results.sort(key=lambda x: x["score"], reverse=True)

76

+

return results[:top_k]

77

+

78

+

79

+

class Reranker:

80

+

def __init__(self, model_name: str = MODELO_RERANKER, use_fp16: bool = True):

81

+

from sentence_transformers import CrossEncoder

82

+

self.model = CrossEncoder(model_name, max_length=512, device="cpu")

83

+

84

+

def rerank(self, query: str, candidates: list[dict], top_k: int = 5) -> list[dict]:

85

+

if not candidates:

86

+

return []

87

+

pairs = [(query, c["contenido"]) for c in candidates]

88

+

scores = self.model.predict(pairs)

89

+

for i, score in enumerate(scores):

90

+

candidates[i]["rerank_score"] = float(score)

91

+

candidates.sort(key=lambda x: x["rerank_score"], reverse=True)

92

+

return candidates[:top_k]

93

+

94

+

95

+

def search_with_rerank(query: str, alpha: float = 0.4, top_k_hybrid: int = 20, top_k_final: int = 5):

96

+

candidates = hybrid_search(query, alpha=alpha, top_k=top_k_hybrid)

97

+

if not candidates:

98

+

return []

99

+

reranker = Reranker()

100

+

return reranker.rerank(query, candidates, top_k=top_k_final)

101

+

102

+

103

+

def main():

104

+

parser = argparse.ArgumentParser(description="Reranking con cross-encoder")

105

+

parser.add_argument("consulta", help="Texto de la consulta")

106

+

parser.add_argument("-k", "--top-k", type=int, default=5, help="Resultados finales")

107

+

parser.add_argument("-c", "--candidatos", type=int, default=20, help="Candidatos para reranking")

108

+

args = parser.parse_args()

109

+

110

+

results = search_with_rerank(query=args.consulta, top_k_hybrid=args.candidatos, top_k_final=args.top_k)

111

+

if not results:

112

+

print("Sin resultados.")

113

+

return

114

+

print(f"\nResultados rerankeados (top {len(results)}):")

115

+

for i, r in enumerate(results, 1):

116

+

preview = r["contenido"].replace("\n", " ")[:150]

117

+

print(f" [{i:2d}] Rerank: {r.get(rerank_score, 0):.4f}")

118

+

print(f" {preview}\n")

119

+

120

+

121

+

if __name__ == "__main__":

122

+

main()

atareao / rag_reranker.py — Reranking con cross-encoder bge-reranker-v2-m3

atareao zrewidował ten Gist 3 weeks ago. Przejdź do rewizji

		@@ -0,0 +1,122 @@
1	+	#!/usr/bin/env python3
2	+	"""
3	+	rag_reranker.py — Reranking con cross-encoder (bge-reranker-v2-m3).
4	+
5	+	Pipeline:
6	+	1. Hybrid search → top 20 candidatos
7	+	2. Cross-encoder reranker → top 5 finales
8	+
9	+	Uso:
10	+	python3 rag_reranker.py "cómo configurar certificados SSL"
11	+	"""
12	+
13	+	import argparse
14	+	import os
15	+	import sys
16	+	import sqlite3
17	+	import math
18	+
19	+	import numpy as np
20	+	from numpy.linalg import norm
21	+
22	+	DB_PATH = os.path.expanduser("~/.cerebro/rag_conocimiento.db")
23	+	EMBEDDING_BYTES = 1024 * 4
24	+	MODELO_RERANKER = "BAAI/bge-reranker-v2-m3"
25	+
26	+
27	+	def obtener_embedding(texto: str) -> np.ndarray \| None:
28	+	import requests
29	+	try:
30	+	resp = requests.post(
31	+	"http://localhost:11434/api/embeddings",
32	+	json={"model": "bge-m3", "prompt": texto}, timeout=15,
33	+	)
34	+	resp.raise_for_status()
35	+	return np.array(resp.json()["embedding"], dtype=np.float32)
36	+	except Exception as e:
37	+	print(f"[aviso] Error: {e}", file=sys.stderr)
38	+	return None
39	+
40	+
41	+	def hybrid_search(query: str, alpha: float = 0.4, top_k: int = 20) -> list[dict]:
42	+	if not os.path.exists(DB_PATH):
43	+	print(f"Error: BD no encontrada", file=sys.stderr)
44	+	return []
45	+	query_emb = obtener_embedding(query)
46	+	if query_emb is None:
47	+	return []
48	+	query_emb = query_emb / (norm(query_emb) + 1e-10)
49	+
50	+	conn = sqlite3.connect(DB_PATH)
51	+	conn.row_factory = sqlite3.Row
52	+	cursor = conn.execute(
53	+	"""SELECT c.id, c.content, e.vector, c.doc_id,
54	+	bm25(chunks_fts) as fts_score
55	+	FROM chunks_fts JOIN chunks c ON chunks_fts.rowid = c.id
56	+	JOIN embeddings e ON e.chunk_id = c.id
57	+	WHERE chunks_fts MATCH ?
58	+	ORDER BY fts_score DESC LIMIT ?""",
59	+	(query, top_k * 2),
60	+	)
61	+	rows = cursor.fetchall()
62	+	conn.close()
63	+
64	+	results = []
65	+	for row in rows:
66	+	emb_vec = np.frombuffer(row["vector"], dtype=np.float32).copy()
67	+	emb_vec = emb_vec / (norm(emb_vec) + 1e-10)
68	+	cos_sim = float(np.dot(query_emb, emb_vec))
69	+	fts_norm = 1.0 / (1.0 + math.exp(-row["fts_score"] / 10.0))
70	+	score_total = alpha * fts_norm + (1.0 - alpha) * cos_sim
71	+	results.append({
72	+	"id": row["id"], "contenido": row["content"],
73	+	"doc_id": row["doc_id"], "score": score_total,
74	+	})
75	+	results.sort(key=lambda x: x["score"], reverse=True)
76	+	return results[:top_k]
77	+
78	+
79	+	class Reranker:
80	+	def __init__(self, model_name: str = MODELO_RERANKER, use_fp16: bool = True):
81	+	from sentence_transformers import CrossEncoder
82	+	self.model = CrossEncoder(model_name, max_length=512, device="cpu")
83	+
84	+	def rerank(self, query: str, candidates: list[dict], top_k: int = 5) -> list[dict]:
85	+	if not candidates:
86	+	return []
87	+	pairs = [(query, c["contenido"]) for c in candidates]
88	+	scores = self.model.predict(pairs)
89	+	for i, score in enumerate(scores):
90	+	candidates[i]["rerank_score"] = float(score)
91	+	candidates.sort(key=lambda x: x["rerank_score"], reverse=True)
92	+	return candidates[:top_k]
93	+
94	+
95	+	def search_with_rerank(query: str, alpha: float = 0.4, top_k_hybrid: int = 20, top_k_final: int = 5):
96	+	candidates = hybrid_search(query, alpha=alpha, top_k=top_k_hybrid)
97	+	if not candidates:
98	+	return []
99	+	reranker = Reranker()
100	+	return reranker.rerank(query, candidates, top_k=top_k_final)
101	+
102	+
103	+	def main():
104	+	parser = argparse.ArgumentParser(description="Reranking con cross-encoder")
105	+	parser.add_argument("consulta", help="Texto de la consulta")
106	+	parser.add_argument("-k", "--top-k", type=int, default=5, help="Resultados finales")
107	+	parser.add_argument("-c", "--candidatos", type=int, default=20, help="Candidatos para reranking")
108	+	args = parser.parse_args()
109	+
110	+	results = search_with_rerank(query=args.consulta, top_k_hybrid=args.candidatos, top_k_final=args.top_k)
111	+	if not results:
112	+	print("Sin resultados.")
113	+	return
114	+	print(f"\nResultados rerankeados (top {len(results)}):")
115	+	for i, r in enumerate(results, 1):
116	+	preview = r["contenido"].replace("\n", " ")[:150]
117	+	print(f" [{i:2d}] Rerank: {r.get(rerank_score, 0):.4f}")
118	+	print(f" {preview}\n")
119	+
120	+
121	+	if __name__ == "__main__":
122	+	main()