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

statement_extractor/models/document.py

@@ -0,0 +1,308 @@
+"""
+Document models for document-level processing.
+
+Document: A document with metadata, pages, and optional summary
+DocumentMetadata: Metadata about the document source
+DocumentPage: A single page within a document
+TextChunk: A chunk of text for processing with page tracking
+ChunkingConfig: Configuration for text chunking
+"""
+
+import uuid
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field
+
+
+class DocumentMetadata(BaseModel):
+    """
+    Metadata about a document source.
+
+    Contains information about where the document came from and
+    who authored it, useful for generating citations.
+    """
+    url: Optional[str] = Field(None, description="URL source of the document")
+    title: Optional[str] = Field(None, description="Document title")
+    year: Optional[int] = Field(None, description="Publication year")
+    authors: list[str] = Field(default_factory=list, description="List of authors")
+    source_type: Optional[str] = Field(
+        None,
+        description="Type of source: 'pdf', 'webpage', 'text', etc."
+    )
+    custom: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Custom metadata fields"
+    )
+
+    def format_citation(self, page_number: Optional[int] = None) -> str:
+        """
+        Format a citation string for this document.
+
+        Args:
+            page_number: Optional page number to include
+
+        Returns:
+            Citation string like "Title - Author, 2024, p. 5"
+        """
+        parts = []
+
+        if self.title:
+            parts.append(self.title)
+
+        if self.authors:
+            if len(self.authors) == 1:
+                parts.append(self.authors[0])
+            elif len(self.authors) == 2:
+                parts.append(f"{self.authors[0]} & {self.authors[1]}")
+            else:
+                parts.append(f"{self.authors[0]} et al.")
+
+        if self.year:
+            parts.append(str(self.year))
+
+        if page_number is not None:
+            parts.append(f"p. {page_number}")
+
+        return " - ".join(parts) if parts else ""
+
+
+class DocumentPage(BaseModel):
+    """
+    A single page within a document.
+
+    Tracks the page number and character offset for citation purposes.
+    """
+    page_number: int = Field(..., description="1-indexed page number")
+    text: str = Field(..., description="Text content of the page")
+    char_offset: int = Field(
+        ...,
+        description="Character offset of this page in the full document text"
+    )
+
+    @property
+    def char_end(self) -> int:
+        """Get the ending character offset of this page."""
+        return self.char_offset + len(self.text)
+
+
+class TextChunk(BaseModel):
+    """
+    A chunk of text for processing.
+
+    Contains the text along with position tracking for mapping
+    extracted statements back to their source pages.
+    """
+    chunk_index: int = Field(..., description="0-indexed chunk number")
+    text: str = Field(..., description="Chunk text content")
+    start_char: int = Field(..., description="Starting character offset in full document")
+    end_char: int = Field(..., description="Ending character offset in full document")
+    page_numbers: list[int] = Field(
+        default_factory=list,
+        description="Page numbers this chunk spans (1-indexed)"
+    )
+    token_count: int = Field(..., description="Number of tokens in this chunk")
+    overlap_chars: int = Field(
+        default=0,
+        description="Number of characters of overlap from previous chunk"
+    )
+    document_id: str = Field(..., description="ID of the source document")
+
+    @property
+    def primary_page(self) -> Optional[int]:
+        """Get the primary page number for this chunk (first page)."""
+        return self.page_numbers[0] if self.page_numbers else None
+
+
+class ChunkingConfig(BaseModel):
+    """
+    Configuration for document chunking.
+
+    Controls how documents are split into chunks for processing.
+    """
+    max_tokens: int = Field(
+        default=2000,
+        ge=100,
+        description="Maximum tokens per chunk (hard limit)"
+    )
+    target_tokens: int = Field(
+        default=1000,
+        ge=50,
+        description="Target tokens per chunk (soft limit, prefers to split here)"
+    )
+    overlap_tokens: int = Field(
+        default=100,
+        ge=0,
+        description="Tokens of overlap between consecutive chunks"
+    )
+    respect_page_boundaries: bool = Field(
+        default=True,
+        description="Try to split at page boundaries when possible"
+    )
+    respect_sentence_boundaries: bool = Field(
+        default=True,
+        description="Try to split at sentence boundaries when possible"
+    )
+
+
+class Document(BaseModel):
+    """
+    A document for processing through the extraction pipeline.
+
+    Contains the full text, optional page structure, metadata for citations,
+    and an optional summary for context.
+    """
+    document_id: str = Field(
+        default_factory=lambda: str(uuid.uuid4()),
+        description="Unique identifier for this document"
+    )
+    metadata: DocumentMetadata = Field(
+        default_factory=DocumentMetadata,
+        description="Document metadata for citations"
+    )
+    pages: list[DocumentPage] = Field(
+        default_factory=list,
+        description="List of pages (optional, for PDFs)"
+    )
+    full_text: str = Field(
+        default="",
+        description="Full text content of the document"
+    )
+    summary: Optional[str] = Field(
+        None,
+        description="Generated summary of the document"
+    )
+
+    @classmethod
+    def from_text(
+        cls,
+        text: str,
+        title: Optional[str] = None,
+        url: Optional[str] = None,
+        **metadata_kwargs,
+    ) -> "Document":
+        """
+        Create a document from plain text.
+
+        Args:
+            text: The document text
+            title: Optional document title
+            url: Optional source URL
+            **metadata_kwargs: Additional metadata fields
+
+        Returns:
+            Document instance
+        """
+        metadata = DocumentMetadata(
+            title=title,
+            url=url,
+            source_type="text",
+            **metadata_kwargs,
+        )
+        return cls(
+            metadata=metadata,
+            full_text=text,
+        )
+
+    @classmethod
+    def from_pages(
+        cls,
+        pages: list[str],
+        title: Optional[str] = None,
+        source_type: str = "pdf",
+        **metadata_kwargs,
+    ) -> "Document":
+        """
+        Create a document from a list of page texts.
+
+        Args:
+            pages: List of page text strings (0-indexed input, stored as 1-indexed)
+            title: Optional document title
+            source_type: Source type (default: "pdf")
+            **metadata_kwargs: Additional metadata fields
+
+        Returns:
+            Document instance
+        """
+        metadata = DocumentMetadata(
+            title=title,
+            source_type=source_type,
+            **metadata_kwargs,
+        )
+
+        # Build pages with character offsets
+        doc_pages = []
+        char_offset = 0
+
+        for i, page_text in enumerate(pages):
+            doc_pages.append(DocumentPage(
+                page_number=i + 1,  # 1-indexed
+                text=page_text,
+                char_offset=char_offset,
+            ))
+            char_offset += len(page_text)
+            if i < len(pages) - 1:
+                char_offset += 1  # Account for newline between pages
+
+        # Join pages with newlines for full text
+        full_text = "\n".join(pages)
+
+        return cls(
+            metadata=metadata,
+            pages=doc_pages,
+            full_text=full_text,
+        )
+
+    def get_page_at_char(self, char_offset: int) -> Optional[int]:
+        """
+        Get the page number containing a character offset.
+
+        Args:
+            char_offset: Character offset in full_text
+
+        Returns:
+            1-indexed page number, or None if no pages defined
+        """
+        if not self.pages:
+            return None
+
+        for page in self.pages:
+            if page.char_offset <= char_offset < page.char_end:
+                return page.page_number
+
+        # If past the last page, return last page
+        if char_offset >= self.pages[-1].char_end:
+            return self.pages[-1].page_number
+
+        return None
+
+    def get_pages_in_range(self, start_char: int, end_char: int) -> list[int]:
+        """
+        Get all page numbers that overlap with a character range.
+
+        Args:
+            start_char: Start character offset
+            end_char: End character offset
+
+        Returns:
+            List of 1-indexed page numbers
+        """
+        if not self.pages:
+            return []
+
+        page_numbers = []
+        for page in self.pages:
+            # Check if page overlaps with range
+            if page.char_offset < end_char and page.char_end > start_char:
+                page_numbers.append(page.page_number)
+
+        return page_numbers
+
+    @property
+    def page_count(self) -> int:
+        """Get the number of pages in the document."""
+        return len(self.pages)
+
+    @property
+    def char_count(self) -> int:
+        """Get the total character count."""
+        return len(self.full_text)

statement_extractor/models/labels.py

@@ -72,6 +72,19 @@ class LabeledStatement(BaseModel):
         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."""
@@ -102,28 +115,41 @@ class LabeledStatement(BaseModel):
         """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": {
-                "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
-                ),
-            },
+            "subject": self._build_entity_dict(self.subject_canonical, "subject"),
             "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
-                ),
-            },
+            "object": self._build_entity_dict(self.object_canonical, "object"),
             "source_text": self.statement.source_text,
             "labels": {
                 label.label_type: label.label_value
@@ -137,6 +163,9 @@ class LabeledStatement(BaseModel):
                 }
                 for t in self.taxonomy_results
             ],
+            "document_id": self.document_id,
+            "page_number": self.page_number,
+            "citation": self.citation,
         }
 
     class Config:

statement_extractor/models/qualifiers.py

@@ -3,15 +3,46 @@ 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 Optional
+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.
@@ -22,6 +53,9 @@ class EntityQualifiers(BaseModel):
     - 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")
@@ -38,11 +72,22 @@ class EntityQualifiers(BaseModel):
         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.org or self.role or self.region or self.country or
-            self.city or self.jurisdiction or self.identifiers
+            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":
@@ -53,6 +98,7 @@ class EntityQualifiers(BaseModel):
         """
         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,
@@ -60,6 +106,8 @@ class EntityQualifiers(BaseModel):
             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,
         )
 
 

statement_extractor/models/statement.py

@@ -1,8 +1,8 @@
 """
 Statement models for the extraction pipeline.
 
-RawTriple: Output of Stage 1 (Splitting)
-PipelineStatement: Output of Stage 2 (Extraction) with refined entities
+SplitSentence: Output of Stage 1 (Splitting) - atomic sentences/statements
+PipelineStatement: Output of Stage 2 (Extraction) with subject-predicate-object triples
 """
 
 from typing import Optional
@@ -12,30 +12,41 @@ from pydantic import BaseModel, Field
 from .entity import ExtractedEntity
 
 
-class RawTriple(BaseModel):
+class SplitSentence(BaseModel):
     """
-    A raw triple from Stage 1 (Splitting).
+    An atomic sentence from Stage 1 (Splitting).
 
-    Contains the basic text components before entity refinement.
-    Generated by T5-Gemma or other splitting plugins.
+    Stage 1 splits text into atomic sentences that can each be converted
+    to subject-predicate-object triples in Stage 2. 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")
+    text: str = Field(..., description="The atomic sentence text")
     confidence: float = Field(
         default=1.0,
         ge=0.0,
         le=1.0,
-        description="Extraction confidence from the splitter"
+        description="Confidence that this is a valid atomic statement"
+    )
+    # 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 sentence was extracted (1-indexed)"
+    )
+    chunk_index: Optional[int] = Field(
+        None,
+        description="Index of the chunk this sentence was extracted from (0-indexed)"
     )
 
     def __str__(self) -> str:
-        return f"{self.subject_text} --[{self.predicate_text}]--> {self.object_text}"
+        return self.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)
+
+# Backwards compatibility alias
+RawTriple = SplitSentence
 
 
 class PipelineStatement(BaseModel):
@@ -63,6 +74,19 @@ class PipelineStatement(BaseModel):
         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}"

statement_extractor/models.py

@@ -217,7 +217,7 @@ class ScoringConfig(BaseModel):
     quality_weight: float = Field(
         default=1.0,
         ge=0.0,
-        description="Weight for groundedness/quality scores in beam selection"
+        description="Weight for confidence scores in beam selection"
     )
     coverage_weight: float = Field(
         default=0.5,

statement_extractor/pipeline/config.py

@@ -16,10 +16,10 @@ class PipelineConfig(BaseModel):
     Controls which stages are enabled, which plugins to use,
     and stage-specific options.
     """
-    # Stage selection (1=Splitting, 2=Extraction, 3=Qualification, 4=Canonicalization, 5=Labeling, 6=Taxonomy)
+    # Stage selection (1=Splitting, 2=Extraction, 3=Qualification, 4=Labeling, 5=Taxonomy)
     enabled_stages: set[int] = Field(
-        default={1, 2, 3, 4, 5, 6},
-        description="Set of enabled stage numbers (1-6)"
+        default={1, 2, 3, 4, 5},
+        description="Set of enabled stage numbers (1-5)"
     )
 
     # Plugin selection
@@ -45,11 +45,7 @@ class PipelineConfig(BaseModel):
     )
     qualifier_options: dict[str, Any] = Field(
         default_factory=dict,
-        description="Options passed to qualifier plugins"
-    )
-    canonicalizer_options: dict[str, Any] = Field(
-        default_factory=dict,
-        description="Options passed to canonicalizer plugins"
+        description="Options passed to qualifier plugins (includes canonicalizers)"
     )
     labeler_options: dict[str, Any] = Field(
         default_factory=dict,
@@ -123,9 +119,8 @@ STAGE_NAMES = {
     1: "splitting",
     2: "extraction",
     3: "qualification",
-    4: "canonicalization",
-    5: "labeling",
-    6: "taxonomy",
+    4: "labeling",
+    5: "taxonomy",
 }
 
 

statement_extractor/pipeline/context.py

@@ -2,7 +2,7 @@
 PipelineContext - Data container that flows through all pipeline stages.
 
 The context accumulates outputs from each stage:
-- Stage 1 (Splitting): raw_triples
+- Stage 1 (Splitting): split_sentences
 - Stage 2 (Extraction): statements
 - Stage 3 (Qualification): qualified_entities
 - Stage 4 (Canonicalization): canonical_entities
@@ -14,7 +14,7 @@ from typing import Any, Optional
 from pydantic import BaseModel, Field
 
 from ..models import (
-    RawTriple,
+    SplitSentence,
     PipelineStatement,
     QualifiedEntity,
     CanonicalEntity,
@@ -37,10 +37,10 @@ class PipelineContext(BaseModel):
         description="Metadata about the source (e.g., document ID, URL, timestamp)"
     )
 
-    # Stage 1 output: Raw triples from splitting
-    raw_triples: list[RawTriple] = Field(
+    # Stage 1 output: Split sentences
+    split_sentences: list[SplitSentence] = Field(
         default_factory=list,
-        description="Raw triples from Stage 1 (Splitting)"
+        description="Atomic sentences from Stage 1 (Splitting)"
     )
 
     # Stage 2 output: Statements with extracted entities