scripts/segmenting.py

"""
segmenting.py — Passage segmentation, classification, and fidelity gates for the
tessera-compressor harness.

Extracted from the harness the compressor was accepted under (same functions the
teacher mint used). Pure text processing: no network, no credentials.

Flow: segment -> group_steps -> classify_passage per passage -> model call ->
gate -> rules fallback on failure. A failed passage costs a few dozen tokens of
savings, never content.
"""
import re

CJK = re.compile(r'[一-鿿㐀-䶿]')
NUM = re.compile(r'\d+(?:\.\d+)?')
IDENT = re.compile(r'`[^`\n]+`|\b[A-Za-z]+(?:_[A-Za-z0-9]+)+\b|\b[a-z]+[A-Z][A-Za-z0-9]*\b')
FENCE = re.compile(r'```.*?```', re.DOTALL)
SENT_SPLIT = re.compile(r'(?<=[.!?;])\s+')
_LIST_MARKER = re.compile(r'(?:^|[\n\s(])(\d{1,2})[.)]\s')
_OPS = set('+-*/=<>≤≥≠∈∀∃¬→⇒%^{}[]')


def segment(text):
    """Split a reasoning block into ordered segments; code fences are atomic and marked."""
    segs = []  # (kind, text) kind ∈ {'code','prose'}
    pos = 0
    for m in FENCE.finditer(text):
        before = text[pos:m.start()]
        segs.extend(('prose', s) for s in _split_prose(before))
        segs.append(('code', m.group(0)))
        pos = m.end()
    segs.extend(('prose', s) for s in _split_prose(text[pos:]))
    return [(k, s) for k, s in segs if s.strip()]


def _split_prose(text):
    out = []
    for line in text.split('\n'):
        line = line.strip()
        if not line:
            continue
        out.extend(s.strip() for s in SENT_SPLIT.split(line) if s.strip())
    return out


def group_steps(segs, max_words=160, max_sents=10):
    """Merge consecutive prose sentences into step-sized passages; code stays atomic."""
    out, buf, words = [], [], 0

    def flush():
        nonlocal buf, words
        if buf:
            out.append(('prose', ' '.join(buf)))
            buf, words = [], 0

    for kind, s in segs:
        if kind == 'code':
            flush()
            out.append((kind, s))
            continue
        buf.append(s)
        words += len(s.split())
        if words >= max_words or len(buf) >= max_sents:
            flush()
    flush()
    return out


def facts(s):
    """Numbers + identifiers that must survive compression.
    List-enumeration markers ("1. Load...") are structure, not facts."""
    nums = set(NUM.findall(s)) - set(_LIST_MARKER.findall(s))
    idents = set(i.strip('`') for i in IDENT.findall(s))
    return nums | idents


def facts_preserved(src, out):
    """Substring presence — regex \\b breaks against adjacent CJK chars.
    Returns the list of MISSING facts (empty list = all preserved)."""
    out_n = out.replace(',', '')
    return [f for f in facts(src) if f.replace(',', '') not in out_n]


def classify_passage(seg, seen_facts, ntok):
    """'load' = fact-dense or novel-fact-bearing (step-faithful treatment);
       'narr' = search/narrative (stub treatment).
    ntok is a callable: text -> token count under your target tokenizer."""
    f = facts(seg)
    novel = f - seen_facts
    toks = max(ntok(seg), 1)
    dens = (len(NUM.findall(seg)) + len(IDENT.findall(seg))
            + sum(seg.count(o) for o in _OPS)) / toks
    if novel and (dens >= 0.08 or len(novel) >= 3):
        return 'load'
    if dens >= 0.15:
        return 'load'
    return 'narr'


def gate(src_seg, rules_seg, out, ntok, novel=None):
    """Deterministic per-passage fidelity gate.
    Returns None if the model output is admissible, else a short fail-reason
    string; on failure the caller uses rules_seg instead.

    novel: the passage's facts that are NOT already in the accumulated chain.
    The prompt tells the model never to restate chain content, so only novel
    facts are required to survive (matching the acceptance harness). Pass None
    to require every fact of the passage (stricter, for chainless use)."""
    if not out or not out.strip():
        return "empty"
    if '```' in out:
        return "fence"
    if len(out) > 2 * len(src_seg) + 40:           # explanation/blow-up guard
        return "blowup"
    required = facts(src_seg) if novel is None else novel
    out_n = out.replace(',', '')
    if any(f.replace(',', '') not in out_n for f in required):
        return "facts"
    if ntok(out) > ntok(rules_seg):                # must not exceed the rules-only version
        return "tokens"
    return None