اعجبني 0
اشتقاق 0

آخر نشاط 3 weeks ago

Script completo de búsqueda híbrida para la base de conocimiento RAG. Combina FTS5 (BM25) con cosine similarity sobre embeddings de bge-m3. Parámetro alpha para ajustar el peso entre búsqueda exacta y semántica. Episodio 821 de atareao con Linux.

atareao's Avatar atareao عدّل هذا المقطع 3 weeks ago. الانتقال إلى التعديل

1 file changed, 150 insertions

rag_hybrid.py(تم إنشاء الملف)

@@ -0,0 +1,150 @@
1 + #!/usr/bin/env python3
2 + """
3 + rag_hybrid.py — Búsqueda híbrida (FTS5 + cosine similarity) para la base de conocimiento RAG.
4 +
5 + Combina lo mejor de dos mundos:
6 + - FTS5: búsqueda textual exacta con stemming y diacríticos
7 + - Embeddings: búsqueda semántica por cosine similarity
8 +
9 + La fórmula de combinación es:
10 + score_total = alpha * sigmoid(fts_score/10) + (1 - alpha) * cosine_sim
11 +
12 + Uso:
13 + python3 rag_hybrid.py "cómo configurar un proxy inverso"
14 + python3 rag_hybrid.py "hooks de git" --alpha 0.6 --limite 10
15 + """
16 +
17 + import argparse
18 + import math
19 + import os
20 + import sqlite3
21 + import sys
22 + from typing import Any
23 +
24 + import numpy as np
25 + from numpy.linalg import norm
26 +
27 + DB_PATH = os.path.expanduser("~/.cerebro/rag_conocimiento.db")
28 + EMBEDDING_DIM = 1024
29 + EMBEDDING_BYTES = EMBEDDING_DIM * 4 # 4096 bytes
30 +
31 +
32 + def blob_to_vector(blob: bytes) -> np.ndarray:
33 + if len(blob) != EMBEDDING_BYTES:
34 + raise ValueError(f"BLOB incorrecto: {len(blob)} bytes (esperados {EMBEDDING_BYTES})")
35 + return np.frombuffer(blob, dtype=np.float32).copy()
36 +
37 +
38 + def normalizar(vector: np.ndarray) -> np.ndarray:
39 + norma = norm(vector)
40 + return vector / norma if norma > 0 else vector
41 +
42 +
43 + def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
44 + return float(np.dot(a, b))
45 +
46 +
47 + def sigmoid_fts(fts_score: float) -> float:
48 + return 1.0 / (1.0 + math.exp(-fts_score / 10.0))
49 +
50 +
51 + def obtener_embedding_ollama(texto: str) -> np.ndarray:
52 + import requests
53 + response = requests.post(
54 + "http://localhost:11434/api/embeddings",
55 + json={"model": "bge-m3", "prompt": texto},
56 + timeout=30,
57 + )
58 + response.raise_for_status()
59 + return np.array(response.json()["embedding"], dtype=np.float32)
60 +
61 +
62 + def hybrid_search(query: str, alpha: float = 0.4, top_k: int = 10, tag: str | None = None) -> list[dict]:
63 + query_emb = obtener_embedding_ollama(query)
64 + query_emb = normalizar(query_emb)
65 +
66 + if not os.path.exists(DB_PATH):
67 + print(f"Error: BD no encontrada en {DB_PATH}", file=sys.stderr)
68 + return []
69 +
70 + conn = sqlite3.connect(DB_PATH)
71 + conn.row_factory = sqlite3.Row
72 +
73 + where_clauses: list[str] = ["chunks_fts MATCH ?"]
74 + params: list[Any] = [query]
75 +
76 + if tag:
77 + where_clauses.append("c.tags LIKE ?")
78 + params.append(f"%{tag}%")
79 +
80 + sql = f"""
81 + SELECT c.id, c.content, e.vector, c.doc_id,
82 + bm25(chunks_fts) as fts_score
83 + FROM chunks_fts
84 + JOIN chunks c ON chunks_fts.rowid = c.id
85 + JOIN embeddings e ON e.chunk_id = c.id
86 + WHERE {" AND ".join(where_clauses)}
87 + ORDER BY fts_score DESC
88 + LIMIT ?
89 + """
90 + params.append(top_k * 2)
91 +
92 + cursor = conn.execute(sql, params)
93 + rows = cursor.fetchall()
94 + conn.close()
95 +
96 + if not rows:
97 + print(f"Sin resultados para: {query}")
98 + return []
99 +
100 + results = []
101 + for row in rows:
102 + try:
103 + emb_vec = blob_to_vector(row["vector"])
104 + except ValueError as e:
105 + print(f"[aviso] {e}", file=sys.stderr)
106 + continue
107 +
108 + emb_vec = normalizar(emb_vec)
109 + cos_sim = cosine_similarity(query_emb, emb_vec)
110 + fts_norm = sigmoid_fts(row["fts_score"])
111 + score_total = alpha * fts_norm + (1.0 - alpha) * cos_sim
112 +
113 + results.append({
114 + "contenido": row["content"],
115 + "doc_id": row["doc_id"],
116 + "score": score_total,
117 + "fts_score": round(row["fts_score"], 4),
118 + "cos_sim": round(cos_sim, 4),
119 + })
120 +
121 + results.sort(key=lambda x: x["score"], reverse=True)
122 + return results[:top_k]
123 +
124 +
125 + def main():
126 + parser = argparse.ArgumentParser(description="Búsqueda híbrida FTS5 + embeddings")
127 + parser.add_argument("consulta", help="Texto de la consulta")
128 + parser.add_argument("--alpha", type=float, default=0.4, help="Peso de FTS5 (0-1)")
129 + parser.add_argument("--limite", type=int, default=10, help="Número de resultados")
130 + parser.add_argument("--tag", help="Filtrar por etiqueta")
131 + args = parser.parse_args()
132 +
133 + resultados = hybrid_search(query=args.consulta, alpha=args.alpha, top_k=args.limite, tag=args.tag)
134 + if not resultados:
135 + sys.exit(0)
136 +
137 + print(f"\n{= * 70}")
138 + print(f" Consulta: {args.consulta} | Alpha: {args.alpha} | Resultados: {len(resultados)}")
139 + print(f"{= * 70}\n")
140 +
141 + for i, r in enumerate(resultados, 1):
142 + preview = r["contenido"].replace("\n", " ")[:120]
143 + print(f" [{i:2d}] Score: {r[score]:.4f} (FTS: {r[fts_score]:.4f} | Cos: {r[cos_sim]:.4f})")
144 + print(f" {preview}\n")
145 +
146 + print(f"{= * 70}")
147 +
148 +
149 + if __name__ == "__main__":
150 + main()
أحدث أقدم

مدعوم بواسطة Opengist ⋅ Load: 25ms⋅

العربية