code/self_healing_finetuning/data_pipeline/repair_answer_quality.py

from __future__ import annotations

import re
from typing import Any


BANNED_ANSWER_PATTERNS = (
    re.compile(r"</?think\b", flags=re.IGNORECASE),
    re.compile(r"\bokay,\s*let'?s see\b", flags=re.IGNORECASE),
    re.compile(r"\bthe user wants\b", flags=re.IGNORECASE),
    re.compile(r"\bi need to\b", flags=re.IGNORECASE),
    re.compile(r"\blet me\b", flags=re.IGNORECASE),
)

CHAT_PRELUDE_RE = re.compile(
    r"^\s*(sure|certainly|of course|here(?:'s| is| are)|as an ai|i (?:can|will|would|am)|okay|ok|let me|i'd be happy)\b",
    flags=re.IGNORECASE,
)
NUMBER_RE = re.compile(r"\$?\d+(?:\.\d+)?%?")
NUMERIC_PROMPT_RE = re.compile(r"\bcalculate\b|from \$?\d|to \$?\d|%|\bratio\b|\bmargin\b|\bebitda\b", flags=re.IGNORECASE)
NUMERIC_COMPLETION_RE = re.compile(
    r"\b(from|to|rose|fell|increased?|decreased?|change[ds]?|delta|ratio|margin|x|%)\b|[/=]",
    flags=re.IGNORECASE,
)
SECTION_NAMES = (
    "reported facts:",
    "calculation:",
    "inference:",
    "risk/tradeoff:",
    "conclusion:",
)


def prompt_requires_numeric(prompt: str, task: Any = None) -> bool:
    return str(task) == "quantitative_qa" or (bool(NUMERIC_PROMPT_RE.search(prompt)) and bool(prompt_numbers(prompt)))


def prompt_numbers(prompt: str) -> set[str]:
    return set(NUMBER_RE.findall(prompt))


def answer_numbers(answer: str) -> set[str]:
    return set(NUMBER_RE.findall(answer))


def has_banned_answer_text(answer: str) -> bool:
    return any(pattern.search(answer) for pattern in BANNED_ANSWER_PATTERNS)


def _terminal_fragment(answer: str) -> str:
    sentences = [item.strip() for item in re.split(r"(?<=[.!?])\s+", answer.strip()) if item.strip()]
    if not sentences:
        return answer.strip()
    return " ".join(sentences[-2:])


def repair_quality_checks(
    *,
    prompt: str,
    answer: str,
    task: Any = None,
    require_sections: bool = False,
) -> dict[str, bool]:
    numbers_in_prompt = prompt_numbers(prompt)
    numbers_in_answer = answer_numbers(answer)
    numeric_required = prompt_requires_numeric(prompt, task)
    required_number_hits = min(2, len(numbers_in_prompt))
    terminal = _terminal_fragment(answer)

    return {
        "non_empty": bool(answer.strip()),
        "no_think_tags": "<think" not in answer.lower() and "</think" not in answer.lower(),
        "no_banned_cot_or_chatty_text": not has_banned_answer_text(answer),
        "no_chatty_preamble": not bool(CHAT_PRELUDE_RE.match(answer)),
        "sections_present_when_required": (not require_sections)
        or all(section in answer.lower() for section in SECTION_NAMES),
        "retains_required_prompt_numbers": (not numbers_in_prompt)
        or len(numbers_in_prompt.intersection(numbers_in_answer)) >= required_number_hits,
        "numeric_completion_when_required": (not numeric_required)
        or (bool(numbers_in_answer) and bool(NUMERIC_COMPLETION_RE.search(answer))),
        "terminal_numeric_conclusion_when_required": (not numeric_required) or bool(NUMBER_RE.search(terminal)),
        "fact_inference_signal": bool(
            re.search(r"\b(fact|reported|inference|suggests|implies|while|however|but)\b", answer, re.IGNORECASE)
        ),
        "risk_tradeoff_signal": bool(re.search(r"\b(risk|tradeoff|offset|downside|pressure|monitor)\b", answer, re.IGNORECASE)),
    }


def answer_admitted_to_training(checks: dict[str, bool]) -> bool:
    required = (
        "non_empty",
        "no_think_tags",
        "no_banned_cot_or_chatty_text",
        "no_chatty_preamble",
        "sections_present_when_required",
        "retains_required_prompt_numbers",
        "numeric_completion_when_required",
        "terminal_numeric_conclusion_when_required",
        "fact_inference_signal",
        "risk_tradeoff_signal",
    )
    return all(checks.get(name) is True for name in required)