odin-engine 0.1.0__py3-none-any.whl

benchmarks/__init__.py

@@ -0,0 +1,17 @@
+"""
+Odin Benchmarks: Academic validation against standard KG datasets.
+
+This module provides:
+- Standard dataset loaders (FB15k-237, WN18RR)
+- KG completion metrics (MRR, Hits@K)
+- Benchmark runners for NPLL evaluation
+- Ablation study tools
+"""
+
+from .metrics import mrr, hits_at_k, evaluate_rankings
+from .datasets import load_fb15k237, load_wn18rr, BenchmarkDataset
+
+__all__ = [
+    "mrr", "hits_at_k", "evaluate_rankings",
+    "load_fb15k237", "load_wn18rr", "BenchmarkDataset"
+]

benchmarks/datasets.py

@@ -0,0 +1,284 @@
+"""
+Standard Knowledge Graph Benchmark Datasets
+
+Provides loaders for:
+- FB15k-237: Freebase subset (14,541 entities, 237 relations)
+- WN18RR: WordNet subset (40,943 entities, 11 relations)
+
+Datasets are downloaded from standard sources and cached locally.
+"""
+
+import os
+import urllib.request
+import tarfile
+import zipfile
+from pathlib import Path
+from typing import List, Tuple, Set, Dict, Optional
+from dataclasses import dataclass
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Dataset URLs - using villmow/datasets_knowledge_embedding (reliable raw files)
+DATASET_BASE_URLS = {
+    "fb15k237": "https://raw.githubusercontent.com/villmow/datasets_knowledge_embedding/master/FB15k-237",
+    "wn18rr": "https://raw.githubusercontent.com/villmow/datasets_knowledge_embedding/master/WN18RR",
+}
+
+# Default cache directory
+DEFAULT_CACHE_DIR = Path.home() / ".cache" / "odin_benchmarks"
+
+
+@dataclass
+class BenchmarkDataset:
+    """Container for a benchmark dataset."""
+    name: str
+    train_triples: List[Tuple[str, str, str]]
+    valid_triples: List[Tuple[str, str, str]]
+    test_triples: List[Tuple[str, str, str]]
+    entities: List[str]
+    relations: List[str]
+    
+    @property
+    def num_entities(self) -> int:
+        return len(self.entities)
+    
+    @property
+    def num_relations(self) -> int:
+        return len(self.relations)
+    
+    @property
+    def num_train(self) -> int:
+        return len(self.train_triples)
+    
+    @property
+    def num_valid(self) -> int:
+        return len(self.valid_triples)
+    
+    @property
+    def num_test(self) -> int:
+        return len(self.test_triples)
+    
+    def get_train_set(self) -> Set[Tuple[str, str, str]]:
+        return set(self.train_triples)
+    
+    def get_all_triples(self) -> Set[Tuple[str, str, str]]:
+        return set(self.train_triples + self.valid_triples + self.test_triples)
+    
+    def __repr__(self) -> str:
+        return (
+            f"BenchmarkDataset({self.name})\n"
+            f"  Entities: {self.num_entities:,}\n"
+            f"  Relations: {self.num_relations}\n"
+            f"  Train: {self.num_train:,}\n"
+            f"  Valid: {self.num_valid:,}\n"
+            f"  Test: {self.num_test:,}"
+        )
+
+
+def _ensure_dir(path: Path):
+    """Ensure directory exists."""
+    path.mkdir(parents=True, exist_ok=True)
+
+
+def _download_file(url: str, dest: Path):
+    """Download a file with progress."""
+    import ssl
+    import certifi
+    
+    logger.info(f"Downloading {url}...")
+    
+    # Try with certifi SSL context first, fall back to unverified
+    try:
+        ssl_context = ssl.create_default_context(cafile=certifi.where())
+        with urllib.request.urlopen(url, context=ssl_context) as response:
+            with open(dest, 'wb') as out_file:
+                out_file.write(response.read())
+    except (ImportError, ssl.SSLError):
+        # Fallback: disable SSL verification (for development only)
+        logger.warning("SSL verification disabled - using unverified context")
+        ssl_context = ssl.create_default_context()
+        ssl_context.check_hostname = False
+        ssl_context.verify_mode = ssl.CERT_NONE
+        with urllib.request.urlopen(url, context=ssl_context) as response:
+            with open(dest, 'wb') as out_file:
+                out_file.write(response.read())
+    
+    logger.info(f"Downloaded to {dest}")
+
+
+def _extract_tar_gz(archive: Path, dest_dir: Path):
+    """Extract a tar.gz archive."""
+    logger.info(f"Extracting {archive}...")
+    with tarfile.open(archive, "r:gz") as tar:
+        tar.extractall(dest_dir)
+
+
+def _extract_zip(archive: Path, dest_dir: Path):
+    """Extract a zip archive."""
+    logger.info(f"Extracting {archive}...")
+    with zipfile.ZipFile(archive, 'r') as zip_ref:
+        zip_ref.extractall(dest_dir)
+
+
+def _load_triples(filepath: Path) -> List[Tuple[str, str, str]]:
+    """Load triples from a TSV file (head, relation, tail)."""
+    triples = []
+    with open(filepath, "r", encoding="utf-8") as f:
+        for line in f:
+            parts = line.strip().split("\t")
+            if len(parts) == 3:
+                h, r, t = parts
+                triples.append((h, r, t))
+    return triples
+
+
+def _extract_entities_and_relations(
+    triples: List[Tuple[str, str, str]]
+) -> Tuple[List[str], List[str]]:
+    """Extract unique entities and relations from triples."""
+    entities = set()
+    relations = set()
+    for h, r, t in triples:
+        entities.add(h)
+        entities.add(t)
+        relations.add(r)
+    return sorted(entities), sorted(relations)
+
+
+def _download_dataset_files(base_url: str, dataset_dir: Path):
+    """Download train/valid/test files for a dataset."""
+    _ensure_dir(dataset_dir)
+    
+    for split in ["train", "valid", "test"]:
+        file_path = dataset_dir / f"{split}.txt"
+        if not file_path.exists():
+            url = f"{base_url}/{split}.txt"
+            _download_file(url, file_path)
+
+
+def load_fb15k237(cache_dir: Optional[Path] = None) -> BenchmarkDataset:
+    """
+    Load FB15k-237 dataset.
+    
+    FB15k-237 is a subset of Freebase with:
+    - 14,541 entities
+    - 237 relations
+    - 310,116 triples
+    
+    This version removes inverse relations from FB15k to prevent
+    data leakage during evaluation.
+    
+    Args:
+        cache_dir: Directory to cache downloaded data
+    
+    Returns:
+        BenchmarkDataset object
+    """
+    cache_dir = Path(cache_dir) if cache_dir else DEFAULT_CACHE_DIR
+    _ensure_dir(cache_dir)
+    
+    dataset_dir = cache_dir / "FB15k-237"
+    
+    # Download if not cached
+    if not (dataset_dir / "train.txt").exists():
+        _download_dataset_files(DATASET_BASE_URLS["fb15k237"], dataset_dir)
+    
+    # Load splits
+    train = _load_triples(dataset_dir / "train.txt")
+    valid = _load_triples(dataset_dir / "valid.txt")
+    test = _load_triples(dataset_dir / "test.txt")
+    
+    # Extract vocab
+    all_triples = train + valid + test
+    entities, relations = _extract_entities_and_relations(all_triples)
+    
+    return BenchmarkDataset(
+        name="FB15k-237",
+        train_triples=train,
+        valid_triples=valid,
+        test_triples=test,
+        entities=entities,
+        relations=relations,
+    )
+
+
+def load_wn18rr(cache_dir: Optional[Path] = None) -> BenchmarkDataset:
+    """
+    Load WN18RR dataset.
+    
+    WN18RR is a subset of WordNet with:
+    - 40,943 entities
+    - 11 relations
+    - 93,003 triples
+    
+    This version removes inverse relations from WN18 to prevent
+    data leakage during evaluation.
+    
+    Args:
+        cache_dir: Directory to cache downloaded data
+    
+    Returns:
+        BenchmarkDataset object
+    """
+    cache_dir = Path(cache_dir) if cache_dir else DEFAULT_CACHE_DIR
+    _ensure_dir(cache_dir)
+    
+    dataset_dir = cache_dir / "WN18RR"
+    
+    # Download if not cached
+    if not (dataset_dir / "train.txt").exists():
+        _download_dataset_files(DATASET_BASE_URLS["wn18rr"], dataset_dir)
+    
+    # Load splits
+    train = _load_triples(dataset_dir / "train.txt")
+    valid = _load_triples(dataset_dir / "valid.txt")
+    test = _load_triples(dataset_dir / "test.txt")
+    
+    # Extract vocab
+    all_triples = train + valid + test
+    entities, relations = _extract_entities_and_relations(all_triples)
+    
+    return BenchmarkDataset(
+        name="WN18RR",
+        train_triples=train,
+        valid_triples=valid,
+        test_triples=test,
+        entities=entities,
+        relations=relations,
+    )
+
+
+def dataset_to_kg(dataset: BenchmarkDataset):
+    """
+    Convert BenchmarkDataset to Odin KnowledgeGraph.
+    
+    Args:
+        dataset: BenchmarkDataset object
+    
+    Returns:
+        KnowledgeGraph object suitable for NPLL training
+    """
+    from npll.core import KnowledgeGraph
+    
+    kg = KnowledgeGraph()
+    
+    # Add all training triples as known facts
+    for h, r, t in dataset.train_triples:
+        kg.add_known_fact(h, r, t)
+    
+    return kg
+
+
+# Quick test
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    
+    print("Loading FB15k-237...")
+    fb = load_fb15k237()
+    print(fb)
+    print()
+    
+    print("Loading WN18RR...")
+    wn = load_wn18rr()
+    print(wn)

benchmarks/metrics.py

@@ -0,0 +1,275 @@
+"""
+Standard Knowledge Graph Completion Metrics
+
+Implements metrics used in academic KG completion papers:
+- Mean Reciprocal Rank (MRR)
+- Hits@K (K=1, 3, 10)
+
+These metrics evaluate link prediction quality:
+Given (h, r, ?), rank all entities by predicted score.
+"""
+
+import numpy as np
+from typing import List, Dict, Tuple, Any
+from dataclasses import dataclass
+
+
+@dataclass
+class RankingResult:
+    """Result of ranking evaluation for a single triple."""
+    head: str
+    relation: str
+    tail: str
+    tail_rank: int  # Rank of true tail among all candidates
+    head_rank: int  # Rank of true head among all candidates (for inverse)
+    num_candidates: int
+
+
+def mrr(ranks: List[int]) -> float:
+    """
+    Mean Reciprocal Rank
+    
+    MRR = (1/|Q|) * Σ (1/rank_i)
+    
+    Args:
+        ranks: List of ranks (1-indexed, where 1 is best)
+    
+    Returns:
+        MRR score in [0, 1]
+    """
+    if not ranks:
+        return 0.0
+    return float(np.mean([1.0 / r for r in ranks]))
+
+
+def hits_at_k(ranks: List[int], k: int) -> float:
+    """
+    Hits@K - proportion of ranks <= K
+    
+    Args:
+        ranks: List of ranks (1-indexed)
+        k: Cutoff threshold
+    
+    Returns:
+        Hits@K score in [0, 1]
+    """
+    if not ranks:
+        return 0.0
+    return float(np.mean([1 if r <= k else 0 for r in ranks]))
+
+
+def evaluate_rankings(results: List[RankingResult]) -> Dict[str, float]:
+    """
+    Compute all standard metrics from ranking results.
+    
+    Args:
+        results: List of RankingResult objects
+    
+    Returns:
+        Dictionary with MRR, Hits@1, Hits@3, Hits@10
+    """
+    if not results:
+        return {
+            "mrr": 0.0,
+            "hits@1": 0.0,
+            "hits@3": 0.0,
+            "hits@10": 0.0,
+            "num_queries": 0
+        }
+    
+    # Use tail ranks (standard protocol: predict tail given head+relation)
+    tail_ranks = [r.tail_rank for r in results]
+    
+    return {
+        "mrr": mrr(tail_ranks),
+        "hits@1": hits_at_k(tail_ranks, 1),
+        "hits@3": hits_at_k(tail_ranks, 3),
+        "hits@10": hits_at_k(tail_ranks, 10),
+        "num_queries": len(results),
+        "mean_rank": float(np.mean(tail_ranks)),
+        "median_rank": float(np.median(tail_ranks)),
+    }
+
+
+def filtered_rank(
+    true_entity: str,
+    scores: Dict[str, float],
+    filter_entities: set,
+) -> int:
+    """
+    Compute filtered rank (standard protocol).
+    
+    Filtered ranking removes other valid answers from consideration,
+    so we only penalize for ranking random entities above the true one.
+    
+    Args:
+        true_entity: The correct answer
+        scores: Dict mapping entity -> score
+        filter_entities: Other valid entities to filter out
+    
+    Returns:
+        Filtered rank (1-indexed)
+    """
+    true_score = scores.get(true_entity, float('-inf'))
+    
+    rank = 1
+    for entity, score in scores.items():
+        if entity == true_entity:
+            continue
+        if entity in filter_entities:
+            continue  # Filter out other valid answers
+        if score > true_score:
+            rank += 1
+    
+    return rank
+
+
+class LinkPredictionEvaluator:
+    """
+    Evaluator for link prediction task.
+    
+    Standard protocol:
+    1. For each test triple (h, r, t):
+       a. Corrupt tail: score all (h, r, e) for e in entities
+       b. Corrupt head: score all (e, r, t) for e in entities
+    2. Filter out other valid triples (filtered setting)
+    3. Report filtered MRR, Hits@1/3/10
+    """
+    
+    def __init__(
+        self,
+        all_entities: List[str],
+        train_triples: set,
+        valid_triples: set = None,
+    ):
+        """
+        Args:
+            all_entities: List of all entity IDs
+            train_triples: Set of (h, r, t) tuples from training
+            valid_triples: Set of (h, r, t) tuples from validation (optional)
+        """
+        self.all_entities = all_entities
+        self.train_triples = train_triples
+        self.valid_triples = valid_triples or set()
+        self.known_triples = train_triples | self.valid_triples
+        
+        # Build index for filtering
+        self._build_filter_index()
+    
+    def _build_filter_index(self):
+        """Build indices for efficient filtering."""
+        # For (h, r, ?): which tails are valid?
+        self.hr_to_tails = {}
+        # For (?, r, t): which heads are valid?
+        self.rt_to_heads = {}
+        
+        for h, r, t in self.known_triples:
+            key_hr = (h, r)
+            key_rt = (r, t)
+            
+            if key_hr not in self.hr_to_tails:
+                self.hr_to_tails[key_hr] = set()
+            self.hr_to_tails[key_hr].add(t)
+            
+            if key_rt not in self.rt_to_heads:
+                self.rt_to_heads[key_rt] = set()
+            self.rt_to_heads[key_rt].add(h)
+    
+    def evaluate_triple(
+        self,
+        head: str,
+        relation: str,
+        tail: str,
+        score_fn,
+    ) -> RankingResult:
+        """
+        Evaluate a single test triple.
+        
+        Args:
+            head: Head entity
+            relation: Relation type
+            tail: True tail entity
+            score_fn: Function (h, r, t) -> float
+        
+        Returns:
+            RankingResult with filtered ranks
+        """
+        # Score all possible tails
+        tail_scores = {}
+        for entity in self.all_entities:
+            tail_scores[entity] = score_fn(head, relation, entity)
+        
+        # Get filter set (other valid tails for this h,r pair)
+        filter_tails = self.hr_to_tails.get((head, relation), set()) - {tail}
+        
+        # Compute filtered rank
+        tail_rank = filtered_rank(tail, tail_scores, filter_tails)
+        
+        # Score all possible heads (for bidirectional evaluation)
+        head_scores = {}
+        for entity in self.all_entities:
+            head_scores[entity] = score_fn(entity, relation, tail)
+        
+        # Get filter set for heads
+        filter_heads = self.rt_to_heads.get((relation, tail), set()) - {head}
+        head_rank = filtered_rank(head, head_scores, filter_heads)
+        
+        return RankingResult(
+            head=head,
+            relation=relation,
+            tail=tail,
+            tail_rank=tail_rank,
+            head_rank=head_rank,
+            num_candidates=len(self.all_entities),
+        )
+    
+    def evaluate_batch(
+        self,
+        test_triples: List[Tuple[str, str, str]],
+        score_fn,
+        verbose: bool = True,
+    ) -> Dict[str, float]:
+        """
+        Evaluate on a batch of test triples.
+        
+        Args:
+            test_triples: List of (h, r, t) tuples
+            score_fn: Function (h, r, t) -> float
+            verbose: Print progress
+        
+        Returns:
+            Dictionary with all metrics
+        """
+        results = []
+        
+        for i, (h, r, t) in enumerate(test_triples):
+            if verbose and (i + 1) % 100 == 0:
+                print(f"  Evaluated {i + 1}/{len(test_triples)} triples...")
+            
+            result = self.evaluate_triple(h, r, t, score_fn)
+            results.append(result)
+        
+        return evaluate_rankings(results)
+
+
+# Convenience function for quick evaluation
+def quick_evaluate(
+    test_triples: List[Tuple[str, str, str]],
+    score_fn,
+    all_entities: List[str],
+    train_triples: set,
+) -> Dict[str, float]:
+    """
+    Quick evaluation helper.
+    
+    Args:
+        test_triples: Test set
+        score_fn: Scoring function
+        all_entities: All entities in KG
+        train_triples: Training triples for filtering
+    
+    Returns:
+        Metrics dictionary
+    """
+    evaluator = LinkPredictionEvaluator(all_entities, train_triples)
+    return evaluator.evaluate_batch(test_triples, score_fn)