Embeddings y bases vectoriales: FAISS, Chroma, pgvector y Qdrant.

import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer

st = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2",
                         device="cpu", local_files_only=True)
frases = [
    "La empresa eleva su previsión de ventas para el año",
    "La compañía mejora sus perspectivas de ingresos anuales",
    "La empresa recorta su previsión de ventas para el año",
    "The company raises its full-year sales guidance",
    "El consejo aprueba el pago de un dividendo",
]
V = st.encode(frases, normalize_embeddings=True,
              show_progress_bar=False)
print(V.shape, V.dtype, "· norma:", np.linalg.norm(V, axis=1).round(3))
sim = pd.DataFrame(V @ V.T, index=["eleva", "mejora", "recorta",
                                   "raises", "dividendo"],
                   columns=["eleva", "mejora", "recorta", "raises",
                            "dividendo"])
print(sim.round(2).to_string())

pas = pd.read_parquet("rag6_pasajes.parquet")
docs = pd.read_parquet("rag6_documentos.parquet")
preg = pd.read_csv("rag6_preguntas.csv")
docs["corto"] = (docs["entidad"].str.replace(", S.A.", "")
                 .str.replace(" UCITS ETF", ""))
docs["anio"] = docs["titulo"].str.extract(r"(20\d\d)$")[0].fillna("")
pas = pas.merge(docs[["id_doc", "entidad", "anio"]], on="id_doc")
pas = pas[pas["tipo"] != "titulo"].reset_index(drop=True)
preg = preg[preg["respondible"]].reset_index(drop=True)
oro = [set(x.split(";")) for x in preg["pasajes_oro"]]
print(len(docs), "documentos ·", len(pas), "pasajes ·", len(preg),
      "preguntas respondibles")
print(pas["tipo"].value_counts().to_string())

import faiss

faiss.omp_set_num_threads(1)
P = st.encode(list(pas["texto"]), batch_size=32,
              normalize_embeddings=True, show_progress_bar=False)
Q = st.encode(list(preg["pregunta"]), batch_size=64,
              normalize_embeddings=True, show_progress_bar=False)
indice = faiss.IndexFlatIP(P.shape[1])   # producto interior = coseno
indice.add(P)
print(indice.ntotal, "vectores ·", P.nbytes // 1024, "KB en float32")
ids = pas["id_pasaje"].to_numpy()

def recall(rankings, k):
    """Fracción de los pasajes de oro que están entre los k primeros."""
    return round(float(np.mean([len(set(ids[r[:k]]) & g) / len(g)
                                for r, g in zip(rankings, oro)])), 3)

_, denso = indice.search(Q, 20)
print({f"recall@{k}": recall(denso, k) for k in (1, 5, 10)})

i = 0
print(preg.loc[i, "pregunta"], "· oro:", preg.loc[i, "pasajes_oro"])
for j in denso[i][:4]:
    print(f"  {ids[j]}  {pas.loc[j, 'entidad'][:22]:22s} "
          f"{pas.loc[j, 'texto'][:52]!r}")

import re
import unicodedata
from nltk.corpus import stopwords
from rank_bm25 import BM25Okapi

VACIAS = set(stopwords.words("spanish"))

def tokens(t):
    t = unicodedata.normalize("NFKD", t.lower())
    t = "".join(c for c in t if not unicodedata.combining(c))
    return [w for w in re.findall(r"\w+", t) if w not in VACIAS]

bm25 = BM25Okapi([tokens(t) for t in pas["texto"]])

def rank_bm25(consultas, mascaras=None):
    S = np.array([bm25.get_scores(tokens(c)) for c in consultas])
    if mascaras is not None:
        S[~mascaras] = -np.inf
    return np.argsort(-S, axis=1, kind="stable")[:, :20]

def rrf(*rankings, k0=60, n=20):
    """Reciprocal rank fusion: suma de 1 / (k0 + posición)."""
    salida = []
    for rs in zip(*rankings):
        p = {}
        for r in rs:
            for pos, j in enumerate(r):
                p[j] = p.get(j, 0) + 1 / (k0 + pos + 1)
        salida.append(sorted(p, key=lambda j: (-p[j], j))[:n])
    return np.array(salida)

lexico = rank_bm25(preg["pregunta"])
hibrido = rrf(denso, lexico)
tabla = pd.DataFrame({m: [recall(r, k) for k in (1, 5, 10)]
                      for m, r in [("denso", denso), ("BM25", lexico),
                                   ("híbrido", hibrido)]},
                     index=["recall@1", "recall@5", "recall@10"])
print(tabla.round(3).to_string())

import chromadb

cliente = chromadb.EphemeralClient()
col = cliente.create_collection(
    "pasajes", embedding_function=None,
    configuration={"hnsw": {"space": "cosine"}})
col.add(ids=list(ids), embeddings=P, documents=list(pas["texto"]),
        metadatas=[{"id_doc": a, "entidad": b, "anio": c}
                   for a, b, c in pas[["id_doc", "entidad", "anio"]]
                   .itertuples(index=False)])

def filtro(pregunta):
    """Entidad citada en la pregunta y, si hay, el año más reciente."""
    m = docs[[c in pregunta for c in docs["corto"]]]
    anios = re.findall(r"20\d\d", pregunta)
    if anios and (m["anio"] == max(anios)).any():
        m = m[m["anio"] == max(anios)]
    return m

def where(m):
    if m.empty:
        return None
    return {"id_doc": {"$in": list(m["id_doc"])}}

r = col.query(query_embeddings=Q[:1], n_results=3,
              where=where(filtro(preg.loc[0, "pregunta"])),
              include=["distances"])
print(r["ids"][0], np.round(r["distances"][0], 3))

pos = {p: i for i, p in enumerate(ids)}

def rank_chroma(E, consultas, n=20):
    salida = []
    for e, c in zip(E, consultas):
        w = where(filtro(c))
        k = n if w is None else min(n, int(pas["id_doc"].isin(
            w["id_doc"]["$in"]).sum()))
        r = col.query(query_embeddings=e[None], n_results=k, where=w,
                      include=[])
        fila = [pos[x] for x in r["ids"][0]]
        salida.append(fila + [fila[-1]] * (n - len(fila)))
    return np.array(salida)

mascaras = np.array([
    pas["id_doc"].isin(filtro(c)["id_doc"]).to_numpy()
    if not filtro(c).empty else np.ones(len(pas), bool)
    for c in preg["pregunta"]])
denso_f = rank_chroma(Q, preg["pregunta"])
lexico_f = rank_bm25(preg["pregunta"], mascaras)
hibrido_f = rrf(denso_f, lexico_f)
for m, r in [("denso", denso_f), ("BM25", lexico_f),
             ("híbrido", hibrido_f)]:
    print(f"{m:8s}", {f"@{k}": recall(r, k) for k in (1, 5, 10)})

def limpia(p):
    for c in docs["corto"].unique():
        p = p.replace(c, "")
    return p

limpias = [limpia(p) for p in preg["pregunta"]]
Q2 = st.encode(limpias, batch_size=64, normalize_embeddings=True,
               show_progress_bar=False)
denso_l = rank_chroma(Q2, preg["pregunta"])
lexico_l = rank_bm25(limpias, mascaras)
hibrido_l = rrf(denso_l, lexico_l)
for m, r in [("denso", denso_l), ("BM25", lexico_l),
             ("híbrido", hibrido_l)]:
    print(f"{m:8s}", {f"@{k}": recall(r, k) for k in (1, 5, 10)})
porpas = lambda r: [len(set(ids[x[:5]]) & g) / len(g)
                    for x, g in zip(r, oro)]
tipo = (pd.DataFrame({"tipo": preg["tipo_pregunta"],
                      "denso": porpas(denso_l),
                      "BM25": porpas(lexico_l),
                      "híbrido": porpas(hibrido_l)})
        .groupby("tipo")
        .agg(n=("denso", "size"), denso=("denso", "mean"),
             BM25=("BM25", "mean"), híbrido=("híbrido", "mean")))
print(tipo.round(2).sort_values("n", ascending=False).to_string())

rng = np.random.default_rng(0)
N, dim, grupos = 50_000, 384, 5_000
centros = rng.standard_normal((grupos, dim)).astype("float32")
X = centros[rng.integers(0, grupos, N)] + rng.standard_normal(
    (N, dim)).astype("float32")
X /= np.linalg.norm(X, axis=1, keepdims=True)
Xq = centros[rng.integers(0, grupos, 500)] + rng.standard_normal(
    (500, dim)).astype("float32")
Xq /= np.linalg.norm(Xq, axis=1, keepdims=True)
exacto = faiss.IndexFlatIP(dim)
exacto.add(X)
_, verdad = exacto.search(Xq, 10)

def r10(I):
    return np.mean([len(set(a) & set(b)) / 10
                    for a, b in zip(I, verdad)])

cuantif = faiss.IndexFlatIP(dim)
ivf = faiss.IndexIVFFlat(cuantif, dim, 256, faiss.METRIC_INNER_PRODUCT)
ivf.cp.seed = 0
ivf.train(X[:20_000])
ivf.add(X)
tam = np.array([ivf.invlists.list_size(i) for i in range(256)])
filas = [("exacto", "-", 1.0, N)]
for nprobe in (1, 4, 16, 64):
    ivf.nprobe = nprobe
    _, I = ivf.search(Xq, 10)
    _, listas = cuantif.search(Xq, nprobe)
    filas.append(("IVF", f"nprobe={nprobe}", r10(I),
                  256 + tam[listas].sum(1).mean()))
hnsw = faiss.IndexHNSWFlat(dim, 16, faiss.METRIC_INNER_PRODUCT)
hnsw.hnsw.efConstruction = 40
hnsw.add(X)
for ef in (16, 32, 64, 128, 256):
    hnsw.hnsw.efSearch = ef
    faiss.cvar.hnsw_stats.reset()
    _, I = hnsw.search(Xq, 10)
    filas.append(("HNSW", f"efSearch={ef}", r10(I),
                  faiss.cvar.hnsw_stats.ndis / len(Xq)))
pq = faiss.IndexIVFPQ(faiss.IndexFlatIP(dim), dim, 256, 48, 8,
                      faiss.METRIC_INNER_PRODUCT)
pq.cp.seed = 0
pq.train(X[:20_000])
pq.add(X)
tam_pq = np.array([pq.invlists.list_size(i) for i in range(256)])
for nprobe in (16, 64):
    pq.nprobe = nprobe
    _, I = pq.search(Xq, 10)
    _, listas = pq.quantizer.search(Xq, nprobe)
    filas.append(("IVF-PQ", f"nprobe={nprobe}", r10(I),
                  256 + tam_pq[listas].sum(1).mean()))
ann = pd.DataFrame(filas, columns=["índice", "parámetro", "recall@10",
                                   "distancias por consulta"])
ann["distancias por consulta"] = ann["distancias por consulta"].round()
print(ann.round(3).astype({"distancias por consulta": int})
      .to_string(index=False))
print("bytes por vector: exacto", dim * 4, "· IVF-PQ", pq.code_size)

# pgvector: PostgreSQL con la extensión vector (no se ejecuta)
# SQL, una vez:
#   CREATE EXTENSION vector;
#   CREATE TABLE pasajes (id text PRIMARY KEY, id_doc text,
#                         texto text, embedding vector(384));
#   CREATE INDEX ON pasajes USING hnsw (embedding vector_cosine_ops);
import psycopg
from pgvector.psycopg import register_vector

with psycopg.connect("dbname=research") as conn:
    register_vector(conn)
    for (i, doc, txt), e in zip(pas[["id_pasaje", "id_doc", "texto"]]
                                .itertuples(index=False), P):
        conn.execute("INSERT INTO pasajes VALUES (%s, %s, %s, %s)",
                     (i, doc, txt, e))
    conn.execute("SET hnsw.ef_search = 100")
    conn.execute("SET hnsw.iterative_scan = relaxed_order")
    filas = conn.execute(
        "SELECT id, texto FROM pasajes WHERE id_doc = %s "
        "ORDER BY embedding <=> %s LIMIT 5", ("R01", Q[0])).fetchall()

# Qdrant: servidor propio o gestionado (no se ejecuta en el curso)
from qdrant_client import QdrantClient
from qdrant_client.models import (Distance, FieldCondition, Filter,
                                  MatchValue, PointStruct, VectorParams)

qd = QdrantClient(url="http://localhost:6333")
qd.create_collection(
    collection_name="pasajes",
    vectors_config=VectorParams(size=384, distance=Distance.COSINE))
qd.upsert(collection_name="pasajes", wait=True, points=[
    PointStruct(id=k, vector=e.tolist(),
                payload={"id_pasaje": i, "id_doc": doc})
    for k, (i, doc, e) in enumerate(zip(ids, pas["id_doc"], P))])
res = qd.query_points(
    collection_name="pasajes", query=Q[0].tolist(), limit=5,
    query_filter=Filter(must=[FieldCondition(
        key="id_doc", match=MatchValue(value="R01"))]),
    with_payload=True).points