Skip to content

causalatee.data.utils

utils

Shared helpers for causalatee's HF-datasets-native schema: entity-marker parsing, whole-document -> per-sentence re-splitting, and same-granularity schema derivations -- everything a dataset conversion_script.py or a downstream evaluation harness needs that isn't specific to any one corpus.

Split into submodules by concern (all re-exported here, so from causalatee.data.utils import X keeps working unchanged):

Several source corpora (BioCause, CaTeRS's brat annotations, CREST's own aggregation) annotate causal relations over a whole document or section, not a single sentence. A per-sentence schema cannot represent a relation whose participants straddle a sentence boundary (discourse-level causality really does occur, e.g. BioCause: "X requires Y ... [new sentence] This substitution decreases Z") -- silently re-basing such a span onto the wrong sentence's text is a real bug, not a hypothetical one (previously observed: a marker landing mid-word, "histidine", in an early BioCause converter draft). splitting's job is to make the correct behavior a single, tested, reusable piece of code instead of a duplicated ad-hoc guard in every converter. * batch -- Dataset.map(fn, batched=True, remove_columns=[...])- ready adapters over the above: identification_batch_to_sentences, extraction_batch_to_sentences, identification_batch_to_detection_sentences (re-splitting to sentence level), and identification_batch_to_detection / identification_batch_to_extraction (same-granularity detection/extraction tables derived from an identification table, for datasets like CTB/ESL that only ship identification). * validation -- verify_dataset, structural sanity checks (conflicting/duplicate relation pairs, entity ids with no parseable marker span) run over a converted split. Every conversion_script.py should call this right before writing a split's parquet file and print any returned errors as warnings.

IMPORTANT (see the markers module docstring for the full explanation): the <eN>...</eN> markers only look like XML/HTML tags -- they are not, and do not need to be well-nested. Do not "fix" a non-nested marker sequence, and do not add crossing-span detection/dropping logic on the assumption that it's needed here -- it isn't.

Typical usage, re-splitting an already-loaded whole-document identification split to sentence level (a detection or extraction split works the same way with the matching batch function)::

import spacy
from causalatee.data.utils import identification_batch_to_sentences

nlp = spacy.load("en_core_web_sm", disable=["ner", "lemmatizer", "tagger"])

def sentence_ranges(text):
    return [(s.start_char, s.end_char) for s in nlp(text).sents]

per_sentence = whole_document_ds.map(
    lambda batch: identification_batch_to_sentences(batch, sentence_ranges),
    batched=True,
    remove_columns=whole_document_ds.column_names,
)

Or, deriving a detection/extraction table at the SAME granularity for a dataset that only ships identification (no sentence splitting involved)::

from causalatee.data.utils import identification_batch_to_detection

detection_ds = identification_ds.map(
    identification_batch_to_detection,
    batched=True,
    remove_columns=identification_ds.column_names,
)

See conf-causality-repro/evaluation/data.py for a complete example that dispatches all three sub-tasks this way.

Functions:

extraction_batch_to_sentences 

extraction_batch_to_sentences(batch: Mapping[str, Sequence], sentence_ranges_fn: Callable[[str], Sequence[Span]]) -> dict[str, list]

Batched wrapper around split_extraction_to_sentences.

Source code in causalatee/data/utils/batch.py 
69
70
71
72
73
74
75
76
77
78
79
80
def extraction_batch_to_sentences(
    batch: Mapping[str, Sequence],
    sentence_ranges_fn: Callable[[str], Sequence[Span]],
) -> dict[str, list]:
    """Batched wrapper around [`split_extraction_to_sentences`][causalatee.data.utils.split_extraction_to_sentences]."""
    out_text: list[str] = []
    out_entity: list[list] = []
    for text, entities in zip(batch["text"], batch["entity"]):
        rows, _ = split_extraction_to_sentences(text, entities, sentence_ranges_fn(text))
        out_text.extend(row["text"] for row in rows)
        out_entity.extend(row["entity"] for row in rows)
    return {"text": out_text, "entity": out_entity}

identification_batch_to_detection 

identification_batch_to_detection(batch: Mapping[str, Sequence]) -> dict[str, list]

Derive a causality-DETECTION batch (text + 0/1 label) from an IDENTIFICATION batch, at the SAME granularity (no sentence splitting).

A row is causal iff it has at least one ACTUALLY CAUSAL relation (see _filter_causal_relations -- a merely non-empty relations list is NOT enough, since some converters explicitly list NoRelation entries for non-causal pairs rather than omitting them). This matches how every converter in this project derives a detection label from relation presence (e.g. CCNC's countercausal sentences, which still count as causal for detection since they're about a causal relationship, just negating it). <eN>...</eN> markers are stripped from text.

NOT the right choice for BioCause: its own detection convention counts a document as causal if it has ANY marked entity, even an unpaired Effect-only one with no Cause and hence no relation -- see identification_batch_to_detection_sentences for that case.

Source code in causalatee/data/utils/batch.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def identification_batch_to_detection(batch: Mapping[str, Sequence]) -> dict[str, list]:
    """Derive a causality-DETECTION batch (text + 0/1 label) from an IDENTIFICATION batch, at the SAME granularity (no
    sentence splitting).

    A row is causal iff it has at least one ACTUALLY CAUSAL relation (see ``_filter_causal_relations`` -- a
    merely non-empty ``relations`` list is NOT enough, since some converters explicitly list NoRelation entries for
    non-causal pairs rather than omitting them). This matches how every converter in this project derives a
    detection label from relation presence (e.g. CCNC's countercausal sentences, which still count as causal for
    detection since they're about a causal relationship, just negating it). ``<eN>...</eN>`` markers are stripped
    from ``text``.

    NOT the right choice for BioCause: its own detection convention counts a document as causal if it has ANY marked
    entity, even an unpaired Effect-only one with no Cause and hence no relation -- see
    [`identification_batch_to_detection_sentences`][causalatee.data.utils.identification_batch_to_detection_sentences]
    for that case.
    """
    out_text: list[str] = []
    out_label: list[int] = []
    for text, relations in zip(batch["text"], batch["relations"]):
        clean_text, _ = parse_entity_markers(text)
        out_text.append(clean_text)
        out_label.append(1 if _filter_causal_relations(relations) else 0)
    return {"text": out_text, "label": out_label}

identification_batch_to_detection_sentences 

identification_batch_to_detection_sentences(batch: Mapping[str, Sequence], sentence_ranges_fn: Callable[[str], Sequence[Span]]) -> dict[str, list]

Derive per-sentence causality-DETECTION rows (text + 0/1 label) from a whole-document IDENTIFICATION batch.

Some corpora's own detection table is whole-document/section only (e.g. BioCause) even though their identification table marks every entity -- cause+effect pairs AND unpaired effect-only ones alike. This reuses that richer information to still produce a per-sentence label: a sentence is causal if it contains ANY marked entity at all, matching the common "a document is causal if it contains any causal event, span or no span" convention.

Source code in causalatee/data/utils/batch.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def identification_batch_to_detection_sentences(
    batch: Mapping[str, Sequence],
    sentence_ranges_fn: Callable[[str], Sequence[Span]],
) -> dict[str, list]:
    """Derive per-sentence causality-DETECTION rows (text + 0/1 label) from a whole-document IDENTIFICATION batch.

    Some corpora's own detection table is whole-document/section only (e.g. BioCause) even though their
    identification table marks every entity -- cause+effect pairs AND unpaired effect-only ones alike. This reuses
    that richer information to still produce a per-sentence label: a sentence is causal if it contains ANY marked
    entity at all, matching the common "a document is causal if it contains any causal event, span or no span"
    convention.
    """
    out_text: list[str] = []
    out_label: list[int] = []
    for text in batch["text"]:
        clean_text, entity_segments = parse_entity_markers(text)
        sent_ranges = sentence_ranges_fn(clean_text)
        groups = [(eid,) for eid in entity_segments]
        rows, _ = split_document_by_sentence(sent_ranges, entity_segments, groups)
        for (s, e), row in zip(sent_ranges, rows):
            out_text.append(clean_text[s:e])
            out_label.append(1 if row["entities"] else 0)
    return {"text": out_text, "label": out_label}

identification_batch_to_extraction 

identification_batch_to_extraction(batch: Mapping[str, Sequence]) -> dict[str, list]

Derive a causal-candidate-EXTRACTION batch (text + entity spans) from an IDENTIFICATION batch, at the SAME granularity (no sentence splitting).

Only entities that participate in at least one ACTUALLY CAUSAL relation are kept (see _filter_causal_relations) -- matching this project's "only causal_eids" convention used by other converters throughout. Rows with no causal relation at all are dropped from the output entirely (nothing to extract), so the output can have fewer rows than the input, same as this module's sentence-splitting adapters.

Source code in causalatee/data/utils/batch.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def identification_batch_to_extraction(batch: Mapping[str, Sequence]) -> dict[str, list]:
    """Derive a causal-candidate-EXTRACTION batch (text + entity spans) from an IDENTIFICATION batch, at the SAME
    granularity (no sentence splitting).

    Only entities that participate in at least one ACTUALLY CAUSAL relation are kept (see
    ``_filter_causal_relations``) -- matching this project's "only causal_eids" convention used by other
    converters throughout. Rows with no causal relation at all are dropped from the output entirely (nothing to
    extract), so the output can have fewer rows than the input, same as this module's sentence-splitting adapters.
    """
    out_text: list[str] = []
    out_entity: list[list] = []
    for text, relations in zip(batch["text"], batch["relations"]):
        causal = _filter_causal_relations(relations)
        if not causal:
            continue
        clean_text, segments = parse_entity_markers(text)
        involved = sorted({eid for rel in causal for eid in (rel["first"], rel["second"])})
        entity = [[x for seg in segments[eid] for x in seg] for eid in involved]
        out_text.append(clean_text)
        out_entity.append(entity)
    return {"text": out_text, "entity": out_entity}

identification_batch_to_sentences 

identification_batch_to_sentences(batch: Mapping[str, Sequence], sentence_ranges_fn: Callable[[str], Sequence[Span]]) -> dict[str, list]

Batched wrapper around split_identification_to_sentences.

sentence_ranges_fn(clean_text) -> [(start, end), ...] supplies sentence boundaries for the marker-stripped text (any segmenter -- spaCy, nltk, a regex splitter).

Source code in causalatee/data/utils/batch.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def identification_batch_to_sentences(
    batch: Mapping[str, Sequence],
    sentence_ranges_fn: Callable[[str], Sequence[Span]],
) -> dict[str, list]:
    """Batched wrapper around
    [`split_identification_to_sentences`][causalatee.data.utils.split_identification_to_sentences].

    ``sentence_ranges_fn(clean_text) -> [(start, end), ...]`` supplies sentence boundaries for the marker-stripped
    text (any segmenter -- spaCy, nltk, a regex splitter).
    """
    out_text: list[str] = []
    out_relations: list[list] = []
    for text, relations in zip(batch["text"], batch["relations"]):
        clean_text, _ = parse_entity_markers(text)
        rows, _ = split_identification_to_sentences(text, relations, sentence_ranges_fn(clean_text))
        out_text.extend(row["text"] for row in rows)
        out_relations.extend(row["relations"] for row in rows)
    return {"text": out_text, "relations": out_relations}

insert_entity_markers 

insert_entity_markers(text: str, segments_by_eid: Mapping[str, Sequence[Span]]) -> str

Insert <eN>...</eN> tags around each id's segment(s).

Left-to-right, stack-based: closes innermost-first and opens outermost (longest span) first at each boundary, so nested and disjoint spans come out correctly nested. Output need not be nested when spans cross -- see the module docstring. Inverse of parse_entity_markers.

Source code in causalatee/data/utils/markers.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def insert_entity_markers(text: str, segments_by_eid: Mapping[str, Sequence[Span]]) -> str:
    """Insert ``<eN>...</eN>`` tags around each id's segment(s).

    Left-to-right, stack-based: closes innermost-first and opens outermost (longest span) first at each boundary, so
    nested and disjoint spans come out correctly nested. Output need not be nested when spans cross -- see the
    module docstring. Inverse of [`parse_entity_markers`][causalatee.data.utils.parse_entity_markers].
    """
    opens: dict[int, list[str]] = {}
    closes: dict[int, list[str]] = {}
    for eid, segments in segments_by_eid.items():
        for start, end in segments:
            opens.setdefault(start, []).append(eid)
            closes.setdefault(end, []).append(eid)

    pieces: list[str] = []
    stack: list[str] = []
    pos = 0
    for boundary in sorted(set(opens) | set(closes)):
        pieces.append(text[pos:boundary])
        pos = boundary
        for eid in reversed(stack.copy()):
            if eid in closes.get(boundary, []):
                pieces.append(f"</{eid}>")
                stack.remove(eid)
        # Open outer (longer) spans first among spans starting here, so
        # nesting stays valid when several ids open at the same boundary.
        for eid in sorted(opens.get(boundary, []), key=lambda e: -_span_length(segments_by_eid[e])):
            pieces.append(f"<{eid}>")
            stack.append(eid)
    pieces.append(text[pos:])
    return "".join(pieces)

parse_entity_markers 

parse_entity_markers(text: str) -> tuple[str, dict[str, list[Span]]]

Strip <eN>/</eN> tags; return (clean_text, segments_by_eid).

An id may open/close more than once -- a discontinuous mention, rendered as repeated <eN>...</eN> occurrences sharing the same id -- every occurrence accumulates onto that id's segment list, in left-to-right order. Inverse of insert_entity_markers.

NOT an XML parser: matching is per entity id (module docstring has the details), so text does not need to be well-nested -- markers for two genuinely crossing spans parse correctly.

Source code in causalatee/data/utils/markers.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def parse_entity_markers(text: str) -> tuple[str, dict[str, list[Span]]]:
    """Strip ``<eN>``/``</eN>`` tags; return (clean_text, segments_by_eid).

    An id may open/close more than once -- a discontinuous mention, rendered as repeated ``<eN>...</eN>`` occurrences
    sharing the same id -- every occurrence accumulates onto that id's segment list, in left-to-right order. Inverse
    of [`insert_entity_markers`][causalatee.data.utils.insert_entity_markers].

    NOT an XML parser: matching is per entity id (module docstring has the details), so ``text`` does not need to be
    well-nested -- markers for two genuinely crossing spans parse correctly.
    """
    spans: dict[str, list[Span]] = {}
    open_positions: dict[str, int] = {}
    clean: list[str] = []
    pos = 0
    for m in _MARKER_RE.finditer(text):
        clean.append(text[pos : m.start()])
        offset = sum(len(c) for c in clean)
        closing, eid = m.group(1), m.group(2)
        if closing:
            spans.setdefault(eid, []).append((open_positions.get(eid, offset), offset))
        else:
            open_positions[eid] = offset
        pos = m.end()
    clean.append(text[pos:])
    return "".join(clean), spans

split_document_by_sentence 

split_document_by_sentence(sent_ranges: Sequence[Span], entity_segments: Mapping[str, Sequence[Span]], groups: Sequence[Sequence[str]]) -> tuple[list[dict], int]

Assign each entity group to the one sentence it fits wholly inside.

A group is a list of entity ids that must be kept together (e.g. the two participants of a relation, or a single unpaired entity). A group survives only if every entity id it references has ALL of its segments wholly inside the SAME sentence's character range; groups that straddle a sentence boundary, or whose entities don't overlap any known sentence at all, are dropped and counted rather than silently corrupted or discarded without a trace.

Returns (rows, num_dropped): rows has one dict per input sentence range, in order, {"entities": {eid: [(local_start, local_end), ...]}, "group_indices": [i, ...]} -- entity offsets already re-based to be local to that sentence, and group_indices the indices into the input groups list that landed in this sentence (empty for sentences with no surviving group, so callers can always align rows 1:1 with sent_ranges).

Source code in causalatee/data/utils/splitting.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def split_document_by_sentence(
    sent_ranges: Sequence[Span],
    entity_segments: Mapping[str, Sequence[Span]],
    groups: Sequence[Sequence[str]],
) -> tuple[list[dict], int]:
    """Assign each entity group to the one sentence it fits wholly inside.

    A ``group`` is a list of entity ids that must be kept together (e.g. the two participants of a relation, or a
    single unpaired entity). A group survives only if every entity id it references has ALL of its segments wholly
    inside the SAME sentence's character range; groups that straddle a sentence boundary, or whose entities don't
    overlap any known sentence at all, are dropped and counted rather than silently corrupted or discarded without a
    trace.

    Returns ``(rows, num_dropped)``: ``rows`` has one dict per input sentence range, in order, ``{"entities": {eid:
    [(local_start, local_end), ...]}, "group_indices": [i, ...]}`` -- entity offsets already re-based to be local to
    that sentence, and ``group_indices`` the indices into the input ``groups`` list that landed in this sentence
    (empty for sentences with no surviving group, so callers can always align rows 1:1 with ``sent_ranges``).
    """
    rows: list[dict] = [{"entities": {}, "group_indices": []} for _ in sent_ranges]
    num_dropped = 0
    for group_index, group in enumerate(groups):
        all_segments = [seg for eid in group for seg in entity_segments[eid]]
        anchor = min((s for s, _ in all_segments), default=None)
        si = _sentence_index(anchor, sent_ranges) if anchor is not None else None
        if si is None:
            num_dropped += 1
            continue
        s, e = sent_ranges[si]
        if not all(_wholly_inside(entity_segments[eid], s, e) for eid in group):
            num_dropped += 1
            continue
        for eid in group:
            if eid not in rows[si]["entities"]:
                rows[si]["entities"][eid] = [(a - s, b - s) for a, b in entity_segments[eid]]
        rows[si]["group_indices"].append(group_index)
    return rows, num_dropped

split_extraction_to_sentences 

split_extraction_to_sentences(text: str, entities: Sequence[Sequence[int]], sent_ranges: Sequence[Span]) -> tuple[list[dict], int]

Whole-document causal-candidate-EXTRACTION example -> per-sentence rows.

entities is causalatee's entity field: a list of flat, even-length document-level offset lists ([s, e] contiguous, [s1, e1, s2, e2, ...] discontinuous). Returns (rows, num_dropped) where each row is {"text": ..., "entity": [[...], ...]}, one per sentence, and num_dropped counts entities that don't fit wholly inside a single sentence.

Source code in causalatee/data/utils/splitting.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def split_extraction_to_sentences(
    text: str,
    entities: Sequence[Sequence[int]],
    sent_ranges: Sequence[Span],
) -> tuple[list[dict], int]:
    """Whole-document causal-candidate-EXTRACTION example -> per-sentence rows.

    ``entities`` is causalatee's ``entity`` field: a list of flat, even-length document-level offset lists (``[s,
    e]`` contiguous, ``[s1, e1, s2, e2, ...]`` discontinuous). Returns ``(rows, num_dropped)`` where each row is
    ``{"text": ..., "entity": [[...], ...]}``, one per sentence, and ``num_dropped`` counts entities that don't fit
    wholly inside a single sentence.
    """
    entity_segments = {
        str(i): [(flat[j], flat[j + 1]) for j in range(0, len(flat), 2)] for i, flat in enumerate(entities)
    }
    groups = [(str(i),) for i in range(len(entities))]
    split_rows, num_dropped = split_document_by_sentence(sent_ranges, entity_segments, groups)

    rows = []
    for (s, e), split_row in zip(sent_ranges, split_rows):
        local_entities = [
            [x for seg in split_row["entities"][group_id] for x in seg]
            for gi in split_row["group_indices"]
            for group_id in [groups[gi][0]]
        ]
        rows.append({"text": text[s:e], "entity": local_entities})
    return rows, num_dropped

split_identification_to_sentences 

split_identification_to_sentences(text: str, relations: Sequence[Mapping], sent_ranges: Sequence[Span]) -> tuple[list[dict], int]

Whole-document causality-IDENTIFICATION example -> per-sentence rows.

text carries <eN>...</eN> entity markers (causalatee's identification schema); relations is the matching list of {"relationship", "first": eid, "second": eid} dicts. sent_ranges are character offsets into the MARKER-STRIPPED text (i.e. what parse_entity_markers returns as clean_text), not the raw text -- markers shift offsets, so a sentence segmenter should run on the clean text first. Returns (rows, num_dropped) where each row is {"text": ..., "relations": [...]}, one per sentence, and num_dropped counts relations whose two entities don't both fit inside a single sentence.

Source code in causalatee/data/utils/splitting.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def split_identification_to_sentences(
    text: str,
    relations: Sequence[Mapping],
    sent_ranges: Sequence[Span],
) -> tuple[list[dict], int]:
    """Whole-document causality-IDENTIFICATION example -> per-sentence rows.

    ``text`` carries ``<eN>...</eN>`` entity markers (causalatee's identification schema); ``relations`` is the
    matching list of ``{"relationship", "first": eid, "second": eid}`` dicts. ``sent_ranges`` are character offsets
    into the MARKER-STRIPPED text (i.e. what [`parse_entity_markers`][causalatee.data.utils.parse_entity_markers]
    returns as ``clean_text``), not the raw
    ``text`` -- markers shift offsets, so a sentence segmenter should run on the clean text first. Returns ``(rows,
    num_dropped)`` where each row is ``{"text": ..., "relations": [...]}``, one per sentence, and ``num_dropped``
    counts relations whose two entities don't both fit inside a single sentence.
    """
    clean_text, entity_segments = parse_entity_markers(text)
    groups = [(rel["first"], rel["second"]) for rel in relations]
    split_rows, num_dropped = split_document_by_sentence(sent_ranges, entity_segments, groups)

    rows = []
    for (s, e), split_row in zip(sent_ranges, split_rows):
        local_relations = [relations[gi] for gi in split_row["group_indices"]]
        local_text = (
            insert_entity_markers(clean_text[s:e], split_row["entities"]) if split_row["entities"] else clean_text[s:e]
        )
        rows.append({"text": local_text, "relations": local_relations})
    return rows, num_dropped

verify_dataset 

verify_dataset(batch: Mapping[str, Sequence] | DataFrame, task: Task) -> list[str]

Run every applicable structural check over one converted split.

batch is column-oriented (batch["text"], batch["relations"], ...) -- same convention as the batch module's Dataset.map-ready functions, so a pandas.DataFrame works directly by column name right before .to_parquet(), with no conversion needed. Returns a flat list of human-readable error descriptions (one entry per issue found), or [] if the split is clean.

Source code in causalatee/data/utils/validation.py 
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def verify_dataset(batch: Mapping[str, Sequence] | pd.DataFrame, task: Task) -> list[str]:
    """Run every applicable structural check over one converted split.

    ``batch`` is column-oriented (``batch["text"]``, ``batch["relations"]``, ...) -- same convention as the
    ``batch`` module's ``Dataset.map``-ready functions, so a ``pandas.DataFrame`` works directly by column name right
    before ``.to_parquet()``, with no conversion needed. Returns a flat list of human-readable error descriptions
    (one entry per issue found), or ``[]`` if the split is clean.
    """
    if task != Task.CausalityIdentification or "relations" not in batch:
        return []
    errors: list[str] = []
    errors.extend(_check_relation_pairs(batch["relations"]))
    errors.extend(_check_entity_markup(batch["text"], batch["relations"]))
    return errors