"""Búsqueda híbrida para un RAG: embeddings (significado) + BM25 (palabras clave), mezclados con RRF.

Es el método de la comparación del vídeo «Cómo lee una IA tus documentos — RAG por dentro» (Frogames).
Uso:  uv run --with sentence-transformers --with rank-bm25 python busqueda_hibrida.py
Cambia TROZOS por los tuyos (trozos de ~1.000 caracteres) y PREGUNTA por la tuya.
"""
import re
import unicodedata

import numpy as np
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

TROZOS = [
    "RAG o ajuste fino: usa RAG cuando el modelo necesita datos que no tiene y ajuste fino cuando quieres cambiar cómo responde.",
    "Para visualizar incrustaciones en dos dimensiones se usan técnicas como t-SNE, que agrupan los textos parecidos.",
    "La cuantización reduce el tamaño del modelo usando menos bits por parámetro.",
    "Weights & Biases sirve para seguir el entrenamiento con gráficos de pérdida y métricas.",
]
PREGUNTA = "¿Cuándo uso RAG y cuándo ajuste fino?"
K = 5          # cuántos trozos le pasas a la IA
RRF_K = 60     # suaviza la fusión: con 60, ir 1.º o 3.º en una lista pesa parecido


def palabras(texto):
    """Minúsculas, sin tildes, solo letras y números: lo que compara BM25."""
    t = unicodedata.normalize("NFKD", texto.lower()).encode("ascii", "ignore").decode()
    return re.findall(r"[a-z0-9]+", t)


modelo = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
emb = modelo.encode(TROZOS, normalize_embeddings=True)            # un vector por trozo (384 números)
q = modelo.encode([PREGUNTA], normalize_embeddings=True)[0]
por_significado = np.argsort(-(emb @ q))                           # coseno = producto escalar (vectores normalizados)

bm25 = BM25Okapi([palabras(t) for t in TROZOS])
por_palabras = np.argsort(-bm25.get_scores(palabras(PREGUNTA)))

puntos = np.zeros(len(TROZOS))
for ranking in (por_significado, por_palabras):                    # RRF: 1 / (RRF_K + puesto) en cada lista
    for puesto, i in enumerate(ranking, start=1):
        puntos[i] += 1 / (RRF_K + puesto)
hibrida = np.argsort(-puntos)

for nombre, ranking in [("embeddings", por_significado), ("BM25", por_palabras), ("híbrida", hibrida)]:
    print(f"\n{nombre}:")
    for puesto, i in enumerate(ranking[:K], start=1):
        print(f"  {puesto}. {TROZOS[i][:70]}")