link-prediction-api

Imported from dice-group/dice-embeddings (.claude/skills/link-prediction-api/SKILL.md). Install upstream with npx skills add dice-group/dice-embeddings --skill link-prediction-api. Copyright stays with the author.

Use a Pre-Trained KGE Model for Inference

When to Use

• Loading a trained model to predict missing triples
• Running 1p / 2p / 3p / intersection / union / negation queries
• Extracting entity or relation embeddings
• Predicting literal values (numerical attributes) with train_literals
• Deploying the model as a web UI via index_serve.py

Reference Implementation

• dicee/knowledge_graph_embeddings.py — KGE class
• dicee/abstracts.py — mixins BaseInteractiveKGE, InteractiveQueryDecomposition, BaseInteractiveTrainKGE (source of is_seen, sample_entity, train_literals, etc.)
• dicee/scripts/index_serve.py — Gradio deployment

1. Load a Pre-Trained Model

from dicee import KGE

# Load from an experiment folder (must contain model.pt + configuration.json)
model = KGE(path="Experiments/2024-01-01_12-00/")

# Or load from a custom path
model = KGE(path="MyRun/")

The path directory must contain:

• model.pt — model weights
• configuration.json — training config
• entity_to_idx.csv — entity index mapping
• relation_to_idx.csv — relation index mapping

2. Predict Missing Links — predict_topk

predict_topk infers the missing element in a triple. Supply exactly two of h, r, t.

# (h, r, ?) — predict missing TAIL entity
results = model.predict_topk(h=["Alice"], r=["knows"], topk=5)
# Returns: [[ ("Bob", 0.95), ("Carol", 0.88), ... ]]

# (?, r, t) — predict missing HEAD entity
results = model.predict_topk(r=["knows"], t=["Bob"], topk=5)

# (h, ?, t) — predict missing RELATION
results = model.predict_topk(h=["Alice"], t=["Bob"], topk=5)

# Batch query — pass lists of strings
results = model.predict_topk(h=["Alice", "Dave"], r=["knows", "friendOf"], topk=3)
# Returns list of B result lists (one per h-r pair)

# Restrict candidates with `within`
results = model.predict_topk(h=["Alice"], r=["knows"],
                              within=["Bob", "Carol", "Eve"], topk=3)

Notes:

• Scores are sigmoid-normalised (0–1)
• batch_size controls memory during inference (default: 1024)
• Input strings must match the KG's entity/relation vocabulary exactly

3. Check Vocabulary

# Is an entity / relation in the model's vocabulary?
model.is_seen(entity="Alice")        # True / False
model.is_seen(relation="knows")      # True / False

# Sample random entities
model.sample_entity(n=10)             # List[str]

# Vocabulary sizes
len(model.entity_to_idx)              # number of entities
len(model.relation_to_idx)            # number of relations

4. Access Raw Embeddings

# Get entity embeddings as list of floats
vecs = model.get_transductive_entity_embeddings(
    indices=["Alice", "Bob"],
    as_list=True       # default — List[List[float]]
)

# As numpy array
vecs = model.get_transductive_entity_embeddings(
    indices=["Alice", "Bob"],
    as_numpy=True
)

# As PyTorch tensor
vecs = model.get_transductive_entity_embeddings(
    indices=["Alice", "Bob"],
    as_pytorch=True
)

5. Multi-Hop Query Answering — answer_multi_hop_query

Answers EPFO (Existential Positive First-Order) queries including conjunctions, disjunctions, and negations.

Supported Query Types

Examples

# 1p: Who does Alice know? (?x : knows(Alice, x))
results = model.answer_multi_hop_query(
    query_type="1p",
    query=("Alice", ("knows",)),
    k=10, tnorm="min"
)

# 2p: Who are the friends-of-friends of Alice?
results = model.answer_multi_hop_query(
    query_type="2p",
    query=("Alice", ("knows", "knows")),
    k=10, tnorm="prod"
)

# 2i: Who is known by BOTH Alice and Bob?
results = model.answer_multi_hop_query(
    query_type="2i",
    query=(("Alice", ("knows",)), ("Bob", ("knows",))),
    k=10, tnorm="min"
)

# Batch of queries at once
results = model.answer_multi_hop_query(
    query_type="1p",
    queries=[
        ("Alice", ("knows",)),
        ("Bob", ("friendOf",)),
    ],
    k=5, tnorm="prod"
)

Parameters

6. Literal Prediction

Train a linear regression head on top of entity embeddings to predict numerical attribute values.

# Train the literal model (requires a TSV file: entity \t attribute \t value)
model.train_literals(train_file_path="KGs/DBpedia/literals_train.tsv")

# Predict attribute values
predictions = model.predict_literals(
    entity=["Berlin", "Paris"],
    attribute=["population", "area"]
)
# Returns: {'Berlin': {'population': 3_500_000, 'area': 891.8}, ...}

7. Evaluate on a Custom Dataset

from dicee import KGE
model = KGE(path="MyRun/")

# Evaluate filtered link prediction metrics (MRR, HITS@k)
triples = [
    ("Alice", "knows", "Bob"),
    ("Carol", "livesIn", "Berlin"),
]
metrics = model.eval_lp_performance(dataset=triples, filtered=True)
print(metrics)  # {'MRR': ..., 'MR': ..., 'HITS@1': ..., 'HITS@3': ..., 'HITS@10': ...}

8. Move Model to Device

model.to("cpu")
model.to("cuda")
model.to("cuda:1")

9. Deploy as a Web UI

Use the built-in Gradio server:

# Serve from an experiment folder
python -m dicee.scripts.index_serve --path "MyRun/"

This launches a Gradio interface at http://localhost:7860 with:

• Subject / Predicate / Object fields
• 1vsAll scoring (fill two fields, leave one blank)
• Top-k entity scores displayed

10. Common Errors