extraction/sectioning.py

"""Structured abstract section detection and conservative sentence splitting."""
from __future__ import annotations

import re
from dataclasses import dataclass


@dataclass(frozen=True)
class AbstractSentence:
    text: str
    source_location: str
    ordinal: int


_HEADING_TO_LOCATION = {
    "BACKGROUND": "abstract_background",
    "BACKGROUND AND OBJECTIVE": "abstract_background",
    "OBJECTIVE": "abstract_background",
    "OBJECTIVES": "abstract_background",
    "AIM": "abstract_background",
    "AIMS": "abstract_background",
    "INTRODUCTION": "abstract_background",
    "METHOD": "abstract_methods",
    "METHODS": "abstract_methods",
    "MATERIAL AND METHODS": "abstract_methods",
    "MATERIALS AND METHODS": "abstract_methods",
    "DESIGN": "abstract_methods",
    "PARTICIPANTS": "abstract_methods",
    "RESULT": "abstract_results",
    "RESULTS": "abstract_results",
    "FINDINGS": "abstract_results",
    "CONCLUSION": "abstract_conclusion",
    "CONCLUSIONS": "abstract_conclusion",
    "INTERPRETATION": "abstract_conclusion",
}

# INLINE_STRUCTURED_ABSTRACT_HEADINGS_V423
_HEADING_RE = re.compile(
    r"(?im)(?<!\S)(BACKGROUND(?: AND OBJECTIVE)?|OBJECTIVES?|AIMS?|INTRODUCTION|METHODS?|MATERIALS? AND METHODS|DESIGN|PARTICIPANTS|RESULTS?|FINDINGS|CONCLUSIONS?|INTERPRETATION)\s*[:.]\s*"
)

_SENTENCE_RE = re.compile(r"(?<=[.!?])\s+(?=(?:[A-Z0-9]|\())")


def _clean(text: str) -> str:
    return re.sub(r"\s+", " ", text).strip()


def split_sentences(text: str) -> list[str]:
    cleaned = _clean(text)
    if not cleaned:
        return []
    parts = [_clean(part) for part in _SENTENCE_RE.split(cleaned)]
    return [part for part in parts if len(part) >= 10]


def section_abstract(abstract: str) -> list[AbstractSentence]:
    matches = list(_HEADING_RE.finditer(abstract))
    out: list[AbstractSentence] = []
    ordinal = 0

    if not matches:
        for sentence in split_sentences(abstract):
            ordinal += 1
            out.append(AbstractSentence(sentence, "abstract_unspecified", ordinal))
        return out

    prefix = abstract[: matches[0].start()].strip()
    for sentence in split_sentences(prefix):
        ordinal += 1
        out.append(AbstractSentence(sentence, "abstract_unspecified", ordinal))

    for index, match in enumerate(matches):
        heading = match.group(1).upper().strip()
        start = match.end()
        end = matches[index + 1].start() if index + 1 < len(matches) else len(abstract)
        section_text = abstract[start:end]
        location = _HEADING_TO_LOCATION.get(heading, "abstract_unspecified")
        for sentence in split_sentences(section_text):
            ordinal += 1
            out.append(AbstractSentence(sentence, location, ordinal))

    return out