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

statement_extractor/models/canonical.py

@@ -0,0 +1,139 @@
+"""
+Canonical models for the extraction pipeline.
+
+CanonicalMatch: Result of matching to a canonical form
+CanonicalEntity: Entity with canonical form from Stage 4
+"""
+
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from .qualifiers import QualifiedEntity
+
+
+class CanonicalMatch(BaseModel):
+    """
+    Result of matching an entity to its canonical form in Stage 4.
+
+    Contains information about how the match was made and confidence level.
+    """
+    canonical_id: Optional[str] = Field(
+        None,
+        description="ID in canonical database (e.g., LEI, Wikidata QID)"
+    )
+    canonical_name: Optional[str] = Field(
+        None,
+        description="Canonical name/label"
+    )
+    match_method: str = Field(
+        ...,
+        description="How the match was made: 'identifier', 'name_exact', 'name_fuzzy', 'llm_verified'"
+    )
+    match_confidence: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        description="Confidence in the canonical match"
+    )
+    match_details: Optional[dict] = Field(
+        None,
+        description="Additional details about the match (e.g., fuzzy score, LLM reasoning)"
+    )
+
+    def is_high_confidence(self, threshold: float = 0.85) -> bool:
+        """Check if this is a high-confidence match."""
+        return self.match_confidence >= threshold
+
+
+class CanonicalEntity(BaseModel):
+    """
+    An entity with canonical form from Stage 4 (Canonicalization).
+
+    Contains the qualified entity plus its canonical match (if found)
+    and a fully qualified name (FQN) for display.
+    """
+    entity_ref: str = Field(..., description="Reference to the original ExtractedEntity")
+    qualified_entity: QualifiedEntity = Field(
+        ...,
+        description="The qualified entity from Stage 3"
+    )
+    canonical_match: Optional[CanonicalMatch] = Field(
+        None,
+        description="Canonical match if found"
+    )
+    fqn: str = Field(
+        ...,
+        description="Fully qualified name, e.g., 'Tim Cook (CEO, Apple Inc)'"
+    )
+
+    @classmethod
+    def from_qualified(
+        cls,
+        qualified: QualifiedEntity,
+        canonical_match: Optional[CanonicalMatch] = None,
+        fqn: Optional[str] = None,
+    ) -> "CanonicalEntity":
+        """Create a CanonicalEntity from a QualifiedEntity."""
+        if fqn is None:
+            # Generate default FQN from qualifiers
+            fqn = cls._generate_fqn(qualified, canonical_match)
+
+        return cls(
+            entity_ref=qualified.entity_ref,
+            qualified_entity=qualified,
+            canonical_match=canonical_match,
+            fqn=fqn,
+        )
+
+    @staticmethod
+    def _generate_fqn(
+        qualified: QualifiedEntity,
+        canonical_match: Optional[CanonicalMatch] = None
+    ) -> str:
+        """
+        Generate a fully qualified name from qualifiers.
+
+        Examples:
+        - PERSON with role+org: "Tim Cook (CEO, Apple Inc)"
+        - ORG with canonical: "Apple Inc (AAPL)"
+        - PERSON with no qualifiers: "Tim Cook"
+        """
+        # Use canonical name if available, otherwise fall back to original text
+        if canonical_match and canonical_match.canonical_name:
+            base_name = canonical_match.canonical_name
+        else:
+            base_name = qualified.original_text
+
+        qualifiers = qualified.qualifiers
+        parts = []
+        seen = set()  # Track seen values to avoid duplicates
+
+        def add_part(value: str) -> None:
+            """Add a part if not already seen (case-insensitive)."""
+            if value and value.lower() not in seen:
+                parts.append(value)
+                seen.add(value.lower())
+
+        # Add role for PERSON entities
+        if qualifiers.role:
+            add_part(qualifiers.role)
+
+        # Add organization for PERSON entities
+        if qualifiers.org:
+            add_part(qualifiers.org)
+
+        # Add ticker for ORG entities
+        if "ticker" in qualifiers.identifiers:
+            add_part(qualifiers.identifiers["ticker"])
+
+        # Add jurisdiction if relevant
+        if qualifiers.jurisdiction and not qualifiers.org:
+            add_part(qualifiers.jurisdiction)
+
+        if parts:
+            return f"{base_name} ({', '.join(parts)})"
+        return base_name
+
+    class Config:
+        frozen = False  # Allow modification during pipeline stages

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

@@ -26,10 +26,9 @@ class EntityType(str, Enum):
 
 class ExtractionMethod(str, Enum):
     """Method used to extract the triple components."""
-    HYBRID = "hybrid"  # Model subject/object + spaCy predicate
-    SPACY = "spacy"  # All components from spaCy dependency parsing
-    SPLIT = "split"  # Subject/object from splitting source text around predicate
-    MODEL = "model"  # All components from T5-Gemma model (when spaCy disabled)
+    HYBRID = "hybrid"  # Model subject/object + GLiNER2 predicate
+    GLINER = "gliner"  # All components from GLiNER2 extraction
+    MODEL = "model"  # All components from T5-Gemma model (when GLiNER2 disabled)
 
 
 class Entity(BaseModel):
@@ -295,9 +294,19 @@ class ExtractionOptions(BaseModel):
         default=True,
         description="Use embedding similarity for predicate deduplication"
     )
-    use_spacy_extraction: bool = Field(
+    use_gliner_extraction: bool = Field(
         default=True,
-        description="Use spaCy for predicate/subject/object extraction (model provides structure + coreference)"
+        description="Use GLiNER2 for predicate/subject/object extraction (model provides structure + coreference)"
+    )
+
+    # GLiNER2 predicate configuration
+    predicates: Optional[list[str]] = Field(
+        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

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",
+]