atareao / rag_samples.py — Hybrid search, reranker y HyDE en Python

#!/usr/bin/env python3 """rag_hybrid.py — Hybrid search: FTS5 + cosine similarity""" import sqlite3, numpy as np, math, os from numpy.linalg import norm DB_PATH = os.path.expanduser("~/.cerebro/rag_conocimiento.db") def blob_to_vector(blob: bytes) -> np.ndarray: return np.frombuffer(blob, dtype=np.float32).copy() def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: return float(np.dot(a, b)) def get_embedding_ollama(text: str) -> np.ndarray: import requests resp = requests.post("http://localhost:11434/api/embeddings", json={"model": "bge-m3", "input": text}) return np.array(resp.json()["embeddings"][0], dtype=np.float32) def hybrid_search(query: str, alpha: float = 0.4, top_k: int = 10): """ alpha: peso para FTS5, (1-alpha) para cosine similarity. 0.6 = nombres propios, 0.2 = conceptual, 0.4 = equilibrio """ conn = sqlite3.connect(DB_PATH) query_emb = get_embedding_ollama(query) query_emb = query_emb / (norm(query_emb) + 1e-10) cursor = conn.execute(""" SELECT c.id, c.contenido, c.embedding, bm25(chunks_fts, 0.0, 0.0, 5.0, 5.0) as fts_score FROM chunks_fts JOIN chunks c ON chunks_fts.rowid = c.id WHERE chunks_fts MATCH ? ORDER BY fts_score DESC LIMIT ? """, (query, top_k * 2)) results = [] for row in cursor: chunk_id, contenido, emb_blob, fts_score = row fts_norm = 1.0 / (1.0 + math.exp(-fts_score / 10.0)) emb_vec = blob_to_vector(emb_blob) emb_vec = emb_vec / (norm(emb_vec) + 1e-10) cos_sim = cosine_similarity(query_emb, emb_vec) score_total = alpha * fts_norm + (1.0 - alpha) * cos_sim results.append((score_total, contenido[:200], fts_norm, cos_sim)) results.sort(key=lambda x: x[0], reverse=True) conn.close() return results[:top_k]

#!/usr/bin/env python3 """rag_hyde.py — HyDE: Hypothetical Document Embeddings""" import requests, numpy as np, sqlite3, os, math from numpy.linalg import norm DB_PATH = os.path.expanduser("~/.cerebro/rag_conocimiento.db") OLLAMA_URL = "http://localhost:11434/api" def generate_hypothetical_document(query: str, model: str ="llama3.2") -> str: prompt = f"""Genera un texto breve y factual que responda a la siguiente pregunta. El texto debe ser informativo, objetivo y escrito en español. Pregunta: {query} Texto informativo:""" resp = requests.post(f"{OLLAMA_URL}/generate", json={"model": model, "prompt": prompt, "stream": False}) return resp.json()["response"] def get_embedding(text: str, model: str ="bge-m3") -> np.ndarray: resp = requests.post(f"{OLLAMA_URL}/embed", json={"model": model, "input": text}) return np.array(resp.json()["embeddings"][0], dtype=np.float32) def blob_to_vector(blob: bytes) -> np.ndarray: return np.frombuffer(blob, dtype=np.float32).copy() def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: return float(np.dot(a, b)) def hyde_search(query: str, top_k: int = 5): hypothetic_doc = generate_hypothetical_document(query) print(f"[HyDE] Documento generado: {hypothetic_doc[:100]}...") hyde_emb = get_embedding(hypothetic_doc) hyde_emb = hyde_emb / (norm(hyde_emb) + 1e-10) conn = sqlite3.connect(DB_PATH) cursor = conn.execute("SELECT id, contenido, embedding FROM chunks") results = [] for row in cursor: contenido, emb_blob = row[1], row[2] emb_vec = blob_to_vector(emb_blob) emb_vec = emb_vec / (norm(emb_vec) + 1e-10) cos_sim = cosine_similarity(hyde_emb, emb_vec) results.append((cos_sim, contenido[:200])) results.sort(key=lambda x: x[0], reverse=True) conn.close() return results[:top_k] def hyde_hybrid_search(query: str, alpha: float = 0.3, top_k: int = 5): """HyDE + Hybrid search combinado. Alpha bajo porque HyDE ya captura semántica.""" hyde_doc = generate_hypothetical_document(query) hyde_emb = get_embedding(hyde_doc) hyde_emb = hyde_emb / (norm(hyde_emb) + 1e-10) conn = sqlite3.connect(DB_PATH) cursor = conn.execute(""" SELECT c.contenido, c.embedding, bm25(chunks_fts, 0, 0, 5, 5) as fts_score FROM chunks_fts JOIN chunks c ON chunks_fts.rowid = c.id WHERE chunks_fts MATCH ? """, (query,)) results = [] for row in cursor: contenido, emb_blob, fts_score = row emb_vec = blob_to_vector(emb_blob) emb_vec = emb_vec / (norm(emb_vec) + 1e-10) cos_sim = cosine_similarity(hyde_emb, emb_vec) fts_norm = 1.0 / (1.0 + math.exp(-fts_score / 10.0)) score_total = alpha * fts_norm + (1.0 - alpha) * cos_sim results.append((score_total, contenido[:200])) results.sort(reverse=True) conn.close() return results[:top_k]

#!/usr/bin/env python3 """rag_reranker.py — Re-ranking con cross-encoder""" from sentence_transformers import CrossEncoder class Reranker: def __init__(self, model_name="BAAI/bge-reranker-v2-m3", use_fp16=True): self.model = CrossEncoder(model_name, max_length=512, device="cpu") def rerank(self, query: str, candidates: list[dict], top_k: int = 5) -> list[dict]: pairs = [(query, c["contenido"]) for c in candidates] scores = self.model.predict(pairs) for i, score in enumerate(scores): candidates[i]["rerank_score"] = float(score) candidates.sort(key=lambda x: x["rerank_score"], reverse=True) return candidates[:top_k] def search_with_rerank(query: str, alpha: float = 0.4, top_k_hybrid: int = 20, top_k_final: int = 5): from rag_hybrid import hybrid_search candidates = hybrid_search(query, alpha=alpha, top_k=top_k_hybrid) candidates_dict = [{"contenido": c[1], "score": c[0]} for c in candidates] reranker = Reranker() return reranker.rerank(query, candidates_dict, top_k=top_k_final)