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

statement_extractor/document/pipeline.py

@@ -0,0 +1,388 @@
+"""
+DocumentPipeline - Orchestrates document-level extraction with chunking and citations.
+
+Wraps ExtractionPipeline to provide document-level features:
+- Text chunking with page awareness
+- Batch processing through pipeline stages
+- Statement deduplication across chunks
+- Citation generation from document metadata
+"""
+
+import logging
+import time
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from ..models.document import ChunkingConfig, Document
+from ..pipeline import ExtractionPipeline, PipelineConfig
+from .chunker import DocumentChunker
+from .context import DocumentContext
+from .deduplicator import StatementDeduplicator
+from .summarizer import DocumentSummarizer
+
+logger = logging.getLogger(__name__)
+
+
+class DocumentPipelineConfig(BaseModel):
+    """Configuration for document pipeline processing."""
+
+    chunking: ChunkingConfig = Field(
+        default_factory=ChunkingConfig,
+        description="Configuration for text chunking"
+    )
+    generate_summary: bool = Field(
+        default=True,
+        description="Whether to generate a document summary"
+    )
+    deduplicate_across_chunks: bool = Field(
+        default=True,
+        description="Whether to deduplicate statements across chunks"
+    )
+    batch_size: int = Field(
+        default=10,
+        ge=1,
+        description="Number of items to process in each batch"
+    )
+    pipeline_config: Optional[PipelineConfig] = Field(
+        default=None,
+        description="Configuration for the underlying extraction pipeline"
+    )
+
+
+class DocumentPipeline:
+    """
+    Document-level extraction pipeline.
+
+    Processes documents through:
+    1. Summary generation (optional)
+    2. Chunking with page awareness
+    3. Batch extraction through all pipeline stages
+    4. Deduplication across chunks
+    5. Citation generation
+
+    Example:
+        >>> pipeline = DocumentPipeline()
+        >>> document = Document.from_text("Long document text...", title="Report")
+        >>> ctx = pipeline.process(document)
+        >>> for stmt in ctx.labeled_statements:
+        ...     print(f"{stmt}: {stmt.citation}")
+    """
+
+    def __init__(
+        self,
+        config: Optional[DocumentPipelineConfig] = None,
+    ):
+        """
+        Initialize the document pipeline.
+
+        Args:
+            config: Document pipeline configuration
+        """
+        self.config = config or DocumentPipelineConfig()
+
+        # Initialize components
+        self._chunker = DocumentChunker(self.config.chunking)
+        self._deduplicator = StatementDeduplicator()
+        self._summarizer = DocumentSummarizer() if self.config.generate_summary else None
+        self._pipeline = ExtractionPipeline(self.config.pipeline_config)
+
+    def process(self, document: Document) -> DocumentContext:
+        """
+        Process a document through the pipeline.
+
+        Args:
+            document: Document to process
+
+        Returns:
+            DocumentContext with all extraction results
+        """
+        logger.info(f"Starting document pipeline: {document.document_id}")
+        start_time = time.time()
+
+        ctx = DocumentContext(document=document)
+
+        try:
+            # Step 1: Generate summary (if enabled)
+            if self.config.generate_summary and self._summarizer:
+                self._generate_summary(document, ctx)
+
+            # Step 2: Chunk the document
+            chunks = self._chunker.chunk_document(document)
+            ctx.chunks = chunks
+            logger.info(f"Created {len(chunks)} chunks")
+
+            if not chunks:
+                logger.warning("No chunks created from document")
+                return ctx
+
+            # Step 3: Process all chunks through Stage 1 (Splitting)
+            self._process_stage1(ctx)
+
+            # Step 4: Deduplicate raw triples
+            if self.config.deduplicate_across_chunks:
+                self._deduplicate_triples(ctx)
+
+            # Step 5: Process through remaining stages (2-6)
+            self._process_remaining_stages(ctx)
+
+            # Step 6: Add citations to statements
+            self._add_citations(ctx)
+
+        except Exception as e:
+            logger.exception("Document pipeline failed")
+            ctx.add_error(f"Pipeline error: {str(e)}")
+            raise
+
+        total_time = time.time() - start_time
+        ctx.record_timing("total", total_time)
+
+        logger.info(
+            f"Document pipeline complete: {ctx.statement_count} statements, "
+            f"{ctx.duplicates_removed} duplicates removed, "
+            f"{total_time:.2f}s"
+        )
+
+        return ctx
+
+    def _generate_summary(self, document: Document, ctx: DocumentContext) -> None:
+        """Generate document summary."""
+        logger.info("Generating document summary")
+        start_time = time.time()
+
+        try:
+            summary = self._summarizer.summarize(document)
+            document.summary = summary
+            logger.info(f"Generated summary: {len(summary)} chars")
+        except Exception as e:
+            logger.error(f"Summary generation failed: {e}")
+            ctx.add_warning(f"Summary generation failed: {e}")
+
+        ctx.record_timing("summarization", time.time() - start_time)
+
+    def _process_stage1(self, ctx: DocumentContext) -> None:
+        """Process all chunks through Stage 1 (Splitting) using batch processing."""
+        from ..pipeline.registry import PluginRegistry
+        from ..plugins.base import PluginCapability
+
+        logger.info(f"Processing {len(ctx.chunks)} chunks through Stage 1 (batch mode)")
+        start_time = time.time()
+
+        # Get the splitter plugin
+        splitters = PluginRegistry.get_splitters()
+        if not splitters:
+            logger.warning("No splitter plugins registered")
+            return
+
+        # Use first enabled splitter
+        splitter = None
+        for s in splitters:
+            plugin_enabled = (
+                self.config.pipeline_config is None or
+                self.config.pipeline_config.is_plugin_enabled(s.name)
+            )
+            if plugin_enabled:
+                splitter = s
+                break
+
+        if not splitter:
+            logger.warning("No enabled splitter plugin found")
+            return
+
+        # Extract all chunk texts
+        chunk_texts = [chunk.text for chunk in ctx.chunks]
+
+        # Create a dummy context for the splitter
+        from ..pipeline.context import PipelineContext
+        dummy_ctx = PipelineContext(
+            source_text="",  # Not used for batch splitting
+            source_metadata=self.config.pipeline_config.model_dump() if self.config.pipeline_config else {},
+        )
+
+        all_triples = []
+
+        # Require batch processing capability
+        if PluginCapability.BATCH_PROCESSING not in splitter.capabilities:
+            raise RuntimeError(
+                f"Splitter plugin '{splitter.name}' does not support batch processing. "
+                "Document pipeline requires BATCH_PROCESSING capability for efficient GPU utilization."
+            )
+
+        logger.info(f"Using batch splitting with {splitter.name}")
+        batch_results = splitter.split_batch(chunk_texts, dummy_ctx)
+
+        # Annotate triples with document/chunk info
+        for chunk, triples in zip(ctx.chunks, batch_results):
+            for triple in triples:
+                triple.document_id = ctx.document.document_id
+                triple.chunk_index = chunk.chunk_index
+                triple.page_number = chunk.primary_page
+                all_triples.append(triple)
+
+        ctx.raw_triples = all_triples
+        ctx.pre_dedup_count = len(all_triples)
+
+        ctx.record_timing("stage1_batch", time.time() - start_time)
+        logger.info(f"Stage 1 produced {len(all_triples)} raw triples from {len(ctx.chunks)} chunks")
+
+    def _deduplicate_triples(self, ctx: DocumentContext) -> None:
+        """Deduplicate raw triples across chunks."""
+        logger.info(f"Deduplicating {len(ctx.raw_triples)} triples")
+        start_time = time.time()
+
+        original_count = len(ctx.raw_triples)
+        ctx.raw_triples = self._deduplicator.deduplicate_batch(ctx.raw_triples)
+        ctx.post_dedup_count = len(ctx.raw_triples)
+
+        removed = original_count - len(ctx.raw_triples)
+        ctx.record_timing("deduplication", time.time() - start_time)
+        logger.info(f"Removed {removed} duplicate triples")
+
+    def _process_remaining_stages(self, ctx: DocumentContext) -> None:
+        """Process through stages 2-6."""
+        logger.info(f"Processing {len(ctx.raw_triples)} triples through stages 2-6")
+        start_time = time.time()
+
+        # Create a pipeline config for stages 2-6
+        # Exclude enabled_stages from base config to avoid duplicate keyword argument
+        base_config = {}
+        if self.config.pipeline_config:
+            base_config = self.config.pipeline_config.model_dump(exclude={"enabled_stages"})
+        stages_config = PipelineConfig(
+            enabled_stages={2, 3, 4, 5, 6},
+            **base_config
+        )
+
+        # Create a combined context with all raw triples
+        from ..pipeline.context import PipelineContext
+
+        combined_ctx = PipelineContext(
+            source_text=ctx.document.full_text,
+            source_metadata={
+                "document_id": ctx.document.document_id,
+                "title": ctx.document.metadata.title,
+            },
+        )
+        combined_ctx.raw_triples = ctx.raw_triples
+
+        # Run stages 2-6
+        pipeline = ExtractionPipeline(stages_config)
+
+        # We need to manually run stages since we're providing pre-existing triples
+        # Stage 2: Extraction
+        if stages_config.is_stage_enabled(2):
+            combined_ctx = pipeline._run_extraction(combined_ctx)
+
+        # Propagate document info to statements
+        for stmt in combined_ctx.statements:
+            # Find the source triple to get document info
+            for triple in ctx.raw_triples:
+                if triple.source_sentence in stmt.source_text:
+                    stmt.document_id = triple.document_id
+                    stmt.chunk_index = triple.chunk_index
+                    stmt.page_number = triple.page_number
+                    break
+
+        # Stage 3: Qualification
+        if stages_config.is_stage_enabled(3):
+            combined_ctx = pipeline._run_qualification(combined_ctx)
+
+        # Stage 4: Canonicalization
+        if stages_config.is_stage_enabled(4):
+            combined_ctx = pipeline._run_canonicalization(combined_ctx)
+
+        # Stage 5: Labeling
+        if stages_config.is_stage_enabled(5):
+            combined_ctx = pipeline._run_labeling(combined_ctx)
+
+        # Stage 6: Taxonomy
+        if stages_config.is_stage_enabled(6):
+            combined_ctx = pipeline._run_taxonomy(combined_ctx)
+
+        # Propagate document info to labeled statements
+        for labeled_stmt in combined_ctx.labeled_statements:
+            labeled_stmt.document_id = labeled_stmt.statement.document_id
+            labeled_stmt.page_number = labeled_stmt.statement.page_number
+
+        ctx.statements = combined_ctx.statements
+        ctx.labeled_statements = combined_ctx.labeled_statements
+
+        # Merge timings
+        for stage, duration in combined_ctx.stage_timings.items():
+            ctx.record_timing(stage, duration)
+
+        ctx.record_timing("stages_2_6_batch", time.time() - start_time)
+        logger.info(f"Stages 2-6 produced {len(ctx.labeled_statements)} labeled statements")
+
+    def _add_citations(self, ctx: DocumentContext) -> None:
+        """Add citations to all labeled statements."""
+        logger.info("Adding citations to statements")
+
+        for stmt in ctx.labeled_statements:
+            citation = ctx.document.metadata.format_citation(stmt.page_number)
+            stmt.citation = citation if citation else None
+
+    def process_text(
+        self,
+        text: str,
+        title: Optional[str] = None,
+        **metadata_kwargs,
+    ) -> DocumentContext:
+        """
+        Process plain text through the document pipeline.
+
+        Convenience method that creates a Document from text.
+
+        Args:
+            text: Text to process
+            title: Optional document title
+            **metadata_kwargs: Additional document metadata
+
+        Returns:
+            DocumentContext with extraction results
+        """
+        document = Document.from_text(text, title=title, **metadata_kwargs)
+        return self.process(document)
+
+    async def process_url(
+        self,
+        url: str,
+        loader_config: Optional["URLLoaderConfig"] = None,
+    ) -> DocumentContext:
+        """
+        Process a URL through the document pipeline.
+
+        Fetches the URL, extracts content (HTML or PDF), and processes it.
+
+        Args:
+            url: URL to process
+            loader_config: Optional loader configuration
+
+        Returns:
+            DocumentContext with extraction results
+        """
+        from .loader import URLLoader, URLLoaderConfig
+
+        loader = URLLoader(loader_config or URLLoaderConfig())
+        document = await loader.load(url)
+        return self.process(document)
+
+    def process_url_sync(
+        self,
+        url: str,
+        loader_config: Optional["URLLoaderConfig"] = None,
+    ) -> DocumentContext:
+        """
+        Process a URL through the document pipeline (synchronous).
+
+        Fetches the URL, extracts content (HTML or PDF), and processes it.
+
+        Args:
+            url: URL to process
+            loader_config: Optional loader configuration
+
+        Returns:
+            DocumentContext with extraction results
+        """
+        import asyncio
+        return asyncio.run(self.process_url(url, loader_config))

statement_extractor/document/summarizer.py

@@ -0,0 +1,195 @@
+"""
+DocumentSummarizer - Generate document summaries using Gemma3.
+
+Creates concise summaries focused on entities, events, and relationships
+that are useful for providing context during extraction.
+"""
+
+import logging
+from typing import Optional
+
+from ..models.document import Document
+
+logger = logging.getLogger(__name__)
+
+
+class DocumentSummarizer:
+    """
+    Generates document summaries using the Gemma3 LLM.
+
+    Summaries focus on:
+    - Key entities mentioned
+    - Important events and actions
+    - Relationships between entities
+    """
+
+    MAX_INPUT_TOKENS = 10_000
+    DEFAULT_MAX_OUTPUT_TOKENS = 300
+
+    def __init__(
+        self,
+        max_input_tokens: int = MAX_INPUT_TOKENS,
+        max_output_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
+    ):
+        """
+        Initialize the summarizer.
+
+        Args:
+            max_input_tokens: Maximum tokens of input to send to the LLM
+            max_output_tokens: Maximum tokens for the summary output
+        """
+        self._max_input_tokens = max_input_tokens
+        self._max_output_tokens = max_output_tokens
+        self._llm = None
+        self._tokenizer = None
+
+    @property
+    def llm(self):
+        """Lazy-load the LLM."""
+        if self._llm is None:
+            from ..llm import get_llm
+            logger.debug("Loading LLM for summarization")
+            self._llm = get_llm()
+        return self._llm
+
+    @property
+    def tokenizer(self):
+        """Lazy-load tokenizer for token counting."""
+        if self._tokenizer is None:
+            from transformers import AutoTokenizer
+            self._tokenizer = AutoTokenizer.from_pretrained(
+                "Corp-o-Rate-Community/statement-extractor",
+                trust_remote_code=True,
+            )
+        return self._tokenizer
+
+    def _count_tokens(self, text: str) -> int:
+        """Count tokens in text."""
+        return len(self.tokenizer.encode(text, add_special_tokens=False))
+
+    def _truncate_to_tokens(self, text: str, max_tokens: int) -> str:
+        """
+        Truncate text to a maximum number of tokens.
+
+        Tries to truncate at sentence boundaries when possible.
+        """
+        token_count = self._count_tokens(text)
+
+        if token_count <= max_tokens:
+            return text
+
+        # Estimate chars per token
+        chars_per_token = len(text) / token_count
+        target_chars = int(max_tokens * chars_per_token * 0.95)  # 5% buffer
+
+        # Truncate
+        truncated = text[:target_chars]
+
+        # Try to end at a sentence boundary
+        last_period = truncated.rfind(". ")
+        last_newline = truncated.rfind("\n")
+        split_pos = max(last_period, last_newline)
+
+        if split_pos > target_chars * 0.7:  # Don't lose too much text
+            truncated = truncated[:split_pos + 1]
+
+        logger.debug(f"Truncated text from {len(text)} to {len(truncated)} chars")
+        return truncated
+
+    def summarize(
+        self,
+        document: Document,
+        custom_prompt: Optional[str] = None,
+    ) -> str:
+        """
+        Generate a summary of the document.
+
+        Args:
+            document: Document to summarize
+            custom_prompt: Optional custom prompt (uses default if not provided)
+
+        Returns:
+            Summary string
+        """
+        if not document.full_text.strip():
+            logger.warning("Cannot summarize empty document")
+            return ""
+
+        logger.info(f"Generating summary for document {document.document_id}")
+
+        # Truncate text to max input tokens
+        text = self._truncate_to_tokens(document.full_text, self._max_input_tokens)
+
+        # Build prompt
+        if custom_prompt:
+            prompt = f"{custom_prompt}\n\n{text}"
+        else:
+            prompt = self._build_prompt(text, document)
+
+        # Generate summary
+        try:
+            summary = self.llm.generate(
+                prompt=prompt,
+                max_tokens=self._max_output_tokens,
+                stop=["\n\n\n", "---"],
+            )
+            summary = summary.strip()
+            logger.info(f"Generated summary ({len(summary)} chars):")
+            # Log summary with indentation for readability
+            for line in summary.split("\n"):
+                logger.info(f"  {line}")
+            return summary
+
+        except Exception as e:
+            logger.error(f"Summary generation failed: {e}")
+            raise
+
+    def _build_prompt(self, text: str, document: Document) -> str:
+        """Build the summarization prompt."""
+        # Include document metadata context if available
+        context_parts = []
+        if document.metadata.title:
+            context_parts.append(f"Title: {document.metadata.title}")
+        if document.metadata.authors:
+            context_parts.append(f"Authors: {', '.join(document.metadata.authors)}")
+        if document.metadata.source_type:
+            context_parts.append(f"Source type: {document.metadata.source_type}")
+
+        context = "\n".join(context_parts) if context_parts else ""
+
+        prompt = f"""Summarize the following document, focusing on:
+1. Key entities (companies, people, locations) mentioned
+2. Important events, actions, and decisions
+3. Relationships between entities
+4. Main topics and themes
+
+Keep the summary concise (2-3 paragraphs) and factual.
+
+{context}
+
+Document text:
+{text}
+
+Summary:"""
+
+        return prompt
+
+    def summarize_text(
+        self,
+        text: str,
+        title: Optional[str] = None,
+    ) -> str:
+        """
+        Generate a summary from plain text.
+
+        Convenience method that creates a temporary Document.
+
+        Args:
+            text: Text to summarize
+            title: Optional document title for context
+
+        Returns:
+            Summary string
+        """
+        document = Document.from_text(text, title=title)
+        return self.summarize(document)

statement_extractor/extractor.py

@@ -392,7 +392,7 @@ class StatementExtractor:
         This is the new extraction pipeline that:
         1. Generates multiple candidates via DBS
         2. Parses each to statements
-        3. Scores each triple for groundedness
+        3. Scores each triple for quality (semantic + entity)
         4. Merges top beams or selects best beam
         5. Deduplicates using embeddings (if enabled)
         """

statement_extractor/models/__init__.py

@@ -43,10 +43,17 @@ else:
 
 # New pipeline models
 from .entity import ExtractedEntity
-from .statement import RawTriple, PipelineStatement
-from .qualifiers import EntityQualifiers, QualifiedEntity
+from .statement import SplitSentence, RawTriple, PipelineStatement
+from .qualifiers import EntityQualifiers, QualifiedEntity, ResolvedRole, ResolvedOrganization
 from .canonical import CanonicalMatch, CanonicalEntity
 from .labels import StatementLabel, LabeledStatement, TaxonomyResult
+from .document import (
+    Document,
+    DocumentMetadata,
+    DocumentPage,
+    TextChunk,
+    ChunkingConfig,
+)
 
 __all__ = [
     # Re-exported from original models.py (backward compatibility)
@@ -62,13 +69,22 @@ __all__ = [
     "ExtractionOptions",
     # New pipeline models
     "ExtractedEntity",
-    "RawTriple",
+    "SplitSentence",
+    "RawTriple",  # Backwards compatibility alias for SplitSentence
     "PipelineStatement",
     "EntityQualifiers",
     "QualifiedEntity",
+    "ResolvedRole",
+    "ResolvedOrganization",
     "CanonicalMatch",
     "CanonicalEntity",
     "StatementLabel",
     "LabeledStatement",
     "TaxonomyResult",
+    # Document models
+    "Document",
+    "DocumentMetadata",
+    "DocumentPage",
+    "TextChunk",
+    "ChunkingConfig",
 ]

statement_extractor/models/canonical.py

@@ -64,9 +64,52 @@ class CanonicalEntity(BaseModel):
     )
     fqn: str = Field(
         ...,
-        description="Fully qualified name, e.g., 'Tim Cook (CEO, Apple Inc)'"
+        description="Fully qualified name, e.g., 'AMAZON CORP INC (SEC-CIK,USA)'"
     )
 
+    @property
+    def name(self) -> Optional[str]:
+        """Get the canonical/legal name if available."""
+        # Prefer legal_name from qualifiers (set by embedding qualifier)
+        if self.qualified_entity.qualifiers.legal_name:
+            return self.qualified_entity.qualifiers.legal_name
+        # Fall back to canonical match name
+        if self.canonical_match and self.canonical_match.canonical_name:
+            return self.canonical_match.canonical_name
+        return None
+
+    @property
+    def qualifiers_dict(self) -> Optional[dict[str, str]]:
+        """
+        Get qualifiers as a dict for serialization.
+
+        Returns a dict with keys like: legal_name, region, source, source_id
+        Only returns non-None values. Returns None if no qualifiers are set.
+        """
+        qualifiers = self.qualified_entity.qualifiers
+        identifiers = qualifiers.identifiers
+        result = {}
+
+        # Add legal name
+        if qualifiers.legal_name:
+            result["legal_name"] = qualifiers.legal_name
+
+        # Add region (prefer region, fall back to jurisdiction/country)
+        if qualifiers.region:
+            result["region"] = qualifiers.region
+        elif qualifiers.jurisdiction:
+            result["region"] = qualifiers.jurisdiction
+        elif qualifiers.country:
+            result["region"] = qualifiers.country
+
+        # Add source and source_id from identifiers
+        if "source" in identifiers:
+            result["source"] = identifiers["source"]
+        if "source_id" in identifiers:
+            result["source_id"] = identifiers["source_id"]
+
+        return result if result else None
+
     @classmethod
     def from_qualified(
         cls,