causalatee.models

Structural (typing.Protocol) interfaces for causality models — one per task, plus PairwiseIdentification (what nearly every trained relation classifier actually implements) and Extraction (the end-to-end composition of all three subtasks). A class satisfies a protocol purely by having a matching __call__ signature — no inheritance required. See each task page's "Models" section for which protocol applies there.

models

Classes

CandidateExtraction

Bases: Protocol

Task.CausalCandidateExtraction: extracts candidate cause/effect spans.

Matches integrations.huggingface.CausalCandidateExtractionPipeline as-is.

Detection

Bases: Protocol

Task.CausalityDetection: classifies whether a sentence is causal at all.

Matches integrations.huggingface.CausalityDetectionPipeline as-is (batch calling is inherited for free from transformers.Pipeline).

ExtractedRelation

Bases: TypedDict

One relationship extracted from raw, unmarked text (see the glossary) -- relationship names which relation TYPE this e1/e2 pair was classified into.

Extraction

Bases: Protocol

The end-to-end task: raw, unmarked text straight to structured relations.

No current model in this repo implements this directly; see compose_extraction to build one from a Detection, a CandidateExtraction, and a PairwiseIdentification model.

Identification

Bases: Protocol

Task.CausalityIdentification in full generality.

Classifies every relation among an arbitrary number of pre-marked spans (<e1>...<eN>) in one call, matching the relations column of the identification dataset schema directly -- a sentence may carry more than two marked entities and more than one relation record. No current model in this repo implements this directly; see lift_pairwise_identification to build one from any PairwiseIdentification model.

IdentifiedRelation

Bases: TypedDict

One relationship among an arbitrary number of marked spans (see the glossary) -- relationship names which relation TYPE this first/second pair was classified into.

PairwiseIdentification

Bases: Protocol

Classifies the relation between exactly the two spans already marked <e1>/<e2> in text.

This is what every current identification model in this repo implements (integrations.huggingface.CausalityIdentificationPipeline), and what nearly every relation-classification model in the literature implements: one relation in, one relation out, entity ids fixed to e1/e2 by convention because the caller already knows which two spans it's asking about.

PairwiseRelation

Bases: TypedDict

One relationship -- a single pair's classification, not the relation itself (see the glossary).

Functions:

compose_extraction

compose_extraction(detection: Detection, candidate_extraction: CandidateExtraction, identification: PairwiseIdentification) -> Extraction

Compose a Detection, a CandidateExtraction, and a PairwiseIdentification model into the end-to-end Extraction interface.

Mirrors lift_pairwise_identification's shape (narrower model(s) in, a model satisfying a broader protocol out), but across three DIFFERENT protocols rather than widening one. Backend-agnostic like every function in this module: each argument can be an HF transformers.Pipeline, a rule-based detector, a custom biaffine model, or any other object satisfying the matching protocol.

See _ComposedExtraction (this module) for the batching and gating strategy.

Source code in causalatee/models/_protocols.py

def compose_extraction(
    detection: Detection,
    candidate_extraction: CandidateExtraction,
    identification: PairwiseIdentification,
) -> Extraction:
    """Compose a ``Detection``, a ``CandidateExtraction``, and a ``PairwiseIdentification`` model into the end-to-end
    ``Extraction`` interface.

    Mirrors ``lift_pairwise_identification``'s shape (narrower model(s) in, a model satisfying a broader protocol
    out), but across three DIFFERENT protocols rather than widening one. Backend-agnostic like every function in
    this module: each argument can be an HF ``transformers.Pipeline``, a rule-based detector, a custom biaffine
    model, or any other object satisfying the matching protocol.

    See ``_ComposedExtraction`` (this module) for the batching and gating strategy.
    """

    return _ComposedExtraction(detection, candidate_extraction, identification)

identify_candidates

identify_candidates(texts: list[str], spans_per_text: list[list[CandidateSpan]], identification: Identification) -> list[list[ExtractedRelation]]

Mark every text's candidate spans at once (<e0>, <e1>, ... one per span, not just two) and identify relations among them via identification, remapping the resulting entity ids back to actual span text.

texts and spans_per_text must be the same length and pairwise-aligned (spans_per_text[i] are texts[i]'s candidate spans). A text with fewer than two spans contributes an empty relation list without ever reaching identification -- callers that need to gate on e.g. Detection first should pass an empty span list for texts that shouldn't be identified at all (see _ComposedExtraction for the causality-gating example).

identification can be any Identification model, including one built by lift_pairwise_identification from a PairwiseIdentification model, or a genuinely joint/N-ary model. Shared by compose_extraction and causalatee.mining's identification stage so the marking/enumeration/remapping logic exists once, not twice.

Source code in causalatee/models/_protocols.py

def identify_candidates(
    texts: list[str],
    spans_per_text: list[list[CandidateSpan]],
    identification: Identification,
) -> list[list[ExtractedRelation]]:
    """Mark every text's candidate spans at once (``<e0>``, ``<e1>``, ... one per span, not just two) and identify
    relations among them via ``identification``, remapping the resulting entity ids back to actual span text.

    ``texts`` and ``spans_per_text`` must be the same length and pairwise-aligned (``spans_per_text[i]`` are
    ``texts[i]``'s candidate spans). A text with fewer than two spans contributes an empty relation list without
    ever reaching ``identification`` -- callers that need to gate on e.g. ``Detection`` first should pass an empty
    span list for texts that shouldn't be identified at all (see ``_ComposedExtraction`` for the causality-gating
    example).

    ``identification`` can be any ``Identification`` model, including one built by
    [`lift_pairwise_identification`][causalatee.models.lift_pairwise_identification] from a
    ``PairwiseIdentification`` model, or a genuinely joint/N-ary model. Shared by ``compose_extraction`` and
    ``causalatee.mining``'s identification stage so the marking/enumeration/remapping logic exists once, not twice.
    """

    marked_texts: list[str] = []
    marked_source_index: list[int] = []
    for i, spans in enumerate(spans_per_text):
        if len(spans) < 2:
            continue
        segments = {f"e{j}": [(span["start"], span["end"])] for j, span in enumerate(spans)}
        marked_texts.append(insert_entity_markers(texts[i], segments))
        marked_source_index.append(i)

    identified = identification(marked_texts) if marked_texts else []

    results: list[list[ExtractedRelation]] = [[] for _ in texts]
    for source_index, relations in zip(marked_source_index, identified):
        text = texts[source_index]
        spans = spans_per_text[source_index]
        extracted: list[ExtractedRelation] = []
        for rel in relations:
            first_span = spans[int(rel["first"][1:])]
            second_span = spans[int(rel["second"][1:])]
            extracted.append(
                {
                    "e1": text[first_span["start"] : first_span["end"]],
                    "e2": text[second_span["start"] : second_span["end"]],
                    "relationship": rel["relationship"].name,
                    "score": rel["score"],
                }
            )
        results[source_index] = extracted
    return results

lift_pairwise_identification

lift_pairwise_identification(model: PairwiseIdentification) -> Identification

Lift any PairwiseIdentification model into the general Identification interface via exhaustive pair enumeration.

See _LiftedPairwiseIdentification (this module) for the enumeration strategy, its batched underlying calls, and its O(n^2)-pairs-per-text cost caveat.

Source code in causalatee/models/_protocols.py

def lift_pairwise_identification(model: PairwiseIdentification) -> Identification:
    """Lift any ``PairwiseIdentification`` model into the general ``Identification`` interface via exhaustive pair
    enumeration.

    See ``_LiftedPairwiseIdentification`` (this module) for the enumeration strategy, its batched underlying calls,
    and its O(n^2)-pairs-per-text cost caveat.
    """

    return _LiftedPairwiseIdentification(model)