Probabilistic Circuits for Knowledge Graph Completion (Practical Guide)

This project is a self-contained reference and runnable demo of a simple, explainable pipeline for Knowledge Graph (KG) completion using mined rules, context discovery, and probabilistic scoring inspired by Probabilistic Circuits (PCs).

It includes:

• A toy KG (family domain) and utilities to use Neo4j optionally.
• A lightweight rule miner (path-based Horn rules and symmetry/subrelation checks).
• Context discovery by grouping co-firing rules that explain training triples.
• Probabilistic context weighting via a simple Noisy-OR model with negative sampling.
• Inference to answer triple queries with probabilities and explanations.
• Optional LLM interface to narrate results or answer natural-language prompts.
• Professional integrations:
  • Neo4j export/import for real graphs
  • AnyBURL runner + rule parsing
  • PyClause adapter (requires external install) for abductive context discovery
  • Juice/pyjuice trainer interface for EM over contexts (with fallback)

Quickstart is below. A full write-up is in this README: see the “Full Write-Up” section for a simple, detailed explanation with examples.

Quickstart

Prereqs: Python 3.9+ recommended.

python -m venv .venv
source .venv/bin/activate  # Windows: .venv\Scripts\activate
pip install -r requirements.txt

# Run the end-to-end demo on a toy family KG
python -m kgpc.toy_demo

# Or use the CLI
python -m kgpc.cli demo
python -m kgpc.cli query "John Grandparent Alice"

Neo4j and LLM hooks are optional. If you want to try them:

• Neo4j: set NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD and load the toy KG via python -m kgpc.cli load-neo4j.
• LLM: set OPENAI_API_KEY and run python -m kgpc.cli ask "Is John Alice's grandparent?".

Pure-Python Pipeline (No Neo4j/No Java)

If you don't have Neo4j or Java, use the RDFlib-backed pipeline:

# 1) Download a dataset
python -m kgpc.cli dataset-download --name WN18RR --out data

# 2) Run the pure-Python pipeline on the train split
python -m kgpc.cli py-pipeline --name WN18RR --split train --samples 10

This will:

• Load triples into an in-memory RDFlib store
• Mine simple Horn rules (Python)
• Discover contexts by co-firing rules (Python)
• Train context weights with EM (Python)
• Print sample triples with probabilities

Train/Save and Query a Model

# Train and save a model on WN18RR train split
python -m kgpc.cli dataset-download --name WN18RR --out data
python -m kgpc.cli py-train --name WN18RR --split train --out data/py_model.json

# Query a triple with the saved model (use exact strings from dataset)
python -m kgpc.cli py-query --name WN18RR --split train --model data/py_model.json --triple "_synset_domain_region_of _hypernym _member_of_domain_region"

# Natural language (only if the relation/entity labels are human-readable)
export OPENAI_API_KEY=...
python -m kgpc.cli py-ask --name WN18RR --split train --model data/py_model.json "Is h r of t?"

Notes:

• Datasets like WN18RR use symbolic identifiers (not human-readable). For best NL behavior, use a dataset with readable labels or provide a mapping layer.
• The context discovery here is Python-only (fallback). If you later provide a Python-installable PyClause, we can swap it in for richer abduction.

Natural Language on WN18RR via Labels

# Build human-friendly labels from WN18RR using WordNet lemmas
python -m kgpc.cli labels-build --name WN18RR --out data/wn18rr_labels.json

# Ask a natural question (maps words to synsets/relations)
python -m kgpc.cli py-ask-nl --name WN18RR --split train --model data/py_model.json \
  --labels data/wn18rr_labels.json "Is dog a hypernym of canine?"

This tries to map the words “dog” / “canine” to WordNet synsets present in WN18RR and uses the trained model’s probabilities. For best results, use lemmas that exist in the dataset splits.

Real Pipeline (Neo4j + AnyBURL + PyClause + Juice)

1. Export triples from Neo4j

export NEO4J_URI=bolt://localhost:7687
export NEO4J_USER=neo4j
export NEO4J_PASSWORD=...
python -m kgpc.cli neo4j-export --out data/triples.tsv

Or import a benchmark dataset and load it to Neo4j:

# Download dataset (FB15k-237 or WN18RR)
python -m kgpc.cli dataset-download --name WN18RR --out data

# Load to Neo4j (train split or all splits)
python -m kgpc.cli dataset-to-neo4j --name WN18RR --split train
# or: --split all

2. Run AnyBURL on exported triples

# Download AnyBURL jar from its repository
python -m kgpc.cli anyburl-run --jar /path/to/AnyBURL.jar --input data/triples.tsv --out data/anyburl_rules.txt --time 120

3. Discover contexts (PyClause recommended)

Install PyClause from source if not on PyPI (adapter expects a pyclause.explain(facts, rules) API). As a fallback, you can run our internal discovery:

python -m kgpc.cli contexts --triples data/triples.tsv --rules data/anyburl_rules.txt --out data/contexts.tsv

4. Train probabilistic circuit with Juice/pyjuice

pip install pyjuice  # if not already installed
python -m kgpc.cli train --contexts data/contexts.tsv --out data/model.tsv

5. Inference

For end-to-end KG QA with explanations using the in-memory pipeline, use ask/query. For models trained via contexts TSV, see infer (weights inspection). A full integration tying Juice-trained weights back into the runtime is scaffolded; we can wire it next based on your dataset scale and preferences.

Security note: Never hardcode secrets (API keys, DB passwords). Use environment variables such as OPENAI_API_KEY, NEO4J_PASSWORD. Rotate any key shared in chat.

Project Structure

kgpc/
  README.md                 # This file
  __init__.py
  cli.py                    # CLI entry points
  toy_demo.py               # One-shot end-to-end demo
  kg.py                     # In-memory KG and helpers
  rule_mining.py            # Simple rule mining
  contexts.py               # Context discovery + scoring data prep
  pc.py                     # Simple probabilistic context model
  inference.py              # Query answering with explanations
  neo4j_utils.py            # Optional Neo4j loader/reader
  llm_interface.py          # Optional LLM interface

Full Write-Up (Simple & Detailed)

1. The Problem: Knowledge Graph Completion

A knowledge graph (KG) stores facts as triples (subject, relation, object).

Example (Family KG):

• (John, Parent, Mary)
• (Mary, Parent, Alice)
• (Alice, Parent, Bob)

KGs are always incomplete. We might be missing:

• (John, Grandparent, Alice)
• (Mary, Ancestor, Bob)

Knowledge graph completion = predicting missing triples.

Two common approaches:

• Embeddings (black box): accurate, but not easily explainable.
• Rules (explainable): interpretable, but can explode to thousands of rules.

This project demonstrates how to tame rule explosion by grouping rules into a small number of probabilistic contexts.

2. Step One: Mining Rules from KG

We mine simple Horn rules from the KG. Examples:

• R1: Parent(x,y) ∧ Parent(y,z) → Grandparent(x,z)
• R2: Parent(x,y) → Ancestor(x,y)
• R3: Grandparent(x,z) → Ancestor(x,z)
• R4: Sibling(x,y) → Sibling(y,x) (symmetry)

Our miner uses observed patterns to propose rules and computes simple confidences.

3. Step Two: Group Rules into Contexts

Rules that co-fire to explain the same training facts are grouped together into contexts. For each known fact (h, r, t) in training, we collect rules whose bodies are satisfied in the KG and whose head predicts (h, r, t). The set of those rules is the “context” explaining that fact. Identical sets across facts define the context clusters.

4. Step Three: Assign Probabilities to Contexts

We estimate how reliable each context is. Our simple model uses a Noisy-OR combination over contexts that support a fact and learns a probability per context using positive facts plus sampled negatives with Laplace smoothing. This approximates a probabilistic circuit and is efficient for a demo.

5. Step Four: Inference with Probabilities

To answer a query like “Is John the grandparent of Alice?”, we:

1. Identify contexts whose rules entail the query.
2. Combine their learned reliabilities with Noisy-OR: p = 1 - ∏(1 - θ_c).
3. Return the probability and the top supporting contexts and rules as explanation.

6. Why This Works

Instead of depending on thousands of individual rules, we operate with a smaller set of probabilistic contexts. This preserves interpretability, reduces redundancy, and often retains much of the predictive power.

7. Pipeline Summary

1. Rule mining (lightweight miner here; AnyBURL/AMIE in production).
2. Context discovery (group co-firing rules per fact).
3. Probabilistic scoring (learn context reliabilities with negatives).
4. Inference (answer queries with probabilities and explanations).

8. Tools and Options

• In-memory KG for speed; optional Neo4j integration.
• Noisy-OR model over contexts; extendable to richer PCs/SPNs.
• Optional LLM interface to narrate or explain outputs.

9. Worked Example (Toy)

Run python -m kgpc.toy_demo to:

• Build a small family KG.
• Mine rules and show their supports/confidences.
• Discover contexts from training triples.
• Learn context probabilities.
• Answer sample queries with probabilities and explanations.

10. Next Steps

• Swap in a real rule miner (AnyBURL/AMIE) and import rules.
• Use a proper abductive engine (e.g., PyClause) for richer explanations.
• Replace Noisy-OR with a full probabilistic circuit/SPN if desired.
• Connect to Neo4j for larger graphs and LLMs for richer interfaces.

Why This Beats A Plain KG

Traditional KGs store raw facts only. This project layers explainable, probabilistic reasoning on top:

• Evidence clustering: We mine generic Horn rules (paths, symmetry, subrelations) and cluster co‑firing rules into a small set of interpretable contexts. This avoids “rule explosion” and keeps reasoning compact and auditible.
• Calibrated probabilities: Each context is given a learned reliability (θ) via an EM procedure. Queries combine evidence with a Noisy‑OR model to produce an explicit probability instead of a binary yes/no.
• NL → triples → explanations: Natural-language prompts are mapped to KG triples using labels/synonyms. We return a probability and an explanation with the relevant contexts and rules, not just a raw edge existence.
• Extensible: Swap or extend miners (AnyBURL/AMIE), abductive engines (PyClause), and proper probabilistic circuits (pyjuice/SPNs) as needed for scale and accuracy.

Example: Capital Of (FB15k‑237)

• Question: “Is Delhi the capital of India?”
• Mapping: (India, /location/country/capital, Delhi)
• Contexts: Capital rules (e.g., inverse consistency) cluster into a “Capital Context” with θ≈0.99.
• Answer: Probability ≈0.99. Explanation cites capital relation and (optionally) inverse‑style rule; includes example bindings where available.

Example: Place of Birth (FB15k‑237)

• Question: “Was Michelle Rodriguez born in San Antonio?”
• Mapping: (Michelle_Rodriguez, /people/person/place_of_birth, San_Antonio)
• Contexts: “Direct PoB Context” dominates; nearby path rules (e.g., contained_by city→state) support auxiliary birthplace inference.
• Answer: Probability ≈0.99 with context θ shown and example witnesses (e.g., San_Antonio→Texas) if available.

Process Summary

1. Mine rules from KG (paths, symmetry, subrelations) with thresholds (support/confidence).
2. Discover contexts by grouping rules that co‑fire to explain facts.
3. Learn θ per context via EM; combine via Noisy‑OR at inference.
4. For NL input, map to a triple via label/synonym resolution, score it, and build an LLM context including probability, contexts, rules, and witnesses.