atareao revidoval tento gist 3 weeks ago. Přejít na revizi
1 file changed, 109 insertions
rag_hyde.py(vytvořil soubor)
| @@ -0,0 +1,109 @@ | |||
| 1 | + | #!/usr/bin/env python3 | |
| 2 | + | """ | |
| 3 | + | rag_hyde.py - HyDE (Hypothetical Document Embeddings) para RAG | |
| 4 | + | ||
| 5 | + | Genera un documento hipotético con Ollama a partir de una consulta, | |
| 6 | + | lo embeddea con bge-m3, y busca por similitud coseno en la BD. | |
| 7 | + | ||
| 8 | + | Uso: | |
| 9 | + | python rag_hyde.py "¿qué es una variable de entorno?" | |
| 10 | + | python rag_hyde.py "explica la configuración de red" --top 10 | |
| 11 | + | """ | |
| 12 | + | ||
| 13 | + | import argparse | |
| 14 | + | import sqlite3 | |
| 15 | + | import sys | |
| 16 | + | import numpy as np | |
| 17 | + | from pathlib import Path | |
| 18 | + | ||
| 19 | + | DB_PATH = Path.home() / ".cerebro" / "rag_conocimiento.db" | |
| 20 | + | OLLAMA_MODEL = "llama3.2" | |
| 21 | + | EMBED_MODEL = "bge-m3" | |
| 22 | + | ||
| 23 | + | HYDE_PROMPT = """Eres un asistente experto en Linux y administración de sistemas. | |
| 24 | + | Dada una pregunta, genera un documento hipotético de 3-5 párrafos que RESPONDA | |
| 25 | + | a la pregunta como si fuera un artículo del blog o la documentación oficial. | |
| 26 | + | ||
| 27 | + | Sé técnico pero claro. Usa markdown si es necesario. | |
| 28 | + | No digas "basado en" ni "hipotético". Responde DIRECTAMENTE. | |
| 29 | + | ||
| 30 | + | Pregunta: {consulta} | |
| 31 | + | ||
| 32 | + | Documento:""" | |
| 33 | + | ||
| 34 | + | ||
| 35 | + | def _ollama(): | |
| 36 | + | try: | |
| 37 | + | import ollama | |
| 38 | + | except ImportError: | |
| 39 | + | print("Necesitas instalar ollama: pip install ollama") | |
| 40 | + | sys.exit(1) | |
| 41 | + | return ollama | |
| 42 | + | ||
| 43 | + | ||
| 44 | + | def generar_doc_hipotetico(consulta: str) -> str: | |
| 45 | + | ollama = _ollama() | |
| 46 | + | prompt = HYDE_PROMPT.format(consulta=consulta) | |
| 47 | + | respuesta = ollama.generate(model=OLLAMA_MODEL, prompt=prompt) | |
| 48 | + | return respuesta["response"] | |
| 49 | + | ||
| 50 | + | ||
| 51 | + | def embed_text(texto: str) -> list[float]: | |
| 52 | + | ollama = _ollama() | |
| 53 | + | respuesta = ollama.embed(model=EMBED_MODEL, input=texto) | |
| 54 | + | return respuesta["embeddings"][0] | |
| 55 | + | ||
| 56 | + | ||
| 57 | + | def cosine_sim(a: np.ndarray, b: np.ndarray) -> float: | |
| 58 | + | return float(np.dot(a, b)) | |
| 59 | + | ||
| 60 | + | ||
| 61 | + | def cargar_embeddings_db() -> list[tuple[int, str, np.ndarray]]: | |
| 62 | + | if not DB_PATH.exists(): | |
| 63 | + | print(f"BD no encontrada: {DB_PATH}") | |
| 64 | + | sys.exit(1) | |
| 65 | + | conn = sqlite3.connect(str(DB_PATH)) | |
| 66 | + | cursor = conn.cursor() | |
| 67 | + | cursor.execute("""SELECT c.id, c.content, e.vector | |
| 68 | + | FROM chunks c JOIN embeddings e ON e.chunk_id = c.id""") | |
| 69 | + | filas = cursor.fetchall() | |
| 70 | + | conn.close() | |
| 71 | + | if not filas: | |
| 72 | + | print("No hay chunks con embeddings en la BD.") | |
| 73 | + | sys.exit(1) | |
| 74 | + | return [(chunk_id, contenido, np.frombuffer(blob, dtype=np.float32).copy()) | |
| 75 | + | for chunk_id, contenido, blob in filas] | |
| 76 | + | ||
| 77 | + | ||
| 78 | + | def buscar_semantica(query_emb: np.ndarray, chunks: list, top_k: int = 5) -> list: | |
| 79 | + | query_norm = query_emb / (np.linalg.norm(query_emb) + 1e-10) | |
| 80 | + | scores = [] | |
| 81 | + | for chunk_id, contenido, vector in chunks: | |
| 82 | + | vec_norm = vector / (np.linalg.norm(vector) + 1e-10) | |
| 83 | + | sim = cosine_sim(query_norm, vec_norm) | |
| 84 | + | scores.append((chunk_id, contenido, sim)) | |
| 85 | + | scores.sort(key=lambda x: x[2], reverse=True) | |
| 86 | + | return scores[:top_k] | |
| 87 | + | ||
| 88 | + | ||
| 89 | + | def main(): | |
| 90 | + | parser = argparse.ArgumentParser(description="HyDE: genera un doc hipotético y busca en la BD") | |
| 91 | + | parser.add_argument("consulta", help="Consulta en lenguaje natural") | |
| 92 | + | parser.add_argument("--top", type=int, default=5, help="Resultados a mostrar") | |
| 93 | + | args = parser.parse_args() | |
| 94 | + | ||
| 95 | + | doc_hipotetico = generar_doc_hipotetico(args.consulta) | |
| 96 | + | print(f"Documento generado:\n{─ * 50}\n{doc_hipotetico[:500]}...\n{─ * 50}\n") | |
| 97 | + | ||
| 98 | + | query_emb = np.array(embed_text(doc_hipotetico), dtype=np.float32) | |
| 99 | + | chunks = cargar_embeddings_db() | |
| 100 | + | resultados = buscar_semantica(query_emb, chunks, args.top) | |
| 101 | + | ||
| 102 | + | for i, (chunk_id, contenido, score) in enumerate(resultados, 1): | |
| 103 | + | preview = contenido[:150].replace("\n", " ") | |
| 104 | + | print(f" {i}. [{score:.4f}] (chunk #{chunk_id})") | |
| 105 | + | print(f" {preview}...") | |
| 106 | + | ||
| 107 | + | ||
| 108 | + | if __name__ == "__main__": | |
| 109 | + | main() | |
Novější
Starší