corp-extractor 0.4.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/extractor.py

@@ -783,7 +783,7 @@ class StatementExtractor:
 
                 if use_gliner_extraction and source_text:
                     try:
-                        from .gliner_extraction import extract_triple_from_text, extract_triple_by_predicate_split
+                        from .gliner_extraction import extract_triple_from_text
 
                         # Get model predicate for fallback/refinement
                         predicate_elem = stmt_elem.find('predicate')
@@ -826,28 +826,6 @@ class StatementExtractor:
                                             source_text=source_text,
                                             extraction_method=ExtractionMethod.GLINER,
                                         ))
-
-                                # Candidate 3: Predicate-split (split source text around predicate)
-                                split_result = extract_triple_by_predicate_split(
-                                    source_text=source_text,
-                                    predicate=gliner_pred,
-                                )
-                                if split_result:
-                                    split_subj, split_pred, split_obj = split_result
-                                    # Only add if different from previous candidates
-                                    is_different_from_hybrid = (split_subj != subject_text or split_obj != object_text)
-                                    is_different_from_gliner = (split_subj != gliner_subj or split_obj != gliner_obj)
-                                    if is_different_from_hybrid and is_different_from_gliner:
-                                        logger.debug(
-                                            f"Adding predicate-split candidate: '{split_subj}' --[{split_pred}]--> '{split_obj}'"
-                                        )
-                                        statements.append(Statement(
-                                            subject=Entity(text=split_subj, type=subject_type),
-                                            predicate=split_pred,
-                                            object=Entity(text=split_obj, type=object_type),
-                                            source_text=source_text,
-                                            extraction_method=ExtractionMethod.SPLIT,
-                                        ))
                             else:
                                 logger.debug(
                                     f"GLiNER2 found no predicate for: '{subject_text}' --> '{object_text}'"

statement_extractor/gliner_extraction.py

@@ -132,18 +132,12 @@ def extract_triple_from_text(
                         if len(entity) >= len(refined_object):
                             refined_object = entity
 
-            # Extract predicate from source text using predicate split
-            predicate_result = extract_triple_by_predicate_split(source_text, model_predicate)
-            if predicate_result:
-                _, extracted_predicate, _ = predicate_result
-            else:
-                extracted_predicate = model_predicate
-
-            if extracted_predicate:
+            # Use model predicate directly (T5-Gemma provides the predicate)
+            if model_predicate:
                 logger.debug(
-                    f"GLiNER2 extracted (entity-refined): subj='{refined_subject}', pred='{extracted_predicate}', obj='{refined_object}'"
+                    f"GLiNER2 extracted (entity-refined): subj='{refined_subject}', pred='{model_predicate}', obj='{refined_object}'"
                 )
-                return (refined_subject, extracted_predicate, refined_object)
+                return (refined_subject, model_predicate, refined_object)
 
         return None
 
@@ -155,70 +149,6 @@ def extract_triple_from_text(
         return None
 
 
-def extract_triple_by_predicate_split(
-    source_text: str,
-    predicate: str,
-) -> tuple[str, str, str] | None:
-    """
-    Extract subject and object by splitting the source text around the predicate.
-
-    This is useful when the predicate is known but subject/object boundaries
-    are uncertain. Uses the predicate as an anchor point.
-
-    Args:
-        source_text: The source sentence
-        predicate: The predicate (verb phrase) to split on
-
-    Returns:
-        Tuple of (subject, predicate, object) or None if split fails
-    """
-    if not source_text or not predicate:
-        return None
-
-    # Find the predicate in the source text (case-insensitive)
-    source_lower = source_text.lower()
-    pred_lower = predicate.lower()
-
-    pred_pos = source_lower.find(pred_lower)
-    if pred_pos < 0:
-        # Try finding just the main verb (first word of predicate)
-        main_verb = pred_lower.split()[0] if pred_lower.split() else ""
-        if main_verb and len(main_verb) > 2:
-            pred_pos = source_lower.find(main_verb)
-            if pred_pos >= 0:
-                # Adjust to use the actual predicate length for splitting
-                predicate = main_verb
-
-    if pred_pos < 0:
-        return None
-
-    # Extract subject (text before predicate, trimmed)
-    subject = source_text[:pred_pos].strip()
-
-    # Extract object (text after predicate, trimmed)
-    pred_end = pred_pos + len(predicate)
-    obj = source_text[pred_end:].strip()
-
-    # Clean up: remove trailing punctuation from object
-    obj = obj.rstrip('.,;:!?')
-
-    # Clean up: remove leading articles/prepositions from object if very short
-    obj_words = obj.split()
-    if obj_words and obj_words[0].lower() in ('a', 'an', 'the', 'to', 'of', 'for'):
-        if len(obj_words) > 1:
-            obj = ' '.join(obj_words[1:])
-
-    # Validate: both subject and object should have meaningful content
-    if len(subject) < 2 or len(obj) < 2:
-        return None
-
-    logger.debug(
-        f"Predicate-split extracted: subj='{subject}', pred='{predicate}', obj='{obj}'"
-    )
-
-    return (subject, predicate, obj)
-
-
 def score_entity_content(text: str) -> float:
     """
     Score how entity-like a text is using GLiNER2 entity recognition.

statement_extractor/llm.py

@@ -0,0 +1,255 @@
+"""
+LLM module for text generation using local models.
+
+Supports:
+- GGUF models via llama-cpp-python (efficient quantized inference)
+- Transformers models via HuggingFace
+
+Usage:
+    from statement_extractor.llm import LLM
+
+    llm = LLM()  # Uses default Gemma3 12B GGUF
+    response = llm.generate("Your prompt here")
+"""
+
+import logging
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+class LLM:
+    """
+    LLM wrapper for text generation.
+
+    Automatically selects the best backend:
+    - GGUF models use llama-cpp-python (efficient, no de-quantization)
+    - Other models use HuggingFace transformers
+    """
+
+    def __init__(
+        self,
+        model_id: str = "google/gemma-3-12b-it-qat-q4_0-gguf",
+        gguf_file: Optional[str] = None,
+        n_ctx: int = 8192,
+        use_4bit: bool = True,
+    ):
+        """
+        Initialize the LLM.
+
+        Args:
+            model_id: HuggingFace model ID
+            gguf_file: GGUF filename (auto-detected if model_id ends with -gguf)
+            n_ctx: Context size for GGUF models
+            use_4bit: Use 4-bit quantization for transformers models
+        """
+        self._model_id = model_id
+        self._gguf_file = gguf_file
+        self._n_ctx = n_ctx
+        self._use_4bit = use_4bit
+
+        # Model instances (lazy loaded)
+        self._llama_model = None  # llama-cpp-python
+        self._transformers_model = None  # HuggingFace transformers
+        self._tokenizer = None
+
+        self._load_failed = False
+
+    @property
+    def is_loaded(self) -> bool:
+        """Check if the model is loaded."""
+        return self._llama_model is not None or self._transformers_model is not None
+
+    def _is_gguf_model(self) -> bool:
+        """Check if the model ID is a GGUF model."""
+        return self._model_id.endswith("-gguf") or self._gguf_file is not None
+
+    def _get_gguf_filename(self) -> str:
+        """Get the GGUF filename from the model ID."""
+        if self._gguf_file:
+            return self._gguf_file
+        # Extract filename from model ID like "google/gemma-3-12b-it-qat-q4_0-gguf"
+        # The actual file is "gemma-3-12b-it-q4_0.gguf" (note: "qat" is removed)
+        model_name = self._model_id.split("/")[-1]
+        if model_name.endswith("-gguf"):
+            model_name = model_name[:-5]  # Remove "-gguf" suffix
+        # Remove "-qat" from the name (it's not in the actual filename)
+        model_name = model_name.replace("-qat", "")
+        return model_name + ".gguf"
+
+    def load(self) -> None:
+        """
+        Load the model.
+
+        Raises:
+            RuntimeError: If the model fails to load
+        """
+        if self.is_loaded or self._load_failed:
+            return
+
+        try:
+            logger.debug(f"Loading LLM: {self._model_id}")
+
+            if self._is_gguf_model():
+                self._load_gguf_model()
+            else:
+                self._load_transformers_model()
+
+            logger.debug("LLM loaded successfully")
+
+        except Exception as e:
+            self._load_failed = True
+            error_msg = f"Failed to load LLM ({self._model_id}): {e}"
+            if "llama_cpp" in str(e).lower() or "llama-cpp" in str(e).lower():
+                error_msg += "\n  Install with: pip install llama-cpp-python"
+            if "accelerate" in str(e):
+                error_msg += "\n  Install with: pip install accelerate"
+            raise RuntimeError(error_msg) from e
+
+    def _load_gguf_model(self) -> None:
+        """Load GGUF model using llama-cpp-python."""
+        try:
+            from llama_cpp import Llama
+            from huggingface_hub import hf_hub_download
+        except ImportError as e:
+            raise ImportError(
+                "llama-cpp-python is required for GGUF models. "
+                "Install with: pip install llama-cpp-python"
+            ) from e
+
+        gguf_file = self._get_gguf_filename()
+        logger.debug(f"Loading GGUF model with file: {gguf_file}")
+
+        # Download the GGUF file from HuggingFace
+        model_path = hf_hub_download(
+            repo_id=self._model_id,
+            filename=gguf_file,
+        )
+
+        # Load with llama-cpp-python
+        self._llama_model = Llama(
+            model_path=model_path,
+            n_ctx=self._n_ctx,
+            n_gpu_layers=-1,  # Use all GPU layers (Metal on Mac, CUDA on Linux)
+            verbose=False,
+        )
+
+    def _load_transformers_model(self) -> None:
+        """Load model using HuggingFace transformers."""
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        self._tokenizer = AutoTokenizer.from_pretrained(self._model_id)
+
+        if self._use_4bit:
+            try:
+                from transformers import BitsAndBytesConfig
+                quantization_config = BitsAndBytesConfig(
+                    load_in_4bit=True,
+                    bnb_4bit_compute_dtype=torch.float16,
+                )
+                self._transformers_model = AutoModelForCausalLM.from_pretrained(
+                    self._model_id,
+                    quantization_config=quantization_config,
+                    device_map="auto",
+                )
+            except ImportError:
+                logger.debug("bitsandbytes not available, loading full precision")
+                self._transformers_model = AutoModelForCausalLM.from_pretrained(
+                    self._model_id,
+                    device_map="auto",
+                    torch_dtype=torch.float16,
+                )
+        else:
+            self._transformers_model = AutoModelForCausalLM.from_pretrained(
+                self._model_id,
+                device_map="auto",
+                torch_dtype=torch.float16,
+            )
+
+    def generate(
+        self,
+        prompt: str,
+        max_tokens: int = 100,
+        stop: Optional[list[str]] = None,
+    ) -> str:
+        """
+        Generate text from a prompt.
+
+        Args:
+            prompt: The input prompt
+            max_tokens: Maximum tokens to generate
+            stop: Stop sequences
+
+        Returns:
+            Generated text (not including the prompt)
+        """
+        self.load()
+
+        if self._llama_model is not None:
+            return self._generate_with_llama(prompt, max_tokens, stop)
+        else:
+            return self._generate_with_transformers(prompt, max_tokens)
+
+    def _generate_with_llama(
+        self,
+        prompt: str,
+        max_tokens: int,
+        stop: Optional[list[str]],
+    ) -> str:
+        """Generate response using llama-cpp-python."""
+        output = self._llama_model(
+            prompt,
+            max_tokens=max_tokens,
+            stop=stop or ["\n\n", "</s>"],
+            echo=False,
+        )
+        return output["choices"][0]["text"]
+
+    def _generate_with_transformers(
+        self,
+        prompt: str,
+        max_tokens: int,
+    ) -> str:
+        """Generate response using transformers."""
+        import torch
+
+        inputs = self._tokenizer(prompt, return_tensors="pt").to(self._transformers_model.device)
+
+        with torch.no_grad():
+            outputs = self._transformers_model.generate(
+                **inputs,
+                max_new_tokens=max_tokens,
+                do_sample=False,
+                pad_token_id=self._tokenizer.pad_token_id,
+            )
+
+        return self._tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+# Singleton instance for shared use
+_default_llm: Optional[LLM] = None
+
+
+def get_llm(
+    model_id: str = "google/gemma-3-12b-it-qat-q4_0-gguf",
+    **kwargs,
+) -> LLM:
+    """
+    Get or create a shared LLM instance.
+
+    Uses a singleton pattern to avoid loading the model multiple times.
+
+    Args:
+        model_id: HuggingFace model ID
+        **kwargs: Additional arguments passed to LLM constructor
+
+    Returns:
+        LLM instance
+    """
+    global _default_llm
+
+    if _default_llm is None or _default_llm._model_id != model_id:
+        _default_llm = LLM(model_id=model_id, **kwargs)
+
+    return _default_llm

statement_extractor/models/__init__.py

@@ -0,0 +1,89 @@
+"""
+Data models for the extraction pipeline.
+
+This module contains all Pydantic models used throughout the pipeline stages:
+- Stage 1 (Splitting): RawTriple
+- Stage 2 (Extraction): ExtractedEntity, PipelineStatement
+- Stage 3 (Qualification): EntityQualifiers, QualifiedEntity
+- Stage 4 (Canonicalization): CanonicalMatch, CanonicalEntity
+- Stage 5 (Labeling): StatementLabel, LabeledStatement
+
+It also re-exports all models from the original models.py for backward compatibility.
+"""
+
+# Import from the original models.py file (now a sibling at the same level)
+# We need to import these BEFORE the local modules to avoid circular imports
+import sys
+import importlib.util
+from pathlib import Path
+
+# Manually load the old models.py to avoid conflict with this package
+_models_py_path = Path(__file__).parent.parent / "models.py"
+if _models_py_path.exists():
+    _spec = importlib.util.spec_from_file_location("_old_models", _models_py_path)
+    _old_models = importlib.util.module_from_spec(_spec)
+    _spec.loader.exec_module(_old_models)
+
+    # Re-export everything from the old models
+    Entity = _old_models.Entity
+    ExtractionMethod = _old_models.ExtractionMethod
+    Statement = _old_models.Statement
+    ExtractionResult = _old_models.ExtractionResult
+    PredicateMatch = _old_models.PredicateMatch
+    PredicateTaxonomy = _old_models.PredicateTaxonomy
+    PredicateComparisonConfig = _old_models.PredicateComparisonConfig
+    ScoringConfig = _old_models.ScoringConfig
+    ExtractionOptions = _old_models.ExtractionOptions
+
+    # Use EntityType from old models
+    EntityType = _old_models.EntityType
+else:
+    # Fallback: define locally if old models.py doesn't exist
+    from .entity import EntityType
+
+# New pipeline models
+from .entity import ExtractedEntity
+from .statement import 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)
+    "Entity",
+    "EntityType",
+    "ExtractionMethod",
+    "Statement",
+    "ExtractionResult",
+    "PredicateMatch",
+    "PredicateTaxonomy",
+    "PredicateComparisonConfig",
+    "ScoringConfig",
+    "ExtractionOptions",
+    # New pipeline models
+    "ExtractedEntity",
+    "RawTriple",
+    "PipelineStatement",
+    "EntityQualifiers",
+    "QualifiedEntity",
+    "ResolvedRole",
+    "ResolvedOrganization",
+    "CanonicalMatch",
+    "CanonicalEntity",
+    "StatementLabel",
+    "LabeledStatement",
+    "TaxonomyResult",
+    # Document models
+    "Document",
+    "DocumentMetadata",
+    "DocumentPage",
+    "TextChunk",
+    "ChunkingConfig",
+]