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

statement_extractor/plugins/labelers/confidence.py

@@ -0,0 +1,138 @@
+"""
+ConfidenceLabeler - Aggregates confidence scores from all pipeline stages.
+
+Combines:
+- Statement extraction confidence
+- Entity extraction confidence
+- Canonical match confidence
+"""
+
+import logging
+from typing import Optional
+
+from ..base import BaseLabelerPlugin, PluginCapability
+from ...pipeline.context import PipelineContext
+from ...pipeline.registry import PluginRegistry
+from ...models import (
+    PipelineStatement,
+    CanonicalEntity,
+    StatementLabel,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@PluginRegistry.labeler
+class ConfidenceLabeler(BaseLabelerPlugin):
+    """
+    Labeler that aggregates confidence scores from all pipeline stages.
+
+    Produces an overall confidence score for each statement.
+    """
+
+    def __init__(
+        self,
+        statement_weight: float = 0.4,
+        subject_weight: float = 0.2,
+        object_weight: float = 0.2,
+        canonical_weight: float = 0.2,
+    ):
+        """
+        Initialize the confidence labeler.
+
+        Args:
+            statement_weight: Weight for statement extraction confidence
+            subject_weight: Weight for subject entity confidence
+            object_weight: Weight for object entity confidence
+            canonical_weight: Weight for canonical match confidence
+        """
+        self._statement_weight = statement_weight
+        self._subject_weight = subject_weight
+        self._object_weight = object_weight
+        self._canonical_weight = canonical_weight
+
+    @property
+    def name(self) -> str:
+        return "confidence_labeler"
+
+    @property
+    def priority(self) -> int:
+        return 100  # Run after other labelers
+
+    @property
+    def capabilities(self) -> PluginCapability:
+        return PluginCapability.NONE
+
+    @property
+    def description(self) -> str:
+        return "Aggregates confidence scores from all pipeline stages"
+
+    @property
+    def label_type(self) -> str:
+        return "confidence"
+
+    def label(
+        self,
+        statement: PipelineStatement,
+        subject_canonical: CanonicalEntity,
+        object_canonical: CanonicalEntity,
+        context: PipelineContext,
+    ) -> Optional[StatementLabel]:
+        """
+        Calculate aggregate confidence for a statement.
+
+        Args:
+            statement: The statement to label
+            subject_canonical: Canonicalized subject
+            object_canonical: Canonicalized object
+            context: Pipeline context
+
+        Returns:
+            StatementLabel with aggregate confidence
+        """
+        scores = []
+        weights = []
+
+        # Statement confidence
+        if statement.confidence_score is not None:
+            scores.append(statement.confidence_score)
+            weights.append(self._statement_weight)
+
+        # Subject entity confidence
+        scores.append(statement.subject.confidence)
+        weights.append(self._subject_weight)
+
+        # Object entity confidence
+        scores.append(statement.object.confidence)
+        weights.append(self._object_weight)
+
+        # Canonical match confidence
+        subj_canon_conf = (
+            subject_canonical.canonical_match.match_confidence
+            if subject_canonical.canonical_match else 0.5
+        )
+        obj_canon_conf = (
+            object_canonical.canonical_match.match_confidence
+            if object_canonical.canonical_match else 0.5
+        )
+        avg_canon_conf = (subj_canon_conf + obj_canon_conf) / 2
+        scores.append(avg_canon_conf)
+        weights.append(self._canonical_weight)
+
+        # Calculate weighted average
+        total_weight = sum(weights)
+        if total_weight > 0:
+            aggregate_confidence = sum(s * w for s, w in zip(scores, weights)) / total_weight
+        else:
+            aggregate_confidence = 0.5
+
+        return StatementLabel(
+            label_type=self.label_type,
+            label_value=round(aggregate_confidence, 3),
+            confidence=1.0,  # High confidence in our calculation
+            labeler=self.name,
+        )
+
+
+# Allow importing without decorator for testing
+ConfidenceLabelerClass = ConfidenceLabeler

statement_extractor/plugins/labelers/relation_type.py

@@ -0,0 +1,87 @@
+"""
+RelationTypeLabeler - Uses predicate category from GLiNER2 extraction.
+
+The relation type comes from the predicate category assigned during
+Stage 2 extraction (GLiNER2). If no category is available, logs an error.
+"""
+
+import logging
+from typing import Optional
+
+from ..base import BaseLabelerPlugin, PluginCapability
+from ...pipeline.context import PipelineContext
+from ...pipeline.registry import PluginRegistry
+from ...models import (
+    PipelineStatement,
+    CanonicalEntity,
+    StatementLabel,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@PluginRegistry.labeler
+class RelationTypeLabeler(BaseLabelerPlugin):
+    """
+    Labeler that uses predicate category from GLiNER2 as relation type.
+
+    The category is set during Stage 2 extraction when GLiNER2 matches
+    a predicate from default_predicates.json (organized by category).
+    """
+
+    @property
+    def name(self) -> str:
+        return "relation_type_labeler"
+
+    @property
+    def priority(self) -> int:
+        return 20
+
+    @property
+    def capabilities(self) -> PluginCapability:
+        return PluginCapability.NONE
+
+    @property
+    def description(self) -> str:
+        return "Uses predicate category from GLiNER2 as relation type"
+
+    @property
+    def label_type(self) -> str:
+        return "relation_type"
+
+    def label(
+        self,
+        statement: PipelineStatement,
+        subject_canonical: CanonicalEntity,
+        object_canonical: CanonicalEntity,
+        context: PipelineContext,
+    ) -> Optional[StatementLabel]:
+        """
+        Get relation type from statement's predicate category.
+
+        Args:
+            statement: The statement to label
+            subject_canonical: Canonicalized subject
+            object_canonical: Canonicalized object
+            context: Pipeline context
+
+        Returns:
+            StatementLabel with relation type, or None if no category
+        """
+        if not statement.predicate_category:
+            logger.error(
+                f"No predicate_category for statement: "
+                f"'{statement.subject.text}' --[{statement.predicate}]--> '{statement.object.text}'"
+            )
+            return None
+
+        return StatementLabel(
+            label_type=self.label_type,
+            label_value=statement.predicate_category,
+            confidence=statement.confidence_score,  # Use statement's confidence
+            labeler=self.name,
+        )
+
+
+# Allow importing without decorator for testing
+RelationTypeLabelerClass = RelationTypeLabeler

statement_extractor/plugins/labelers/sentiment.py

@@ -0,0 +1,159 @@
+"""
+SentimentLabeler - Classifies statement sentiment.
+
+Uses GLiNER2 classification when available, falls back to pattern matching.
+"""
+
+import logging
+import re
+from typing import Optional
+
+from ..base import BaseLabelerPlugin, ClassificationSchema, PluginCapability
+from ...pipeline.context import PipelineContext
+from ...pipeline.registry import PluginRegistry
+from ...models import (
+    PipelineStatement,
+    CanonicalEntity,
+    StatementLabel,
+)
+
+logger = logging.getLogger(__name__)
+
+# Positive predicates and patterns
+POSITIVE_PATTERNS = [
+    r'\b(acquired|announced|launched|released|expanded|grew|increased|improved|won|awarded)\b',
+    r'\b(partnered|collaborated|joined|signed|agreed|approved|completed|achieved)\b',
+    r'\b(invested|funded|raised|promoted|hired|appointed)\b',
+]
+
+# Negative predicates and patterns
+NEGATIVE_PATTERNS = [
+    r'\b(lost|declined|decreased|dropped|fell|failed|fired|laid off|resigned)\b',
+    r'\b(sued|accused|charged|investigated|fined|penalized|rejected|denied)\b',
+    r'\b(closed|shut down|cancelled|terminated|withdrew|abandoned)\b',
+]
+
+# Neutral predicates (don't include words that are also in positive/negative lists)
+NEUTRAL_PATTERNS = [
+    r'\b(said|stated|reported|confirmed|disclosed)\b',
+    r'\b(is|was|are|were|has|have|had)\b',
+    r'\b(located|based|headquartered|operates|employs)\b',
+]
+
+
+def classify_sentiment(text: str) -> tuple[str, float]:
+    """
+    Classify sentiment of text using pattern matching.
+
+    Returns (sentiment, confidence) where sentiment is 'positive', 'negative', or 'neutral'.
+    """
+    text_lower = text.lower()
+
+    positive_matches = sum(
+        len(re.findall(pattern, text_lower, re.IGNORECASE))
+        for pattern in POSITIVE_PATTERNS
+    )
+    negative_matches = sum(
+        len(re.findall(pattern, text_lower, re.IGNORECASE))
+        for pattern in NEGATIVE_PATTERNS
+    )
+    neutral_matches = sum(
+        len(re.findall(pattern, text_lower, re.IGNORECASE))
+        for pattern in NEUTRAL_PATTERNS
+    )
+
+    total_matches = positive_matches + negative_matches + neutral_matches
+
+    if total_matches == 0:
+        return "neutral", 0.5
+
+    if positive_matches > negative_matches and positive_matches > neutral_matches:
+        confidence = min(0.6 + (positive_matches / total_matches) * 0.3, 0.9)
+        return "positive", confidence
+
+    if negative_matches > positive_matches and negative_matches > neutral_matches:
+        confidence = min(0.6 + (negative_matches / total_matches) * 0.3, 0.9)
+        return "negative", confidence
+
+    return "neutral", 0.6
+
+
+@PluginRegistry.labeler
+class SentimentLabeler(BaseLabelerPlugin):
+    """
+    Labeler that classifies statement sentiment.
+
+    Provides a ClassificationSchema so GLiNER2 can run classification.
+    Falls back to pattern matching if no pre-computed result available.
+    """
+
+    SENTIMENT_CHOICES = ["positive", "negative", "neutral"]
+
+    @property
+    def name(self) -> str:
+        return "sentiment_labeler"
+
+    @property
+    def priority(self) -> int:
+        return 10
+
+    @property
+    def capabilities(self) -> PluginCapability:
+        return PluginCapability.NONE
+
+    @property
+    def description(self) -> str:
+        return "Classifies statement sentiment (positive/negative/neutral)"
+
+    @property
+    def label_type(self) -> str:
+        return "sentiment"
+
+    @property
+    def classification_schema(self) -> ClassificationSchema:
+        """Provide schema for GLiNER2 to run classification."""
+        return ClassificationSchema(
+            label_type=self.label_type,
+            choices=self.SENTIMENT_CHOICES,
+            description="Classify the sentiment of this statement",
+            scope="statement",
+        )
+
+    def label(
+        self,
+        statement: PipelineStatement,
+        subject_canonical: CanonicalEntity,
+        object_canonical: CanonicalEntity,
+        context: PipelineContext,
+    ) -> Optional[StatementLabel]:
+        """
+        Classify sentiment of a statement.
+
+        First checks for pre-computed result from extractor (GLiNER2).
+        Falls back to pattern matching if not available.
+        """
+        # Check for pre-computed classification from extractor
+        result = context.get_classification(statement.source_text, self.label_type)
+        if result:
+            label_value, confidence = result
+            return StatementLabel(
+                label_type=self.label_type,
+                label_value=label_value,
+                confidence=confidence,
+                labeler=self.name,
+            )
+
+        # Fallback: pattern-based classification
+        text_to_analyze = f"{statement.predicate} {statement.source_text}"
+        sentiment, confidence = classify_sentiment(text_to_analyze)
+
+        return StatementLabel(
+            label_type=self.label_type,
+            label_value=sentiment,
+            confidence=confidence,
+            labeler=self.name,
+        )
+
+
+# Allow importing without decorator for testing
+SentimentLabelerClass = SentimentLabeler