llm_query/span_marking/utils.py

import re
from collections.abc import Callable

SPAN_PATTERN = re.compile(r"\{\{(.*?)\}\}")

TokenizeFn = Callable[[str], tuple[str, ...]]


def is_pass_response(content: str) -> bool:
    stripped = content.strip()
    candidate = stripped
    if len(candidate) >= 2 and candidate[0] == candidate[-1] and candidate[0] in "\"'":
        candidate = candidate[1:-1].strip()
    return candidate.casefold() == "pass"


def strip_markers(marked_text: str) -> tuple[str, list[tuple[int, int]]]:
    """
    Remove {{ }} markers and return cleaned text plus dissimilar character ranges
    in the cleaned text.
    """
    cleaned_parts: list[str] = []
    dissimilar_ranges: list[tuple[int, int]] = []
    cleaned_offset = 0
    last_end = 0
    for match in SPAN_PATTERN.finditer(marked_text):
        before = marked_text[last_end:match.start()]
        cleaned_parts.append(before)
        cleaned_offset += len(before)
        span_text = match.group(1)
        span_start = cleaned_offset
        span_end = span_start + len(span_text)
        dissimilar_ranges.append((span_start, span_end))
        cleaned_parts.append(span_text)
        cleaned_offset = span_end
        last_end = match.end()
    cleaned_parts.append(marked_text[last_end:])
    return "".join(cleaned_parts), dissimilar_ranges


def token_character_spans(text: str, tokens: tuple[str, ...]) -> list[tuple[int, int]] | None:
    """Locate each token in order within text. Returns None if a token cannot be found."""
    spans: list[tuple[int, int]] = []
    position = 0
    for token in tokens:
        index = text.find(token, position)
        if index < 0:
            return None
        spans.append((index, index + len(token)))
        position = index + len(token)
    return spans


def extract_labels_from_marked(
        original_text: str,
        marked_text: str,
        tokenize: TokenizeFn,
) -> tuple[int, ...]:
    """
    Convert a {{marked}} copy of original_text into per-token dissimilar labels.

    Requires only that cleaned and original token counts match (token strings may
    differ, e.g. apostrophe variants). Length mismatch yields all zeros.
    """
    original_tokens = tokenize(original_text)
    if is_pass_response(marked_text):
        return tuple(0 for _ in original_tokens)

    cleaned_text, dissimilar_ranges = strip_markers(marked_text)
    cleaned_tokens = tokenize(cleaned_text)
    if len(cleaned_tokens) != len(original_tokens):
        return tuple(0 for _ in original_tokens)

    character_spans = token_character_spans(cleaned_text, cleaned_tokens)
    if character_spans is None:
        return tuple(0 for _ in original_tokens)

    labels: list[int] = []
    for start, end in character_spans:
        is_dissimilar = any(
            start < span_end and end > span_start
            for span_start, span_end in dissimilar_ranges
        )
        labels.append(1 if is_dissimilar else 0)
    return tuple(labels)