corp-extractor 0.2.7__py3-none-any.whl

statement_extractor/canonicalization.py

@@ -0,0 +1,196 @@
+"""
+Entity canonicalization for statement deduplication.
+
+Provides default canonicalization functions and a Canonicalizer class
+for normalizing entity text before comparison.
+"""
+
+import re
+from typing import Callable, Optional
+
+from .models import Statement
+
+
+# Common determiners to remove from the start of entity text
+DETERMINERS = frozenset(["the", "a", "an", "this", "that", "these", "those"])
+
+
+def default_entity_canonicalizer(text: str) -> str:
+    """
+    Default entity canonicalization function.
+
+    Transformations:
+    - Trim leading/trailing whitespace
+    - Convert to lowercase
+    - Remove leading determiners (the, a, an, etc.)
+    - Normalize internal whitespace (multiple spaces -> single)
+
+    Args:
+        text: The entity text to canonicalize
+
+    Returns:
+        Canonicalized text
+
+    Example:
+        >>> default_entity_canonicalizer("  The  Apple Inc.  ")
+        'apple inc.'
+        >>> default_entity_canonicalizer("A new product")
+        'new product'
+    """
+    # Trim and lowercase
+    result = text.strip().lower()
+
+    # Normalize internal whitespace
+    result = re.sub(r'\s+', ' ', result)
+
+    # Remove leading determiners
+    words = result.split()
+    if words and words[0] in DETERMINERS:
+        result = ' '.join(words[1:])
+
+    return result.strip()
+
+
+class Canonicalizer:
+    """
+    Canonicalize entities for deduplication.
+
+    Supports custom canonicalization functions for entities.
+    Predicate comparison uses embeddings (see PredicateComparer).
+
+    Example:
+        >>> canon = Canonicalizer()
+        >>> canon.canonicalize_entity("The Apple Inc.")
+        'apple inc.'
+
+        >>> # With custom function
+        >>> canon = Canonicalizer(entity_fn=lambda x: x.upper())
+        >>> canon.canonicalize_entity("Apple Inc.")
+        'APPLE INC.'
+    """
+
+    def __init__(
+        self,
+        entity_fn: Optional[Callable[[str], str]] = None,
+    ):
+        """
+        Initialize the canonicalizer.
+
+        Args:
+            entity_fn: Custom function to canonicalize entity text.
+                       If None, uses default_entity_canonicalizer.
+        """
+        self.entity_fn = entity_fn or default_entity_canonicalizer
+
+    def canonicalize_entity(self, text: str) -> str:
+        """
+        Canonicalize an entity string.
+
+        Args:
+            text: Entity text to canonicalize
+
+        Returns:
+            Canonicalized text
+        """
+        return self.entity_fn(text)
+
+    def canonicalize_statement_entities(
+        self,
+        statement: Statement
+    ) -> tuple[str, str]:
+        """
+        Return canonicalized (subject, object) tuple.
+
+        Note: Predicate comparison uses embeddings, not text canonicalization.
+
+        Args:
+            statement: Statement to canonicalize
+
+        Returns:
+            Tuple of (canonicalized_subject, canonicalized_object)
+        """
+        return (
+            self.canonicalize_entity(statement.subject.text),
+            self.canonicalize_entity(statement.object.text),
+        )
+
+    def create_dedup_key(
+        self,
+        statement: Statement,
+        predicate_canonical: Optional[str] = None
+    ) -> tuple[str, str, str]:
+        """
+        Create a deduplication key for a statement.
+
+        For exact-match deduplication (when not using embedding-based comparison).
+
+        Args:
+            statement: Statement to create key for
+            predicate_canonical: Optional canonical predicate (if taxonomy was used)
+
+        Returns:
+            Tuple of (subject, predicate, object) for deduplication
+        """
+        subj = self.canonicalize_entity(statement.subject.text)
+        obj = self.canonicalize_entity(statement.object.text)
+        pred = predicate_canonical or statement.predicate.lower().strip()
+        return (subj, pred, obj)
+
+
+def deduplicate_statements_exact(
+    statements: list[Statement],
+    entity_canonicalizer: Optional[Callable[[str], str]] = None,
+    detect_reversals: bool = True,
+) -> list[Statement]:
+    """
+    Deduplicate statements using exact text matching.
+
+    Use this when embedding-based deduplication is disabled.
+    When duplicates are found, entity types are merged - specific types
+    (ORG, PERSON, etc.) take precedence over UNKNOWN.
+
+    When detect_reversals=True, also detects reversed duplicates where
+    subject and object are swapped. The first occurrence determines the
+    canonical orientation.
+
+    Args:
+        statements: List of statements to deduplicate
+        entity_canonicalizer: Optional custom canonicalization function
+        detect_reversals: Whether to detect reversed duplicates (default True)
+
+    Returns:
+        Deduplicated list with merged entity types
+    """
+    if len(statements) <= 1:
+        return statements
+
+    canonicalizer = Canonicalizer(entity_fn=entity_canonicalizer)
+
+    # Map from dedup key to index in unique list
+    seen: dict[tuple[str, str, str], int] = {}
+    unique: list[Statement] = []
+
+    for stmt in statements:
+        key = canonicalizer.create_dedup_key(stmt)
+        # Also compute reversed key (object, predicate, subject)
+        reversed_key = (key[2], key[1], key[0])
+
+        if key in seen:
+            # Direct duplicate found - merge entity types
+            existing_idx = seen[key]
+            existing_stmt = unique[existing_idx]
+            merged_stmt = existing_stmt.merge_entity_types_from(stmt)
+            unique[existing_idx] = merged_stmt
+        elif detect_reversals and reversed_key in seen:
+            # Reversed duplicate found - merge entity types (accounting for reversal)
+            existing_idx = seen[reversed_key]
+            existing_stmt = unique[existing_idx]
+            # Merge types from the reversed statement
+            merged_stmt = existing_stmt.merge_entity_types_from(stmt.reversed())
+            unique[existing_idx] = merged_stmt
+        else:
+            # New unique statement
+            seen[key] = len(unique)
+            unique.append(stmt)
+
+    return unique

statement_extractor/cli.py

@@ -0,0 +1,215 @@
+"""
+Command-line interface for statement extraction.
+
+Usage:
+    corp-extractor "Your text here"
+    corp-extractor -f input.txt
+    cat input.txt | corp-extractor -
+"""
+
+import sys
+from typing import Optional
+
+import click
+
+from . import __version__
+from .models import (
+    ExtractionOptions,
+    PredicateComparisonConfig,
+    PredicateTaxonomy,
+    ScoringConfig,
+)
+
+
+@click.command()
+@click.argument("text", required=False)
+@click.option("-f", "--file", "input_file", type=click.Path(exists=True), help="Read input from file")
+@click.option(
+    "-o", "--output",
+    type=click.Choice(["table", "json", "xml"], case_sensitive=False),
+    default="table",
+    help="Output format (default: table)"
+)
+@click.option("--json", "output_json", is_flag=True, help="Output as JSON (shortcut for -o json)")
+@click.option("--xml", "output_xml", is_flag=True, help="Output as XML (shortcut for -o xml)")
+# Beam search options
+@click.option("-b", "--beams", type=int, default=4, help="Number of beams for diverse beam search (default: 4)")
+@click.option("--diversity", type=float, default=1.0, help="Diversity penalty for beam search (default: 1.0)")
+@click.option("--max-tokens", type=int, default=2048, help="Maximum tokens to generate (default: 2048)")
+# Deduplication options
+@click.option("--no-dedup", is_flag=True, help="Disable deduplication")
+@click.option("--no-embeddings", is_flag=True, help="Disable embedding-based deduplication (faster)")
+@click.option("--no-merge", is_flag=True, help="Disable beam merging (select single best beam)")
+@click.option("--dedup-threshold", type=float, default=0.65, help="Similarity threshold for deduplication (default: 0.65)")
+# Quality options
+@click.option("--min-confidence", type=float, default=0.0, help="Minimum confidence threshold 0-1 (default: 0)")
+# Taxonomy options
+@click.option("--taxonomy", type=click.Path(exists=True), help="Load predicate taxonomy from file (one per line)")
+@click.option("--taxonomy-threshold", type=float, default=0.5, help="Similarity threshold for taxonomy matching (default: 0.5)")
+# Device options
+@click.option("--device", type=click.Choice(["auto", "cuda", "cpu"]), default="auto", help="Device to use (default: auto)")
+# Output options
+@click.option("-v", "--verbose", is_flag=True, help="Show verbose output with confidence scores")
+@click.option("-q", "--quiet", is_flag=True, help="Suppress progress messages")
+@click.version_option(version=__version__)
+def main(
+    text: Optional[str],
+    input_file: Optional[str],
+    output: str,
+    output_json: bool,
+    output_xml: bool,
+    beams: int,
+    diversity: float,
+    max_tokens: int,
+    no_dedup: bool,
+    no_embeddings: bool,
+    no_merge: bool,
+    dedup_threshold: float,
+    min_confidence: float,
+    taxonomy: Optional[str],
+    taxonomy_threshold: float,
+    device: str,
+    verbose: bool,
+    quiet: bool,
+):
+    """
+    Extract structured statements from text.
+
+    TEXT can be provided as an argument, read from a file with -f, or piped via stdin.
+
+    \b
+    Examples:
+        corp-extractor "Apple announced a new iPhone."
+        corp-extractor -f article.txt --json
+        corp-extractor -f article.txt -o json --beams 8
+        cat article.txt | corp-extractor -
+        echo "Tim Cook is CEO of Apple." | corp-extractor - --verbose
+
+    \b
+    Output formats:
+        table  Human-readable table (default)
+        json   JSON with full metadata
+        xml    Raw XML from model
+    """
+    # Determine output format
+    if output_json:
+        output = "json"
+    elif output_xml:
+        output = "xml"
+
+    # Get input text
+    input_text = _get_input_text(text, input_file)
+    if not input_text:
+        raise click.UsageError(
+            "No input provided. Use: statement-extractor \"text\", "
+            "statement-extractor -f file.txt, or pipe via stdin."
+        )
+
+    if not quiet:
+        click.echo(f"Processing {len(input_text)} characters...", err=True)
+
+    # Load taxonomy if provided
+    predicate_taxonomy = None
+    if taxonomy:
+        predicate_taxonomy = PredicateTaxonomy.from_file(taxonomy)
+        if not quiet:
+            click.echo(f"Loaded taxonomy with {len(predicate_taxonomy.predicates)} predicates", err=True)
+
+    # Configure predicate comparison
+    predicate_config = PredicateComparisonConfig(
+        similarity_threshold=taxonomy_threshold,
+        dedup_threshold=dedup_threshold,
+    )
+
+    # Configure scoring
+    scoring_config = ScoringConfig(min_confidence=min_confidence)
+
+    # Configure extraction options
+    options = ExtractionOptions(
+        num_beams=beams,
+        diversity_penalty=diversity,
+        max_new_tokens=max_tokens,
+        deduplicate=not no_dedup,
+        embedding_dedup=not no_embeddings,
+        merge_beams=not no_merge,
+        predicate_taxonomy=predicate_taxonomy,
+        predicate_config=predicate_config,
+        scoring_config=scoring_config,
+    )
+
+    # Import here to allow --help without loading torch
+    from .extractor import StatementExtractor
+
+    # Create extractor with specified device
+    device_arg = None if device == "auto" else device
+    extractor = StatementExtractor(device=device_arg)
+
+    if not quiet:
+        click.echo(f"Using device: {extractor.device}", err=True)
+
+    # Run extraction
+    try:
+        if output == "xml":
+            result = extractor.extract_as_xml(input_text, options)
+            click.echo(result)
+        elif output == "json":
+            result = extractor.extract_as_json(input_text, options)
+            click.echo(result)
+        else:
+            # Table format
+            result = extractor.extract(input_text, options)
+            _print_table(result, verbose)
+    except Exception as e:
+        raise click.ClickException(f"Extraction failed: {e}")
+
+
+def _get_input_text(text: Optional[str], input_file: Optional[str]) -> Optional[str]:
+    """Get input text from argument, file, or stdin."""
+    if text == "-" or (text is None and input_file is None and not sys.stdin.isatty()):
+        # Read from stdin
+        return sys.stdin.read().strip()
+    elif input_file:
+        # Read from file
+        with open(input_file, "r", encoding="utf-8") as f:
+            return f.read().strip()
+    elif text:
+        return text.strip()
+    return None
+
+
+def _print_table(result, verbose: bool):
+    """Print statements in a human-readable table format."""
+    if not result.statements:
+        click.echo("No statements extracted.")
+        return
+
+    click.echo(f"\nExtracted {len(result.statements)} statement(s):\n")
+    click.echo("-" * 80)
+
+    for i, stmt in enumerate(result.statements, 1):
+        subject_type = f" ({stmt.subject.type.value})" if stmt.subject.type.value != "UNKNOWN" else ""
+        object_type = f" ({stmt.object.type.value})" if stmt.object.type.value != "UNKNOWN" else ""
+
+        click.echo(f"{i}. {stmt.subject.text}{subject_type}")
+        click.echo(f"   --[{stmt.predicate}]-->")
+        click.echo(f"   {stmt.object.text}{object_type}")
+
+        if verbose:
+            if stmt.confidence_score is not None:
+                click.echo(f"   Confidence: {stmt.confidence_score:.2f}")
+
+            if stmt.canonical_predicate:
+                click.echo(f"   Canonical: {stmt.canonical_predicate}")
+
+            if stmt.was_reversed:
+                click.echo(f"   (subject/object were swapped)")
+
+            if stmt.source_text:
+                source = stmt.source_text[:60] + "..." if len(stmt.source_text) > 60 else stmt.source_text
+                click.echo(f"   Source: \"{source}\"")
+
+        click.echo("-" * 80)
+
+
+if __name__ == "__main__":
+    main()