Spaces:

alezsd
/

NSF-RAG-Codex

Running

NSF-RAG-Codex / rag_corrector.py

Alexander Sanchez

updated with Paloma's errata report

2c07e6b 6 months ago

9.06 kB

"""
rag_corrector.py
────────────────
Núcleo del sistema RAG.
  1. Detecta posibles errores HTR y grafías modernas en el texto de entrada.
  2. Recupera ejemplos similares del vector store (few-shot dinámico).
  3. Construye el prompt con reglas, ejemplos y alertas.
  4. Llama al LLM y devuelve el texto corregido + trazabilidad.
Uso:
    from rag_corrector import RAGCorrector
    corrector = RAGCorrector(vector_store)
    result = corrector.correct("texto htr aqui")
    print(result["corrected"])
    print(result["prompt"])          # para depuración
"""

import os
import re
from typing import List, Dict, Tuple
from openai import OpenAI
from dotenv import load_dotenv

from knowledge_base import HTR_ERROR_PATTERNS, GRAFIA_PATTERNS

load_dotenv()

MODEL = os.getenv("OPENAI_MODEL", "gpt-4o")
TOP_K = int(os.getenv("TOP_K", 5))

SYSTEM_PROMPT = """Eres un corrector especializado en documentos notariales y judiciales \
españoles del siglo XVI (castellano antiguo).
Tu ÚNICA tarea es corregir los errores introducidos por el proceso automático de \
reconocimiento de texto manuscrito (HTR). NO debes modernizar el texto bajo ninguna \
circunstancia.
REGLAS ABSOLUTAS — incumplirlas invalida la corrección:
1. Conserva SIEMPRE las grafías propias del s.XVI:
   fizo, fazer, hazer, merçed, vezino, mesmo, çibdad, escriuano,
   dho (=dicho), q̃ (=que), nro (=nuestro), vn/vna, etc.
2. NO conviertas f→h inicial latina (fizo ≠ hizo, fazer ≠ hacer).
3. Conserva abreviaturas y tildes voladas (q̃, nro, dho, sr).
4. Corrige SOLO lo que claramente sea un error HTR (grafema confundido visualmente).
5. Si no estás seguro de si algo es error HTR o forma s.XVI válida → conserva el original.
6. Presta especial atención a los errores marcados como [CRÍTICO] — son los de mayor \
impacto en la comprensión del texto.
7. NO corrijas ningún patrón marcado con ⚠ NO_CORREGIR — esas formas son válidas en s.XVI.
8. Responde ÚNICAMENTE con el texto corregido. Sin explicaciones, sin comillas."""


class RAGCorrector:
    def __init__(self, vector_store):
        self.vs = vector_store
        self.client = OpenAI(
            api_key=os.getenv("OPENAI_API_KEY"),
            base_url=os.getenv("OPENAI_BASE_URL", "https://api.x.ai/v1"),
        )

    # ── API pública ──────────────────────────────────────────────────────────

    def correct(self, htr_text: str, top_k: int = TOP_K, model: str = None) -> Dict:
        """
        Corrige un texto HTR usando RAG.
        Retorna dict con:
          corrected     : str   — texto corregido
          prompt        : str   — prompt completo enviado al LLM
          retrieved     : list  — documentos recuperados del vector store
          htr_errors    : list  — patrones HTR detectados
          grafia_warns  : list  — grafías modernas detectadas (alertas)
          model         : str   — modelo usado
        """
        retrieved    = self.vs.retrieve(htr_text, k=top_k)
        htr_errors   = self._detect_htr_errors(htr_text)
        grafia_warns = self._detect_grafias(htr_text)

        prompt = self._build_prompt(htr_text, retrieved, htr_errors, grafia_warns)
        corrected = self._call_llm(prompt, model=model or MODEL)

        return {
            "corrected":    corrected,
            "prompt":       prompt,
            "retrieved":    retrieved,
            "htr_errors":   htr_errors,
            "grafia_warns": grafia_warns,
            "model":        model or MODEL,
        }

    # ── Detección de patrones ────────────────────────────────────────────────

    def _detect_htr_errors(self, text: str) -> List[Dict]:
        """Detecta patrones de error HTR en el texto usando word boundaries."""
        found = []
        for p in HTR_ERROR_PATTERNS:
            htr_val = p.get("htr", "")
            # Ignorar patrones descriptivos (fusion_palabras, omision_*, etc.)
            if "_" in htr_val or len(htr_val) > 10:
                continue
            pattern = r'\b' + re.escape(htr_val) + r'\b'
            if re.search(pattern, text, re.IGNORECASE):
                found.append(p)
        return found

    def _detect_grafias(self, text: str) -> List[Dict]:
        """Detecta formas modernas que NO deberían modernizarse."""
        found = []
        lower = text.lower()
        for p in GRAFIA_PATTERNS:
            pattern = r'\b' + re.escape(p["modern"].lower()) + r'\b'
            if re.search(pattern, lower):
                found.append(p)
        return found

    # ── Constructor de prompt ────────────────────────────────────────────────

    def _build_prompt(
        self,
        htr_text:     str,
        retrieved:    List[Dict],
        htr_errors:   List[Dict],
        grafia_warns: List[Dict],
    ) -> str:

        sections = []

        # ── Bloque 1: Few-shot dinámico (ejemplos recuperados) ────────────────
        if retrieved:
            examples = []
            for i, doc in enumerate(retrieved, 1):
                corr = "; ".join(doc["corrections"]) if doc["corrections"] else "—"
                examples.append(
                    f"Ejemplo {i} [{doc['type']}, {doc['region']}, {doc['date']}]"
                    f" caligrafia: {doc.get('caligrafia', '?')}"
                    f" (similitud={doc['score']}):\n"
                    f"  HTR:  \"{doc['htr']}\"\n"
                    f"  GT:   \"{doc['gt']}\"\n"
                    f"  Correcciones aplicadas: {corr}"
                )
            sections.append(
                "EJEMPLOS DEL CORPUS (similares al texto a corregir):\n"
                + "\n\n".join(examples)
            )

        # ── Bloque 2: Alertas de errores HTR detectados ───────────────────────
        if htr_errors:
            # Ordenar: críticos primero
            SEV_ORDER = {"alta": 0, "media": 1, "baja": 2}
            sorted_errors = sorted(
                htr_errors,
                key=lambda p: SEV_ORDER.get(p.get("severity", "media"), 1)
            )

            hints = []
            for p in sorted_errors:
                sev      = p.get("severity", "media").upper()
                htr_c    = p.get("htr", "")
                gt_c     = p.get("gt", "")
                context  = p.get("context", "")
                tendencia= p.get("tendencia", "")
                warning  = p.get("warning")

                # Ejemplos: hasta 3 de la lista completa
                ejemplos_raw = p.get("examples", p.get("example", ""))
                if isinstance(ejemplos_raw, list):
                    ejemplos_str = "; ".join(ejemplos_raw[:3])
                else:
                    ejemplos_str = ejemplos_raw

                linea = (
                    f"  • [{sev}] '{htr_c}' puede ser '{gt_c}'\n"
                    f"    Causa    : {context}"
                )
                if tendencia:
                    linea += f"\n    Tendencia: {tendencia}"
                if ejemplos_str:
                    linea += f"\n    Ejemplos : {ejemplos_str}"
                if warning:
                    linea += f"\n    ⚠ {warning}"

                hints.append(linea)

            sections.append(
                "POSIBLES ERRORES HTR DETECTADOS EN ESTE TEXTO:\n"
                + "\n\n".join(hints)
            )

        # ── Bloque 3: Alertas de grafías modernas ─────────────────────────────
        if grafia_warns:
            warns = "\n".join(
                f"  • '{p['modern']}' → mantener como '{p['ancient']}': {p['rule']}"
                for p in grafia_warns
            )
            sections.append(
                "ALERTA — GRAFÍAS QUE NO DEBEN MODERNIZARSE:\n" + warns
            )

        context_block = "\n\n".join(sections)

        return (
            f"{context_block}\n\nTEXTO HTR A CORREGIR:\n\"{htr_text}\""
            if context_block
            else f"TEXTO HTR A CORREGIR:\n\"{htr_text}\""
        )

    # ── Llamada al LLM ───────────────────────────────────────────────────────

    def _call_llm(self, user_prompt: str, model: str = MODEL) -> str:
        response = self.client.chat.completions.create(
            model=model,
            temperature=0.1,
            max_tokens=1024,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user",   "content": user_prompt},
            ],
        )
        return response.choices[0].message.content.strip().strip('"').strip("'")

	"""
	rag_corrector.py
	────────────────
	Núcleo del sistema RAG.
	1. Detecta posibles errores HTR y grafías modernas en el texto de entrada.
	2. Recupera ejemplos similares del vector store (few-shot dinámico).
	3. Construye el prompt con reglas, ejemplos y alertas.
	4. Llama al LLM y devuelve el texto corregido + trazabilidad.

	Uso:
	from rag_corrector import RAGCorrector
	corrector = RAGCorrector(vector_store)
	result = corrector.correct("texto htr aqui")
	print(result["corrected"])
	print(result["prompt"]) # para depuración
	"""

	import os
	import re
	from typing import List, Dict, Tuple
	from openai import OpenAI
	from dotenv import load_dotenv

	from knowledge_base import HTR_ERROR_PATTERNS, GRAFIA_PATTERNS

	load_dotenv()

	MODEL = os.getenv("OPENAI_MODEL", "gpt-4o")
	TOP_K = int(os.getenv("TOP_K", 5))

	SYSTEM_PROMPT = """Eres un corrector especializado en documentos notariales y judiciales \
	españoles del siglo XVI (castellano antiguo).
	Tu ÚNICA tarea es corregir los errores introducidos por el proceso automático de \
	reconocimiento de texto manuscrito (HTR). NO debes modernizar el texto bajo ninguna \
	circunstancia.
	REGLAS ABSOLUTAS — incumplirlas invalida la corrección:
	1. Conserva SIEMPRE las grafías propias del s.XVI:
	fizo, fazer, hazer, merçed, vezino, mesmo, çibdad, escriuano,
	dho (=dicho), q̃ (=que), nro (=nuestro), vn/vna, etc.
	2. NO conviertas f→h inicial latina (fizo ≠ hizo, fazer ≠ hacer).
	3. Conserva abreviaturas y tildes voladas (q̃, nro, dho, sr).
	4. Corrige SOLO lo que claramente sea un error HTR (grafema confundido visualmente).
	5. Si no estás seguro de si algo es error HTR o forma s.XVI válida → conserva el original.
	6. Presta especial atención a los errores marcados como [CRÍTICO] — son los de mayor \
	impacto en la comprensión del texto.
	7. NO corrijas ningún patrón marcado con ⚠ NO_CORREGIR — esas formas son válidas en s.XVI.
	8. Responde ÚNICAMENTE con el texto corregido. Sin explicaciones, sin comillas."""


	class RAGCorrector:
	def __init__(self, vector_store):
	self.vs = vector_store
	self.client = OpenAI(
	api_key=os.getenv("OPENAI_API_KEY"),
	base_url=os.getenv("OPENAI_BASE_URL", "https://api.x.ai/v1"),
	)

	# ── API pública ──────────────────────────────────────────────────────────

	def correct(self, htr_text: str, top_k: int = TOP_K, model: str = None) -> Dict:
	"""
	Corrige un texto HTR usando RAG.
	Retorna dict con:
	corrected : str — texto corregido
	prompt : str — prompt completo enviado al LLM
	retrieved : list — documentos recuperados del vector store
	htr_errors : list — patrones HTR detectados
	grafia_warns : list — grafías modernas detectadas (alertas)
	model : str — modelo usado
	"""
	retrieved = self.vs.retrieve(htr_text, k=top_k)
	htr_errors = self._detect_htr_errors(htr_text)
	grafia_warns = self._detect_grafias(htr_text)

	prompt = self._build_prompt(htr_text, retrieved, htr_errors, grafia_warns)
	corrected = self._call_llm(prompt, model=model or MODEL)

	return {
	"corrected": corrected,
	"prompt": prompt,
	"retrieved": retrieved,
	"htr_errors": htr_errors,
	"grafia_warns": grafia_warns,
	"model": model or MODEL,
	}

	# ── Detección de patrones ────────────────────────────────────────────────

	def _detect_htr_errors(self, text: str) -> List[Dict]:
	"""Detecta patrones de error HTR en el texto usando word boundaries."""
	found = []
	for p in HTR_ERROR_PATTERNS:
	htr_val = p.get("htr", "")
	# Ignorar patrones descriptivos (fusion_palabras, omision_*, etc.)
	if "_" in htr_val or len(htr_val) > 10:
	continue
	pattern = r'\b' + re.escape(htr_val) + r'\b'
	if re.search(pattern, text, re.IGNORECASE):
	found.append(p)
	return found

	def _detect_grafias(self, text: str) -> List[Dict]:
	"""Detecta formas modernas que NO deberían modernizarse."""
	found = []
	lower = text.lower()
	for p in GRAFIA_PATTERNS:
	pattern = r'\b' + re.escape(p["modern"].lower()) + r'\b'
	if re.search(pattern, lower):
	found.append(p)
	return found

	# ── Constructor de prompt ────────────────────────────────────────────────

	def _build_prompt(
	self,
	htr_text: str,
	retrieved: List[Dict],
	htr_errors: List[Dict],
	grafia_warns: List[Dict],
	) -> str:

	sections = []

	# ── Bloque 1: Few-shot dinámico (ejemplos recuperados) ────────────────
	if retrieved:
	examples = []
	for i, doc in enumerate(retrieved, 1):
	corr = "; ".join(doc["corrections"]) if doc["corrections"] else "—"
	examples.append(
	f"Ejemplo {i} [{doc['type']}, {doc['region']}, {doc['date']}]"
	f" caligrafia: {doc.get('caligrafia', '?')}"
	f" (similitud={doc['score']}):\n"
	f" HTR: \"{doc['htr']}\"\n"
	f" GT: \"{doc['gt']}\"\n"
	f" Correcciones aplicadas: {corr}"
	)
	sections.append(
	"EJEMPLOS DEL CORPUS (similares al texto a corregir):\n"
	+ "\n\n".join(examples)
	)

	# ── Bloque 2: Alertas de errores HTR detectados ───────────────────────
	if htr_errors:
	# Ordenar: críticos primero
	SEV_ORDER = {"alta": 0, "media": 1, "baja": 2}
	sorted_errors = sorted(
	htr_errors,
	key=lambda p: SEV_ORDER.get(p.get("severity", "media"), 1)
	)

	hints = []
	for p in sorted_errors:
	sev = p.get("severity", "media").upper()
	htr_c = p.get("htr", "")
	gt_c = p.get("gt", "")
	context = p.get("context", "")
	tendencia= p.get("tendencia", "")
	warning = p.get("warning")

	# Ejemplos: hasta 3 de la lista completa
	ejemplos_raw = p.get("examples", p.get("example", ""))
	if isinstance(ejemplos_raw, list):
	ejemplos_str = "; ".join(ejemplos_raw[:3])
	else:
	ejemplos_str = ejemplos_raw

	linea = (
	f" • [{sev}] '{htr_c}' puede ser '{gt_c}'\n"
	f" Causa : {context}"
	)
	if tendencia:
	linea += f"\n Tendencia: {tendencia}"
	if ejemplos_str:
	linea += f"\n Ejemplos : {ejemplos_str}"
	if warning:
	linea += f"\n ⚠ {warning}"

	hints.append(linea)

	sections.append(
	"POSIBLES ERRORES HTR DETECTADOS EN ESTE TEXTO:\n"
	+ "\n\n".join(hints)
	)

	# ── Bloque 3: Alertas de grafías modernas ─────────────────────────────
	if grafia_warns:
	warns = "\n".join(
	f" • '{p['modern']}' → mantener como '{p['ancient']}': {p['rule']}"
	for p in grafia_warns
	)
	sections.append(
	"ALERTA — GRAFÍAS QUE NO DEBEN MODERNIZARSE:\n" + warns
	)

	context_block = "\n\n".join(sections)

	return (
	f"{context_block}\n\nTEXTO HTR A CORREGIR:\n\"{htr_text}\""
	if context_block
	else f"TEXTO HTR A CORREGIR:\n\"{htr_text}\""
	)

	# ── Llamada al LLM ───────────────────────────────────────────────────────

	def _call_llm(self, user_prompt: str, model: str = MODEL) -> str:
	response = self.client.chat.completions.create(
	model=model,
	temperature=0.1,
	max_tokens=1024,
	messages=[
	{"role": "system", "content": SYSTEM_PROMPT},
	{"role": "user", "content": user_prompt},
	],
	)
	return response.choices[0].message.content.strip().strip('"').strip("'")