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

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/models/__init__.py

@@ -44,9 +44,16 @@ else:
 # New pipeline models
 from .entity import ExtractedEntity
 from .statement import RawTriple, PipelineStatement
-from .qualifiers import EntityQualifiers, QualifiedEntity
+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)
@@ -66,9 +73,17 @@ __all__ = [
     "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,

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: