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

statement_extractor/models/labels.py

@@ -0,0 +1,220 @@
+"""
+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"
+    )
+    # Document tracking fields
+    document_id: Optional[str] = Field(
+        None,
+        description="ID of the source document (for document pipeline)"
+    )
+    page_number: Optional[int] = Field(
+        None,
+        description="Page number where this statement was extracted (1-indexed)"
+    )
+    citation: Optional[str] = Field(
+        None,
+        description="Formatted citation string (e.g., 'Title - Author, 2024, p. 5')"
+    )
+
+    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 _build_entity_dict(self, canonical: CanonicalEntity, entity_type: str) -> dict:
+        """Build entity dict for serialization."""
+        statement_entity = self.statement.subject if entity_type == "subject" else self.statement.object
+        fqn = self.subject_fqn if entity_type == "subject" else self.object_fqn
+
+        # Get canonical_id from identifiers or canonical_match
+        identifiers = canonical.qualified_entity.qualifiers.identifiers
+        canonical_id = identifiers.get("canonical_id")
+        if not canonical_id and canonical.canonical_match:
+            canonical_id = canonical.canonical_match.canonical_id
+
+        result = {
+            "text": statement_entity.text,
+            "type": statement_entity.type.value,
+            "fqn": fqn,
+            "canonical_id": canonical_id,
+        }
+
+        # Add name if available
+        if canonical.name:
+            result["name"] = canonical.name
+
+        # Add qualifiers if available
+        qualifiers_dict = canonical.qualifiers_dict
+        if qualifiers_dict:
+            result["qualifiers"] = qualifiers_dict
+
+        return result
+
+    def as_dict(self) -> dict:
+        """Convert to a simplified dictionary representation."""
+        return {
+            "subject": self._build_entity_dict(self.subject_canonical, "subject"),
+            "predicate": self.statement.predicate,
+            "object": self._build_entity_dict(self.object_canonical, "object"),
+            "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
+            ],
+            "document_id": self.document_id,
+            "page_number": self.page_number,
+            "citation": self.citation,
+        }
+
+    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,139 @@
+"""
+Qualifier models for the extraction pipeline.
+
+EntityQualifiers: Semantic qualifiers and external identifiers
+QualifiedEntity: Entity with qualification information from Stage 3
+ResolvedRole: Canonical role information from database
+ResolvedOrganization: Canonical organization information from database
+"""
+
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field
+
+from .entity import EntityType
+
+
+class ResolvedRole(BaseModel):
+    """
+    Resolved/canonical role information for a person.
+
+    Populated when matching a person against the database,
+    capturing the canonical role from Wikidata or other sources.
+    """
+    canonical_name: str = Field(..., description="Canonical role name (e.g., 'Chief Executive Officer')")
+    canonical_id: Optional[str] = Field(None, description="Full canonical ID (e.g., 'wikidata:Q484876')")
+    source: str = Field(..., description="Source of resolution (e.g., 'wikidata')")
+    source_id: Optional[str] = Field(None, description="ID in the source (e.g., 'Q484876' for Wikidata)")
+
+
+class ResolvedOrganization(BaseModel):
+    """
+    Resolved/canonical organization information.
+
+    Populated when resolving an organization mentioned in context
+    against the organization database (GLEIF, SEC, Companies House, Wikidata).
+    """
+    canonical_name: str = Field(..., description="Canonical organization name")
+    canonical_id: str = Field(..., description="Full canonical ID (e.g., 'LEI:549300XYZ', 'SEC-CIK:1234567')")
+    source: str = Field(..., description="Source of resolution (e.g., 'gleif', 'sec_edgar', 'wikidata')")
+    source_id: str = Field(..., description="ID in the source")
+    region: Optional[str] = Field(None, description="Organization's region/jurisdiction")
+    match_confidence: float = Field(default=1.0, description="Confidence in the match (0-1)")
+    match_details: Optional[dict[str, Any]] = Field(None, description="Additional match details")
+
+
+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
+    """
+    # Canonical name from database (for ORG entities)
+    legal_name: Optional[str] = Field(None, description="Canonical legal name from database")
+
+    # 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."
+    )
+
+    # Resolved canonical information (for PERSON entities)
+    resolved_role: Optional[ResolvedRole] = Field(
+        None,
+        description="Canonical role information from database lookup"
+    )
+    resolved_org: Optional[ResolvedOrganization] = Field(
+        None,
+        description="Canonical organization information from database lookup"
+    )
+
+    def has_any_qualifier(self) -> bool:
+        """Check if any qualifier or identifier is set."""
+        return bool(
+            self.legal_name or self.org or self.role or self.region or self.country or
+            self.city or self.jurisdiction or self.identifiers or
+            self.resolved_role or self.resolved_org
+        )
+
+    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(
+            legal_name=other.legal_name or self.legal_name,
+            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,
+            resolved_role=other.resolved_role or self.resolved_role,
+            resolved_org=other.resolved_org or self.resolved_org,
+        )
+
+
+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,101 @@
+"""
+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"
+    )
+    # Document tracking fields
+    document_id: Optional[str] = Field(
+        None,
+        description="ID of the source document (for document pipeline)"
+    )
+    page_number: Optional[int] = Field(
+        None,
+        description="Page number where this triple was extracted (1-indexed)"
+    )
+    chunk_index: Optional[int] = Field(
+        None,
+        description="Index of the chunk this triple was extracted from (0-indexed)"
+    )
+
+    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')"
+    )
+    # Document tracking fields
+    document_id: Optional[str] = Field(
+        None,
+        description="ID of the source document (for document pipeline)"
+    )
+    page_number: Optional[int] = Field(
+        None,
+        description="Page number where this statement was extracted (1-indexed)"
+    )
+    chunk_index: Optional[int] = Field(
+        None,
+        description="Index of the chunk this statement was extracted from (0-indexed)"
+    )
+
+    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,129 @@
+"""
+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=Labeling, 5=Taxonomy)
+    enabled_stages: set[int] = Field(
+        default={1, 2, 3, 4, 5},
+        description="Set of enabled stage numbers (1-5)"
+    )
+
+    # 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 (includes canonicalizers)"
+    )
+    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: "labeling",
+    5: "taxonomy",
+}
+
+
+def get_stage_name(stage: int) -> str:
+    """Get the human-readable name for a stage."""
+    return STAGE_NAMES.get(stage, f"stage_{stage}")