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

statement_extractor/extractor.py

@@ -721,16 +721,17 @@ class StatementExtractor:
         Parse XML output into Statement objects.
 
         Uses model for subject, object, entity types, and source_text.
-        Always uses spaCy for predicate extraction (model predicates are unreliable).
+        Always uses GLiNER2 for predicate extraction (model predicates are unreliable).
 
         Produces two candidates for each statement:
-        1. Hybrid: model subject/object + spaCy predicate
-        2. spaCy-only: all components from spaCy
+        1. Hybrid: model subject/object + GLiNER2 predicate
+        2. GLiNER2-only: all components from GLiNER2
 
         Both go into the candidate pool; scoring/dedup picks the best.
         """
         statements: list[Statement] = []
-        use_spacy_extraction = options.use_spacy_extraction if options else True
+        use_gliner_extraction = options.use_gliner_extraction if options else True
+        predicates = options.predicates if options else None
 
         try:
             root = ET.fromstring(xml_output)
@@ -780,75 +781,59 @@ class StatementExtractor:
                     logger.debug(f"Skipping statement: missing subject or object from model")
                     continue
 
-                if use_spacy_extraction and source_text:
+                if use_gliner_extraction and source_text:
                     try:
-                        from .spacy_extraction import extract_triple_from_text, extract_triple_by_predicate_split
-                        spacy_result = extract_triple_from_text(
+                        from .gliner_extraction import extract_triple_from_text
+
+                        # Get model predicate for fallback/refinement
+                        predicate_elem = stmt_elem.find('predicate')
+                        model_predicate = predicate_elem.text.strip() if predicate_elem is not None and predicate_elem.text else ""
+
+                        gliner_result = extract_triple_from_text(
                             source_text=source_text,
                             model_subject=subject_text,
                             model_object=object_text,
-                            model_predicate="",  # Don't pass model predicate
+                            model_predicate=model_predicate,
+                            predicates=predicates,
                         )
-                        if spacy_result:
-                            spacy_subj, spacy_pred, spacy_obj = spacy_result
+                        if gliner_result:
+                            gliner_subj, gliner_pred, gliner_obj = gliner_result
 
-                            if spacy_pred:
-                                # Candidate 1: Hybrid (model subject/object + spaCy predicate)
+                            if gliner_pred:
+                                # Candidate 1: Hybrid (model subject/object + GLiNER2 predicate)
                                 logger.debug(
-                                    f"Adding hybrid candidate: '{subject_text}' --[{spacy_pred}]--> '{object_text}'"
+                                    f"Adding hybrid candidate: '{subject_text}' --[{gliner_pred}]--> '{object_text}'"
                                 )
                                 statements.append(Statement(
                                     subject=Entity(text=subject_text, type=subject_type),
-                                    predicate=spacy_pred,
+                                    predicate=gliner_pred,
                                     object=Entity(text=object_text, type=object_type),
                                     source_text=source_text,
                                     extraction_method=ExtractionMethod.HYBRID,
                                 ))
 
-                                # Candidate 2: spaCy-only (if different from hybrid)
-                                if spacy_subj and spacy_obj:
-                                    is_different = (spacy_subj != subject_text or spacy_obj != object_text)
+                                # Candidate 2: GLiNER2-only (if different from hybrid)
+                                if gliner_subj and gliner_obj:
+                                    is_different = (gliner_subj != subject_text or gliner_obj != object_text)
                                     if is_different:
                                         logger.debug(
-                                            f"Adding spaCy-only candidate: '{spacy_subj}' --[{spacy_pred}]--> '{spacy_obj}'"
-                                        )
-                                        statements.append(Statement(
-                                            subject=Entity(text=spacy_subj, type=subject_type),
-                                            predicate=spacy_pred,
-                                            object=Entity(text=spacy_obj, type=object_type),
-                                            source_text=source_text,
-                                            extraction_method=ExtractionMethod.SPACY,
-                                        ))
-
-                                # Candidate 3: Predicate-split (split source text around predicate)
-                                split_result = extract_triple_by_predicate_split(
-                                    source_text=source_text,
-                                    predicate=spacy_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_spacy = (split_subj != spacy_subj or split_obj != spacy_obj)
-                                    if is_different_from_hybrid and is_different_from_spacy:
-                                        logger.debug(
-                                            f"Adding predicate-split candidate: '{split_subj}' --[{split_pred}]--> '{split_obj}'"
+                                            f"Adding GLiNER2-only candidate: '{gliner_subj}' --[{gliner_pred}]--> '{gliner_obj}'"
                                         )
                                         statements.append(Statement(
-                                            subject=Entity(text=split_subj, type=subject_type),
-                                            predicate=split_pred,
-                                            object=Entity(text=split_obj, type=object_type),
+                                            subject=Entity(text=gliner_subj, type=subject_type),
+                                            predicate=gliner_pred,
+                                            object=Entity(text=gliner_obj, type=object_type),
                                             source_text=source_text,
-                                            extraction_method=ExtractionMethod.SPLIT,
+                                            extraction_method=ExtractionMethod.GLINER,
                                         ))
                             else:
                                 logger.debug(
-                                    f"spaCy found no predicate for: '{subject_text}' --> '{object_text}'"
+                                    f"GLiNER2 found no predicate for: '{subject_text}' --> '{object_text}'"
                                 )
                     except Exception as e:
-                        logger.debug(f"spaCy extraction failed: {e}")
+                        logger.debug(f"GLiNER2 extraction failed: {e}")
                 else:
-                    # spaCy disabled - fall back to model predicate
+                    # GLiNER2 disabled - fall back to model predicate
                     predicate_elem = stmt_elem.find('predicate')
                     model_predicate = predicate_elem.text.strip() if predicate_elem is not None and predicate_elem.text else ""
 

statement_extractor/gliner_extraction.py

@@ -0,0 +1,218 @@
+"""
+GLiNER2-based triple extraction.
+
+Uses GLiNER2 for relation extraction and entity recognition to extract
+subject, predicate, and object from source text. T5-Gemma model provides
+triple structure and coreference resolution, while GLiNER2 handles
+linguistic analysis.
+
+The GLiNER2 model is loaded automatically on first use.
+"""
+
+import logging
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+# Lazy-loaded GLiNER2 model
+_model = None
+
+
+def _get_model():
+    """
+    Lazy-load the GLiNER2 model.
+
+    Uses the base model (205M parameters) which is CPU-optimized.
+    """
+    global _model
+    if _model is None:
+        from gliner2 import GLiNER2
+
+        logger.info("Loading GLiNER2 model 'fastino/gliner2-base-v1'...")
+        _model = GLiNER2.from_pretrained("fastino/gliner2-base-v1")
+        logger.debug("GLiNER2 model loaded")
+    return _model
+
+
+def extract_triple_from_text(
+    source_text: str,
+    model_subject: str,
+    model_object: str,
+    model_predicate: str,
+    predicates: Optional[list[str]] = None,
+) -> tuple[str, str, str] | None:
+    """
+    Extract subject, predicate, object from source text using GLiNER2.
+
+    Returns a GLiNER2-based triple that can be added to the candidate pool
+    alongside the model's triple. The existing scoring/dedup logic will
+    pick the best one.
+
+    Args:
+        source_text: The source sentence to analyze
+        model_subject: Subject from T5-Gemma (used for matching and fallback)
+        model_object: Object from T5-Gemma (used for matching and fallback)
+        model_predicate: Predicate from T5-Gemma (used when no predicates provided)
+        predicates: Optional list of predefined relation types to extract
+
+    Returns:
+        Tuple of (subject, predicate, object) from GLiNER2, or None if extraction fails
+    """
+    if not source_text:
+        return None
+
+    try:
+        model = _get_model()
+
+        if predicates:
+            # Use relation extraction with predefined predicates
+            result = model.extract_relations(source_text, predicates)
+
+            # Find best matching relation
+            relation_data = result.get("relation_extraction", {})
+            best_match = None
+            best_confidence = 0.0
+
+            for rel_type, relations in relation_data.items():
+                for rel in relations:
+                    # Handle both tuple format and dict format
+                    if isinstance(rel, tuple):
+                        head, tail = rel
+                        confidence = 1.0
+                    else:
+                        head = rel.get("head", {}).get("text", "")
+                        tail = rel.get("tail", {}).get("text", "")
+                        confidence = min(
+                            rel.get("head", {}).get("confidence", 0.5),
+                            rel.get("tail", {}).get("confidence", 0.5)
+                        )
+
+                    # Score based on match with model hints
+                    score = confidence
+                    if model_subject.lower() in head.lower() or head.lower() in model_subject.lower():
+                        score += 0.2
+                    if model_object.lower() in tail.lower() or tail.lower() in model_object.lower():
+                        score += 0.2
+
+                    if score > best_confidence:
+                        best_confidence = score
+                        best_match = (head, rel_type, tail)
+
+            if best_match:
+                logger.debug(
+                    f"GLiNER2 extracted (relation): subj='{best_match[0]}', pred='{best_match[1]}', obj='{best_match[2]}'"
+                )
+                return best_match
+
+        else:
+            # No predicate list provided - use GLiNER2 for entity extraction
+            # and extract predicate from source text using the model's hint
+
+            # Extract entities to refine subject/object boundaries
+            entity_types = [
+                "person", "organization", "company", "location", "city", "country",
+                "product", "event", "date", "money", "quantity"
+            ]
+            result = model.extract_entities(source_text, entity_types)
+            entities = result.get("entities", {})
+
+            # Find entities that match model subject/object
+            refined_subject = model_subject
+            refined_object = model_object
+
+            for entity_type, entity_list in entities.items():
+                for entity in entity_list:
+                    entity_lower = entity.lower()
+                    # Check if this entity matches or contains the model's subject/object
+                    if model_subject.lower() in entity_lower or entity_lower in model_subject.lower():
+                        # Use the entity text if it's more complete
+                        if len(entity) >= len(refined_subject):
+                            refined_subject = entity
+                    if model_object.lower() in entity_lower or entity_lower in model_object.lower():
+                        if len(entity) >= len(refined_object):
+                            refined_object = entity
+
+            # Use model predicate directly (T5-Gemma provides the predicate)
+            if model_predicate:
+                logger.debug(
+                    f"GLiNER2 extracted (entity-refined): subj='{refined_subject}', pred='{model_predicate}', obj='{refined_object}'"
+                )
+                return (refined_subject, model_predicate, refined_object)
+
+        return None
+
+    except ImportError as e:
+        logger.warning(f"GLiNER2 not installed: {e}")
+        return None
+    except Exception as e:
+        logger.debug(f"GLiNER2 extraction failed: {e}")
+        return None
+
+
+def score_entity_content(text: str) -> float:
+    """
+    Score how entity-like a text is using GLiNER2 entity recognition.
+
+    Returns:
+        1.0 - Recognized as a named entity with high confidence
+        0.8 - Recognized as an entity with moderate confidence
+        0.6 - Partially recognized or contains entity-like content
+        0.2 - Not recognized as any entity type
+    """
+    if not text or not text.strip():
+        return 0.2
+
+    try:
+        model = _get_model()
+
+        # Check if text is recognized as common entity types
+        entity_types = [
+            "person", "organization", "company", "location", "city", "country",
+            "product", "event", "date", "money", "quantity"
+        ]
+
+        result = model.extract_entities(
+            text,
+            entity_types,
+            include_confidence=True
+        )
+
+        # Result format: {'entities': {'person': [{'text': '...', 'confidence': 0.99}], ...}}
+        entities_dict = result.get("entities", {})
+
+        # Find best matching entity across all types
+        best_confidence = 0.0
+        text_lower = text.lower().strip()
+
+        for entity_type, entity_list in entities_dict.items():
+            for entity in entity_list:
+                if isinstance(entity, dict):
+                    entity_text = entity.get("text", "").lower().strip()
+                    confidence = entity.get("confidence", 0.5)
+                else:
+                    # Fallback for string format
+                    entity_text = str(entity).lower().strip()
+                    confidence = 0.8
+
+                # Check if entity covers most of the input text
+                if entity_text == text_lower:
+                    # Exact match
+                    best_confidence = max(best_confidence, confidence)
+                elif entity_text in text_lower or text_lower in entity_text:
+                    # Partial match - reduce confidence
+                    best_confidence = max(best_confidence, confidence * 0.8)
+
+        if best_confidence >= 0.9:
+            return 1.0
+        elif best_confidence >= 0.7:
+            return 0.8
+        elif best_confidence >= 0.5:
+            return 0.6
+        elif best_confidence > 0:
+            return 0.4
+        else:
+            return 0.2
+
+    except Exception as e:
+        logger.debug(f"Entity scoring failed for '{text}': {e}")
+        return 0.5  # Neutral score on error

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,74 @@
+"""
+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
+from .canonical import CanonicalMatch, CanonicalEntity
+from .labels import StatementLabel, LabeledStatement, TaxonomyResult
+
+__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",
+    "CanonicalMatch",
+    "CanonicalEntity",
+    "StatementLabel",
+    "LabeledStatement",
+    "TaxonomyResult",
+]