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

statement_extractor/plugins/__init__.py

@@ -4,10 +4,9 @@ Plugins module for the extraction pipeline.
 Contains all plugin implementations organized by stage:
 - splitters/: Stage 1 - Text to atomic triples
 - extractors/: Stage 2 - Refine entities and relations
-- qualifiers/: Stage 3 - Add qualifiers and identifiers
-- canonicalizers/: Stage 4 - Resolve canonical forms
-- labelers/: Stage 5 - Classify statements
-- taxonomy/: Stage 6 - Taxonomy classification
+- qualifiers/: Stage 3 - Qualify entities (add identifiers, canonical names, FQN)
+- labelers/: Stage 4 - Classify statements
+- taxonomy/: Stage 5 - Taxonomy classification
 """
 
 from .base import (
@@ -16,13 +15,20 @@ from .base import (
     BaseSplitterPlugin,
     BaseExtractorPlugin,
     BaseQualifierPlugin,
-    BaseCanonicalizerPlugin,
     BaseLabelerPlugin,
     BaseTaxonomyPlugin,
+    # Content acquisition plugins
+    ContentType,
+    ScraperResult,
+    PDFParseResult,
+    BaseScraperPlugin,
+    BasePDFParserPlugin,
 )
 
 # Import plugin modules for auto-registration
-from . import splitters, extractors, qualifiers, canonicalizers, labelers, taxonomy
+from . import splitters, extractors, qualifiers, labelers, taxonomy
+# Content acquisition plugins
+from . import scrapers, pdf
 
 __all__ = [
     "PluginCapability",
@@ -30,14 +36,20 @@ __all__ = [
     "BaseSplitterPlugin",
     "BaseExtractorPlugin",
     "BaseQualifierPlugin",
-    "BaseCanonicalizerPlugin",
     "BaseLabelerPlugin",
     "BaseTaxonomyPlugin",
+    # Content acquisition plugins
+    "ContentType",
+    "ScraperResult",
+    "PDFParseResult",
+    "BaseScraperPlugin",
+    "BasePDFParserPlugin",
     # Plugin modules
     "splitters",
     "extractors",
     "qualifiers",
-    "canonicalizers",
     "labelers",
     "taxonomy",
+    "scrapers",
+    "pdf",
 ]

statement_extractor/plugins/base.py

@@ -4,15 +4,20 @@ Base plugin classes for the extraction pipeline.
 Defines the abstract interfaces for each pipeline stage:
 - BaseSplitterPlugin: Stage 1 - Text → RawTriple
 - BaseExtractorPlugin: Stage 2 - RawTriple → PipelineStatement
-- BaseQualifierPlugin: Stage 3 - Entity → EntityQualifiers
-- BaseCanonicalizerPlugin: Stage 4 - QualifiedEntity → CanonicalMatch
-- BaseLabelerPlugin: Stage 5 - Statement → StatementLabel
-- BaseTaxonomyPlugin: Stage 6 - Statement → TaxonomyResult
+- BaseQualifierPlugin: Stage 3 - Entity → CanonicalEntity
+- BaseLabelerPlugin: Stage 4 - Statement → StatementLabel
+- BaseTaxonomyPlugin: Stage 5 - Statement → TaxonomyResult
+
+Content acquisition plugins (for URL processing):
+- BaseScraperPlugin: Fetch content from URLs
+- BasePDFParserPlugin: Extract text from PDFs
 """
 
 from abc import ABC, abstractmethod
-from enum import Flag, auto
-from typing import TYPE_CHECKING
+from enum import Enum, Flag, auto
+from typing import TYPE_CHECKING, Any, Optional
+
+from pydantic import BaseModel, Field
 
 if TYPE_CHECKING:
     from ..pipeline.context import PipelineContext
@@ -20,9 +25,6 @@ if TYPE_CHECKING:
         RawTriple,
         PipelineStatement,
         ExtractedEntity,
-        EntityQualifiers,
-        QualifiedEntity,
-        CanonicalMatch,
         CanonicalEntity,
         StatementLabel,
         TaxonomyResult,
@@ -40,6 +42,56 @@ class PluginCapability(Flag):
     CACHING = auto()            # Supports result caching
 
 
+def get_available_vram_gb() -> float:
+    """
+    Get available GPU VRAM in GB.
+
+    Returns 0.0 if no GPU is available or VRAM cannot be determined.
+    """
+    try:
+        import torch
+        if torch.cuda.is_available():
+            device = torch.cuda.current_device()
+            total = torch.cuda.get_device_properties(device).total_memory
+            allocated = torch.cuda.memory_allocated(device)
+            available = (total - allocated) / (1024 ** 3)  # Convert to GB
+            return available
+        elif torch.backends.mps.is_available():
+            # MPS doesn't expose VRAM info; estimate based on typical Apple Silicon
+            # Return a conservative estimate
+            return 8.0
+    except ImportError:
+        pass
+    return 0.0
+
+
+def calculate_batch_size(
+    per_item_vram_gb: float,
+    overhead_gb: float = 2.0,
+    min_batch: int = 1,
+    max_batch: int = 32,
+) -> int:
+    """
+    Calculate optimal batch size based on available VRAM.
+
+    Args:
+        per_item_vram_gb: VRAM required per item in GB
+        overhead_gb: Reserved VRAM for model weights and system overhead
+        min_batch: Minimum batch size
+        max_batch: Maximum batch size cap
+
+    Returns:
+        Optimal batch size for the current GPU
+    """
+    available = get_available_vram_gb()
+    if available <= 0 or per_item_vram_gb <= 0:
+        return min_batch
+
+    usable = max(0, available - overhead_gb)
+    batch_size = int(usable / per_item_vram_gb)
+    return max(min_batch, min(batch_size, max_batch))
+
+
 class BasePlugin(ABC):
     """
     Base class for all pipeline plugins.
@@ -74,6 +126,50 @@ class BasePlugin(ABC):
         """Human-readable description of this plugin."""
         return ""
 
+    @property
+    def model_vram_gb(self) -> float:
+        """
+        Estimated VRAM required for model weights in GB.
+
+        Override this if the plugin loads a GPU model. This is used to
+        reserve memory overhead when calculating batch sizes.
+
+        Default is 0.0 (no GPU model).
+        """
+        return 0.0
+
+    @property
+    def per_item_vram_gb(self) -> float:
+        """
+        Estimated VRAM required per item during batch processing in GB.
+
+        Override this for plugins with BATCH_PROCESSING capability.
+        Used to calculate optimal batch size: batch = (available - overhead) / per_item
+
+        Default is 0.1 GB (100MB) as a conservative estimate.
+        """
+        return 0.1
+
+    def get_optimal_batch_size(self, max_batch: int = 32) -> int:
+        """
+        Calculate optimal batch size based on available VRAM and plugin requirements.
+
+        Args:
+            max_batch: Maximum batch size cap
+
+        Returns:
+            Optimal batch size for current GPU state
+        """
+        if not (PluginCapability.BATCH_PROCESSING in self.capabilities):
+            return 1
+
+        return calculate_batch_size(
+            per_item_vram_gb=self.per_item_vram_gb,
+            overhead_gb=self.model_vram_gb + 1.0,  # Add 1GB system overhead
+            min_batch=1,
+            max_batch=max_batch,
+        )
+
 
 class BaseSplitterPlugin(BasePlugin):
     """
@@ -101,6 +197,27 @@ class BaseSplitterPlugin(BasePlugin):
         """
         ...
 
+    def split_batch(
+        self,
+        texts: list[str],
+        context: "PipelineContext",
+    ) -> list[list["RawTriple"]]:
+        """
+        Split multiple texts into atomic triples in a single batch.
+
+        Default implementation calls split() for each text sequentially.
+        Plugins with BATCH_PROCESSING capability should override this
+        for efficient GPU batching using get_optimal_batch_size().
+
+        Args:
+            texts: List of input texts to split
+            context: Pipeline context for accessing metadata and config
+
+        Returns:
+            List of RawTriple lists, one per input text
+        """
+        return [self.split(text, context) for text in texts]
+
 
 class BaseExtractorPlugin(BasePlugin):
     """
@@ -132,10 +249,14 @@ class BaseExtractorPlugin(BasePlugin):
 
 class BaseQualifierPlugin(BasePlugin):
     """
-    Stage 3 plugin: Add qualifiers and identifiers to entities.
+    Stage 3 plugin: Qualify entities with identifiers and canonical forms.
+
+    Processes entities of specific types and adds:
+    - Semantic qualifiers (role, org for PERSON entities)
+    - External identifiers (LEI, company number, SEC CIK)
+    - Canonical name and FQN (fully qualified name)
 
-    Processes entities of specific types and adds semantic qualifiers
-    (role, org) or external identifiers (LEI, company number).
+    Returns a CanonicalEntity ready for use in labeled statements.
     """
 
     @property
@@ -167,67 +288,26 @@ class BaseQualifierPlugin(BasePlugin):
         self,
         entity: "ExtractedEntity",
         context: "PipelineContext",
-    ) -> "EntityQualifiers | None":
+    ) -> "CanonicalEntity | None":
         """
-        Add qualifiers to an entity.
+        Qualify an entity and return its canonical form.
+
+        This method should:
+        1. Look up identifiers (LEI, CIK, company number, etc.)
+        2. Find the canonical name if available
+        3. Generate the FQN (fully qualified name)
+        4. Return a CanonicalEntity with all information
 
         Args:
             entity: The entity to qualify
             context: Pipeline context (for accessing source text, other entities)
 
         Returns:
-            EntityQualifiers with added information, or None if nothing to add
+            CanonicalEntity with qualifiers and FQN, or None if entity not found
         """
         ...
 
 
-class BaseCanonicalizerPlugin(BasePlugin):
-    """
-    Stage 4 plugin: Resolve entities to canonical forms.
-
-    Takes qualified entities and finds their canonical representations
-    using various matching strategies (identifier, name, fuzzy, LLM).
-    """
-
-    @property
-    @abstractmethod
-    def supported_entity_types(self) -> set["EntityType"]:
-        """Entity types this plugin can canonicalize."""
-        ...
-
-    @abstractmethod
-    def find_canonical(
-        self,
-        entity: "QualifiedEntity",
-        context: "PipelineContext",
-    ) -> "CanonicalMatch | None":
-        """
-        Find canonical form for an entity.
-
-        Args:
-            entity: Qualified entity to canonicalize
-            context: Pipeline context
-
-        Returns:
-            CanonicalMatch if found, None otherwise
-        """
-        ...
-
-    def format_fqn(
-        self,
-        entity: "QualifiedEntity",
-        match: "CanonicalMatch | None",
-    ) -> str:
-        """
-        Format the fully qualified name for display.
-
-        Can be overridden by subclasses for custom formatting.
-        Default implementation uses CanonicalEntity._generate_fqn.
-        """
-        from ..models import CanonicalEntity
-        return CanonicalEntity._generate_fqn(entity, match)
-
-
 class ClassificationSchema:
     """
     Schema for simple multi-choice classification (2-20 choices).
@@ -309,7 +389,7 @@ class TaxonomySchema:
 
 class BaseLabelerPlugin(BasePlugin):
     """
-    Stage 5 plugin: Apply labels to statements.
+    Stage 4 plugin: Apply labels to statements.
 
     Adds classification labels (sentiment, relation type, confidence)
     to the final labeled statements.
@@ -380,7 +460,7 @@ class BaseLabelerPlugin(BasePlugin):
 
 class BaseTaxonomyPlugin(BasePlugin):
     """
-    Stage 6 plugin: Classify statements against a taxonomy.
+    Stage 5 plugin: Classify statements against a taxonomy.
 
     Taxonomy classification is separate from labeling because:
     - It operates on large taxonomies (100s-1000s of labels)
@@ -444,3 +524,193 @@ class BaseTaxonomyPlugin(BasePlugin):
             List of TaxonomyResult objects (empty if none above threshold)
         """
         ...
+
+    def classify_batch(
+        self,
+        items: list[tuple["PipelineStatement", "CanonicalEntity", "CanonicalEntity"]],
+        context: "PipelineContext",
+    ) -> list[list["TaxonomyResult"]]:
+        """
+        Classify multiple statements against the taxonomy in a single batch.
+
+        Default implementation calls classify() for each statement sequentially.
+        Plugins with BATCH_PROCESSING capability should override this
+        for efficient GPU batching using get_optimal_batch_size().
+
+        Args:
+            items: List of (statement, subject_canonical, object_canonical) tuples
+            context: Pipeline context
+
+        Returns:
+            List of TaxonomyResult lists, one per input statement
+        """
+        return [
+            self.classify(stmt, subj, obj, context)
+            for stmt, subj, obj in items
+        ]
+
+
+# =============================================================================
+# Content Acquisition Plugins (for URL processing)
+# =============================================================================
+
+
+class ContentType(str, Enum):
+    """Content type detected from URL or HTTP response."""
+    HTML = "html"
+    PDF = "pdf"
+    BINARY = "binary"
+    UNKNOWN = "unknown"
+
+
+class ScraperResult(BaseModel):
+    """Result from a scraper plugin."""
+    url: str = Field(description="Original URL requested")
+    final_url: str = Field(description="Final URL after redirects")
+    content: bytes = Field(description="Raw content bytes")
+    content_type: ContentType = Field(description="Detected content type")
+    headers: dict[str, str] = Field(default_factory=dict, description="Response headers")
+    error: Optional[str] = Field(default=None, description="Error message if fetch failed")
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    @property
+    def ok(self) -> bool:
+        """Check if the fetch was successful."""
+        return self.error is None and len(self.content) > 0
+
+
+class PDFParseResult(BaseModel):
+    """Result from a PDF parser plugin."""
+    pages: list[str] = Field(description="Extracted text for each page")
+    page_count: int = Field(description="Total number of pages in PDF")
+    metadata: dict[str, Any] = Field(default_factory=dict, description="PDF metadata (title, author, etc)")
+    error: Optional[str] = Field(default=None, description="Error message if parsing failed")
+
+    @property
+    def ok(self) -> bool:
+        """Check if parsing was successful."""
+        return self.error is None
+
+    @property
+    def full_text(self) -> str:
+        """Get concatenated text from all pages."""
+        return "\n\n".join(self.pages)
+
+
+class BaseScraperPlugin(BasePlugin):
+    """
+    Plugin for fetching content from URLs.
+
+    Scrapers handle HTTP requests, redirects, retries, and content type detection.
+    They return raw bytes that can be processed by appropriate parsers (HTML, PDF, etc).
+
+    Example implementation:
+        @PluginRegistry.scraper
+        class MyScraperPlugin(BaseScraperPlugin):
+            @property
+            def name(self) -> str:
+                return "my_scraper"
+
+            async def fetch(self, url: str, timeout: float = 30.0) -> ScraperResult:
+                # Implement fetching logic
+                ...
+    """
+
+    @property
+    def capabilities(self) -> PluginCapability:
+        """Scrapers support async processing by default."""
+        return PluginCapability.ASYNC_PROCESSING | PluginCapability.EXTERNAL_API
+
+    @abstractmethod
+    async def fetch(self, url: str, timeout: float = 30.0) -> ScraperResult:
+        """
+        Fetch content from a URL.
+
+        Args:
+            url: The URL to fetch
+            timeout: Request timeout in seconds
+
+        Returns:
+            ScraperResult with content, content type, and any errors
+        """
+        ...
+
+    async def head(self, url: str, timeout: float = 10.0) -> ScraperResult:
+        """
+        Check content type without downloading the full body.
+
+        Default implementation does a full fetch. Override for efficiency.
+
+        Args:
+            url: The URL to check
+            timeout: Request timeout in seconds
+
+        Returns:
+            ScraperResult with content_type populated (content may be empty)
+        """
+        return await self.fetch(url, timeout)
+
+    def is_supported_url(self, url: str) -> bool:
+        """
+        Check if this scraper can handle the URL.
+
+        Override to restrict to specific URL patterns or domains.
+
+        Args:
+            url: The URL to check
+
+        Returns:
+            True if this scraper can handle the URL
+        """
+        return True
+
+
+class BasePDFParserPlugin(BasePlugin):
+    """
+    Plugin for extracting text from PDF files.
+
+    PDF parsers take raw PDF bytes and extract text content page by page.
+    They may support OCR for image-heavy PDFs.
+
+    Example implementation:
+        @PluginRegistry.pdf_parser
+        class MyPDFParserPlugin(BasePDFParserPlugin):
+            @property
+            def name(self) -> str:
+                return "my_pdf_parser"
+
+            def parse(self, pdf_bytes: bytes, ...) -> PDFParseResult:
+                # Implement parsing logic
+                ...
+    """
+
+    @abstractmethod
+    def parse(
+        self,
+        pdf_bytes: bytes,
+        max_pages: int = 500,
+        use_ocr: bool = False,
+    ) -> PDFParseResult:
+        """
+        Extract text from PDF bytes.
+
+        Args:
+            pdf_bytes: Raw PDF file content
+            max_pages: Maximum number of pages to process
+            use_ocr: Force OCR even for text-extractable PDFs
+
+        Returns:
+            PDFParseResult with extracted text for each page
+        """
+        ...
+
+    @property
+    def supports_ocr(self) -> bool:
+        """
+        Whether this parser supports OCR for image-heavy PDFs.
+
+        Returns:
+            True if OCR is available
+        """
+        return False

statement_extractor/plugins/extractors/gliner2.py

@@ -180,6 +180,16 @@ class GLiNER2Extractor(BaseExtractorPlugin):
     def description(self) -> str:
         return "GLiNER2 model for entity and relation extraction"
 
+    @property
+    def model_vram_gb(self) -> float:
+        """GLiNER2 model weights ~0.8GB."""
+        return 0.8
+
+    @property
+    def per_item_vram_gb(self) -> float:
+        """Each triple during batch processing ~0.1GB."""
+        return 0.1
+
     def _get_model(self):
         """Lazy-load the GLiNER2 model."""
         if self._model is None:

statement_extractor/plugins/labelers/taxonomy.py

@@ -8,9 +8,19 @@ there are too many possible values for simple multi-choice classification.
 import json
 import logging
 from pathlib import Path
-from typing import Optional
+from typing import Optional, TypedDict
 
 from ..base import BaseLabelerPlugin, TaxonomySchema, PluginCapability
+
+
+class TaxonomyEntry(TypedDict):
+    """Structure for each taxonomy label entry."""
+    description: str
+    id: int
+    mnli_label: str
+    embedding_label: str
+
+
 from ...pipeline.context import PipelineContext
 from ...models import (
     PipelineStatement,
@@ -214,7 +224,7 @@ class TaxonomyLabeler(BaseLabelerPlugin):
         self._top_k_categories = top_k_categories
         self._min_confidence = min_confidence
 
-        self._taxonomy: Optional[dict[str, dict[str, int]]] = None
+        self._taxonomy: Optional[dict[str, dict[str, TaxonomyEntry]]] = None
         self._classifier: Optional[TaxonomyClassifier] = None
 
     @property
@@ -250,7 +260,7 @@ class TaxonomyLabeler(BaseLabelerPlugin):
             scope="statement",
         )
 
-    def _load_taxonomy(self) -> dict[str, dict[str, int]]:
+    def _load_taxonomy(self) -> dict[str, dict[str, TaxonomyEntry]]:
         """Load taxonomy from JSON file."""
         if self._taxonomy is not None:
             return self._taxonomy
@@ -358,12 +368,15 @@ class TaxonomyLabeler(BaseLabelerPlugin):
         taxonomy = self._load_taxonomy()
 
         if category and category in taxonomy:
-            return taxonomy[category].get(label)
+            entry = taxonomy[category].get(label)
+            if entry:
+                return entry.get("id")
 
         # Search all categories for flat classification
         for cat_labels in taxonomy.values():
             if label in cat_labels:
-                return cat_labels[label]
+                entry = cat_labels[label]
+                return entry.get("id")
 
         return None
 

statement_extractor/plugins/labelers/taxonomy_embedding.py

@@ -11,10 +11,19 @@ import json
 import logging
 import time
 from pathlib import Path
-from typing import Optional
+from typing import Optional, TypedDict
 
 import numpy as np
 
+
+class TaxonomyEntry(TypedDict):
+    """Structure for each taxonomy label entry."""
+    description: str
+    id: int
+    mnli_label: str
+    embedding_label: str
+
+
 from ..base import BaseLabelerPlugin, TaxonomySchema, PluginCapability
 from ...pipeline.context import PipelineContext
 from ...models import (
@@ -106,14 +115,14 @@ class EmbeddingClassifier:
 
     def precompute_label_embeddings(
         self,
-        taxonomy: dict[str, dict[str, int]],
+        taxonomy: dict[str, dict[str, TaxonomyEntry]],
         categories: Optional[list[str]] = None,
     ) -> None:
         """
         Pre-compute embeddings for all label names.
 
         Args:
-            taxonomy: Taxonomy dict {category: {label: id, ...}, ...}
+            taxonomy: Taxonomy dict {category: {label: TaxonomyEntry, ...}, ...}
             categories: Categories to include (default: all)
         """
         self._load_model()
@@ -314,7 +323,7 @@ class EmbeddingTaxonomyLabeler(BaseLabelerPlugin):
         self._top_k_categories = top_k_categories
         self._min_confidence = min_confidence
 
-        self._taxonomy: Optional[dict[str, dict[str, int]]] = None
+        self._taxonomy: Optional[dict[str, dict[str, TaxonomyEntry]]] = None
         self._classifier: Optional[EmbeddingClassifier] = None
         self._embeddings_computed = False
 
@@ -350,7 +359,7 @@ class EmbeddingTaxonomyLabeler(BaseLabelerPlugin):
             scope="statement",
         )
 
-    def _load_taxonomy(self) -> dict[str, dict[str, int]]:
+    def _load_taxonomy(self) -> dict[str, dict[str, TaxonomyEntry]]:
         """Load taxonomy from JSON file."""
         if self._taxonomy is not None:
             return self._taxonomy
@@ -456,7 +465,9 @@ class EmbeddingTaxonomyLabeler(BaseLabelerPlugin):
         taxonomy = self._load_taxonomy()
 
         if category in taxonomy:
-            return taxonomy[category].get(label)
+            entry = taxonomy[category].get(label)
+            if entry:
+                return entry.get("id")
 
         return None
 

statement_extractor/plugins/pdf/__init__.py

@@ -0,0 +1,10 @@
+"""
+PDF parser plugins for extracting text from PDF files.
+
+Built-in parsers:
+- pypdf_parser: Default PDF parser using PyMuPDF with optional OCR
+"""
+
+from .pypdf import PyPDFParserPlugin
+
+__all__ = ["PyPDFParserPlugin"]