kernels/training_utils.py

import html
import re
from typing import Any, Dict, List


_HTML_TAG_RE = re.compile(r"<[^>]+>")
_ESCAPED_TAG_RE = re.compile(r"\\<[^>]*\\>")   # backslash-escaped: \</ref\>
_SELF_CLOSING_RE = re.compile(r"/>")            # stray self-closing fragments
_CASCADE_RE = re.compile(r"[}\]\"]{3,}")        # }}}}, ]]]], """ cascades
_BRACKET_RE = re.compile(r"[\{\}\[\]]")
_PUNCT_RE = re.compile(r"[|\\/]+")
_WHITESPACE_RE = re.compile(r"\s+")


def sanitize_text(text: str) -> str:
    """Remove HTML, layout, and bracket-cascade artifacts from training text.

    Handles:
    - Standard HTML tags: <tag>, </tag>
    - Backslash-escaped angle brackets from Docling: \\</ref\\>, \\</span\\>
    - Unicode-encoded angle brackets: \\u003c / \\u003e
    - Stray self-closing fragments: />
    - HTML entities: &gt; &lt; &amp; &nbsp; etc.
    - Cascading closing-bracket artifacts: }}}}, ]]]]
    """
    if not text:
        return ""

    # 1. Decode unicode-escaped angle brackets (\u003c → <, \u003e → >)
    #    These appear in post-Docling JSON that has been serialised with
    #    ensure_ascii=True and then re-read as a Python string.
    cleaned = text.replace("\\u003c", "<").replace("\\u003e", ">")
    cleaned = cleaned.replace("\u003c", "<").replace("\u003e", ">")

    # 2. Strip backslash-escaped tags before the HTML regex sees them
    cleaned = _ESCAPED_TAG_RE.sub(" ", cleaned)

    # 3. Unescape remaining HTML entities (&amp; &nbsp; &#160; etc.) so that
    #    the HTML-tag regex can then cleanly remove any residual tags.
    cleaned = html.unescape(cleaned)

    # 4. Remove all remaining HTML tags
    cleaned = _HTML_TAG_RE.sub(" ", cleaned)

    # 5. Remove stray self-closing fragment "/>", leftover bare ">" / "<"
    cleaned = _SELF_CLOSING_RE.sub(" ", cleaned)
    cleaned = cleaned.replace(">", " ").replace("<", " ")

    # 6. Remove cascading closing-bracket artifacts (}}}}}, ]]]]])
    cleaned = _CASCADE_RE.sub(" ", cleaned)

    # 7. Remove remaining individual bracket noise
    cleaned = _BRACKET_RE.sub(" ", cleaned)

    # 8. Remove pipe / backslash / forward-slash noise
    cleaned = _PUNCT_RE.sub(" ", cleaned)

    # 9. Collapse all whitespace
    cleaned = _WHITESPACE_RE.sub(" ", cleaned)
    return cleaned.strip()


def sanitize_example(example: Dict[str, Any]) -> Dict[str, Any]:
    """Sanitize instruction, input, and output fields in a training example."""
    sanitized = {}
    for key in ("instruction", "input", "output"):
        value = example.get(key, "")
        if isinstance(value, str):
            sanitized[key] = sanitize_text(value)
        else:
            sanitized[key] = value
    return sanitized


_THINK_BLOCK_RE = re.compile(r"<think>(.*?)</think>", re.DOTALL)
_STEP_LABEL_RE = re.compile(r"\bStep\s+\d+", re.IGNORECASE)


def _thinking_quality_check(output_text: str) -> bool:
    """Return True when a CoT output passes minimum quality standards.

    A reasoning example is rejected when any of the following hold:
    - No <think>...</think> block is present.
    - The think block body is fewer than 30 words (too terse to be useful).
    - The think block has fewer than 3 labelled step markers ("Step N").
    - The final JSON answer appears verbatim inside the think block (answer
      leakage — the model would learn to write the answer before reasoning).

    Non-reasoning outputs (those without a <think> block at all but also without
    the 'think' instruction) are passed through unchanged; this check only fires
    when a <think> block is present.
    """
    think_match = _THINK_BLOCK_RE.search(output_text)
    if think_match is None:
        # Not a CoT example — no think block expected, let it through.
        return True

    think_body = think_match.group(1).strip()

    # Minimum word count inside the think block.
    if len(think_body.split()) < 30:
        return False

    # Check for templated Step N patterns — these indicate generated reasoning
    # rather than genuine analytical thinking.  Allow up to 2 step labels
    # (natural in analytical writing) but reject heavily templated outputs.
    step_labels = _STEP_LABEL_RE.findall(think_body)
    if len(step_labels) > 5:
        return False

    # Check for answer leakage: find the text after </think> and verify none
    # of its JSON values appear verbatim inside the think block.
    after_think = output_text[think_match.end():].strip()
    try:
        answer_obj = __import__("json").loads(after_think)
        for v in answer_obj.values():
            if isinstance(v, str) and len(v) > 8 and v in think_body:
                # Tolerate short strings that appear as normal reasoning words;
                # flag only long values that indicate the answer was pre-written.
                if len(v) > 30:
                    return False
    except Exception:
        pass  # Non-JSON final answer — leakage check not applicable.

    return True

    think_body = think_match.group(1).strip()

    # Minimum word count inside the think block.
    if len(think_body.split()) < 30:
        return False

    # Must have at least 3 step labels.
    step_labels = _STEP_LABEL_RE.findall(think_body)
    if len(step_labels) < 3:
        return False

    # Check for answer leakage: find the text after </think> and verify none
    # of its JSON values appear verbatim inside the think block.
    after_think = output_text[think_match.end():].strip()
    try:
        answer_obj = __import__("json").loads(after_think)
        for v in answer_obj.values():
            if isinstance(v, str) and len(v) > 8 and v in think_body:
                # Tolerate short strings that appear as normal reasoning words;
                # flag only long values that indicate the answer was pre-written.
                if len(v) > 30:
                    return False
    except Exception:
        pass  # Non-JSON final answer — leakage check not applicable.

    return True


def _near_duplicate(a: str, b: str, threshold: float = 0.88) -> bool:
    """Return True when two strings are suspiciously similar (ratio >= threshold).

    Uses SequenceMatcher which is fast enough for the dataset sizes we see
    (hundreds to low thousands of examples).  The threshold of 0.88 catches
    single-entity paraphrases (e.g. one node name swapped in an otherwise
    identical sentence, ratio ~0.90) while staying well above the ratio for
    genuinely distinct examples (unrelated edges typically score <0.3).
    """
    from difflib import SequenceMatcher
    return SequenceMatcher(None, a, b, autojunk=False).ratio() >= threshold


def curate_examples(examples: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """Remove duplicates, near-duplicates, and low-signal examples before training.

    Pass 1 — Exact dedup: drop any example whose (instruction, input, output)
              triple has already been seen.
    Pass 2 — Near-dedup: for each surviving example, check its output against
              the last N accepted outputs using SequenceMatcher.  If similarity
              >= 0.92 the example is skipped.  We only compare against a sliding
              window of the most recent 200 outputs to keep runtime linear.
    Pass 3 — Quality floor: skip examples whose output has fewer than 4 words.
    Pass 4 — Thinking quality: for CoT/reasoning examples (those containing a
              <think> block), verify the block meets minimum structural standards
              so the model learns genuine reasoning rather than degenerate CoT.
    """
    curated: List[Dict[str, Any]] = []
    seen_exact: set = set()
    recent_outputs: List[str] = []   # sliding window for near-dup check
    WINDOW = 200

    for example in examples:
        if not isinstance(example, dict):
            continue

        sanitized = sanitize_example(example)
        instruction = str(sanitized.get("instruction", "") or "").strip()
        input_text = str(sanitized.get("input", "") or "").strip()
        output_text = str(sanitized.get("output", "") or "").strip()

        if not instruction or not input_text or not output_text:
            continue
        if len(output_text.split()) < 4:
            continue

        # Pass 1 — exact dedup
        signature = (instruction.lower(), input_text.lower(), output_text.lower())
        if signature in seen_exact:
            continue
        seen_exact.add(signature)

        # Pass 2 — near-dup check against recent outputs
        out_lower = output_text.lower()
        if any(_near_duplicate(out_lower, prev) for prev in recent_outputs[-WINDOW:]):
            continue

        # Pass 4 — CoT/thinking quality gate (no-op for non-reasoning examples)
        if not _thinking_quality_check(output_text):
            continue

        recent_outputs.append(out_lower)
        # Preserve task_type so format_prompt can detect reasoning examples
        # and route them through the thinking-safe formatter.
        entry: Dict[str, Any] = {
            "instruction": instruction,
            "input": input_text,
            "output": output_text,
        }
        if example.get("task_type"):
            entry["task_type"] = example["task_type"]
        curated.append(entry)

    return curated