corp-extractor 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

statement_extractor/models/entity.py

@@ -0,0 +1,102 @@
+"""
+Entity models for the extraction pipeline.
+
+ExtractedEntity represents entities identified during extraction with
+confidence scores and span information.
+
+Note: EntityType is imported from the original models.py for consistency.
+"""
+
+from typing import Optional, TYPE_CHECKING
+import uuid
+
+from pydantic import BaseModel, Field
+
+# Import EntityType from parent module to avoid duplication
+# This will be populated by __init__.py which loads from old models.py
+if TYPE_CHECKING:
+    from enum import Enum
+
+    class EntityType(str, Enum):
+        """Supported entity types for subjects and objects."""
+        ORG = "ORG"
+        PERSON = "PERSON"
+        GPE = "GPE"
+        LOC = "LOC"
+        PRODUCT = "PRODUCT"
+        EVENT = "EVENT"
+        WORK_OF_ART = "WORK_OF_ART"
+        LAW = "LAW"
+        DATE = "DATE"
+        MONEY = "MONEY"
+        PERCENT = "PERCENT"
+        QUANTITY = "QUANTITY"
+        UNKNOWN = "UNKNOWN"
+else:
+    # At runtime, we need to import it from somewhere
+    # Try the old models.py location first
+    try:
+        import importlib.util
+        from pathlib import Path
+        _models_py_path = Path(__file__).parent.parent / "models.py"
+        _spec = importlib.util.spec_from_file_location("_old_models", _models_py_path)
+        _old_models = importlib.util.module_from_spec(_spec)
+        _spec.loader.exec_module(_old_models)
+        EntityType = _old_models.EntityType
+    except Exception:
+        # Fallback to defining it here
+        from enum import Enum
+
+        class EntityType(str, Enum):
+            """Supported entity types for subjects and objects."""
+            ORG = "ORG"
+            PERSON = "PERSON"
+            GPE = "GPE"
+            LOC = "LOC"
+            PRODUCT = "PRODUCT"
+            EVENT = "EVENT"
+            WORK_OF_ART = "WORK_OF_ART"
+            LAW = "LAW"
+            DATE = "DATE"
+            MONEY = "MONEY"
+            PERCENT = "PERCENT"
+            QUANTITY = "QUANTITY"
+            UNKNOWN = "UNKNOWN"
+
+
+class ExtractedEntity(BaseModel):
+    """
+    An entity extracted from text with type and confidence information.
+
+    Used in Stage 2 (Extraction) and flows through subsequent stages.
+    """
+    text: str = Field(..., description="The entity text as extracted")
+    type: EntityType = Field(default=EntityType.UNKNOWN, description="The entity type")
+    span: Optional[tuple[int, int]] = Field(
+        None,
+        description="Character offsets (start, end) in source text"
+    )
+    confidence: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        description="Confidence score for this entity extraction"
+    )
+    entity_ref: str = Field(
+        default_factory=lambda: str(uuid.uuid4()),
+        description="Unique reference ID for tracking this entity through the pipeline"
+    )
+
+    def __str__(self) -> str:
+        return f"{self.text} ({self.type.value})"
+
+    def __hash__(self) -> int:
+        return hash(self.entity_ref)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, ExtractedEntity):
+            return False
+        return self.entity_ref == other.entity_ref
+
+    class Config:
+        frozen = False  # Allow modification during pipeline stages

statement_extractor/models/labels.py

@@ -0,0 +1,191 @@
+"""
+Label models for the extraction pipeline.
+
+StatementLabel: A label applied to a statement
+LabeledStatement: Final output from Stage 5 with all labels
+TaxonomyResult: Taxonomy classification from Stage 6
+"""
+
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel, Field
+
+from .statement import PipelineStatement
+from .canonical import CanonicalEntity
+
+
+class StatementLabel(BaseModel):
+    """
+    A label applied to a statement in Stage 5 (Labeling).
+
+    Labels can represent sentiment, relation type, confidence, or
+    any other classification applied by labeler plugins.
+    """
+    label_type: str = Field(
+        ...,
+        description="Type of label: 'sentiment', 'relation_type', 'confidence', etc."
+    )
+    label_value: Union[str, float, bool] = Field(
+        ...,
+        description="The label value (string for classification, float for scores)"
+    )
+    confidence: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        description="Confidence in this label"
+    )
+    labeler: Optional[str] = Field(
+        None,
+        description="Name of the labeler plugin that produced this label"
+    )
+
+    def is_high_confidence(self, threshold: float = 0.8) -> bool:
+        """Check if this is a high-confidence label."""
+        return self.confidence >= threshold
+
+
+class LabeledStatement(BaseModel):
+    """
+    Final output from Stage 5 (Labeling) with taxonomy from Stage 6.
+
+    Contains the original statement, canonicalized subject and object,
+    all labels applied by labeler plugins, and taxonomy classifications.
+    """
+    statement: PipelineStatement = Field(
+        ...,
+        description="The original statement from Stage 2"
+    )
+    subject_canonical: CanonicalEntity = Field(
+        ...,
+        description="Canonicalized subject entity"
+    )
+    object_canonical: CanonicalEntity = Field(
+        ...,
+        description="Canonicalized object entity"
+    )
+    labels: list[StatementLabel] = Field(
+        default_factory=list,
+        description="Labels applied to this statement"
+    )
+    taxonomy_results: list["TaxonomyResult"] = Field(
+        default_factory=list,
+        description="Taxonomy classifications from Stage 6"
+    )
+
+    def get_label(self, label_type: str) -> Optional[StatementLabel]:
+        """Get a label by type, or None if not found."""
+        for label in self.labels:
+            if label.label_type == label_type:
+                return label
+        return None
+
+    def get_labels_by_type(self, label_type: str) -> list[StatementLabel]:
+        """Get all labels of a specific type."""
+        return [label for label in self.labels if label.label_type == label_type]
+
+    def add_label(self, label: StatementLabel) -> None:
+        """Add a label to this statement."""
+        self.labels.append(label)
+
+    @property
+    def subject_fqn(self) -> str:
+        """Get the subject's fully qualified name."""
+        return self.subject_canonical.fqn
+
+    @property
+    def object_fqn(self) -> str:
+        """Get the object's fully qualified name."""
+        return self.object_canonical.fqn
+
+    def __str__(self) -> str:
+        """Format as FQN triple."""
+        return f"{self.subject_fqn} --[{self.statement.predicate}]--> {self.object_fqn}"
+
+    def as_dict(self) -> dict:
+        """Convert to a simplified dictionary representation."""
+        return {
+            "subject": {
+                "text": self.statement.subject.text,
+                "type": self.statement.subject.type.value,
+                "fqn": self.subject_fqn,
+                "canonical_id": (
+                    self.subject_canonical.canonical_match.canonical_id
+                    if self.subject_canonical.canonical_match else None
+                ),
+            },
+            "predicate": self.statement.predicate,
+            "object": {
+                "text": self.statement.object.text,
+                "type": self.statement.object.type.value,
+                "fqn": self.object_fqn,
+                "canonical_id": (
+                    self.object_canonical.canonical_match.canonical_id
+                    if self.object_canonical.canonical_match else None
+                ),
+            },
+            "source_text": self.statement.source_text,
+            "labels": {
+                label.label_type: label.label_value
+                for label in self.labels
+            },
+            "taxonomy": [
+                {
+                    "category": t.category,
+                    "label": t.label,
+                    "confidence": t.confidence,
+                }
+                for t in self.taxonomy_results
+            ],
+        }
+
+    class Config:
+        frozen = False  # Allow modification during pipeline stages
+
+
+class TaxonomyResult(BaseModel):
+    """
+    Result of taxonomy classification from Stage 6.
+
+    Represents a classification of a statement against a taxonomy,
+    typically with a category (top-level) and label (specific topic).
+    """
+    taxonomy_name: str = Field(
+        ...,
+        description="Name of the taxonomy (e.g., 'esg_topics', 'industry_codes')"
+    )
+    category: str = Field(
+        ...,
+        description="Top-level category (e.g., 'environment', 'governance')"
+    )
+    label: str = Field(
+        ...,
+        description="Specific label within the category (e.g., 'carbon emissions')"
+    )
+    label_id: Optional[int] = Field(
+        None,
+        description="Numeric ID for reproducibility"
+    )
+    confidence: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        description="Classification confidence"
+    )
+    classifier: Optional[str] = Field(
+        None,
+        description="Name of the taxonomy plugin that produced this result"
+    )
+    metadata: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional metadata (e.g., runner-up labels, scores)"
+    )
+
+    @property
+    def full_label(self) -> str:
+        """Get the full label in category:label format."""
+        return f"{self.category}:{self.label}"
+
+    def is_high_confidence(self, threshold: float = 0.7) -> bool:
+        """Check if this is a high-confidence classification."""
+        return self.confidence >= threshold

statement_extractor/models/qualifiers.py

@@ -0,0 +1,91 @@
+"""
+Qualifier models for the extraction pipeline.
+
+EntityQualifiers: Semantic qualifiers and external identifiers
+QualifiedEntity: Entity with qualification information from Stage 3
+"""
+
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from .entity import EntityType
+
+
+class EntityQualifiers(BaseModel):
+    """
+    Qualifiers that provide context and identifiers for an entity.
+
+    Populated by Stage 3 (Qualification) plugins such as:
+    - PersonQualifierPlugin: Adds role, org for PERSON entities
+    - GLEIFQualifierPlugin: Adds LEI for ORG entities
+    - CompaniesHouseQualifierPlugin: Adds UK company number
+    - SECEdgarQualifierPlugin: Adds SEC CIK, ticker
+    """
+    # Semantic qualifiers (for PERSON entities)
+    org: Optional[str] = Field(None, description="Organization/employer name")
+    role: Optional[str] = Field(None, description="Job title/position/role")
+
+    # Location qualifiers
+    region: Optional[str] = Field(None, description="State/province/region")
+    country: Optional[str] = Field(None, description="Country name or ISO code")
+    city: Optional[str] = Field(None, description="City name")
+    jurisdiction: Optional[str] = Field(None, description="Legal jurisdiction (e.g., 'UK', 'US-DE')")
+
+    # External identifiers (keyed by identifier type)
+    identifiers: dict[str, str] = Field(
+        default_factory=dict,
+        description="External identifiers: lei, ch_number, sec_cik, ticker, wikidata_qid, etc."
+    )
+
+    def has_any_qualifier(self) -> bool:
+        """Check if any qualifier or identifier is set."""
+        return bool(
+            self.org or self.role or self.region or self.country or
+            self.city or self.jurisdiction or self.identifiers
+        )
+
+    def merge_with(self, other: "EntityQualifiers") -> "EntityQualifiers":
+        """
+        Merge qualifiers from another instance, preferring non-None values.
+
+        Returns a new EntityQualifiers with merged values.
+        """
+        merged_identifiers = {**self.identifiers, **other.identifiers}
+        return EntityQualifiers(
+            org=other.org or self.org,
+            role=other.role or self.role,
+            region=other.region or self.region,
+            country=other.country or self.country,
+            city=other.city or self.city,
+            jurisdiction=other.jurisdiction or self.jurisdiction,
+            identifiers=merged_identifiers,
+        )
+
+
+class QualifiedEntity(BaseModel):
+    """
+    An entity with qualification information from Stage 3.
+
+    Links back to the original ExtractedEntity via entity_ref and
+    adds qualifiers from various qualification plugins.
+    """
+    entity_ref: str = Field(..., description="Reference to the original ExtractedEntity")
+    original_text: str = Field(..., description="Original entity text")
+    entity_type: EntityType = Field(..., description="Entity type")
+    qualifiers: EntityQualifiers = Field(
+        default_factory=EntityQualifiers,
+        description="Qualifiers and identifiers for this entity"
+    )
+    qualification_sources: list[str] = Field(
+        default_factory=list,
+        description="List of plugins that contributed qualifiers"
+    )
+
+    def add_qualifier_source(self, source: str) -> None:
+        """Add a qualification source to the list."""
+        if source not in self.qualification_sources:
+            self.qualification_sources.append(source)
+
+    class Config:
+        frozen = False  # Allow modification during pipeline stages

statement_extractor/models/statement.py

@@ -0,0 +1,75 @@
+"""
+Statement models for the extraction pipeline.
+
+RawTriple: Output of Stage 1 (Splitting)
+PipelineStatement: Output of Stage 2 (Extraction) with refined entities
+"""
+
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from .entity import ExtractedEntity
+
+
+class RawTriple(BaseModel):
+    """
+    A raw triple from Stage 1 (Splitting).
+
+    Contains the basic text components before entity refinement.
+    Generated by T5-Gemma or other splitting plugins.
+    """
+    subject_text: str = Field(..., description="Raw subject text")
+    predicate_text: str = Field(..., description="Raw predicate text")
+    object_text: str = Field(..., description="Raw object text")
+    source_sentence: str = Field(..., description="The source sentence this triple was extracted from")
+    confidence: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        description="Extraction confidence from the splitter"
+    )
+
+    def __str__(self) -> str:
+        return f"{self.subject_text} --[{self.predicate_text}]--> {self.object_text}"
+
+    def as_tuple(self) -> tuple[str, str, str]:
+        """Return as a simple (subject, predicate, object) tuple."""
+        return (self.subject_text, self.predicate_text, self.object_text)
+
+
+class PipelineStatement(BaseModel):
+    """
+    A statement with extracted entities from Stage 2 (Extraction).
+
+    Contains refined subject/object entities with types, spans, and confidence.
+    This is the main statement type that flows through stages 2-5.
+    """
+    subject: ExtractedEntity = Field(..., description="The subject entity")
+    predicate: str = Field(..., description="The relationship/predicate text")
+    predicate_category: Optional[str] = Field(
+        None,
+        description="Category/domain of the predicate (e.g., 'ownership_control', 'employment_leadership')"
+    )
+    object: ExtractedEntity = Field(..., description="The object entity")
+    source_text: str = Field(..., description="The source text this statement was extracted from")
+    confidence_score: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        description="Overall confidence score for this statement"
+    )
+    extraction_method: Optional[str] = Field(
+        None,
+        description="Method used to extract this statement (e.g., 'hybrid', 'gliner', 'model')"
+    )
+
+    def __str__(self) -> str:
+        return f"{self.subject.text} --[{self.predicate}]--> {self.object.text}"
+
+    def as_triple(self) -> tuple[str, str, str]:
+        """Return as a simple (subject, predicate, object) tuple."""
+        return (self.subject.text, self.predicate, self.object.text)
+
+    class Config:
+        frozen = False  # Allow modification during pipeline stages

statement_extractor/models.py

@@ -28,7 +28,6 @@ class ExtractionMethod(str, Enum):
     """Method used to extract the triple components."""
     HYBRID = "hybrid"  # Model subject/object + GLiNER2 predicate
     GLINER = "gliner"  # All components from GLiNER2 extraction
-    SPLIT = "split"  # Subject/object from splitting source text around predicate
     MODEL = "model"  # All components from T5-Gemma model (when GLiNER2 disabled)
 
 
@@ -305,6 +304,10 @@ class ExtractionOptions(BaseModel):
         default=None,
         description="Optional list of predefined predicate types for GLiNER2 relation extraction (e.g., ['works_for', 'founded'])"
     )
+    use_default_predicates: bool = Field(
+        default=True,
+        description="Use default predicate taxonomy when no custom predicates provided (enables GLiNER2 relation extraction)"
+    )
 
     # Verbose logging
     verbose: bool = Field(

statement_extractor/pipeline/__init__.py

@@ -0,0 +1,39 @@
+"""
+Pipeline module for the extraction pipeline.
+
+This module provides the core pipeline infrastructure:
+- PipelineContext: Data container that flows through all stages
+- PipelineConfig: Configuration for stage/plugin selection
+- PluginRegistry: Registration and discovery of plugins
+- ExtractionPipeline: Main orchestrator class
+
+Plugins are auto-loaded when this module is imported.
+"""
+
+from .context import PipelineContext
+from .config import PipelineConfig
+from .registry import PluginRegistry
+from .orchestrator import ExtractionPipeline
+
+
+def _load_plugins():
+    """Load all plugins by importing their modules."""
+    import logging
+
+    try:
+        from ..plugins import splitters, extractors, qualifiers, canonicalizers, labelers, taxonomy
+        # The @PluginRegistry decorators register plugins on import
+    except ImportError as e:
+        logging.debug(f"Some plugins failed to load: {e}")
+
+
+# Auto-load plugins on module import
+_load_plugins()
+
+
+__all__ = [
+    "PipelineContext",
+    "PipelineConfig",
+    "PluginRegistry",
+    "ExtractionPipeline",
+]

statement_extractor/pipeline/config.py

@@ -0,0 +1,134 @@
+"""
+PipelineConfig - Configuration for stage/plugin selection.
+
+Controls which stages are enabled and which plugins to use.
+"""
+
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field
+
+
+class PipelineConfig(BaseModel):
+    """
+    Configuration for the extraction pipeline.
+
+    Controls which stages are enabled, which plugins to use,
+    and stage-specific options.
+    """
+    # Stage selection (1=Splitting, 2=Extraction, 3=Qualification, 4=Canonicalization, 5=Labeling, 6=Taxonomy)
+    enabled_stages: set[int] = Field(
+        default={1, 2, 3, 4, 5, 6},
+        description="Set of enabled stage numbers (1-6)"
+    )
+
+    # Plugin selection
+    enabled_plugins: Optional[set[str]] = Field(
+        None,
+        description="Set of enabled plugin names (None = all enabled)"
+    )
+    disabled_plugins: set[str] = Field(
+        default_factory=lambda: {
+            "mnli_taxonomy_classifier",  # Disabled by default - use embedding_taxonomy_classifier instead (faster)
+        },
+        description="Set of disabled plugin names"
+    )
+
+    # Stage-specific options
+    splitter_options: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Options passed to splitter plugins"
+    )
+    extractor_options: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Options passed to extractor plugins"
+    )
+    qualifier_options: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Options passed to qualifier plugins"
+    )
+    canonicalizer_options: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Options passed to canonicalizer plugins"
+    )
+    labeler_options: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Options passed to labeler plugins"
+    )
+    taxonomy_options: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Options passed to taxonomy plugins"
+    )
+
+    # General options
+    fail_fast: bool = Field(
+        default=True,
+        description="Stop processing on first error (otherwise continue and collect errors)"
+    )
+    parallel_processing: bool = Field(
+        default=False,
+        description="Enable parallel processing where possible"
+    )
+    max_statements: Optional[int] = Field(
+        None,
+        description="Maximum number of statements to process (None = unlimited)"
+    )
+
+    def is_stage_enabled(self, stage: int) -> bool:
+        """Check if a stage is enabled."""
+        return stage in self.enabled_stages
+
+    def is_plugin_enabled(self, plugin_name: str) -> bool:
+        """Check if a plugin is enabled."""
+        if plugin_name in self.disabled_plugins:
+            return False
+        if self.enabled_plugins is None:
+            return True
+        return plugin_name in self.enabled_plugins
+
+    @classmethod
+    def from_stage_string(cls, stages: str, **kwargs) -> "PipelineConfig":
+        """
+        Create config from a stage string.
+
+        Examples:
+            "1,2,3" -> stages 1, 2, 3
+            "1-3" -> stages 1, 2, 3
+            "1-5" -> all stages
+        """
+        enabled = set()
+        for part in stages.split(","):
+            part = part.strip()
+            if "-" in part:
+                start, end = part.split("-", 1)
+                for i in range(int(start), int(end) + 1):
+                    enabled.add(i)
+            else:
+                enabled.add(int(part))
+        return cls(enabled_stages=enabled, **kwargs)
+
+    @classmethod
+    def default(cls) -> "PipelineConfig":
+        """Create a default configuration with all stages enabled."""
+        return cls()
+
+    @classmethod
+    def minimal(cls) -> "PipelineConfig":
+        """Create a minimal configuration with only splitting and extraction."""
+        return cls(enabled_stages={1, 2})
+
+
+# Stage name mapping
+STAGE_NAMES = {
+    1: "splitting",
+    2: "extraction",
+    3: "qualification",
+    4: "canonicalization",
+    5: "labeling",
+    6: "taxonomy",
+}
+
+
+def get_stage_name(stage: int) -> str:
+    """Get the human-readable name for a stage."""
+    return STAGE_NAMES.get(stage, f"stage_{stage}")