Comparar modelos de embeddings: dimensión, latencia y recall

# Datos y preparación que el video da por hechos: ejecuta esto antes del resto del código.
import json
from collections import Counter
from urllib.request import urlopen
BASE = "https://pypro-assets.sfo3.digitaloceanspaces.com/cursos/rag-con-ollama/datos"

url = f"{BASE}/preguntas.jsonl"
texto = urlopen(url).read().decode()
preguntas = [json.loads(l) for l in texto.splitlines()]
evaluables = [p for p in preguntas
              if p["tipo"] != "sin_respuesta"]
print(len(preguntas), len(evaluables))
tipos = Counter(p["tipo"] for p in evaluables)
for tipo, n in tipos.items():
    print(tipo, n)

# Paso que el video ejecuta sin mostrarlo en pantalla:
q22 = {"pregunta": "¿Cuánto tarda un envío internacional a Canadá?",
       "doc": "envios-internacionales.md", "secciones": ["Países con envío"]}
nomic = [{"doc": "envios-internacionales.md", "seccion": "Introducción"}, {"doc": "envios-internacionales.md", "seccion": "Costos, impuestos y aduanas"}, {"doc": "envios-internacionales.md", "seccion": "Países con envío"}, {"doc": "envios-internacionales.md", "seccion": "Devoluciones internacionales"}, {"doc": "envios-nacionales.md", "seccion": "Introducción"}]
bge = [{"doc": "envios-internacionales.md", "seccion": "Países con envío"}, {"doc": "faq-clientes.md", "seccion": "Pedidos y envíos"}, {"doc": "envios-internacionales.md", "seccion": "Tiempos de preparación"}, {"doc": "envios-nacionales.md", "seccion": "Tiempos de entrega"}, {"doc": "envios-internacionales.md", "seccion": "Devoluciones internacionales"}]

def acierto(c, p):
    return (c["doc"] == p["doc"]
            and c["seccion"] in p["secciones"])
def mrr(ranking, p):
    for i, c in enumerate(ranking, start=1):
        if acierto(c, p):
            return 1 / i
    return 0.0
print(mrr(nomic, q22), mrr(bge, q22))

# Paso que el video ejecuta sin mostrarlo en pantalla:
import io, re, time, zipfile
import chromadb
import ollama

def cargar_corpus(url=f"{BASE}/pyshop-kb.zip"):
    docs = []
    with zipfile.ZipFile(io.BytesIO(urlopen(url).read())) as z:
        for nombre in sorted(n for n in z.namelist() if n.endswith(".md")):
            crudo = z.read(nombre).decode("utf-8")
            meta, texto = re.match(r"---\n(.*?)\n---\n(.*)", crudo, re.S).groups()
            campos = dict(linea.split(": ", 1) for linea in meta.splitlines())
            docs.append({"doc": nombre.split("/")[-1], "titulo": campos["titulo"],
                         "tipo": campos["tipo"], "area": campos["area"], "texto": texto.strip()})
    return docs

def secciones_de(doc):
    salida = []
    for parte in re.split(r"^(?=## )", doc["texto"], flags=re.M):
        if parte.strip():
            m = re.match(r"## (.+)", parte)
            salida.append((m.group(1).strip() if m else "Introducción", parte.strip()))
    return salida

def chunk_markdown(doc, tam=800, solape=100):
    chunks = []
    for seccion, texto in secciones_de(doc):
        trozos = [texto] if len(texto) <= tam else [
            texto[i:i + tam] for i in range(0, max(len(texto) - solape, 1), tam - solape)]
        for trozo in trozos:
            chunks.append({"id": f"{doc['doc']}#{len(chunks)}", "texto": trozo, "doc": doc["doc"],
                           "titulo": doc["titulo"], "tipo": doc["tipo"], "area": doc["area"],
                           "seccion": seccion})
    return chunks

def embed(textos, modelo, tipo="documento", lote=64):
    if isinstance(textos, str):
        textos = [textos]
    if modelo.startswith("nomic-embed-text"):
        prefijo = "search_document: " if tipo == "documento" else "search_query: "
        textos = [prefijo + t for t in textos]
    vectores = []
    for i in range(0, len(textos), lote):
        vectores += ollama.embed(model=modelo, input=textos[i:i + lote]).embeddings
    return vectores

cliente = chromadb.PersistentClient(path="pyshop_chroma_744")
METADATOS = ("doc", "titulo", "tipo", "area", "seccion")

def indexar(chunks, modelo, nombre):
    """Una colección por modelo (coseno), con los embeddings calculados por Ollama."""
    if nombre in [c.name for c in cliente.list_collections()]:
        cliente.delete_collection(nombre)
    col = cliente.create_collection(nombre, metadata={"hnsw:space": "cosine"})
    col.add(ids=[c["id"] for c in chunks], documents=[c["texto"] for c in chunks],
            metadatas=[{k: c[k] for k in METADATOS} for c in chunks],
            embeddings=embed([c["texto"] for c in chunks], modelo, "documento"))
    return col

def buscar(col, pregunta, modelo, k=10):
    r = col.query(query_embeddings=embed(pregunta, modelo, "consulta"), n_results=k)
    return [{"id": i, "distancia": d, **m}
            for i, d, m in zip(r["ids"][0], r["distances"][0], r["metadatas"][0])]

MODELOS = ["nomic-embed-text", "mxbai-embed-large", "bge-m3"]
docs = cargar_corpus()
chunks = [c for doc in docs for c in chunk_markdown(doc)]

filas = []
for m in MODELOS:
    col = indexar(chunks, m, f"kb_{m}")
    for p in evaluables:
        top10 = buscar(col, p["pregunta"], m, k=10)
        filas.append({"modelo": m, "tipo": p["tipo"],
                      "mrr": mrr(top10, p)})
print(len(filas), "filas")

# Paso que el video ejecuta sin mostrarlo en pantalla:
import pandas as pd

df = pd.DataFrame(filas)
df["hit@1"] = df["mrr"] == 1
df["hit@5"] = df["mrr"] >= 0.2
cols = ["hit@1", "hit@5", "mrr"]
tabla = df.groupby("modelo")[cols].mean()
print(tabla.round(2))

t = df.pivot_table("mrr", index="tipo",
                   columns="modelo")
t.columns = ["bge-m3", "mxbai", "nomic"]
print(t.round(2))

for m in MODELOS:
    embed(["hola"], m)  # calentar
    t0 = time.perf_counter()
    col = indexar(chunks, m, f"kb_{m}")
    t1 = time.perf_counter()
    for p in evaluables:
        buscar(col, p["pregunta"], m)
    ms = (time.perf_counter() - t1) / 54 * 1000
    print(m, f"{t1 - t0:.1f} s", f"{ms:.0f} ms")

for m, dim in [("nomic", 768), ("bge-m3", 1024)]:
    kb = 160 * dim * 4 / 1024
    gb = 1_000_000 * dim * 4 / 1e9
    print(m, f"{kb:.0f} KB", f"{gb:.1f} GB/millón")