add-model

name	add-model
description	Add a new Knowledge Graph Embedding model to the dicee framework. Use when implementing a new KGE model, adding a custom scoring function, creating a new algebra-based embedding, or extending BaseKGE with a new architecture.

Add a New KGE Model to dicee

When to Use

• Implementing a new KGE model (bilinear, translational, convolutional, algebra-based)
• Adding a custom scoring function to the framework
• Extending an existing model with new components

Reference Implementations

• Simple bilinear model: dicee/models/real.py → DistMult, TransE
• Clifford algebra model: dicee/models/clifford.py → Keci
• Convolutional model: dicee/models/quaternion.py → ConvQ
• Transformer model: dicee/models/transformers.py → BytE, or real.py → CoKE
• BaseKGE contract: dicee/models/base_model.py

---

Step-by-Step Procedure

Step 1 — Create the model file

Create dicee/models/mymodel.py. Every model must extend BaseKGE and implement two forward methods:

from .base_model import BaseKGE
import torch

class MyModel(BaseKGE):
    def __init__(self, args: dict):
        super().__init__(args)
        self.name = 'MyModel'
        # Access config params via self.args dict:
        #   self.embedding_dim  — already set by BaseKGE
        #   self.num_entities   — already set by BaseKGE
        #   self.num_relations  — already set by BaseKGE
        #   self.args.get("p", 0)  — Clifford p param
        #   self.args.get("q", 1)  — Clifford q param
        # Define additional parameters:
        self.W = torch.nn.Parameter(torch.empty(self.embedding_dim, self.embedding_dim))
        torch.nn.init.xavier_normal_(self.W)

    def forward_triples(self, x: torch.LongTensor) -> torch.FloatTensor:
        """
        Score individual triples.
        x: (batch_size, 3) — columns are [head_idx, relation_idx, tail_idx]
        Returns: (batch_size,) — one score per triple
        """
        head_ent_emb = self.entity_embeddings(x[:, 0])      # (B, D)
        rel_emb = self.relation_embeddings(x[:, 1])          # (B, D)
        tail_ent_emb = self.entity_embeddings(x[:, 2])       # (B, D)
        # Implement scoring: e.g., for DistMult:
        return (head_ent_emb * rel_emb * tail_ent_emb).sum(dim=1)

    def forward_k_vs_all(self, x: torch.LongTensor) -> torch.FloatTensor:
        """
        Score (head, relation) pairs against ALL entities.
        x: (batch_size, 2) — columns are [head_idx, relation_idx]
        Returns: (batch_size, num_entities) — score matrix
        """
        head_ent_emb = self.entity_embeddings(x[:, 0])      # (B, D)
        rel_emb = self.relation_embeddings(x[:, 1])          # (B, D)
        # score against all entity embeddings:
        return (head_ent_emb * rel_emb) @ self.entity_embeddings.weight.T

Important constraints:

• args is a plain dict — use self.args.get("key", default) for custom params
• BaseKGE.__init__ initialises self.entity_embeddings, self.relation_embeddings, the loss function, dropout layers, and normalisation — do not redefine them
• For Clifford models: embedding_dim / (p + q + 1) must be an integer; assert this in __init__
• Use self.apply_layer_norm, self.apply_batch_norm, self.input_dp, self.hidden_dp from BaseKGE for consistency

Step 2 — Register the model

Edit dicee/models/__init__.py to export the new model:

# Add to existing imports:
from .mymodel import MyModel  # noqa

Note: this file uses from .real import * style star-imports for most modules, plus explicit imports for clifford.py and dualE.py. transformers.py, literal.py, ensemble.py, and fsdp_models.py are deliberately NOT re-exported here — follow the existing pattern for the file you're adding to.

The model is then selectable via --model MyModel on the CLI or args.model = "MyModel" in Python.

Also add MyModel to the --model list in README.md (search for the line starting with * ```--model under "Knowledge Graph Embedding Models").

Step 3 — (Optional) Add custom config params

If the model needs new hyperparameters, add them to dicee/config.py inside the Namespace.__init__ method:

self.my_param: float = 0.5
"""Description of my_param"""

Access them in the model as self.args.get("my_param", 0.5).

Step 4 — Write tests

Add a test in tests/ (or extend an existing test file). Minimum test:

from dicee.executer import Execute
from dicee.config import Namespace

def test_my_model():
    args = Namespace()
    args.model = 'MyModel'
    args.dataset_dir = 'KGs/UMLS'
    args.scoring_technique = 'KvsAll'
    args.num_epochs = 1
    args.embedding_dim = 32
    args.eval_model = 'train_val_test'
    result = Execute(args).start()
    assert result['Train']['MRR'] > 0

---

Common Pitfalls

Mistake	Fix
Redefining entity_embeddings in __init__	Don't — BaseKGE creates it
Using args.embedding_dim (namespace-style)	Use self.embedding_dim (already parsed)
Clifford r not integer	Assert (embedding_dim / (p+q+1)).is_integer() early
Missing forward_k_vs_all	Required for KvsAll / 1vsAll / AllvsAll / KvsSample scoring
Missing forward_triples	Required for NegSample / FixedNegSample scoring
Forgetting # noqa in __init__.py import	Linter removes "unused" imports

---

BaseKGE Attributes Available to All Models

self.embedding_dim     # int — embedding vector size
self.num_entities      # int — vocabulary size (entities)
self.num_relations     # int — vocabulary size (relations)
self.entity_embeddings # nn.Embedding(num_entities, embedding_dim)
self.relation_embeddings # nn.Embedding(num_relations, embedding_dim)
self.input_dp          # nn.Dropout(input_dropout_rate)
self.hidden_dp         # nn.Dropout(hidden_dropout_rate)
self.feature_map_dp    # nn.Dropout(feature_map_dropout_rate)
self.loss              # loss function (BCEWithLogitsLoss or CrossEntropyLoss)
self.loss_history      # list of per-epoch losses
self.args              # dict of full config

---

Quick Reference: Existing Models by Complexity

Complexity	Model	File	Key idea
Simplest	DistMult	real.py	Hadamard product h⊙r⊙t
Translational	TransE	real.py	Distance ‖h+r−t‖
Complex	ComplEx	complex.py	Complex vectors
Quaternion	QMult	quaternion.py	Hamilton product
Clifford	Keci	clifford.py	Cl(p,q) algebra
Convolutional	ConvQ	quaternion.py	Conv over quaternion
Transformer	CoKE	real.py	Attention over triples
Function space	LFMult	function_space.py	Learned function embeddings