RAG sobre informes, folletos y KID, y cómo evaluarlo.

import re
import unicodedata

import numpy as np
import pandas as pd
from rank_bm25 import BM25Okapi

docs = pd.read_parquet("rag6_documentos.parquet")
pas = pd.read_parquet("rag6_pasajes.parquet")
preg = pd.read_csv("rag6_preguntas.csv")
TEXTO = dict(zip(docs.id_doc, docs.texto))
TITULO = dict(zip(docs.id_doc, docs.titulo))
P = pas.set_index("id_pasaje")
NOTA = "¹²³⁴⁵⁶⁷⁸⁹"


def seccion(doc, i):
    hs = re.findall(r"^##+ (.*)$", TEXTO[doc][:i + 1], re.M)
    return hs[-1] if hs else ""


def trocear(doc, n=800):
    t = TEXTO[doc]
    bl = [(m.start(), m.end())
          for m in re.finditer(r"[^\n](?:.|\n(?!\n))*", t)]
    notas = {t[a]: (a, b) for a, b in bl if t[a] in NOTA}
    trozos = []
    for a, b in bl:
        s = t[a:b]
        if s.startswith("#") or s[0] in NOTA:
            continue
        marcas = [notas[c] for c in dict.fromkeys(s) if c in notas]
        fija = s.startswith("|") or bool(marcas)
        ult = trozos[-1] if trozos else None
        if (ult and not fija and not ult["fija"]
                and seccion(doc, a) == seccion(doc, ult["tramos"][0][0])
                and b - ult["tramos"][0][0] <= n):
            ult["tramos"][0] = (ult["tramos"][0][0], b)
            continue
        trozos.append(dict(doc=doc, tramos=[(a, b)] + marcas,
                           fija=fija))
    for c in trozos:
        cuerpo = "\n\n".join(t[i:f] for i, f in c["tramos"])
        sec = seccion(doc, c["tramos"][0][0])
        c["texto"] = f"{TITULO[doc]} > {sec}\n\n{cuerpo}"
    return trozos


VACIAS = set("""de la el en y a los las del que se por con un una
para es al lo su sus o como fue cual cuanto cuantos cuales
segun""".split())


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


def cubierta(trozos, pid):
    """Fracción de los caracteres del pasaje dentro de los trozos."""
    r = P.loc[pid]
    iv = sorted((max(i, r.inicio), min(f, r.fin)) for c in trozos
                if c["doc"] == r.id_doc for i, f in c["tramos"])
    total, hasta = 0, r.inicio
    for a, b in iv:
        a = max(a, hasta)
        if b > a:
            total, hasta = total + b - a, b
    return total / (r.fin - r.inicio)


TR = [c for d in TEXTO for c in trocear(d)]
bm = BM25Okapi([palabras(c["texto"]) for c in TR])


def buscar(pregunta, k=3):
    sc = bm.get_scores(palabras(pregunta))
    orden = np.argsort(-sc, kind="stable")[:k]
    return [TR[j] for j in orden], float(sc[orden[0]])


print(len(TR), "trozos;", len(preg), "preguntas,",
      int((~preg.respondible).sum()), "sin respuesta")

NUM = re.compile(r"\d[\d.,]*\d|\d")
CITA = re.compile(r"\[(\d+)\]")


def cifras(s):
    """Conjunto de cifras de un texto, sin las citas [n]."""
    out = set()
    for t in NUM.findall(CITA.sub(" ", s)):
        if "," in t:
            t = t.replace(".", "").replace(",", ".")
        elif re.fullmatch(r"\d{1,3}(\.\d{3})+", t):
            t = t.replace(".", "")
        out.add(round(float(t.rstrip(".")), 4))
    return out


def calla(resp):
    return "no consta" in resp.lower()


def exacta(resp, q):
    """Cifras de oro presentes y ninguna ajena (salvo las de la
    pregunta); sin cifras en el oro, sus palabras."""
    if not q.respondible:
        return calla(resp)
    if calla(resp):
        return False
    oro, dichas = cifras(q.respuesta_oro), cifras(resp)
    if oro:
        return oro <= dichas and dichas <= oro | cifras(q.pregunta)
    return set(palabras(q.respuesta_oro)) <= set(palabras(resp))


def contiene(resp, q):
    """Versión laxa: las cifras (o palabras) de oro están dentro."""
    oro = cifras(q.respuesta_oro)
    if oro:
        return oro <= cifras(resp)
    return set(palabras(q.respuesta_oro)) <= set(palabras(resp))


def fiel(resp, contexto, q):
    """Toda cifra de la respuesta está en el contexto o la pregunta."""
    return cifras(resp) <= cifras(contexto) | cifras(q.pregunta)


def citas_correctas(citados, q):
    """Fracción de trozos citados que cubren un pasaje de oro."""
    if not citados or not q.respondible:
        return np.nan
    oro = q.pasajes_oro.split(";")
    return np.mean([any(cubierta([c], p) >= 0.5 for p in oro)
                    for c in citados])


print(sorted(cifras("Ventas de 13.668,9 millones, un 2,2 % más [2]")))
print(sorted(cifras("40.3% de 5.505 en 2023")))

def mejor_linea(trozo, pregunta):
    cuerpo = trozo["texto"].split("\n\n", 1)[1]
    lineas = [x for x in re.split(r"\n|(?<=\.) ", cuerpo) if x.strip()]
    qp = set(palabras(pregunta))
    return max(lineas, key=lambda x: len(qp & set(palabras(x))))


base = []
for q in preg.itertuples():
    top, s1 = buscar(q.pregunta)
    oro = q.pasajes_oro.split(";") if q.respondible else []
    cob = (np.mean([cubierta(top, p) >= 0.5 for p in oro])
           if oro else np.nan)
    trozo = top[0]["texto"].split("\n\n", 1)[1]
    linea = mejor_linea(top[0], q.pregunta)
    base.append(dict(
        id=q.id_pregunta, tipo=q.tipo_pregunta, resp=q.respondible,
        score=s1, cobertura=cob,
        trozo=contiene(trozo, q) if q.respondible else np.nan,
        linea=contiene(linea, q) if q.respondible else np.nan,
        linea_exacta=exacta(linea, q) if q.respondible else np.nan,
        cita=citas_correctas(top[:1], q)))
base = pd.DataFrame(base)
r = base[base.resp]
print(r[["cobertura", "trozo", "linea", "linea_exacta", "cita"]]
      .astype(float).mean().round(3).to_dict())
cols = ["cobertura", "trozo", "linea", "cita"]
r.groupby("tipo")[cols].mean().round(2)

from sklearn.metrics import roc_auc_score

print(base.groupby("resp").score.describe()[["mean", "50%", "min",
                                             "max"]].round(2))
print("AUC:", round(roc_auc_score(~base.resp, -base.score), 3))
u = base[~base.resp].score.median()        # calla en la mitad
calla_bien = (~base.resp & (base.score < u)).sum()
calla_mal = (base.resp & (base.score < u)).sum()
print(f"umbral {u:.2f}: calla en {calla_bien} de 17 sin respuesta"
      f" y en {calla_mal} de 246 respondibles")

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

torch.manual_seed(0)
torch.use_deterministic_algorithms(True)
CUOTA = {"texto": 3, "tabla": 3, "nota_pie": 1, "texto_y_nota": 1,
         "tabla_y_nota": 1, "no_respondible": 3}
muestra = pd.concat([preg[preg.tipo_pregunta == t]
                     .sample(n, random_state=0)
                     for t, n in CUOTA.items()])

SISTEMA = (
    "Eres un analista financiero. Responde en español con una frase "
    "corta, usando solo los fragmentos. Termina la frase con el número "
    "del fragmento que la apoya entre corchetes. Si los fragmentos no "
    "contienen la respuesta, responde exactamente: No consta.\n\n"
    "Ejemplo:\n[1] La plantilla media fue de 1.200 personas.\n"
    "[2] El consejo tiene 9 miembros.\n"
    "Pregunta: ¿Cuántos miembros tiene el consejo?\n"
    "Respuesta: El consejo tiene 9 miembros [2].")


def mensajes(pregunta, top):
    ctx = "\n\n".join(f"[{i + 1}] {c['texto']}"
                      for i, c in enumerate(top))
    return [{"role": "system", "content": SISTEMA},
            {"role": "user",
             "content": f"{ctx}\n\nPregunta: {pregunta}\nRespuesta:"}]


modelo = "Qwen/Qwen2.5-0.5B-Instruct"
tok = AutoTokenizer.from_pretrained(modelo)
llm = AutoModelForCausalLM.from_pretrained(modelo, dtype=torch.float32)
print(muestra.tipo_pregunta.value_counts().to_dict())

filas = []
for q in muestra.itertuples():
    top, _ = buscar(q.pregunta)
    e = tok.apply_chat_template(mensajes(q.pregunta, top),
                                add_generation_prompt=True,
                                return_tensors="pt", return_dict=True)
    with torch.inference_mode():
        out = llm.generate(**e, max_new_tokens=40, do_sample=False,
                           pad_token_id=tok.eos_token_id)
    n = e["input_ids"].shape[1]
    resp = tok.decode(out[0, n:], skip_special_tokens=True).strip()
    citados = [top[int(i) - 1] for i in CITA.findall(resp)
               if 1 <= int(i) <= len(top)]
    contexto = "\n".join(c["texto"] for c in top)
    oro = q.pasajes_oro.split(";") if q.respondible else []
    filas.append(dict(
        id=q.id_pregunta, tipo=q.tipo_pregunta, resp=q.respondible,
        tokens=n, respuesta=resp.replace("\n", " ")[:60],
        exacta=exacta(resp, q), contiene=contiene(resp, q),
        fiel=fiel(resp, contexto, q) if not calla(resp) else np.nan,
        cita=citas_correctas(citados, q), cita_alguna=bool(citados),
        cobertura=(np.mean([cubierta(top, p) >= 0.5 for p in oro])
                   if oro else np.nan)))
llmr = pd.DataFrame(filas)
llmr[["id", "tipo", "respuesta", "exacta", "fiel", "cita"]]

bm_ = base.set_index("id").loc[muestra.id_pregunta].reset_index()
calla_base = bm_.score < u
exacta_base = np.where(bm_.resp, ~calla_base & (bm_.linea_exacta == 1),
                       calla_base)


def ic(x):
    """Proporción con su intervalo de Wilson al 95 %."""
    x = np.asarray(x, float)
    x = x[~np.isnan(x)]
    n, p, z = len(x), x.mean(), 1.96
    c = (p + z**2 / (2 * n)) / (1 + z**2 / n)
    h = (z * np.sqrt(p * (1 - p) / n + z**2 / (4 * n**2))
         / (1 + z**2 / n))
    return f"{p:.2f} [{c - h:.2f}, {c + h:.2f}] n={n}"


r_llm, r_base = llmr[llmr.resp], bm_[bm_.resp]
tabla = pd.DataFrame({
    "LLM 0,5B": [ic(llmr.exacta), ic(r_llm.contiene), "-",
                 ic(r_llm.fiel), ic(r_llm.cita), ic(r_llm.cita_alguna),
                 ic(llmr[~llmr.resp].exacta)],
    "extracción": [ic(exacta_base), ic(r_base.linea), ic(r_base.trozo),
                   "1.00 (no genera)", ic(r_base.cita),
                   "1.00 (siempre [1])", ic(calla_base[~bm_.resp])]},
    index=["exacta (12)", "contiene el oro", "  (trozo entero)",
           "fiel al contexto", "citas correctas", "cita algo",
           "calla si no consta"])
print("cobertura del buscador en la muestra:",
      round(r_llm.cobertura.mean(), 3))
tabla

import anthropic

cliente = anthropic.Anthropic()          # ANTHROPIC_API_KEY del entorno
doc = "R10"
respuesta = cliente.messages.create(
    model="claude-opus-5-5",
    max_tokens=1024,
    messages=[{"role": "user", "content": [
        {"type": "document",
         "source": {"type": "text", "media_type": "text/plain",
                    "data": TEXTO[doc]},
         "title": TITULO[doc],
         "citations": {"enabled": True}},
        {"type": "text",
         "text": "¿Cuál es el indicador resumido de riesgo?"}]}])
for bloque in respuesta.content:
    for c in getattr(bloque, "citations", None) or []:
        print(c.start_char_index, c.end_char_index, c.cited_text)

import json

import anthropic
import httpx2


def simulador(request):
    """Responde como la API de Messages con citas char_location."""
    cuerpo = json.loads(request.content)
    partes = cuerpo["messages"][0]["content"]
    doc, pregunta = partes[0], partes[1]["text"]
    t = doc["source"]["data"]
    frases = [(m.start(), m.end())
              for m in re.finditer(r"[^\n.]+(?:\.(?=\s|$))?", t)
              if len(m.group().strip()) > 20]
    bm_f = BM25Okapi([palabras(t[a:b]) for a, b in frases])
    a, b = frases[int(np.argmax(bm_f.get_scores(palabras(pregunta))))]
    cita = {"type": "char_location", "cited_text": t[a:b],
            "document_index": 0, "document_title": doc["title"],
            "start_char_index": a, "end_char_index": b}
    return httpx2.Response(200, json={
        "id": "msg_simulado", "type": "message", "role": "assistant",
        "model": cuerpo["model"], "stop_reason": "end_turn",
        "stop_sequence": None,
        "usage": {"input_tokens": 0, "output_tokens": 0},
        "content": [
            {"type": "text", "text": "Según el documento, "},
            {"type": "text", "text": t[a:b].strip(),
             "citations": [cita]}]})


cliente = anthropic.Anthropic(
    api_key="sk-simulado", base_url="http://simulado",
    http_client=httpx2.Client(
        transport=httpx2.MockTransport(simulador)))


def pregunta_con_citas(q):
    doc = q.id_doc
    r = cliente.messages.create(
        model="simulado", max_tokens=200,
        messages=[{"role": "user", "content": [
            {"type": "document",
             "source": {"type": "text", "media_type": "text/plain",
                        "data": TEXTO[doc]},
             "title": TITULO[doc], "citations": {"enabled": True}},
            {"type": "text", "text": q.pregunta}]}])
    citas = [c for b in r.content for c in (b.citations or [])]
    tramo = [dict(doc=doc, tramos=[(c.start_char_index,
                                    c.end_char_index)])
             for c in citas]
    pas_doc = pas[pas.id_doc == doc].id_pasaje
    tocados = [p for p in pas_doc if cubierta(tramo, p) > 0]
    return citas, tocados


Q = preg.set_index("id_pregunta")
for qid in ["QA196", "QA105", "QA230"]:
    q = Q.loc[qid]
    citas, tocados = pregunta_con_citas(q)
    c = citas[0]
    print(qid, q.pregunta[:58])
    print("   cita", c.start_char_index, c.end_char_index,
          repr(c.cited_text[:50]))
    ok = bool(set(tocados) & set(q.pasajes_oro.split(";")))
    print("   pasajes citados:", tocados, "| oro:", q.pasajes_oro,
          "| correcta:", ok)