Algorithms

base

Abstract interface implemented by every linking algorithm.

BaseLinker

Bases: ABC

Common interface for Entity Alignment, Entity Linking, WSD and WSA.

Subclasses implement a single task by consuming two entity datasets (whose shape depends on the task), an optional supporting graph, and a blocking strategy used to restrict the candidate search space.

link abstractmethod

link(
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    graph: Graph = None,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> list[AlignmentResult]

Link entities in dataset1 to entities in dataset2.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities (e.g. mentions, source KG entities).	required
dataset2	list[Entity] | EntitySource	Target entities (e.g. KB entries, target KG entities), or an EntitySource wrapping a target too large to materialize as a list[Entity]. Support depends on blocking -- see candidate_pairs.	required
graph	Graph	Optional supporting graph as a triple list, a networkx.Graph, or an rdflib.Graph.	None
blocking	BlockingStrategy	Strategy used to generate candidate pairs before scoring.	DEFAULT_BLOCKING

Returns:

Type	Description
list[AlignmentResult]	A list of predicted links.

Source code in src/linkingtk/algorithms/base.py

@abstractmethod
def link(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    graph: Graph = None,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> list[AlignmentResult]:
    """Link entities in ``dataset1`` to entities in ``dataset2``.

    Args:
        dataset1: Source entities (e.g. mentions, source KG entities).
        dataset2: Target entities (e.g. KB entries, target KG entities),
            or an [EntitySource][linkingtk.core.source.EntitySource]
            wrapping a target too large to materialize as a
            ``list[Entity]``. Support depends on ``blocking`` -- see
            [candidate_pairs][linkingtk.blocking.base.BlockingStrategy.candidate_pairs].
        graph: Optional supporting graph as a triple list, a
            ``networkx.Graph``, or an ``rdflib.Graph``.
        blocking: Strategy used to generate candidate pairs before
            scoring.

    Returns:
        A list of predicted links.
    """

String-similarity baseline

string_similarity

General string-similarity baseline linker.

Scores each blocked candidate pair by the similarity of a configurable text field from either side (e.g. a mention's context against a candidate's gloss, or two entities' labels), then ranks candidates per source entity. This is a general-purpose baseline for Entity Alignment, Entity Linking and Word Sense Alignment; the classic Lesk algorithm for WSD (linkingtk.algorithms.wsd.lesk.LeskLinker) is implemented as a preconfigured instance of this class.

StringSimilarityLinker

StringSimilarityLinker(
    source_field: Field = "label",
    target_field: Field = "label",
    metric: Metric = "jaccard",
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

Baseline linker that scores candidates by string similarity.

For each blocked (entity1, entity2) pair, extracts source_field text from entity1 and target_field text from entity2, and scores the pair with metric (higher is always better, regardless of metric). For each source entity, the highest-scoring candidate is returned as the predicted link, with the rest kept as ranked alternatives.

Parameters:

Name	Type	Description	Default
source_field	Field	Which text to extract from each dataset1 entity: "label" (all labels, space-joined), "description", or "context". A callable taking an Entity and returning str may be passed instead, for fields not covered above.	'label'
target_field	Field	Which text to extract from each dataset2 entity; same options as source_field.	'label'
metric	Metric	"word_overlap" (raw shared-token count), "jaccard" (normalized token overlap), or "levenshtein" (normalized edit-distance similarity). A callable taking two strings and returning a similarity score (higher is better) may be passed instead, for metrics not covered above.	'jaccard'
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher (each source entity's highest-scoring candidate, independently — multiple sources may map to the same target). Pass OptimalMatcher for a globally optimal one-to-one assignment instead.	DEFAULT_MATCHER

Source code in src/linkingtk/algorithms/string_similarity.py

def __init__(
    self,
    source_field: Field = "label",
    target_field: Field = "label",
    metric: Metric = "jaccard",
    matching: Matcher = DEFAULT_MATCHER,
) -> None:
    self.source_field = source_field
    self.target_field = target_field
    self.metric = metric
    self.matching = matching

jaccard

jaccard(text1: str, text2: str) -> float

Token-set Jaccard similarity, normalized to [0, 1].

Source code in src/linkingtk/algorithms/string_similarity.py

def jaccard(text1: str, text2: str) -> float:
    """Token-set Jaccard similarity, normalized to ``[0, 1]``."""
    tokens1, tokens2 = _tokenize(text1), _tokenize(text2)
    union = tokens1 | tokens2
    return len(tokens1 & tokens2) / len(union) if union else 0.0

levenshtein_similarity

levenshtein_similarity(text1: str, text2: str) -> float

Character-level similarity derived from normalized edit distance, in [0, 1].

Source code in src/linkingtk/algorithms/string_similarity.py

def levenshtein_similarity(text1: str, text2: str) -> float:
    """Character-level similarity derived from normalized edit distance, in ``[0, 1]``."""
    text1, text2 = text1.lower(), text2.lower()
    denom = len(text1) + len(text2)
    if denom == 0:
        return 1.0
    return 1.0 - edit_distance(text1, text2) / denom

word_overlap

word_overlap(text1: str, text2: str) -> float

Count of shared tokens, as used by the classic Lesk algorithm.

Source code in src/linkingtk/algorithms/string_similarity.py

def word_overlap(text1: str, text2: str) -> float:
    """Count of shared tokens, as used by the classic Lesk algorithm."""
    return float(len(_tokenize(text1) & _tokenize(text2)))

Feature-based classifier

feature_classifier

Feature-based classifier linker for Entity Alignment and Word Sense Alignment.

Scores each blocked candidate pair with a classical ML classifier trained on hand-crafted similarity features, rather than a single fixed metric — in the style of EntMatcher (https://github.com/DexterZeng/EntMatcher). EntMatcher itself is a research repo, not a dependency of this project; the idea reused here is that a globally optimal one-to-one assignment (OptimalMatcher) can outperform independent per-source argmax matching, regardless of what produces the underlying similarity score.

Classifier

Bases: Protocol

Structural type for a fit/predict_proba classifier (e.g. scikit-learn's).

FeatureClassifierLinker

FeatureClassifierLinker(
    features: list[tuple[str, FeatureFn]] | None = None,
    vectorizer: Vectorizer | None = None,
    classifier: Classifier | None = None,
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

Scores candidate pairs with a classical ML classifier over hand-crafted features.

Unlike StringSimilarityLinker (a single fixed metric), this combines several similarity signals — DEFAULT_FEATURES (string-overlap/edit-distance metrics on the label and description fields, reused from linkingtk.algorithms.string_similarity) plus a TF-IDF cosine similarity feature over combined label+description text — into a feature vector scored by a trainable classifier. Must be fit before link can be called.

Parameters:

Name	Type	Description	Default
features	list[tuple[str, FeatureFn]] | None	(name, fn) pairs, each fn scoring one (entity1, entity2) pair. Defaults to DEFAULT_FEATURES.	None
vectorizer	Vectorizer | None	Any object exposing scikit-learn's fit_transform/transform interface, used for the TF-IDF cosine similarity feature. Defaults to TfidfVectorizer(stop_words="english"). Fit once in fit, over both datasets' combined label+description text, and reused via transform in link — unlike EmbeddingSimilarityBlocker, which has no explicit fit step and refits per call.	None
classifier	Classifier | None	Any object exposing scikit-learn's fit/ predict_proba interface, expected to return probabilities in [0, 1] for its positive class. Defaults to StandardScaler + LogisticRegression in a sklearn.pipeline.make_pipeline.	None
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher (each source entity's highest-scoring candidate, independently, like StringSimilarityLinker — multiple source entities may map to the same target). Pass OptimalMatcher for a globally optimal one-to-one assignment instead, which can outperform greedy matching when two source entities' individually-best candidate is the same target. See EntMatcherLinker for a preconfigured instance using it.	DEFAULT_MATCHER

Source code in src/linkingtk/algorithms/feature_classifier.py

def __init__(
    self,
    features: list[tuple[str, FeatureFn]] | None = None,
    vectorizer: Vectorizer | None = None,
    classifier: Classifier | None = None,
    matching: Matcher = DEFAULT_MATCHER,
) -> None:
    self.features = features if features is not None else DEFAULT_FEATURES
    self.vectorizer: Vectorizer = (
        vectorizer if vectorizer is not None else TfidfVectorizer(stop_words="english")
    )
    self.classifier: Classifier = (
        classifier
        if classifier is not None
        else make_pipeline(StandardScaler(), LogisticRegression())
    )
    self.matching = matching
    self._tfidf_fitted = False
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
    negatives: list[tuple[Entity, Entity]] | None = None,
    negative_ratio: int = 1,
    random_state: int | None = None,
) -> FeatureClassifierLinker

Train the classifier on candidate pairs from blocking.

Positive examples are the candidate pairs found in ground_truth. Negative examples default to a simple, self-contained strategy: the remaining candidate pairs, sampled uniformly at random (up to negative_ratio times the number of positives). Pass negatives explicitly — e.g. from sample_hard_negatives — for more deliberate negative mining instead; when given, negative_ratio and random_state are ignored.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities.	required
dataset2	list[Entity]	Target entities.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) true pairs.	required
blocking	BlockingStrategy	Strategy used to generate training candidate pairs.	DEFAULT_BLOCKING
negatives	list[tuple[Entity, Entity]] | None	Explicit negative (entity1, entity2) pairs to train on, overriding the default random-sampling strategy.	None
negative_ratio	int	Maximum number of negative examples sampled per positive example. Ignored if negatives is given.	1
random_state	int | None	Seed for negative-example sampling. Ignored if negatives is given.	None

Returns:

Type	Description
FeatureClassifierLinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If no ground_truth pair survives blocking (no positive examples to learn from), or if negatives is empty (given explicitly, or after sampling) — a classifier can't train on a single class.

Source code in src/linkingtk/algorithms/feature_classifier.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
    negatives: list[tuple[Entity, Entity]] | None = None,
    negative_ratio: int = 1,
    random_state: int | None = None,
) -> FeatureClassifierLinker:
    """Train the classifier on candidate pairs from ``blocking``.

    Positive examples are the candidate pairs found in
    ``ground_truth``. Negative examples default to a simple,
    self-contained strategy: the remaining candidate pairs, sampled
    uniformly at random (up to ``negative_ratio`` times the number of
    positives). Pass ``negatives`` explicitly — e.g. from
    [sample_hard_negatives][linkingtk.blocking.negative_sampling.sample_hard_negatives]
    — for more deliberate negative mining instead; when given,
    ``negative_ratio`` and ``random_state`` are ignored.

    Args:
        dataset1: Source entities.
        dataset2: Target entities.
        ground_truth: List of ``(source_id, target_id)`` true pairs.
        blocking: Strategy used to generate training candidate pairs.
        negatives: Explicit negative ``(entity1, entity2)`` pairs to
            train on, overriding the default random-sampling strategy.
        negative_ratio: Maximum number of negative examples sampled
            per positive example. Ignored if ``negatives`` is given.
        random_state: Seed for negative-example sampling. Ignored if
            ``negatives`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If no ``ground_truth`` pair survives
            ``blocking`` (no positive examples to learn from), or if
            ``negatives`` is empty (given explicitly, or after
            sampling) — a classifier can't train on a single class.
    """
    candidates = list(blocking.candidate_pairs(dataset1, dataset2))
    ground_truth_set = set(ground_truth)
    positives = [(e1, e2) for e1, e2 in candidates if (e1.id, e2.id) in ground_truth_set]
    if not positives:
        raise LinkingTKError(
            "No ground-truth pairs survived blocking; fit() has nothing to learn "
            "from. Check that `blocking` actually finds the ground-truth pairs "
            "(see Evaluator.evaluate_blocking) before training a classifier on top "
            "of it."
        )

    if negatives is None:
        negatives = self._sample_random_negatives(
            candidates, ground_truth_set, len(positives), negative_ratio, random_state
        )
    elif not negatives:
        raise LinkingTKError(
            "`negatives` is empty; fit() cannot train a classifier on a single class."
        )

    training_pairs = positives + negatives
    labels = [1] * len(positives) + [0] * len(negatives)

    tfidf1, tfidf2 = self._fit_tfidf(dataset1, dataset2)
    X = self._feature_matrix(training_pairs, tfidf1, tfidf2)
    self.classifier.fit(X, labels)
    self._fitted = True
    return self

LLM-based linker

llm

General-purpose linker that prompts an LLM to decide/score links.

Usable across Entity Alignment, Entity Linking, Word Sense Disambiguation and Word Sense Alignment -- selected via task, the same way StringSimilarityLinker is one general-purpose class rather than four task-specific subclasses. Talks to an LLM through LlmClient (#60), so this module never imports openai/anthropic itself.

LlmBaseLinker

LlmBaseLinker(
    client: LlmClient,
    task: Task = "el",
    matching: Matcher = DEFAULT_MATCHER,
    max_tokens: int = 2048,
)

Bases: BaseLinker

Prompts an LLM to score/rank each source entity's blocked candidates.

Unlike StringSimilarityLinker, which scores each candidate pair independently, this linker makes one LLM call per source entity, showing all of its blocked candidates at once via LlmClient.complete_structured -- cheaper than one call per pair, and gives the model real comparative context instead of siloed yes/no judgments.

Parameters:

Name	Type	Description	Default
client	LlmClient	An already-constructed LlmClient (e.g. from create_client). This class never imports openai/anthropic itself.	required
task	Task	Which PromptTemplate instruction to use -- "ea", "el", "wsd", or "wsa".	'el'
matching	Matcher	Strategy used to resolve scored candidates into final links, same as StringSimilarityLinker's matching argument.	DEFAULT_MATCHER
max_tokens	int	Forwarded to client.complete_structured per call. Default 2048, not complete_structured's own 1024 default -- measured directly against a real WSD benchmark (examples/llm_benchmark.py) with wide (top_k=50) blocking: a source entity with many real candidates can need more than 1024 tokens just to rank all of them, and 1024 produced truncated, unparseable JSON responses for exactly those higher-candidate-count entities (a systematic gap, not isolated noise). Widen further if using an even larger candidate width than 50.	2048

Note

graph is accepted for interface compliance but not used, same as StringSimilarityLinker. LLM calls are made sequentially, one per source entity with at least one blocked candidate.

Source code in src/linkingtk/algorithms/llm.py

def __init__(
    self,
    client: LlmClient,
    task: Task = "el",
    matching: Matcher = DEFAULT_MATCHER,
    max_tokens: int = 2048,
) -> None:
    self.client = client
    self.task = task
    self.matching = matching
    self.max_tokens = max_tokens

PromptTemplate dataclass

PromptTemplate(instruction: str)

A task's instruction to the LLM, prepended to a source entity + its candidates.

Attributes:

Name	Type	Description
instruction	str	The system-message text describing the task and the expected judgment -- the only part that varies per task. Entity rendering itself (_format_entity) is shared across every task, since the per-task Entity-shape differences (labels/description/context populated differently per task) already fall out of "render whichever fields are non-empty" without per-task branching.

_llm_prompting

Shared LLM-prompting building blocks for LlmBaseLinker and LlmRerankerLinker.

Extracted from llm.py (#21) so #23's reranker reuses the exact same per-task instructions, entity rendering, structured-output schema, and hallucinated-id recovery, rather than a parallel copy.

PromptTemplate dataclass

PromptTemplate(instruction: str)

A task's instruction to the LLM, prepended to a source entity + its candidates.

Attributes:

Name	Type	Description
instruction	str	The system-message text describing the task and the expected judgment -- the only part that varies per task. Entity rendering itself (_format_entity) is shared across every task, since the per-task Entity-shape differences (labels/description/context populated differently per task) already fall out of "render whichever fields are non-empty" without per-task branching.

LLM re-ranker

llm_reranker

LLM re-ranking layer over an already-scored bi-encoder/cross-encoder candidate set.

Filed as #23: a cheap, already-trained ReFinEDLinker/ BlinkLinker/ GlossBertLinker/ EwiserLinker/ EscLinker narrows blocked candidates to a scored top-k first; only that narrow window is then shown to an LLM for re-ranking, unlike LlmBaseLinker (#21), which sends an LLM every blocked candidate with no first-stage narrowing. No specific paper is being ported here (unlike ChatEALinker, #22) -- generalizes the same "retrieve cheap, then let an LLM re-rank only the top-k" idea to Entity Linking and Word Sense Disambiguation.

CandidateScorer

Bases: Protocol

What LlmRerankerLinker needs from its base_linker.

ReFinEDLinker, BlinkLinker, GlossBertLinker, EwiserLinker, and EscLinker all implement this -- their own link() calls it directly too, so a reranked source entity's top-k is always built from the same scores link() would otherwise resolve straight into matching.match().

score_candidates

score_candidates(
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]

Blocked candidates per source entity id: (target_id, score), unsorted.

Source code in src/linkingtk/algorithms/llm_reranker.py

def score_candidates(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]:
    """Blocked candidates per source entity id: `(target_id, score)`, unsorted."""

LlmRerankerLinker

LlmRerankerLinker(
    base_linker: CandidateScorer,
    client: LlmClient,
    task: Task = "el",
    top_k: int = 10,
    threshold: float | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    max_tokens: int = 2048,
)

Bases: BaseLinker

Re-ranks a base_linker's top-top_k blocked candidates with an LLM.

No fit() -- like LlmBaseLinker, the LLM needs no training; only base_linker's own model does, via whatever training path it normally uses (e.g. Trainer, CrossEncoderTrainer, or a loaded checkpoint), done by the caller before constructing this.

Parameters:

Name	Type	Description	Default
base_linker	CandidateScorer	An already-fitted/loaded CandidateScorer -- any of ReFinEDLinker/BlinkLinker/GlossBertLinker/ EwiserLinker/EscLinker.	required
client	LlmClient	An already-constructed LlmClient. This class never imports openai/anthropic itself.	required
task	Task	Which PromptTemplate instruction to use -- typically "el" or "wsd", matching base_linker's own task.	'el'
top_k	int	How many of base_linker's own highest-scored candidates per source entity to actually show the LLM. The LLM can only ever re-order within this window, never reach a candidate ranked below it -- the actual cost/quality trade-off this class exists to make, unlike LlmBaseLinker's unbounded candidate count.	10
threshold	float | None	Confidence shortcut -- if base_linker's own top two scores for a source entity differ by more than this, skip the LLM call entirely and keep base_linker's own ranking (a simplified analog of ChatEALinker's same idea, without its iterative window-widening -- #23 isn't tied to a specific paper's algorithm). Defaults to None (disabled): unlike ChatEALinker's own normalized [0, 1] multi-dimension score, base_linker's score scale is arbitrary and method-specific -- a bi-encoder's cosine similarity lands roughly in [-1, 1] (a fixed threshold like 0.5 is meaningful there), but a cross-encoder's raw logit margin (e.g. GlossBertLinker/EwiserLinker/EscLinker) is unbounded and routinely exceeds any such threshold by 10x or more -- measured directly during #23's own WSD benchmark: a threshold=0.5 default silently skipped the LLM for 28/30 sampled entities, masking any real re-ranking effect entirely. Only pass a threshold when you know base_linker's own score scale well enough to pick one.	None
matching	Matcher	Strategy used to resolve scored candidates into final links.	DEFAULT_MATCHER
max_tokens	int	Forwarded to client.complete_structured per call.	2048

Note

graph is accepted for interface compliance but not used, same as LlmBaseLinker. blocking is applied twice per link() call -- once inside base_linker.score_candidates(), once again here to recover real Entity objects for prompt rendering (scores alone aren't enough to build a prompt). BlockingStrategys are cheap index lookups next to a transformer forward pass, so this isn't worth threading Entity objects through CandidateScorer's return type to avoid.

Source code in src/linkingtk/algorithms/llm_reranker.py

def __init__(
    self,
    base_linker: CandidateScorer,
    client: LlmClient,
    task: Task = "el",
    top_k: int = 10,
    threshold: float | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    max_tokens: int = 2048,
) -> None:
    self.base_linker = base_linker
    self.client = client
    self.task = task
    self.top_k = top_k
    self.threshold = threshold
    self.matching = matching
    self.max_tokens = max_tokens

Entity Alignment

ea

Entity Alignment (EA) algorithms.

EA finds matches between two knowledge graphs, often across languages. Entities typically carry labels and descriptions but rarely a context, and a supporting graph is normally available. Related systems: OpenEA, Simple-HHEA, ChatEA, ProLEA, EntMatcher.

AliNetLinker

AliNetLinker(
    layer_dims: list[int] | None = None,
    num_epochs: int = 500,
    learning_rate: float = 0.001,
    batch_size: int = 3000,
    neg_triple_num: int = 10,
    neg_margin: float = 1.5,
    neg_margin_balance: float = 0.1,
    sim_th: float = 0.0,
    bootstrap_start_epoch: int = 20,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via AliNet's gated multi-hop GCN/attention embeddings.

Must be fit before link can be called. See the module docstring for what's ported vs. deliberately deviated from OpenEA.

Parameters:

Name	Type	Description	Default
layer_dims	list[int] | None	Layer width sequence -- layer_dims[0] is the initial embedding's dimensionality, each subsequent entry one more GCN+attention+highway layer's output width. OpenEA's published EN-FR-15K-V1 config uses [500, 400, 300].	None
num_epochs	int	Training epochs. OpenEA's config allows up to 2000 with early stopping.	500
learning_rate	float	Adam's learning rate. OpenEA's published value is 0.001.	0.001
batch_size	int	Positive pairs sampled per training step (with replacement, matching OpenEA). OpenEA's published value is 3000.	3000
neg_triple_num	int	Negatives drawn per positive pair (k, both sides independently). OpenEA's published value is 10.	10
neg_margin	float	Margin for margin_ranking_loss_l2_squared. OpenEA's published value is 1.5.	1.5
neg_margin_balance	float	Weight applied to the negative-pair loss term. OpenEA's published value is 0.1.	0.1
sim_th	float	Minimum similarity for a bootstrapped pair to be accepted -- also the bootstrapping on/off switch: OpenEA's own run() only calls its bootstrapping step if sim_th > 0.0, and its own published config uses 0.0 (off). See the module docstring.	0.0
bootstrap_start_epoch	int	First epoch a bootstrapping round may run, if sim_th > 0. OpenEA's published value (start_augment * eval_freq) is 2 * 10 = 20.	20
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/alinet.py

def __init__(
    self,
    layer_dims: list[int] | None = None,
    num_epochs: int = 500,
    learning_rate: float = 0.001,
    batch_size: int = 3000,
    neg_triple_num: int = 10,
    neg_margin: float = 1.5,
    neg_margin_balance: float = 0.1,
    sim_th: float = 0.0,
    bootstrap_start_epoch: int = 20,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.layer_dims = layer_dims if layer_dims is not None else [500, 400, 300]
    self.num_epochs = num_epochs
    self.learning_rate = learning_rate
    self.batch_size = batch_size
    self.neg_triple_num = neg_triple_num
    self.neg_margin = neg_margin
    self.neg_margin_balance = neg_margin_balance
    self.sim_th = sim_th
    self.bootstrap_start_epoch = bootstrap_start_epoch
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> AliNetLinker

Train the gated multi-hop GCN/attention stack, optionally bootstrapping new pairs.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (same approach IPTransELinker uses), needed for enhance_triples/ generate_2hop_pairs.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs used as positive pairs in the margin loss, and as the initial bootstrapping seed set.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. Also the cadence bootstrapping rounds run on, when enabled. If None (default), trains the full num_epochs unconditionally (bootstrapping still runs on the eval_every cadence if enabled).	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth and (if enabled) run a bootstrapping round.	10

Returns:

Type	Description
AliNetLinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/alinet.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> AliNetLinker:
    """Train the gated multi-hop GCN/attention stack, optionally bootstrapping new pairs.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (same approach
            [IPTransELinker][linkingtk.algorithms.ea.iptranse.IPTransELinker]
            uses), needed for
            [enhance_triples][linkingtk.algorithms.ea._alinet_training.enhance_triples]/
            [generate_2hop_pairs][linkingtk.algorithms.ea._alinet_training.generate_2hop_pairs].
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs used as positive pairs in the margin loss, and as
            the initial bootstrapping seed set.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must already
            be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left unspecified,
            training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. Also the cadence bootstrapping
            rounds run on, when enabled. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally (bootstrapping
            still runs on the ``eval_every`` cadence if enabled).
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``
            and (if enabled) run a bootstrapping round.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples, or if ``device`` is
            invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
    except ImportError as exc:
        raise OptionalDependencyError("AliNetLinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped = map_triples_to_ids(triples, entity_to_id, relation_to_id)
    num_entities = len(entity_to_id)

    triples1 = [t for t in triples if t[0] in ids1]
    triples2 = [t for t in triples if t[0] in ids2]
    mapped1 = map_triples_to_ids(triples1, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2, entity_to_id, relation_to_id)

    seed_pairs = [
        (entity_to_id[s], entity_to_id[t])
        for s, t in ground_truth
        if s in entity_to_id and t in entity_to_id
    ]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train AliNet's margin loss with."
        )

    entities1_ids = [entity_to_id[eid] for eid in ids1 if eid in entity_to_id]
    entities2_ids = [entity_to_id[eid] for eid in ids2 if eid in entity_to_id]

    two_hop_pairs = generate_2hop_pairs(mapped1) | generate_2hop_pairs(mapped2)
    two_indices, two_values = pairs_to_symmetric_adjacency(two_hop_pairs, num_entities)
    two_indices, two_values = normalize_adjacency_coo(two_indices, two_values, num_entities)
    two_hop_adjacency = coo_to_torch_sparse(
        two_indices, two_values, (num_entities, num_entities), device
    )

    def _build_one_hop_adjacency(pairs: list[tuple[int, int]]) -> torch.Tensor:
        enhanced1, enhanced2 = enhance_triples(mapped1, mapped2, pairs)
        all_triples = np.array(
            mapped.tolist() + list(enhanced1) + list(enhanced2), dtype=np.int64
        )
        indices, values = build_primal_adjacency(all_triples, num_entities)
        indices, values = normalize_adjacency_coo(indices, values, num_entities)
        return coo_to_torch_sparse(indices, values, (num_entities, num_entities), device)

    model = build_alinet_model(num_entities, self.layer_dims).to(device)
    _set_adjacency(model, _build_one_hop_adjacency(seed_pairs), two_hop_adjacency)
    optimizer = torch.optim.Adam(model.parameters(), lr=self.learning_rate)

    aligned1 = {source for source, _target in seed_pairs}
    aligned2 = {target for _source, target in seed_pairs}
    unaligned1 = [entity for entity in entities1_ids if entity not in aligned1]
    unaligned2 = [entity for entity in entities2_ids if entity not in aligned2]

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        model.train()
        batch_size = min(self.batch_size, len(seed_pairs))
        batch_idx = rng.integers(0, len(seed_pairs), size=batch_size)
        pos_left = torch.tensor([seed_pairs[i][0] for i in batch_idx], device=device)
        pos_right = torch.tensor([seed_pairs[i][1] for i in batch_idx], device=device)

        neg_left_np, neg_right_np = sample_uniform_cross_kg_negatives(
            entities1_ids,
            entities2_ids,
            batch_size,
            self.neg_triple_num,
            rng,
            exclude_pairs=set(seed_pairs),
        )
        neg_left = torch.from_numpy(neg_left_np).long().to(device)
        neg_right = torch.from_numpy(neg_right_np).long().to(device)

        embeddings = model()
        loss = margin_ranking_loss_l2_squared(
            embeddings,
            pos_left,
            pos_right,
            neg_left,
            neg_right,
            self.neg_margin,
            self.neg_margin_balance,
        )
        optimizer.zero_grad()
        loss.backward()  # type: ignore[no-untyped-call]
        optimizer.step()

        if (epoch + 1) % eval_every == 0:
            model.eval()
            with torch.no_grad():
                current_embeds = model().cpu().numpy()
            model.train()

            if val_pairs:
                hits1 = _validation_hits1(current_embeds, entity_to_id, val_pairs)
                if hits1 <= best_hits1:
                    epochs_without_improvement += 1
                    if epochs_without_improvement >= patience:
                        break
                else:
                    best_hits1 = hits1
                    epochs_without_improvement = 0

            if self.sim_th > 0.0 and epoch + 1 >= self.bootstrap_start_epoch:
                new_seed_pairs, unaligned1, unaligned2 = run_bootstrapping_round(
                    current_embeds, seed_pairs, unaligned1, unaligned2, self.sim_th
                )
                if len(new_seed_pairs) > len(seed_pairs):
                    seed_pairs = new_seed_pairs
                    _set_adjacency(
                        model, _build_one_hop_adjacency(seed_pairs), two_hop_adjacency
                    )

    model.eval()
    with torch.no_grad():
        final_embeds = model().cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/alinet.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/alinet.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

AttrELinker

AttrELinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    literal_len: int = 5,
    char_freq_threshold: float = 0.0001,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a shared TransE space regularized by character embeddings.

Must be fit before link can be called. See the module docstring for how this differs from IPTransELinker and JAPELinker.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of every trained embedding table (entities, relations, character-space entities, attributes, characters). OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each one pass of structural training, one pass of attribute/character training, and one (repeated) pass of joint-loss training. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for both the structural and attribute-triple training (each derives its own step count from its own triple count), and the divisor for the joint loss's repeat count.	5000
learning_rate	float	SGD's learning rate, for every optimizer here. OpenEA's published value is 0.01.	0.01
margin	float	Hinge margin for both the structural and attribute-triple losses. OpenEA's published value is 1.5.	1.5
literal_len	int	Fixed character-sequence length each attribute value is truncated/padded to. OpenEA's published value is 5.	5
char_freq_threshold	float	Minimum share of total (distinct-value) character occurrences required to keep a character in the vocabulary; rarer characters map to the padding id. OpenEA's published value is 0.0001.	0.0001
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/attre.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    literal_len: int = 5,
    char_freq_threshold: float = 0.0001,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.margin = margin
    self.literal_len = literal_len
    self.char_freq_threshold = char_freq_threshold
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> AttrELinker

Train a shared TransE space jointly with a character/attribute embedding space.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (same approach IPTransELinker uses).	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs. Seeds the shared embedding tables (each pair's target is aliased to its source's id, in both the SE and CE spaces).	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint.	required
attribute_triples1	list[Triple]	KG1's (entity_id, predicate, value) attribute triples. Required -- see the module docstring's deviation note.	required
attribute_triples2	list[Triple]	KG2's attribute triples. Required.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
AttrELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If attribute_triples1 and attribute_triples2 are both empty, or if none of ground_truth's pairs have both ids present in dataset1/dataset2, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/attre.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> AttrELinker:
    """Train a shared TransE space jointly with a character/attribute embedding space.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (same approach
            [IPTransELinker][linkingtk.algorithms.ea.iptranse.IPTransELinker]
            uses).
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs. Seeds the shared embedding tables (each pair's
            target is aliased to its source's id, in both the SE and
            CE spaces).
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples. **Required** -- see the module
            docstring's deviation note.
        attribute_triples2: KG2's attribute triples. **Required**.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If ``attribute_triples1`` and
            ``attribute_triples2`` are both empty, or if none of
            ``ground_truth``'s pairs have both ids present in
            ``dataset1``/``dataset2``, or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    if not attribute_triples1 and not attribute_triples2:
        raise LinkingTKError(
            "AttrELinker.fit() needs attribute triples -- `attribute_triples1` "
            "and `attribute_triples2` are both empty, leaving nothing to train "
            "the character/attribute embedding half with."
        )

    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("AttrELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    seed_pairs = [(s, t) for s, t in ground_truth if s in ids1 and t in ids2]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in "
            "`dataset1`/`dataset2`; fit() has no seed pairs to train "
            "AttrE's shared embedding tables with."
        )

    entity_to_id, relation_to_id = build_shared_id_mappings(
        triples1_labels + triples2_labels, seed_pairs
    )
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)
    se_ctx1 = build_kg_context(mapped1)
    se_ctx2 = build_kg_context(mapped2)

    (
        ce_ctx1,
        ce_ctx2,
        attribute_to_id,
        char_to_id,
        value_char_ids_tensor,
    ) = self._build_attribute_context(
        torch, entity_to_id, attribute_triples1, attribute_triples2, device
    )

    entity_embeds = self._init_embedding(torch, functional, len(entity_to_id), device)
    relation_embeds = self._init_embedding(torch, functional, len(relation_to_id), device)
    entity_embeds_ce = self._init_embedding(torch, functional, len(entity_to_id), device)
    attr_embeds = self._init_embedding(torch, functional, max(1, len(attribute_to_id)), device)
    char_embeds = self._init_embedding(torch, functional, len(char_to_id) + 1, device)

    se_optimizer = torch.optim.SGD([entity_embeds, relation_embeds], lr=self.learning_rate)
    ce_optimizer = torch.optim.SGD(
        [entity_embeds_ce, attr_embeds, char_embeds], lr=self.learning_rate
    )
    joint_optimizer = torch.optim.SGD([entity_embeds, entity_embeds_ce], lr=self.learning_rate)

    joint_entity_ids = np.array(sorted(entity_to_id.values()), dtype=np.int64)
    joint_steps = max(1, math.ceil(len(joint_entity_ids) / self.batch_size))

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_structural_epoch(
            entity_embeds,
            relation_embeds,
            se_optimizer,
            se_ctx1,
            se_ctx2,
            rng,
            self.batch_size,
            self.margin,
        )
        train_attr_epoch(
            entity_embeds_ce,
            attr_embeds,
            char_embeds,
            value_char_ids_tensor,
            ce_optimizer,
            ce_ctx1,
            ce_ctx2,
            rng,
            self.batch_size,
            self.margin,
        )
        train_joint_epoch(
            entity_embeds, entity_embeds_ce, joint_optimizer, joint_entity_ids, joint_steps
        )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/attre.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/attre.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

BootEALinker

BootEALinker(
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    sub_epoch: int = 10,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    pos_margin: float = 0.01,
    neg_margin: float = 2.0,
    neg_margin_balance: float = 0.2,
    neg_triple_num: int = 10,
    truncated_epsilon: float = 0.9,
    sim_th: float = 0.7,
    k: int = 10,
    bootstrap_pool_size: int | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a shared TransE space, grown by editing-based bootstrapping.

Must be fit before link can be called. See the module docstring for how this differs from MTransELinker and IPTransELinker.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Total structural-training-epoch cap, split across outer bootstrapping iterations of sub_epoch epochs each. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	2000
sub_epoch	int	Epochs per outer iteration, between bootstrapping rounds. OpenEA's published value is 10.	10
batch_size	int	Mini-batch size for both the structural and alignment losses. OpenEA's published value is 5000.	5000
learning_rate	float	Adagrad's learning rate, for both optimizers. OpenEA's published value is 0.01.	0.01
pos_margin	float	Positive-triple margin in the "limited" loss -- positive scores aren't penalized until they exceed this. OpenEA's published value is 0.01.	0.01
neg_margin	float	Negative-triple margin -- negative scores aren't penalized until they drop below this. OpenEA's published value is 2.0.	2.0
neg_margin_balance	float	Weight applied to the negative-triple loss term relative to the positive term. OpenEA's published value is 0.2.	0.2
neg_triple_num	int	Negative triples sampled per positive. OpenEA's published value is 10.	10
truncated_epsilon	float	Fraction of each KG's entities excluded from an entity's truncated negative-sampling candidate set (only the top 1 - truncated_epsilon most similar entities are eligible). OpenEA's published value is 0.9.	0.9
sim_th	float	Minimum similarity for a bootstrap round's candidate pair. OpenEA's published value is 0.7.	0.7
k	int	Max candidate matches kept per reference-pool row before max-weight bipartite matching. OpenEA's published value is 10.	10
bootstrap_pool_size	int | None	Cap on each side's bootstrap reference pool size (random subsample, fixed for the whole run) -- same deviation-mitigation as IPTransELinker's own parameter: this repo's broader-than-OpenEA reference pool (see the module docstring) means find_mwgm_pairs's scipy.optimize.linear_sum_assignment step runs over a denser candidate set than OpenEA's own val+test-only pool would produce. None (default) uses the full pool.	None
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/bootea.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    sub_epoch: int = 10,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    pos_margin: float = 0.01,
    neg_margin: float = 2.0,
    neg_margin_balance: float = 0.2,
    neg_triple_num: int = 10,
    truncated_epsilon: float = 0.9,
    sim_th: float = 0.7,
    k: int = 10,
    bootstrap_pool_size: int | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.sub_epoch = sub_epoch
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.pos_margin = pos_margin
    self.neg_margin = neg_margin
    self.neg_margin_balance = neg_margin_balance
    self.neg_triple_num = neg_triple_num
    self.truncated_epsilon = truncated_epsilon
    self.sim_th = sim_th
    self.k = k
    self.bootstrap_pool_size = bootstrap_pool_size
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> BootEALinker

Train a shared TransE space with seed pseudo-triples and iterative bootstrapping.

See the module docstring for the three mechanisms that distinguish this from a plain structural baseline.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (same approach IPTransELinker uses) and as the bootstrap reference pool's source.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs. Seeds the main structural triple set (via pseudo- triples) and is excluded from the bootstrap reference pool.	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked, and training stops after patience checks with no improvement (no further bootstrapping rounds run once stopped, matching OpenEA's own early-stop-before- bootstrap ordering).	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
BootEALinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples -- nothing to seed the structural triple set with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/bootea.py

def fit(  # noqa: PLR0915 -- the bootstrapping loop genuinely needs this much sequential setup
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> BootEALinker:
    """Train a shared TransE space with seed pseudo-triples and iterative bootstrapping.

    See the module docstring for the three mechanisms that
    distinguish this from a plain structural baseline.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (same approach
            [IPTransELinker][linkingtk.algorithms.ea.iptranse.IPTransELinker]
            uses) and as the bootstrap reference pool's source.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs. Seeds the main structural triple set (via pseudo-
            triples) and is excluded from the bootstrap reference pool.
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked,
            and training stops after ``patience`` checks with no
            improvement (no further bootstrapping rounds run once
            stopped, matching OpenEA's own early-stop-before-
            bootstrap ordering).
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples -- nothing to seed
            the structural triple set with -- or if ``device`` is
            invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("BootEALinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    entity_to_id, relation_to_id = build_id_mappings(triples1_labels + triples2_labels)
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)

    seed_pairs = [(s, t) for s, t in ground_truth if s in entity_to_id and t in entity_to_id]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train BootEA's structural "
            "triple set with."
        )
    seed_pair_ids = [(entity_to_id[s], entity_to_id[t]) for s, t in seed_pairs]

    real_ctx1 = build_kg_context(mapped1)
    real_ctx2 = build_kg_context(mapped2)
    seed_triples1, seed_triples2 = pseudo_triples_for_pairs(
        seed_pair_ids,
        real_ctx1.by_head,
        real_ctx1.by_tail,
        real_ctx2.by_head,
        real_ctx2.by_tail,
    )
    combined1 = self._augment_triples(mapped1, seed_triples1)
    combined2 = self._augment_triples(mapped2, seed_triples2)
    ctx1 = build_kg_context(combined1)
    ctx2 = build_kg_context(combined2)

    kg1_entity_ids = np.array(
        sorted(entity_to_id[e] for e in ids1 if e in entity_to_id), dtype=np.int64
    )
    kg2_entity_ids = np.array(
        sorted(entity_to_id[e] for e in ids2 if e in entity_to_id), dtype=np.int64
    )
    neighbor_k1 = max(1, int(round((1 - self.truncated_epsilon) * len(kg1_entity_ids))))
    neighbor_k2 = max(1, int(round((1 - self.truncated_epsilon) * len(kg2_entity_ids))))

    entity_embeds, relation_embeds = self._init_embeddings(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    structural_optimizer = torch.optim.Adagrad(
        [entity_embeds, relation_embeds], lr=self.learning_rate
    )
    alignment_optimizer = torch.optim.Adagrad(
        [entity_embeds, relation_embeds], lr=self.learning_rate
    )

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]

    pool1_labels, pool2_labels = reference_pools(sorted(ids1), sorted(ids2), seed_pairs)
    pool1_ids = np.array([entity_to_id[e] for e in pool1_labels], dtype=np.int64)
    pool2_ids = np.array([entity_to_id[e] for e in pool2_labels], dtype=np.int64)
    if self.bootstrap_pool_size is not None:
        if len(pool1_ids) > self.bootstrap_pool_size:
            pool1_ids = rng.choice(pool1_ids, size=self.bootstrap_pool_size, replace=False)
        if len(pool2_ids) > self.bootstrap_pool_size:
            pool2_ids = rng.choice(pool2_ids, size=self.bootstrap_pool_size, replace=False)

    labeled_alignment: dict[int, int] = {}
    neighbors1: dict[int, npt.NDArray[np.int64]] = {}
    neighbors2: dict[int, npt.NDArray[np.int64]] = {}

    num_outer_iters = max(1, self.num_epochs // self.sub_epoch)
    best_hits1 = -1.0
    epochs_without_improvement = 0
    epoch = 0
    stopped_early = False

    for _outer in range(num_outer_iters):
        for _ in range(self.sub_epoch):
            train_structural_epoch(
                entity_embeds,
                relation_embeds,
                structural_optimizer,
                ctx1,
                ctx2,
                neighbors1,
                neighbors2,
                rng,
                self.batch_size,
                self.pos_margin,
                self.neg_margin,
                self.neg_margin_balance,
                self.neg_triple_num,
            )
            epoch += 1

            if val_pairs and epoch % eval_every == 0:
                with torch.no_grad():
                    current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
                hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
                if hits1 <= best_hits1:
                    epochs_without_improvement += 1
                    if epochs_without_improvement >= patience:
                        stopped_early = True
                        break
                else:
                    best_hits1 = hits1
                    epochs_without_improvement = 0

        if stopped_early:
            break

        if len(pool1_ids) and len(pool2_ids):
            with torch.no_grad():
                current = functional.normalize(entity_embeds, dim=1).cpu().numpy()
            sim_mat = current[pool1_ids] @ current[pool2_ids].T
            curr_matches = find_mwgm_pairs(sim_mat, self.sim_th, self.k)
            labeled_alignment = edit_labeled_alignment(
                labeled_alignment, {(i, j) for i, j, _ in curr_matches}, sim_mat
            )

            if labeled_alignment:
                aligned_pairs = [
                    (int(pool1_ids[i]), int(pool2_ids[j])) for i, j in labeled_alignment.items()
                ]
                align_triples1, align_triples2 = pseudo_triples_for_pairs(
                    aligned_pairs, ctx1.by_head, ctx1.by_tail, ctx2.by_head, ctx2.by_tail
                )
                train_alignment_epoch(
                    entity_embeds,
                    relation_embeds,
                    alignment_optimizer,
                    np.array(align_triples1, dtype=np.int64).reshape(-1, 3),
                    np.array(align_triples2, dtype=np.int64).reshape(-1, 3),
                    self.batch_size,
                )

        with torch.no_grad():
            current = functional.normalize(entity_embeds, dim=1).cpu().numpy()
        neighbors1 = compute_truncated_neighbors(current, kg1_entity_ids, neighbor_k1)
        neighbors2 = compute_truncated_neighbors(current, kg2_entity_ids, neighbor_k2)

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/bootea.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/bootea.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

ChatEALinker

ChatEALinker(
    base_linker: EmbeddingLinker,
    client: LlmClient,
    top_k: int = 20,
    neigh_num: int = 5,
    desc_neigh_num: int = 25,
    desc_max_tokens: int = 80,
    threshold: float = 0.5,
    history_len: int = 3,
    relation_names: dict[str, str] | None = None,
    temporal_triples: list[TemporalTriple] | None = None,
    use_name: bool = True,
    use_desc: bool = True,
    use_struct: bool = True,
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

Re-ranks an already-fitted embedding linker's top-K candidates with an LLM.

Unlike most linkers in linkingtk.algorithms.ea, has no fit() -- base_linker must already be fitted before construction (its embeddings are what generates candidates), and the LLM itself needs no training. This mirrors LlmBaseLinker, which is also fit()-free for the same reason.

Parameters:

Name	Type	Description	Default
base_linker	EmbeddingLinker	An already-fitted linker exposing source_embedding/ target_embedding (see EmbeddingLinker) -- e.g. a fitted SimpleHHEALinker.	required
client	LlmClient	An LlmClient (e.g. from create_client). This class never imports openai/anthropic itself.	required
top_k	int	Candidates considered per source entity, ranked by base_linker's own cosine similarity. The reference's own default (20).	20
neigh_num	int	Neighbor triples shown per entity in the reasoning prompt (the reference's own --neigh 5 default for the main re-ranking pass).	5
desc_neigh_num	int	Neighbor triples shown per entity when generating its one-sentence description (the reference's own --neigh 25 default -- wider than neigh_num since description generation happens once per entity, not once per comparison).	25
desc_max_tokens	int	Forwarded to the description-generation LLM call.	80
threshold	float	If the base method's own top-2 candidate scores differ by more than this, accept the top-1 candidate immediately with no LLM call at all (the reference's own confidence shortcut).	0.5
history_len	int	Chat turns (user+assistant pairs) kept in the running conversation history within one source entity's re-ranking.	3
relation_names	dict[str, str] | None	relation_id -> display name, e.g. from IcewsWikiDataset.load_relation_labels. Falls back to the raw id for any relation not present, or when omitted entirely.	None
temporal_triples	list[TemporalTriple] | None	Optional (subject_id, relation_id, object_id, start_label, end_label) temporal facts (e.g. from IcewsWikiDataset.load_temporal_graphs). When given, used instead of graph for neighbor rendering (strictly more informative -- same ids, plus time) and enables the time-scoring dimension. None (default) disables the time dimension entirely, matching SimpleHHEALinker.fit's own convention for the same reference feature-toggle.	None
use_name	bool	Whether to score name similarity.	True
use_desc	bool	Whether to generate and score entity descriptions (requires use_name, see resolve_dimensions).	True
use_struct	bool	Whether to score structural (neighbor-tuple) similarity.	True
matching	Matcher	Strategy used to resolve scored candidates into final links, same as every other linker's matching argument.	DEFAULT_MATCHER

Note

blocking is accepted for interface compliance but not used -- candidates come from base_linker's own embedding similarity, not a BlockingStrategy, same pattern LlmBaseLinker/ StringSimilarityLinker already document for their own unused params. dataset2 must be a list[Entity], not an EntitySource -- candidate generation needs every target's embedding upfront.

Source code in src/linkingtk/algorithms/ea/chatea.py

def __init__(
    self,
    base_linker: EmbeddingLinker,
    client: LlmClient,
    top_k: int = 20,
    neigh_num: int = 5,
    desc_neigh_num: int = 25,
    desc_max_tokens: int = 80,
    threshold: float = 0.5,
    history_len: int = 3,
    relation_names: dict[str, str] | None = None,
    temporal_triples: list[TemporalTriple] | None = None,
    use_name: bool = True,
    use_desc: bool = True,
    use_struct: bool = True,
    matching: Matcher = DEFAULT_MATCHER,
) -> None:
    self.base_linker = base_linker
    self.client = client
    self.top_k = top_k
    self.neigh_num = neigh_num
    self.desc_neigh_num = desc_neigh_num
    self.desc_max_tokens = desc_max_tokens
    self.threshold = threshold
    self.history_len = history_len
    self.relation_names = relation_names or {}
    self.temporal_triples = temporal_triples
    self.use_name = use_name
    self.use_desc = use_desc
    self.use_struct = use_struct
    self.matching = matching

EntMatcherLinker

EntMatcherLinker()

Bases: FeatureClassifierLinker

EA linker combining hand-crafted similarity features with optimal one-to-one matching.

A preconfigured FeatureClassifierLinker (matching=OptimalMatcher()) — see that class for the feature set, training, and matching behavior. Kept as its own named class since it maps to a recognizable reference (EntMatcher's key insight: a globally optimal assignment can outperform independent per-source matching).

Source code in src/linkingtk/algorithms/ea/entmatcher.py

def __init__(self) -> None:
    super().__init__(matching=OptimalMatcher())

GCNAlignLinker

GCNAlignLinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    learning_rate: float = 8.0,
    neg_triple_num: int = 5,
    gamma: float = 3.0,
    min_weight: float = 0.3,
    use_attributes: bool = False,
    attr_dim: int = 100,
    attr_top_fraction: float = 0.7,
    beta: float = 0.9,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a 2-layer GCN's propagated embeddings.

Must be fit before link can be called. See the module docstring for what's ported vs. deliberately deviated from OpenEA, and for the structural-only-vs.-sa (structural+attribute) tradeoff.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the propagated structural embeddings (OpenEA's se_dim). OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs. OpenEA's published value allows up to 2000 with early stopping.	500
learning_rate	float	Plain SGD's learning rate -- see the module docstring for why this is unusually high by this package's conventions. OpenEA's published value is 8.	8.0
neg_triple_num	int	Negatives sampled per seed pair per side (k), refreshed every 10 epochs. OpenEA's published value is 5.	5
gamma	float	Margin for margin_ranking_loss_l1. OpenEA's published value is 3.	3.0
min_weight	float	Floor applied to every adjacency edge weight -- see build_weighted_adjacency. OpenEA's own value (hardcoded, not configurable in the reference) is 0.3.	0.3
use_attributes	bool	Whether to also train the attribute (ae) branch -- needs fit()'s attribute_triples1/attribute_triples2 to be given. Defaults to False (structural-only, needs nothing beyond plain relational triples).	False
attr_dim	int	Dimensionality of the propagated attribute embeddings (OpenEA's ae_dim). Only used if use_attributes. OpenEA's published value is 100.	100
attr_top_fraction	float	Fraction of the combined attribute-predicate vocabulary kept as feature columns -- see build_attribute_features. Only used if use_attributes. OpenEA's own value (hardcoded) is 0.7.	0.7
beta	float	Weight on the structural half of the concatenated [se * beta, ae * (1 - beta)] scoring embedding. Only used if use_attributes. OpenEA's published value is 0.9.	0.9
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/gcn_align.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    learning_rate: float = 8.0,
    neg_triple_num: int = 5,
    gamma: float = 3.0,
    min_weight: float = 0.3,
    use_attributes: bool = False,
    attr_dim: int = 100,
    attr_top_fraction: float = 0.7,
    beta: float = 0.9,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.learning_rate = learning_rate
    self.neg_triple_num = neg_triple_num
    self.gamma = gamma
    self.min_weight = min_weight
    self.use_attributes = use_attributes
    self.attr_dim = attr_dim
    self.attr_top_fraction = attr_top_fraction
    self.beta = beta
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> GCNAlignLinker

Propagate structural (and, if enabled, attribute) embeddings through the graph.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities. Unused beyond mirroring sibling linkers' signatures -- which entities get embeddings is determined entirely by graph.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs used as positive pairs in the margin loss.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader. Also drives the structural adjacency both branches propagate over, even when use_attributes is set.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10
attribute_triples1	list[Triple] | None	KG1's (entity_id, predicate, value) attribute triples, e.g. from EnFr15KAttrDataset.load_attribute_triples. Required if use_attributes (set in __init__); ignored otherwise.	None
attribute_triples2	list[Triple] | None	KG2's own attribute triples. See attribute_triples1.	None

Returns:

Type	Description
GCNAlignLinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples, if use_attributes is set but no attribute triples (or no attribute predicate clears attr_top_fraction's cutoff) are given, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/gcn_align.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> GCNAlignLinker:
    """Propagate structural (and, if enabled, attribute) embeddings through the graph.

    Args:
        dataset1: Source entities. Unused beyond mirroring sibling
            linkers' signatures -- which entities get embeddings is
            determined entirely by ``graph``.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs used as positive pairs in the margin loss.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must already
            be disjoint, as they are from that loader. Also drives the
            structural adjacency both branches propagate over, even
            when ``use_attributes`` is set.
        random_state: Seed for reproducible training. Left unspecified,
            training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples, e.g. from
            [EnFr15KAttrDataset.load_attribute_triples][linkingtk.datasets.openea_native._OpenEANativeDataset.load_attribute_triples].
            Required if ``use_attributes`` (set in ``__init__``);
            ignored otherwise.
        attribute_triples2: KG2's own attribute triples. See
            ``attribute_triples1``.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples, if ``use_attributes``
            is set but no attribute triples (or no attribute predicate
            clears ``attr_top_fraction``'s cutoff) are given, or if
            ``device`` is invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
    except ImportError as exc:
        raise OptionalDependencyError("GCNAlignLinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    triples = to_triples(graph)
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped = map_triples_to_ids(triples, entity_to_id, relation_to_id)
    num_entities = len(entity_to_id)

    seed_pairs = [
        (entity_to_id[s], entity_to_id[t])
        for s, t in ground_truth
        if s in entity_to_id and t in entity_to_id
    ]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train GCN-Align's margin loss with."
        )

    r2f, r2if = compute_relation_functionality(mapped)
    raw_indices, raw_values = build_weighted_adjacency(mapped, r2f, r2if, self.min_weight)
    norm_indices, norm_values = normalize_adjacency_coo(raw_indices, raw_values, num_entities)
    adjacency = coo_to_torch_sparse(
        norm_indices, norm_values, (num_entities, num_entities), device
    )

    model_se = build_gcn_align_branch(num_entities, self.embedding_dim).to(device)
    optimizer_se = torch.optim.SGD(model_se.parameters(), lr=self.learning_rate)

    model_ae: torch.nn.Module | None = None
    optimizer_ae: torch.optim.Optimizer | None = None
    attr_features: torch.Tensor | None = None
    if self.use_attributes:
        if not attribute_triples1 and not attribute_triples2:
            raise LinkingTKError(
                "GCNAlignLinker(use_attributes=True) needs `attribute_triples1`/"
                "`attribute_triples2` -- pass entities and attribute triples from "
                "EnFr15KAttrDataset (or similar), not EnFr15KDataset."
            )
        attr_indices, attr_values, num_attrs = build_attribute_features(
            attribute_triples1 or [],
            attribute_triples2 or [],
            entity_to_id,
            self.attr_top_fraction,
        )
        if num_attrs == 0:
            raise LinkingTKError(
                "GCNAlignLinker(use_attributes=True): no attribute predicate cleared "
                "`attr_top_fraction`'s cutoff -- nothing to train the attribute branch with."
            )
        attr_features = coo_to_torch_sparse(
            attr_indices, attr_values, (num_entities, num_attrs), device
        )
        model_ae = build_gcn_align_attr_branch(num_attrs, self.attr_dim).to(device)
        optimizer_ae = torch.optim.SGD(model_ae.parameters(), lr=self.learning_rate)

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    pos_left = torch.tensor([s for s, _ in seed_pairs], device=device)
    pos_right = torch.tensor([t for _, t in seed_pairs], device=device)

    def _resample_negatives() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        neg_left_np, neg_right_np, neg2_left_np, neg2_right_np = sample_negatives(
            seed_pairs, num_entities, self.neg_triple_num, rng
        )
        return (
            torch.from_numpy(neg_left_np).long().to(device),
            torch.from_numpy(neg_right_np).long().to(device),
            torch.from_numpy(neg2_left_np).long().to(device),
            torch.from_numpy(neg2_right_np).long().to(device),
        )

    def _combined_embeddings(
        se: npt.NDArray[Any], ae: npt.NDArray[Any] | None
    ) -> npt.NDArray[Any]:
        if ae is None:
            return se
        return np.concatenate([se * self.beta, ae * (1.0 - self.beta)], axis=1)

    neg_left, neg_right, neg2_left, neg2_right = _resample_negatives()

    for epoch in range(self.num_epochs):
        if epoch % _NEG_REFRESH_INTERVAL == 0 and epoch > 0:
            neg_left, neg_right, neg2_left, neg2_right = _resample_negatives()

        embeddings_se = model_se(adjacency)
        loss_se = margin_ranking_loss_l1(
            embeddings_se,
            pos_left,
            pos_right,
            neg_left,
            neg_right,
            neg2_left,
            neg2_right,
            self.gamma,
        )
        optimizer_se.zero_grad()
        loss_se.backward()  # type: ignore[no-untyped-call]
        optimizer_se.step()

        if model_ae is not None and optimizer_ae is not None and attr_features is not None:
            embeddings_ae = model_ae(adjacency, attr_features)
            loss_ae = margin_ranking_loss_l1(
                embeddings_ae,
                pos_left,
                pos_right,
                neg_left,
                neg_right,
                neg2_left,
                neg2_right,
                self.gamma,
            )
            optimizer_ae.zero_grad()
            loss_ae.backward()  # type: ignore[no-untyped-call]
            optimizer_ae.step()

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                se_np = model_se(adjacency).cpu().numpy()
                ae_np = (
                    model_ae(adjacency, attr_features).cpu().numpy()
                    if model_ae is not None and attr_features is not None
                    else None
                )
            current_embeds = _combined_embeddings(se_np, ae_np)
            hits1 = _validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_se = model_se(adjacency).cpu().numpy()
        final_ae = (
            model_ae(adjacency, attr_features).cpu().numpy()
            if model_ae is not None and attr_features is not None
            else None
        )
    final_embeds = _combined_embeddings(final_se, final_ae)
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/gcn_align.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/gcn_align.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

IMUSELinker

IMUSELinker(
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    name_sim_threshold: float = 0.6,
    entity_sim_threshold: float = 0.6,
    attr_sim_threshold: float = 0.6,
    top_k_attribute_pairs: int = 10,
    bootstrap_iterations: int = 1,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Unsupervised EA: bootstraps its own seed pairs from attribute-value similarity.

Must be fit before link can be called. Unlike every other linker in this family, fit takes no ground_truth -- see the module docstring.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each pairing one epoch of structural triple-loss training with one epoch of alignment-loss training. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	2000
batch_size	int	Mini-batch size for structural triple-loss training; the alignment-loss step count is derived from it (matching OpenEA).	5000
learning_rate	float	SGD's learning rate, for both optimizers -- OpenEA's published value is 0.01.	0.01
margin	float	Hinge margin for the structural triple loss. OpenEA's published value is 1.5.	1.5
name_sim_threshold	float	Minimum Levenshtein ratio for two attribute predicates' local names to be considered aligned. OpenEA hardcodes 0.6 at this call site.	0.6
entity_sim_threshold	float	Minimum average Levenshtein ratio (across shared aligned-predicate values) for two entities to be considered aligned. OpenEA's published value (sim_thresholds_ent) is 0.6.	0.6
attr_sim_threshold	float	Same as entity_sim_threshold but for re-aligning attribute predicates from aligned entities (sim_thresholds_attr, only used if ``bootstrap_iterations 1). OpenEA's published value is0.6``.	0.6
top_k_attribute_pairs	int	Maximum aligned attribute-predicate pairs kept after name-similarity matching, ranked by combined triple count. OpenEA hardcodes 10.	10
bootstrap_iterations	int	Outer bootstrap rounds (entity alignment, then optionally attribute-predicate re-alignment from those entities). OpenEA's published default (interactive_model_iter_num) is 1, at which the attribute-predicate-realignment step never runs -- see the module docstring.	1
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/imuse.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    name_sim_threshold: float = 0.6,
    entity_sim_threshold: float = 0.6,
    attr_sim_threshold: float = 0.6,
    top_k_attribute_pairs: int = 10,
    bootstrap_iterations: int = 1,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.margin = margin
    self.name_sim_threshold = name_sim_threshold
    self.entity_sim_threshold = entity_sim_threshold
    self.attr_sim_threshold = attr_sim_threshold
    self.top_k_attribute_pairs = top_k_attribute_pairs
    self.bootstrap_iterations = bootstrap_iterations
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> IMUSELinker

Bootstrap an alignment from attribute similarity, then train jointly.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- used to partition graph into each KG's own triples (same approach KDCoELinker uses).	required
dataset2	list[Entity]	Target entities. See dataset1.	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint.	required
attribute_triples1	list[Triple]	KG1's (entity_id, predicate, value) attribute triples -- required, not optional (unlike JAPELinker/ KDCoELinker): IMUSE has no ground truth at all, so bootstrapping from attributes is its only source of cross-KG signal, with no fallback.	required
attribute_triples2	list[Triple]	KG2's attribute triples. See attribute_triples1.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. Unlike ground_truth, this is accepted here -- see the module docstring's note on what "unsupervised" does and doesn't cover. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
IMUSELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If both attribute_triples1 and attribute_triples2 are empty (nothing to bootstrap an alignment with), or if bootstrapping finds no confident aligned entity pairs at all (nothing to train the alignment loss with), or if device is invalid or unavailable.
OptionalDependencyError	If torch or rapidfuzz isn't installed.

Source code in src/linkingtk/algorithms/ea/imuse.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> IMUSELinker:
    """Bootstrap an alignment from attribute similarity, then train jointly.

    Args:
        dataset1: Source entities -- used to partition ``graph`` into
            each KG's own triples (same approach
            [KDCoELinker][linkingtk.algorithms.ea.kdcoe.KDCoELinker]
            uses).
        dataset2: Target entities. See ``dataset1``.
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples -- **required**, not optional (unlike
            [JAPELinker][linkingtk.algorithms.ea.jape.JAPELinker]/
            [KDCoELinker][linkingtk.algorithms.ea.kdcoe.KDCoELinker]):
            IMUSE has no ground truth at all, so bootstrapping from
            attributes is its only source of cross-KG signal, with no
            fallback.
        attribute_triples2: KG2's attribute triples. See
            ``attribute_triples1``.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. Unlike ``ground_truth``, this
            *is* accepted here -- see the module docstring's note on
            what "unsupervised" does and doesn't cover. If ``None``
            (default), trains the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If both ``attribute_triples1`` and
            ``attribute_triples2`` are empty (nothing to bootstrap an
            alignment with), or if bootstrapping finds no confident
            aligned entity pairs at all (nothing to train the
            alignment loss with), or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch or rapidfuzz isn't installed.
    """
    try:
        import rapidfuzz  # noqa: F401 -- import-only check, used inside _imuse_text.py
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("IMUSELinker", "kge") from exc

    if not attribute_triples1 and not attribute_triples2:
        raise LinkingTKError(
            "Both `attribute_triples1` and `attribute_triples2` are empty; "
            "IMUSELinker.fit() has no attribute signal to bootstrap an "
            "alignment with (it has no ground_truth to fall back on)."
        )

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1 = [t for t in triples if t[0] in ids1]
    triples2 = [t for t in triples if t[0] in ids2]

    entity_to_id, relation_to_id = build_id_mappings(triples1 + triples2)
    mapped1 = map_triples_to_ids(triples1, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(mapped1)
    ctx2 = build_kg_context(mapped2)

    bootstrapped_pairs = bootstrap_alignment(
        attribute_triples1,
        attribute_triples2,
        self.name_sim_threshold,
        self.entity_sim_threshold,
        self.attr_sim_threshold,
        self.top_k_attribute_pairs,
        self.bootstrap_iterations,
    )
    seed_pairs = [
        (s, t) for s, t in bootstrapped_pairs if s in entity_to_id and t in entity_to_id
    ]
    if not seed_pairs:
        raise LinkingTKError(
            "Bootstrapping from attribute-value similarity found no confident "
            "aligned entity pairs (or none with both ids present in `graph`'s "
            "own triples); fit() has no seed pairs to train IMUSE's alignment "
            "loss with."
        )
    seed_source_ids = np.array([entity_to_id[s] for s, _ in seed_pairs], dtype=np.int64)
    seed_target_ids = np.array([entity_to_id[t] for _, t in seed_pairs], dtype=np.int64)

    entity_embeds, relation_embeds = self._init_embeddings(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    triple_optimizer = torch.optim.SGD([entity_embeds, relation_embeds], lr=self.learning_rate)
    align_optimizer = torch.optim.SGD([entity_embeds], lr=self.learning_rate)
    align_steps = max(1, math.ceil(len(seed_pairs) / self.batch_size))

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_structural_epoch(
            entity_embeds,
            relation_embeds,
            triple_optimizer,
            ctx1,
            ctx2,
            rng,
            self.batch_size,
            self.margin,
        )
        train_align_epoch(
            entity_embeds, align_optimizer, seed_source_ids, seed_target_ids, align_steps
        )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/imuse.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/imuse.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

IPTransELinker

IPTransELinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    path_parm: float = 0.1,
    sim_th: float = 0.7,
    bootstrap_every: int = 100,
    bootstrap_pool_size: int | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a shared TransE space, path loss, and bootstrapping.

Must be fit before link can be called. See the module docstring for how this differs from MTransELinker.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each one pass over both KGs' triples (plus path-loss minibatches). OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for triple-loss training, split proportionally between the two KGs by triple count each epoch (matching OpenEA).	5000
learning_rate	float	Adagrad's learning rate, for both the main and bootstrap optimizers -- OpenEA's published value is 0.01.	0.01
margin	float	Hinge margin for the triple and path losses. OpenEA's published value is 1.5.	1.5
path_parm	float	Weight applied to the path (relation-composition) loss. OpenEA's published value is 0.1.	0.1
sim_th	float	Minimum embedding similarity for a bootstrap round to accept a newly found pair. OpenEA's published value is 0.7.	0.7
bootstrap_every	int	Run a bootstrapping round every this many epochs. OpenEA's published value (bp_freq) is 100.	100
bootstrap_pool_size	int | None	Cap on each side's bootstrap reference pool size (random subsample, fixed for the whole run) -- see the module docstring's deviation note. None (default) uses the full pool.	None
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/iptranse.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    path_parm: float = 0.1,
    sim_th: float = 0.7,
    bootstrap_every: int = 100,
    bootstrap_pool_size: int | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.margin = margin
    self.path_parm = path_parm
    self.sim_th = sim_th
    self.bootstrap_every = bootstrap_every
    self.bootstrap_pool_size = bootstrap_pool_size
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> IPTransELinker

Train a shared TransE space with margin+path loss and periodic bootstrapping.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (see the module docstring).	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs. Seeds the shared embedding table (each pair's target is aliased to its source's id) and is excluded from the bootstrap reference pool.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
IPTransELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in dataset1/dataset2 -- nothing to seed the shared embedding table with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/iptranse.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> IPTransELinker:
    """Train a shared TransE space with margin+path loss and periodic bootstrapping.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (see the module docstring).
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs. Seeds the shared embedding table (each pair's
            target is aliased to its source's id) and is excluded
            from the bootstrap reference pool.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must
            already be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``dataset1``/``dataset2`` -- nothing to
            seed the shared embedding table with -- or if ``device``
            is invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("IPTransELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    seed_pairs = [(s, t) for s, t in ground_truth if s in ids1 and t in ids2]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in "
            "`dataset1`/`dataset2`; fit() has no seed pairs to train "
            "IPTransE's shared embedding table with."
        )

    entity_to_id, relation_to_id = build_shared_id_mappings(
        triples1_labels + triples2_labels, seed_pairs
    )
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(triples1_labels, mapped1, relation_to_id)
    ctx2 = build_kg_context(triples2_labels, mapped2, relation_to_id)

    pool1_ids, pool2_ids = self._bootstrap_pools(
        dataset1, dataset2, entity_to_id, seed_pairs, rng
    )
    combined_entity_pool = np.union1d(ctx1.entity_pool, ctx2.entity_pool)

    entity_embeds, relation_embeds = self._init_embeddings(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    optimizer = torch.optim.Adagrad([entity_embeds, relation_embeds], lr=self.learning_rate)
    alignment_optimizer = torch.optim.Adagrad(
        [entity_embeds, relation_embeds], lr=self.learning_rate
    )

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_epoch(
            entity_embeds,
            relation_embeds,
            optimizer,
            ctx1,
            ctx2,
            rng,
            self.batch_size,
            self.margin,
            self.path_parm,
        )

        if (epoch + 1) % self.bootstrap_every == 0:
            bootstrap_round(
                entity_embeds,
                relation_embeds,
                alignment_optimizer,
                ctx1,
                ctx2,
                pool1_ids,
                pool2_ids,
                combined_entity_pool,
                rng,
                self.batch_size,
                self.margin,
                self.sim_th,
            )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/iptranse.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/iptranse.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

JAPELinker

JAPELinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    neg_alpha: float = 0.1,
    top_attr_threshold: float = 0.9,
    attr_sim_mat_threshold: float = 0.95,
    attr_max_epoch: int = 200,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a shared TransE space, plus an inspectable attribute similarity.

Must be fit before link can be called. See the module docstring for how this differs from MTransELinker/ IPTransELinker, and in particular why the attribute-similarity matrix built from attribute_triples1/attribute_triples2 no longer perturbs the trained structural embeddings (it's exposed post-fit() as attr_sim_mat_ instead).

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained structural and attribute embeddings. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Structural training epochs. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for both the structural triple loss and (if attribute triples are given) Attr2Vec's training.	5000
learning_rate	float	Adagrad's learning rate, for every optimizer used here -- OpenEA's published value is 0.01.	0.01
neg_alpha	float	Weight on the negative-triple term in the structural loss (pos_loss - neg_alpha * neg_loss). OpenEA's published value is 0.1.	0.1
top_attr_threshold	float	Fraction of each KG's distinct attribute predicates to keep, by descending frequency. OpenEA's published value is 0.9.	0.9
attr_sim_mat_threshold	float	Attribute-similarity values below this are zeroed out in the exposed attr_sim_mat_. OpenEA's published value is 0.95.	0.95
attr_max_epoch	int	Attr2Vec training epochs. OpenEA's published value is 200.	200
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/jape.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    neg_alpha: float = 0.1,
    top_attr_threshold: float = 0.9,
    attr_sim_mat_threshold: float = 0.95,
    attr_max_epoch: int = 200,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.neg_alpha = neg_alpha
    self.top_attr_threshold = top_attr_threshold
    self.attr_sim_mat_threshold = attr_sim_mat_threshold
    self.attr_max_epoch = attr_max_epoch
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False
    self.attr_sim_mat_: npt.NDArray[np.floating[Any]] | None = None
    """Thresholded attribute-similarity matrix from the last `fit()` call, or ``None`` if no
    attribute triples were given. Rows/columns correspond to
    `attr_pool1_ids_`/`attr_pool2_ids_`. Diagnostic only -- see the
    module docstring for why this doesn't affect the trained
    embeddings."""
    self.attr_pool1_ids_: list[str] = []
    """Entity ids indexing `attr_sim_mat_`'s rows."""
    self.attr_pool2_ids_: list[str] = []
    """Entity ids indexing `attr_sim_mat_`'s columns."""

attr_pool1_ids_ instance-attribute

attr_pool1_ids_: list[str] = []

Entity ids indexing attr_sim_mat_'s rows.

attr_pool2_ids_ instance-attribute

attr_pool2_ids_: list[str] = []

Entity ids indexing attr_sim_mat_'s columns.

attr_sim_mat_ instance-attribute

attr_sim_mat_: NDArray[floating[Any]] | None = None

Thresholded attribute-similarity matrix from the last fit() call, or None if no attribute triples were given. Rows/columns correspond to attr_pool1_ids_/attr_pool2_ids_. Diagnostic only -- see the module docstring for why this doesn't affect the trained embeddings.

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> JAPELinker

Train a shared TransE space; optionally build an inspectable attribute-similarity matrix.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (same approach as IPTransELinker.fit, for the same reason: negative sampling needs each KG's own entity pool).	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs. Seeds the shared embedding table (each pair's target is aliased to its source's id) and, if attribute triples are given, merges each pair's attribute vocabularies for Attr2Vec training.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2)) -- entity ids on both sides must already be disjoint.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10
attribute_triples1	list[Triple] | None	KG1's (entity_id, predicate, value) attribute triples, e.g. from a GraphDatasetLoader that supports load_attribute_triples() (see linkingtk.datasets.openea_native). If None (default) or empty, together with attribute_triples2, the attribute pipeline is skipped entirely -- doesn't affect the trained structural embeddings either way (see the module docstring's discovered-upstream-bug note), but skips training Attr2Vec and populating attr_sim_mat_.	None
attribute_triples2	list[Triple] | None	KG2's attribute triples. See attribute_triples1.	None

Returns:

Type	Description
JAPELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in dataset1/dataset2 -- nothing to seed the shared embedding table with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/jape.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> JAPELinker:
    """Train a shared TransE space; optionally build an inspectable attribute-similarity matrix.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (same approach as
            `IPTransELinker.fit`, for the same reason: negative
            sampling needs each KG's own entity pool).
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs. Seeds the shared embedding table (each pair's
            target is aliased to its source's id) and, if attribute
            triples are given, merges each pair's attribute
            vocabularies for Attr2Vec training.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)``) -- entity ids
            on both sides must already be disjoint.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples, e.g. from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]
            that supports
            ``load_attribute_triples()``
            (see [linkingtk.datasets.openea_native][]). If ``None``
            (default) or empty, together with ``attribute_triples2``,
            the attribute pipeline is skipped entirely -- doesn't
            affect the trained structural embeddings either way (see
            the module docstring's discovered-upstream-bug note), but
            skips training Attr2Vec and populating ``attr_sim_mat_``.
        attribute_triples2: KG2's attribute triples. See
            ``attribute_triples1``.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``dataset1``/``dataset2`` -- nothing to
            seed the shared embedding table with -- or if ``device``
            is invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("JAPELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    seed_pairs = [(s, t) for s, t in ground_truth if s in ids1 and t in ids2]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in "
            "`dataset1`/`dataset2`; fit() has no seed pairs to train "
            "JAPE's shared embedding table with."
        )

    entity_to_id, relation_to_id = build_shared_id_mappings(
        triples1_labels + triples2_labels, seed_pairs
    )
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(mapped1)
    ctx2 = build_kg_context(mapped2)

    entity_embeds, relation_embeds = self._init_structural_embeddings(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    optimizer = torch.optim.Adagrad([entity_embeds, relation_embeds], lr=self.learning_rate)

    self.attr_sim_mat_, self.attr_pool1_ids_, self.attr_pool2_ids_ = (
        self._setup_attribute_pipeline(
            dataset1,
            dataset2,
            entity_to_id,
            seed_pairs,
            attribute_triples1,
            attribute_triples2,
            rng,
            device,
        )
    )

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_triple_epoch(
            entity_embeds,
            relation_embeds,
            optimizer,
            ctx1,
            ctx2,
            rng,
            self.batch_size,
            self.neg_alpha,
        )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/jape.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/jape.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

KDCoELinker

KDCoELinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    alpha: float = 5.0,
    new_param: float = 0.1,
    wv_dim: int = 300,
    default_desc_length: int = 4,
    desc_batch_size: int = 512,
    desc_sim_th: float = 0.95,
    sim_th: float = 0.8,
    max_co_training_iters: int = 5,
    word_embed_url: str = _FASTTEXT_URL,
    cache_dir: Path | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Co-trains a structural TransE+mapping space with a description-text encoder.

Must be fit before link can be called. See the module docstring for how this differs from MTransELinker and IPTransELinker.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings and the (square) mapping matrix. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Per-phase training-epoch cap (structural or description), each co-training iteration. OpenEA's config allows up to 2000 per phase with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for structural triple-loss training; the mapping-loss batch size is derived from it (same number of minibatches per epoch for both, matching OpenEA).	5000
learning_rate	float	Adagrad's learning rate, for every optimizer here (structural, mapping, mapping-new, description) -- OpenEA's published value is 0.01.	0.01
margin	float	Hinge margin for the structural triple loss. OpenEA's published value is 1.5.	1.5
alpha	float	Weight applied to the seed-pairs mapping loss. OpenEA's published value is 5.0 (matches MTransELinker's own default).	5.0
new_param	float	Weight applied to the bootstrapped-pairs mapping loss. OpenEA's published value is 0.1.	0.1
wv_dim	int	Pretrained word-vector dimensionality. Must match word_embed_url's file -- OpenEA's published fastText file is 300-dimensional.	300
default_desc_length	int	Fixed word-sequence length each entity's description is padded/truncated to. OpenEA's published value is 4.	4
desc_batch_size	int	Mini-batch size for description-encoder training. OpenEA's published value is 512.	512
desc_sim_th	float	Kept for interface parity with OpenEA's published value (0.95), but not the operative filter for the description pathway's bootstrap rounds -- see the module docstring's "reciprocal best-match" deviation. Description similarities saturate near 1.0 on real data, making an absolute threshold meaningless; bootstrap rounds instead use _iptranse_training.find_mutual_pairs's reciprocal-agreement filter.	0.95
sim_th	float	Minimum (mapped) structural-embedding similarity for a bootstrap round to accept a newly found pair. OpenEA's published value is 0.8.	0.8
max_co_training_iters	int	Cap on outer co-training iterations (each alternating a description-encoder phase and a structural phase). OpenEA's published value (max_iter) is 5.	5
word_embed_url	str	URL of a fastText-format .vec.zip (see _kdcoe_text.py's load_fasttext_vectors). Defaults to the real, published pretrained file; tests override this with a file:// url pointing at a tiny local fixture.	_FASTTEXT_URL
cache_dir	Path | None	Passed through to fetch_cached for both the word-vector download and (indirectly, via callers' attribute_triples) any dataset assets. None uses fetch_cached's own default cache location.	None
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/kdcoe.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    alpha: float = 5.0,
    new_param: float = 0.1,
    wv_dim: int = 300,
    default_desc_length: int = 4,
    desc_batch_size: int = 512,
    desc_sim_th: float = 0.95,
    sim_th: float = 0.8,
    max_co_training_iters: int = 5,
    word_embed_url: str = _FASTTEXT_URL,
    cache_dir: Path | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.margin = margin
    self.alpha = alpha
    self.new_param = new_param
    self.wv_dim = wv_dim
    self.default_desc_length = default_desc_length
    self.desc_batch_size = desc_batch_size
    self.desc_sim_th = desc_sim_th
    self.sim_th = sim_th
    self.max_co_training_iters = max_co_training_iters
    self.word_embed_url = word_embed_url
    self.cache_dir = cache_dir
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._mapping: npt.NDArray[np.floating[Any]] | None = None
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> KDCoELinker

Co-train the structural+mapping space and the description encoder.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (same approach IPTransELinker uses) and as the description pathway's label-fallback source.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs, used as both phases' seed training pairs.	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping in both phases -- every eval_every epochs, Hits@1 is checked, and that phase stops after patience checks with no improvement. If None (default), each phase trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping a phase early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10
attribute_triples1	list[Triple] | None	KG1's (entity_id, predicate, value) attribute triples, if available -- entities with a .../description-like predicate triple use its value as description text; others fall back to their own label (see the module docstring). May be omitted (or partial) -- the description pathway still runs, just with less real description text.	None
attribute_triples2	list[Triple] | None	KG2's attribute triples. See attribute_triples1.	None

Returns:

Type	Description
KDCoELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples -- nothing to seed either phase's training with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/kdcoe.py

def fit(  # noqa: PLR0915 -- co-training loop genuinely needs this much sequential setup
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> KDCoELinker:
    """Co-train the structural+mapping space and the description encoder.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (same approach
            [IPTransELinker][linkingtk.algorithms.ea.iptranse.IPTransELinker]
            uses) and as the description pathway's label-fallback source.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs, used as both phases' seed training pairs.
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping in *both* phases -- every ``eval_every`` epochs,
            Hits@1 is checked, and that phase stops after ``patience``
            checks with no improvement. If ``None`` (default), each
            phase trains the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping a phase early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples, if available -- entities with a
            ``.../description``-like predicate triple use its value as
            description text; others fall back to their own label
            (see the module docstring). May be omitted (or partial) --
            the description pathway still runs, just with less real
            description text.
        attribute_triples2: KG2's attribute triples. See
            ``attribute_triples1``.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples -- nothing to seed
            either phase's training with -- or if ``device`` is
            invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("KDCoELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    entity_to_id, relation_to_id = build_id_mappings(triples1_labels + triples2_labels)
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(mapped1)
    ctx2 = build_kg_context(mapped2)

    seed_pairs = [(s, t) for s, t in ground_truth if s in entity_to_id and t in entity_to_id]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train KDCoE's structural "
            "and description phases with."
        )
    seed_source_ids = np.array([entity_to_id[s] for s, _ in seed_pairs], dtype=np.int64)
    seed_target_ids = np.array([entity_to_id[t] for _, t in seed_pairs], dtype=np.int64)

    word_embeds, word_ids = self._build_description_inputs(
        dataset1, dataset2, attribute_triples1 or [], attribute_triples2 or [], device
    )

    entity_embeds, relation_embeds, mapping_mat, eye = self._init_structural(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    structural_optimizer = torch.optim.Adagrad(
        [entity_embeds, relation_embeds], lr=self.learning_rate
    )
    mapping_optimizer = torch.optim.Adagrad([mapping_mat, entity_embeds], lr=self.learning_rate)
    mapping_optimizer_new = torch.optim.Adagrad(
        [mapping_mat, entity_embeds], lr=self.learning_rate
    )
    desc_encoder = build_description_encoder(self.wv_dim, self.default_desc_length).to(device)
    desc_optimizer = torch.optim.Adagrad(desc_encoder.parameters(), lr=self.learning_rate)

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    pool1, pool2 = reference_pools(sorted(ids1), sorted(ids2), seed_pairs)

    desc_bootstrapped_pairs: set[tuple[str, str]] = set()
    rel_bootstrapped_pairs: set[tuple[str, str]] = set()
    triple_steps = max(1, math.ceil((len(ctx1.triples) + len(ctx2.triples)) / self.batch_size))

    for _iteration in range(self.max_co_training_iters):
        self._train_desc_phase(
            torch,
            desc_encoder,
            desc_optimizer,
            word_embeds,
            word_ids,
            seed_pairs + list(desc_bootstrapped_pairs),
            val_pairs,
            rng,
            patience,
            eval_every,
        )
        desc_found = self._find_new_pairs_via_desc(
            torch, desc_encoder, word_embeds, word_ids, pool1, pool2
        )
        if desc_found and desc_found <= desc_bootstrapped_pairs:
            break
        desc_bootstrapped_pairs |= desc_found

        mapping_new_pairs = rel_bootstrapped_pairs | desc_bootstrapped_pairs
        bootstrap_source_ids = np.array(
            [entity_to_id[s] for s, _ in mapping_new_pairs], dtype=np.int64
        )
        bootstrap_target_ids = np.array(
            [entity_to_id[t] for _, t in mapping_new_pairs], dtype=np.int64
        )
        self._train_structural_phase(
            torch,
            functional,
            entity_embeds,
            relation_embeds,
            mapping_mat,
            eye,
            structural_optimizer,
            mapping_optimizer,
            mapping_optimizer_new,
            ctx1,
            ctx2,
            seed_source_ids,
            seed_target_ids,
            bootstrap_source_ids,
            bootstrap_target_ids,
            entity_to_id,
            val_pairs,
            rng,
            triple_steps,
            patience,
            eval_every,
        )
        rel_found = self._find_new_pairs_via_structural(
            torch, functional, entity_embeds, mapping_mat, entity_to_id, pool1, pool2
        )
        if rel_found and rel_found <= rel_bootstrapped_pairs:
            break
        rel_bootstrapped_pairs |= rel_found

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
        final_mapping = mapping_mat.cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._mapping = final_mapping
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side (projected).

Source code in src/linkingtk/algorithms/ea/kdcoe.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side (projected)."""
    return self._project(self._embedding(entity_id))

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side (unprojected).

Source code in src/linkingtk/algorithms/ea/kdcoe.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side (unprojected)."""
    return self._embedding(entity_id)

KGELinker

KGELinker(
    model: str = "TransE",
    embedding_dim: int = 50,
    num_epochs: int = 50,
    batch_size: int | None = None,
    learning_rate: float = 0.01,
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

Scores candidate pairs by trained knowledge-graph-embedding similarity.

Must be fit before link can be called. Unlike FeatureClassifierLinker, training needs the KGs' relational structure, not just entity text -- see fit's graph argument.

Parameters:

Name	Type	Description	Default
model	str	A model name from pykeen's registry (pykeen.models.model_resolver, ~40 entries), e.g. "TransE", "DistMult", "RotatE". Not restricted to those three -- pykeen already provides its own named-model resolution, so there's nothing to add on top of it here -- but they're what this class is validated against; other models are best-effort, since not every pykeen model shares the same constructor keyword shape (e.g. some need more than embedding_dim).	'TransE'
embedding_dim	int	Dimensionality of the trained entity/relation embeddings.	50
num_epochs	int	Training epochs.	50
batch_size	int | None	Training batch size. Defaults to pykeen's own auto-selection (None).	None
learning_rate	float	Adam's learning rate. pykeen's own default (0.001, via torch.optim.Adam's default) converges too slowly for the seed-alignment triples to pull related entities together within a practical epoch budget at small scale; 0.01 converges reliably in testing.	0.01
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher. OptimalMatcher's dense Hungarian-algorithm cost matrix is explicitly not meant for KG-scale entity counts (see its docstring) -- avoid it here for anything but small graphs.	DEFAULT_MATCHER

Source code in src/linkingtk/algorithms/ea/kge.py

def __init__(
    self,
    model: str = "TransE",
    embedding_dim: int = 50,
    num_epochs: int = 50,
    batch_size: int | None = None,
    learning_rate: float = 0.01,
    matching: Matcher = DEFAULT_MATCHER,
) -> None:
    self.model = model
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.matching = matching
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
) -> KGELinker

Train KGE embeddings jointly over graph plus seed alignment triples.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities.	required
dataset2	list[Entity]	Target entities.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs, added to the training graph as seed alignment triples (see the module docstring).	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, pykeen trains non-deterministically.	None

Returns:

Type	Description
KGELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If no ground_truth pair's ids are actually present in dataset1/dataset2 -- nothing to seed training with.
OptionalDependencyError	If pykeen isn't installed.

Source code in src/linkingtk/algorithms/ea/kge.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
) -> KGELinker:
    """Train KGE embeddings jointly over ``graph`` plus seed alignment triples.

    Args:
        dataset1: Source entities.
        dataset2: Target entities.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs, added to the training graph as seed alignment
            triples (see the module docstring).
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must
            already be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left
            unspecified, pykeen trains non-deterministically.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If no ``ground_truth`` pair's ids are actually
            present in ``dataset1``/``dataset2`` -- nothing to seed
            training with.
        OptionalDependencyError: If pykeen isn't installed.
    """
    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    seed_pairs = [(s, t) for s, t in ground_truth if s in ids1 and t in ids2]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have ids present in `dataset1`/"
            "`dataset2`; fit() has no seed alignment triples to train with."
        )

    try:
        import torch
        from pykeen.models import model_resolver
        from pykeen.training import SLCWATrainingLoop
        from pykeen.triples import TriplesFactory
        from pykeen.utils import set_random_seed
        from torch.optim import Adam
    except ImportError as exc:
        raise OptionalDependencyError("KGELinker", "kge") from exc

    if random_state is not None:
        set_random_seed(random_state)

    triples = to_triples(graph) + [(s, _SEED_RELATION, t) for s, t in seed_pairs]
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped_triples = map_triples_to_ids(triples, entity_to_id, relation_to_id)
    triples_factory = TriplesFactory(
        mapped_triples=mapped_triples,
        entity_to_id=entity_to_id,
        relation_to_id=relation_to_id,
    )
    model: Any = model_resolver.make(
        self.model,
        triples_factory=triples_factory,
        embedding_dim=self.embedding_dim,
        random_seed=random_state,
    )
    training_loop = SLCWATrainingLoop(
        model=model,
        triples_factory=triples_factory,
        optimizer=Adam(params=model.get_grad_params(), lr=self.learning_rate),
    )
    training_loop.train(
        triples_factory=triples_factory,
        num_epochs=self.num_epochs,
        batch_size=self.batch_size,
        use_tqdm=False,
    )

    model.eval()
    with torch.no_grad():
        embeddings = model.entity_representations[0](indices=None)
    if embeddings.is_complex():
        embeddings = torch.view_as_real(embeddings).reshape(embeddings.shape[0], -1)
    embeddings_array = embeddings.detach().numpy()
    self._id_to_vector = {
        entity_id: embeddings_array[index]
        for entity_id, index in triples_factory.entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/kge.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/kge.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

MTransELinker

MTransELinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    alpha: float = 5.0,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a learned cross-KG mapping over shared TransE embeddings.

Must be fit before link can be called.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings and the (square) mapping matrix. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each alternating one epoch of triple-loss training with one epoch of mapping-loss training. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for triple-loss training; the per-epoch mapping-loss batch size is derived from it (matching OpenEA: the same number of mini-batches per epoch for both).	5000
learning_rate	float	Adagrad's learning rate, for both the embedding and mapping optimizers -- OpenEA's published value is 0.01.	0.01
alpha	float	Weight applied to the mapping loss (approximation error plus orthogonality regularization). OpenEA's published value is 5.0.	5.0
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/mtranse.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    alpha: float = 5.0,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.alpha = alpha
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._mapping: npt.NDArray[np.floating[Any]] | None = None
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> MTransELinker

Train shared TransE embeddings plus a learned cross-KG mapping.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities. Unused beyond mirroring KGELinker.fit's signature -- which entities get embeddings is determined entirely by graph.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs used to train the mapping matrix. Unlike KGELinker, these are not injected into graph as extra triples.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
MTransELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples -- nothing to train the mapping with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/mtranse.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> MTransELinker:
    """Train shared TransE embeddings plus a learned cross-KG mapping.

    Args:
        dataset1: Source entities. Unused beyond mirroring
            [KGELinker.fit][linkingtk.algorithms.ea.kge.KGELinker.fit]'s
            signature -- which entities get embeddings is determined
            entirely by ``graph``.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs used to train the mapping matrix. Unlike
            ``KGELinker``, these are *not* injected into ``graph`` as
            extra triples.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must
            already be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples -- nothing to train
            the mapping with -- or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("MTransELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)

    triples = to_triples(graph)
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped_triples = map_triples_to_ids(triples, entity_to_id, relation_to_id)
    heads = torch.from_numpy(mapped_triples[:, 0]).long().to(device)
    rels = torch.from_numpy(mapped_triples[:, 1]).long().to(device)
    tails = torch.from_numpy(mapped_triples[:, 2]).long().to(device)

    seed_pairs = [(s, t) for s, t in ground_truth if s in entity_to_id and t in entity_to_id]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train MTransE's mapping with."
        )
    seed_source = torch.tensor([entity_to_id[s] for s, _ in seed_pairs], device=device)
    seed_target = torch.tensor([entity_to_id[t] for _, t in seed_pairs], device=device)

    entity_embeds = torch.nn.Parameter(
        functional.normalize(
            torch.randn(len(entity_to_id), self.embedding_dim, device=device), dim=1
        )
    )
    relation_embeds = torch.nn.Parameter(
        functional.normalize(
            torch.randn(len(relation_to_id), self.embedding_dim, device=device), dim=1
        )
    )
    mapping_mat = torch.nn.Parameter(
        torch.nn.init.orthogonal_(
            torch.empty(self.embedding_dim, self.embedding_dim, device=device)
        )
    )
    eye = torch.eye(self.embedding_dim, device=device)

    triple_optimizer = torch.optim.Adagrad(
        [entity_embeds, relation_embeds], lr=self.learning_rate
    )
    # `entity_embeds` is included here (not just `mapping_mat`) to match
    # OpenEA's `generate_optimizer(mapping_loss, lr, var_list=None, ...)`,
    # which implicitly differentiates against every trainable variable
    # the loss touches -- see the module docstring.
    mapping_optimizer = torch.optim.Adagrad([mapping_mat, entity_embeds], lr=self.learning_rate)

    triple_steps = max(1, math.ceil(len(heads) / self.batch_size))
    seed_batch_size = max(1, len(seed_source) // triple_steps)

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        perm = torch.randperm(len(heads), device=device)
        for step in range(triple_steps):
            idx = perm[step * self.batch_size : (step + 1) * self.batch_size]
            h = functional.normalize(entity_embeds[heads[idx]], dim=1)
            r = functional.normalize(relation_embeds[rels[idx]], dim=1)
            t = functional.normalize(entity_embeds[tails[idx]], dim=1)
            triple_loss = ((h + r - t) ** 2).sum()
            triple_optimizer.zero_grad()
            triple_loss.backward()  # type: ignore[no-untyped-call]
            triple_optimizer.step()

        for _step in range(triple_steps):
            batch_idx = torch.randint(0, len(seed_source), (seed_batch_size,), device=device)
            s_vecs = functional.normalize(entity_embeds[seed_source[batch_idx]], dim=1)
            t_vecs = functional.normalize(entity_embeds[seed_target[batch_idx]], dim=1)
            map_loss = ((t_vecs - s_vecs @ mapping_mat) ** 2).sum()
            orth_loss = ((mapping_mat @ mapping_mat.T - eye) ** 2).sum()
            mapping_loss = self.alpha * (map_loss + orth_loss)
            mapping_optimizer.zero_grad()
            mapping_loss.backward()  # type: ignore[no-untyped-call]
            mapping_optimizer.step()

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
                current_mapping = mapping_mat.cpu().numpy()
            hits1 = _validation_hits1(current_embeds, current_mapping, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
        final_mapping = mapping_mat.cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._mapping = final_mapping
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side (projected).

Source code in src/linkingtk/algorithms/ea/mtranse.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side (projected)."""
    return self._project(self._embedding(entity_id))

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side (unprojected).

Source code in src/linkingtk/algorithms/ea/mtranse.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side (unprojected)."""
    return self._embedding(entity_id)

MultiKELinker

MultiKELinker(
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    batch_size: int = 5000,
    entity_batch_size: int = 5000,
    attribute_batch_size: int = 5000,
    learning_rate: float = 0.001,
    common_space_learning_rate: float = 0.004,
    neg_triple_num: int = 10,
    cv_weight: float = 1.0,
    predicate_init_sim: float = 0.9,
    predicate_soft_sim: float = 0.8,
    start_predicate_soft_alignment: int = 10,
    min_predicate_triple_count: int = 10,
    literal_encoder_model: str = "distilbert-base-multilingual-cased",
    literal_max_length: int = 16,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a shared embedding table unifying name/relation/attribute views.

Must be fit before link can be called. See the module docstring for the full architecture and the three documented scope cuts.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of every trained embedding table. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each one pass of relation-view, common-space, cross-KG relation-entity (and, once warmed up, cross-KG relation-predicate) training, then the same four steps for the attribute view. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	2000
batch_size	int	Mini-batch size for the relation view's main loss and its two cross-KG inference terms.	5000
entity_batch_size	int	Mini-batch size for common-space learning.	5000
attribute_batch_size	int	Mini-batch size for the attribute view's main loss and its two cross-KG inference terms.	5000
learning_rate	float	SGD's learning rate for every optimizer except common-space learning. OpenEA's published value is 0.001.	0.001
common_space_learning_rate	float	SGD's learning rate for common-space learning specifically (OpenEA's own ITC_learning_rate, deliberately distinct from learning_rate). Published value is 0.004.	0.004
neg_triple_num	int	Negatives sampled per positive relation triple -- only the main relation-view loss uses negative sampling at all (see the module docstring). Published value is 10.	10
cv_weight	float	Weight applied to the common-space-learning loss. Published value is 1.0.	1.0
predicate_init_sim	float	Levenshtein-ratio threshold (of humanized local names) for the one-time mutual-best-match predicate alignment. Published value is 0.9.	0.9
predicate_soft_sim	float	Similarity floor used to rescale an aligned predicate pair's weight (_multike_text.zoom_weight). Published value is 0.80.	0.8
start_predicate_soft_alignment	int	Epoch (0-indexed) at which the two predicate-alignment cross-KG loss terms start contributing -- a warm-up gate. Published value is 10.	10
min_predicate_triple_count	int	Attribute predicates with fewer than this many combined (both-KG) triples are dropped before anything else. Published value is 10.	10
literal_encoder_model	str	Hugging Face model id for the literal encoder (see the module docstring's "The literal encoder" section). Defaults to distilbert-base-multilingual-cased.	'distilbert-base-multilingual-cased'
literal_max_length	int	Token-truncation length passed to the literal encoder for every name/value string.	16
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Also used for the literal encoder's forward pass. Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/multike.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    batch_size: int = 5000,
    entity_batch_size: int = 5000,
    attribute_batch_size: int = 5000,
    learning_rate: float = 0.001,
    common_space_learning_rate: float = 0.004,
    neg_triple_num: int = 10,
    cv_weight: float = 1.0,
    predicate_init_sim: float = 0.9,
    predicate_soft_sim: float = 0.80,
    start_predicate_soft_alignment: int = 10,
    min_predicate_triple_count: int = 10,
    literal_encoder_model: str = "distilbert-base-multilingual-cased",
    literal_max_length: int = 16,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.entity_batch_size = entity_batch_size
    self.attribute_batch_size = attribute_batch_size
    self.learning_rate = learning_rate
    self.common_space_learning_rate = common_space_learning_rate
    self.neg_triple_num = neg_triple_num
    self.cv_weight = cv_weight
    self.predicate_init_sim = predicate_init_sim
    self.predicate_soft_sim = predicate_soft_sim
    self.start_predicate_soft_alignment = start_predicate_soft_alignment
    self.min_predicate_triple_count = min_predicate_triple_count
    self.literal_encoder_model = literal_encoder_model
    self.literal_max_length = literal_max_length
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> MultiKELinker

Train the three-view embedding and a shared space unifying them.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs -- MultiKE's only real supervision channel (see the module docstring's "cross-KG entity inference").	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint.	required
attribute_triples1	list[Triple]	KG1's (entity_id, predicate, value) attribute triples. Required -- see the module docstring's deviation note.	required
attribute_triples2	list[Triple]	KG2's attribute triples. Required.	required
random_state	int | None	Seed for reproducible training (also seeds the literal encoder's random projection). Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
MultiKELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If attribute_triples1 and attribute_triples2 are both empty, or if none of ground_truth's pairs have both ids present in graph's own triples, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/multike.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> MultiKELinker:
    """Train the three-view embedding and a shared space unifying them.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs -- MultiKE's only real supervision channel (see the
            module docstring's "cross-KG entity inference").
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples. **Required** -- see the module
            docstring's deviation note.
        attribute_triples2: KG2's attribute triples. **Required**.
        random_state: Seed for reproducible training (also seeds the
            literal encoder's random projection). Left unspecified,
            training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If ``attribute_triples1`` and
            ``attribute_triples2`` are both empty, or if none of
            ``ground_truth``'s pairs have both ids present in
            ``graph``'s own triples, or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    if not attribute_triples1 and not attribute_triples2:
        raise LinkingTKError(
            "MultiKELinker.fit() needs attribute triples -- `attribute_triples1` "
            "and `attribute_triples2` are both empty, leaving nothing to train "
            "the attribute view or literal encoder with."
        )

    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("MultiKELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1 = [t for t in triples if t[0] in ids1]
    triples2 = [t for t in triples if t[0] in ids2]

    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped1 = map_triples_to_ids(triples1, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(mapped1)
    ctx2 = build_kg_context(mapped2)

    seed_pairs = [(s, t) for s, t in ground_truth if s in entity_to_id and t in entity_to_id]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train MultiKE's cross-KG "
            "entity inference with -- its only real supervision channel."
        )

    attr1_in = [t for t in attribute_triples1 if t[0] in entity_to_id]
    attr2_in = [t for t in attribute_triples2 if t[0] in entity_to_id]
    attr1_filtered, attr2_filtered = filter_frequent_predicates(
        attr1_in, attr2_in, self.min_predicate_triple_count
    )
    attr1_clean = self._clean_triples(attr1_filtered)
    attr2_clean = self._clean_triples(attr2_filtered)

    attribute_to_id = {
        predicate: index
        for index, predicate in enumerate(
            sorted({p for _, p, _ in attr1_clean} | {p for _, p, _ in attr2_clean})
        )
    }
    values = sorted({v for _, _, v in attr1_clean} | {v for _, _, v in attr2_clean})
    value_to_id = {value: index for index, value in enumerate(values)}

    def map_attribute_triples(mapped_triples: list[Triple]) -> npt.NDArray[np.int64]:
        rows = [
            (entity_to_id[e], attribute_to_id[a], value_to_id[v]) for e, a, v in mapped_triples
        ]
        return np.array(rows, dtype=np.int64) if rows else np.empty((0, 3), dtype=np.int64)

    attr_mapped1 = map_attribute_triples(attr1_clean)
    attr_mapped2 = map_attribute_triples(attr2_clean)

    aligned_relations = align_predicates_by_name(
        {p for _, p, _ in triples1}, {p for _, p, _ in triples2}, self.predicate_init_sim
    )
    aligned_attributes = align_predicates_by_name(
        {p for _, p, _ in attr1_clean}, {p for _, p, _ in attr2_clean}, self.predicate_init_sim
    )

    weighted1 = weight_attribute_triples(
        attr1_clean, aligned_attributes, self.predicate_soft_sim, is_kg1=True
    )
    weighted2 = weight_attribute_triples(
        attr2_clean, aligned_attributes, self.predicate_soft_sim, is_kg1=False
    )
    actx1 = build_attribute_context(
        attr_mapped1, np.array([w for *_, w in weighted1], dtype=np.float32)
    )
    actx2 = build_attribute_context(
        attr_mapped2, np.array([w for *_, w in weighted2], dtype=np.float32)
    )

    cross_rel1, cross_rel2 = generate_cross_kg_relation_triples(seed_pairs, triples1, triples2)
    cross_rel_ids = np.concatenate(
        [
            map_triples_to_ids(cross_rel1, entity_to_id, relation_to_id),
            map_triples_to_ids(cross_rel2, entity_to_id, relation_to_id),
        ],
        axis=0,
    )
    cross_attr1, cross_attr2 = generate_cross_kg_attribute_triples(
        seed_pairs, attr1_clean, attr2_clean
    )
    cross_attr_ids = np.concatenate(
        [map_attribute_triples(cross_attr1), map_attribute_triples(cross_attr2)], axis=0
    )

    pred_rel1_ids, pred_rel1_w = self._relation_ids_and_weights(
        substitute_predicate_triples(triples1, aligned_relations, is_kg1=True),
        entity_to_id,
        relation_to_id,
    )
    pred_rel2_ids, pred_rel2_w = self._relation_ids_and_weights(
        substitute_predicate_triples(triples2, aligned_relations, is_kg1=False),
        entity_to_id,
        relation_to_id,
    )
    pred_rel_ids = np.concatenate([pred_rel1_ids, pred_rel2_ids], axis=0)
    pred_rel_w = np.concatenate([pred_rel1_w, pred_rel2_w], axis=0)

    pred_attr1_ids, pred_attr1_w = self._attribute_ids_and_weights(
        substitute_predicate_triples(attr1_clean, aligned_attributes, is_kg1=True),
        entity_to_id,
        attribute_to_id,
        value_to_id,
    )
    pred_attr2_ids, pred_attr2_w = self._attribute_ids_and_weights(
        substitute_predicate_triples(attr2_clean, aligned_attributes, is_kg1=False),
        entity_to_id,
        attribute_to_id,
        value_to_id,
    )
    pred_attr_ids = np.concatenate([pred_attr1_ids, pred_attr2_ids], axis=0)
    pred_attr_w = np.concatenate([pred_attr1_w, pred_attr2_w], axis=0)

    name_embeds, literal_embeds = self._encode_literals(
        dataset1, dataset2, entity_to_id, values, random_state, device
    )
    name_embeds_t = torch.from_numpy(name_embeds).to(device)
    literal_embeds_t = torch.from_numpy(literal_embeds).to(device)

    rv_ent_embeds = self._init_embedding(torch, functional, len(entity_to_id), device)
    rel_embeds = self._init_embedding(torch, functional, max(1, len(relation_to_id)), device)
    av_ent_embeds = self._init_embedding(torch, functional, len(entity_to_id), device)
    attr_embeds = self._init_embedding_raw(torch, max(1, len(attribute_to_id)), device)
    ent_embeds = self._init_embedding(torch, functional, len(entity_to_id), device)
    scorer = build_attribute_scorer(self.embedding_dim).to(device)

    relation_optimizer = torch.optim.SGD(
        [rv_ent_embeds, rel_embeds, ent_embeds], lr=self.learning_rate
    )
    attribute_optimizer = torch.optim.SGD(
        [av_ent_embeds, attr_embeds, ent_embeds, *scorer.parameters()], lr=self.learning_rate
    )
    common_space_optimizer = torch.optim.SGD(
        [ent_embeds, rv_ent_embeds, av_ent_embeds], lr=self.common_space_learning_rate
    )
    cross_kg_relation_entity_optimizer = torch.optim.SGD(
        [rv_ent_embeds, rel_embeds], lr=self.learning_rate
    )
    cross_kg_attribute_entity_optimizer = torch.optim.SGD(
        [av_ent_embeds, attr_embeds, *scorer.parameters()], lr=self.learning_rate
    )
    cross_kg_relation_predicate_optimizer = torch.optim.SGD(
        [rv_ent_embeds, rel_embeds], lr=self.learning_rate
    )
    cross_kg_attribute_predicate_optimizer = torch.optim.SGD(
        [av_ent_embeds, attr_embeds, *scorer.parameters()], lr=self.learning_rate
    )

    entity_ids = np.arange(len(entity_to_id), dtype=np.int64)
    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_relation_view_epoch(
            rv_ent_embeds,
            rel_embeds,
            ent_embeds,
            name_embeds_t,
            relation_optimizer,
            ctx1,
            ctx2,
            rng,
            self.batch_size,
            self.neg_triple_num,
        )
        train_common_space_epoch(
            ent_embeds,
            rv_ent_embeds,
            av_ent_embeds,
            name_embeds_t,
            common_space_optimizer,
            entity_ids,
            rng,
            self.entity_batch_size,
            self.cv_weight,
        )
        train_cross_kg_relation_entity_epoch(
            rv_ent_embeds,
            rel_embeds,
            cross_kg_relation_entity_optimizer,
            cross_rel_ids,
            rng,
            self.batch_size,
        )
        if epoch >= self.start_predicate_soft_alignment:
            train_cross_kg_relation_predicate_epoch(
                rv_ent_embeds,
                rel_embeds,
                cross_kg_relation_predicate_optimizer,
                pred_rel_ids,
                pred_rel_w,
                rng,
                self.batch_size,
            )

        train_attribute_view_epoch(
            av_ent_embeds,
            attr_embeds,
            literal_embeds_t,
            ent_embeds,
            name_embeds_t,
            scorer,
            attribute_optimizer,
            actx1,
            actx2,
            rng,
            self.attribute_batch_size,
        )
        train_common_space_epoch(
            ent_embeds,
            rv_ent_embeds,
            av_ent_embeds,
            name_embeds_t,
            common_space_optimizer,
            entity_ids,
            rng,
            self.entity_batch_size,
            self.cv_weight,
        )
        train_cross_kg_attribute_entity_epoch(
            av_ent_embeds,
            attr_embeds,
            literal_embeds_t,
            scorer,
            cross_kg_attribute_entity_optimizer,
            cross_attr_ids,
            rng,
            self.attribute_batch_size,
        )
        if epoch >= self.start_predicate_soft_alignment:
            train_cross_kg_attribute_predicate_epoch(
                av_ent_embeds,
                attr_embeds,
                literal_embeds_t,
                scorer,
                cross_kg_attribute_predicate_optimizer,
                pred_attr_ids,
                pred_attr_w,
                rng,
                self.attribute_batch_size,
            )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(ent_embeds, dim=1).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(ent_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/multike.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/multike.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

RDGCNLinker

RDGCNLinker(
    embedding_dim: int = 300,
    num_epochs: int = 500,
    learning_rate: float = 0.002,
    neg_triple_num: int = 125,
    gamma: float = 1.0,
    alpha: float = 0.1,
    beta: float = 0.3,
    default_name_length: int = 4,
    word_embed_url: str = _FASTTEXT_URL,
    cache_dir: Path | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via RDGCN's dual/primal-attention-plus-GCN embeddings.

Must be fit before link can be called. See the module docstring for what's ported vs. deliberately deviated from OpenEA, and why dataset1/dataset2 must come from EnFr15KAttrDataset rather than EnFr15KDataset.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Entity-embedding dimensionality -- must equal the pretrained word vectors' own dimensionality (no projection layer sits between them). OpenEA's published EN-FR-15K-V1 config uses 300 (fastText's native size).	300
num_epochs	int	Training epochs. OpenEA's config allows up to 2000 with early stopping.	500
learning_rate	float	Adam's learning rate. OpenEA's published value is 0.002.	0.002
neg_triple_num	int	Hard negatives mined per seed pair per side (k), refreshed every 10 epochs. OpenEA's published value is 125.	125
gamma	float	Margin for margin_ranking_loss_l1. OpenEA's published value is 1.0.	1.0
alpha	float	Round-1 primal-embedding mixing weight. OpenEA's published value is 0.1.	0.1
beta	float	Round-2 primal-embedding mixing weight. OpenEA's published value is 0.3.	0.3
default_name_length	int	Max local-name tokens summed per entity for the initial embedding. OpenEA's published value is 4.	4
word_embed_url	str	URL of a fastText-format .vec.zip (see _kdcoe_text.py's load_fasttext_vectors). Defaults to the real, published pretrained file; tests override this with a file:// url pointing at a tiny local fixture.	_FASTTEXT_URL
cache_dir	Path | None	Passed through to fetch_cached for the word-vector download. None uses fetch_cached's own default cache location.	None
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/rdgcn.py

def __init__(
    self,
    embedding_dim: int = 300,
    num_epochs: int = 500,
    learning_rate: float = 0.002,
    neg_triple_num: int = 125,
    gamma: float = 1.0,
    alpha: float = 0.1,
    beta: float = 0.3,
    default_name_length: int = 4,
    word_embed_url: str = _FASTTEXT_URL,
    cache_dir: Path | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.learning_rate = learning_rate
    self.neg_triple_num = neg_triple_num
    self.gamma = gamma
    self.alpha = alpha
    self.beta = beta
    self.default_name_length = default_name_length
    self.word_embed_url = word_embed_url
    self.cache_dir = cache_dir
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> RDGCNLinker

Train the two-round dual/primal-attention GCN over the combined graph.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- must carry non-empty labels (e.g. from EnFr15KAttrDataset), used to seed each entity's initial embedding from pretrained word vectors of its local name.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs used as positive pairs in the margin loss.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
RDGCNLinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If no entity in dataset1/dataset2 has a non-empty label, if none of ground_truth's pairs have both ids present in graph's own triples, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/rdgcn.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> RDGCNLinker:
    """Train the two-round dual/primal-attention GCN over the combined graph.

    Args:
        dataset1: Source entities -- must carry non-empty `labels`
            (e.g. from
            [EnFr15KAttrDataset][linkingtk.datasets.EnFr15KAttrDataset]),
            used to seed each entity's initial embedding from
            pretrained word vectors of its local name.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs used as positive pairs in the margin loss.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must already
            be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left unspecified,
            training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If no entity in ``dataset1``/``dataset2`` has a
            non-empty label, if none of ``ground_truth``'s pairs have
            both ids present in ``graph``'s own triples, or if
            ``device`` is invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    entities = dataset1 + dataset2
    if not any(label_texts(entity) for entity in entities):
        raise LinkingTKError(
            "RDGCNLinker.fit() needs entities with non-empty `labels` to seed name "
            "embeddings from -- pass entities from EnFr15KAttrDataset (or similar), "
            "not EnFr15KDataset. See the module docstring."
        )

    try:
        import torch
    except ImportError as exc:
        raise OptionalDependencyError("RDGCNLinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)

    triples = to_triples(graph)
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped = map_triples_to_ids(triples, entity_to_id, relation_to_id)
    num_entities = len(entity_to_id)
    num_relations = len(relation_to_id)

    seed_pairs = [
        (entity_to_id[s], entity_to_id[t])
        for s, t in ground_truth
        if s in entity_to_id and t in entity_to_id
    ]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train RDGCN's margin loss with."
        )

    entity_by_id = {entity.id: entity for entity in entities}
    vocabulary = {
        token
        for entity in entities
        for token in tokenize_description(label_texts(entity)[0] if label_texts(entity) else "")
    }
    word_vectors = load_fasttext_vectors(self.word_embed_url, vocabulary, self.cache_dir)
    ordered_entities = [
        entity_by_id.get(entity_id, Entity(id=entity_id, labels=[]))
        for entity_id, _index in sorted(entity_to_id.items(), key=lambda item: item[1])
    ]
    name_embeddings = init_name_embeddings(
        ordered_entities, word_vectors, self.embedding_dim, self.default_name_length
    )
    initial_embeds_np = np.stack(
        [name_embeddings[entity.id] for entity in ordered_entities]
    ).astype(np.float32)
    initial_embeds = torch.from_numpy(initial_embeds_np).to(device)

    primal_indices, primal_values = build_primal_adjacency(mapped, num_entities)
    primal_indices, primal_values = normalize_adjacency_coo(
        primal_indices, primal_values, num_entities
    )
    primal_adjacency = coo_to_torch_sparse(
        primal_indices, primal_values, (num_entities, num_entities), device
    )

    dual_adjacency_np = build_dual_graph(mapped, num_relations, num_entities)
    dual_adjacency = torch.from_numpy(dual_adjacency_np).float().to(device)

    (head_indices, head_values), (tail_indices, tail_values) = build_relation_masks(
        mapped, num_relations
    )
    head_mask = coo_to_torch_sparse(
        head_indices, head_values, (num_relations, num_entities), device
    )
    tail_mask = coo_to_torch_sparse(
        tail_indices, tail_values, (num_relations, num_entities), device
    )

    edge_heads_np, edge_tails_np, edge_relations_np = build_edge_relations(mapped)
    edge_heads = torch.from_numpy(edge_heads_np).long().to(device)
    edge_tails = torch.from_numpy(edge_tails_np).long().to(device)
    edge_relations = torch.from_numpy(edge_relations_np).long().to(device)

    model = build_rdgcn_model(initial_embeds, self.embedding_dim, self.alpha, self.beta).to(
        device
    )
    optimizer = torch.optim.Adam(model.parameters(), lr=self.learning_rate)

    pos_left = torch.tensor([s for s, _ in seed_pairs], device=device)
    pos_right = torch.tensor([t for _, t in seed_pairs], device=device)

    def _forward() -> torch.Tensor:
        return model(  # type: ignore[no-any-return]
            primal_adjacency,
            dual_adjacency,
            head_mask,
            tail_mask,
            edge_heads,
            edge_tails,
            edge_relations,
        )

    with torch.no_grad():
        neg_left, neg_right, neg2_left, neg2_right = mine_hard_negatives(
            _forward(), pos_left, pos_right, self.neg_triple_num
        )

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        if epoch % _NEG_REFRESH_INTERVAL == 0 and epoch > 0:
            with torch.no_grad():
                neg_left, neg_right, neg2_left, neg2_right = mine_hard_negatives(
                    _forward(), pos_left, pos_right, self.neg_triple_num
                )

        embeddings = _forward()
        loss = margin_ranking_loss_l1(
            embeddings,
            pos_left,
            pos_right,
            neg_left,
            neg_right,
            neg2_left,
            neg2_right,
            self.gamma,
        )
        optimizer.zero_grad()
        loss.backward()  # type: ignore[no-untyped-call]
        optimizer.step()

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = _forward().cpu().numpy()
            hits1 = _validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = _forward().cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/rdgcn.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/rdgcn.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

RSN4EALinker

RSN4EALinker(
    hidden_size: int = 100,
    num_layers: int = 2,
    keep_prob: float = 0.6,
    max_length: int = 15,
    alpha: float = 0.7,
    beta: float = 0.7,
    repeat_times: int = 2,
    max_paths: int | None = None,
    num_epochs: int = 30,
    batch_size: int = 512,
    learning_rate: float = 0.0005,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a Recurrent Skipping Network over sampled walk paths.

Must be fit before link can be called. Unlike MTransELinker/ SEALinker, scoring uses raw (unprojected) embeddings on both sides -- RSN4EA never builds a mapping matrix; see the module docstring.

Parameters:

Name	Type	Description	Default
hidden_size	int	Entity/relation embedding dimensionality and LSTM hidden size. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_layers	int	LSTM depth. OpenEA's published value is 2.	2
keep_prob	float	Per-LSTM-layer output dropout keep-probability. OpenEA's published value is 0.6.	0.6
max_length	int	Sampled path length -- must be odd and >= 3 (entities at even positions, relations at odd positions). OpenEA's published value is 15.	15
alpha	float	Path-sampling depth-bias weight (down-weights continuations already directly reachable from a walk's own start entity). OpenEA's published value is 0.7.	0.7
beta	float	Path-sampling cross-KG bias weight (up-weights continuations whose tail is a seed-pair entity). OpenEA's published value is 0.7.	0.7
repeat_times	int	Walks sampled per row of the augmented knowledge base. OpenEA's published value is 2.	2
max_paths	int | None	Optional cap on the number of sampled paths kept for training (random subsample, fixed for the run) -- see the module docstring. None (default) keeps every successfully sampled path.	None
num_epochs	int	Training epochs. OpenEA's published value is 30 (RSN4EA converges in far fewer epochs than this family's TransE-based methods).	30
batch_size	int	Mini-batch size for path training. OpenEA's published value is 512.	512
learning_rate	float	Adam's learning rate (see the module docstring for why Adam, not the config's stated Adagrad). OpenEA's published value is 0.0005.	0.0005
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/rsn4ea.py

def __init__(
    self,
    hidden_size: int = 100,
    num_layers: int = 2,
    keep_prob: float = 0.6,
    max_length: int = 15,
    alpha: float = 0.7,
    beta: float = 0.7,
    repeat_times: int = 2,
    max_paths: int | None = None,
    num_epochs: int = 30,
    batch_size: int = 512,
    learning_rate: float = 0.0005,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    if max_length < 3 or max_length % 2 == 0:
        raise LinkingTKError(f"max_length must be odd and >= 3, got {max_length!r}.")
    self.hidden_size = hidden_size
    self.num_layers = num_layers
    self.keep_prob = keep_prob
    self.max_length = max_length
    self.alpha = alpha
    self.beta = beta
    self.repeat_times = repeat_times
    self.max_paths = max_paths
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> RSN4EALinker

Sample walk paths over an alias-expanded graph and train the RSN.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities. Unused beyond mirroring sibling linkers' signatures -- which entities get embeddings is determined entirely by graph.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs. Used to build the alias-substituted graph (_rsn4ea_training.build_augmented_kb) that lets sampled walks cross between KGs -- RSN4EA's only alignment-training mechanism, no separate loss channel.	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint, as they are from a GraphDatasetLoader's load_graphs().	required
random_state	int | None	Seed for reproducible path sampling and training. Left unspecified, both are non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
RSN4EALinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples, if no walk path could be sampled at all, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/rsn4ea.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> RSN4EALinker:
    """Sample walk paths over an alias-expanded graph and train the RSN.

    Args:
        dataset1: Source entities. Unused beyond mirroring sibling
            linkers' signatures -- which entities get embeddings is
            determined entirely by ``graph``.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs. Used to build the alias-substituted graph
            (`_rsn4ea_training.build_augmented_kb`)
            that lets sampled walks cross between KGs -- RSN4EA's only
            alignment-training mechanism, no separate loss channel.
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint, as they are
            from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``.
        random_state: Seed for reproducible path sampling and training.
            Left unspecified, both are non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples, if no walk path
            could be sampled at all, or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
    except ImportError as exc:
        raise OptionalDependencyError("RSN4EALinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    triples = to_triples(graph)
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped = map_triples_to_ids(triples, entity_to_id, relation_to_id)

    seed_pairs = [
        (entity_to_id[s], entity_to_id[t])
        for s, t in ground_truth
        if s in entity_to_id and t in entity_to_id
    ]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to build RSN4EA's alias-expanded "
            "graph with."
        )

    augmented_kb, num_relations = build_augmented_kb(
        mapped, len(entity_to_id), len(relation_to_id), seed_pairs
    )
    paths = sample_paths(
        augmented_kb,
        seed_pairs,
        rng,
        self.max_length,
        self.alpha,
        self.beta,
        self.repeat_times,
        self.max_paths,
    )
    if len(paths) == 0:
        raise LinkingTKError(
            "No walk paths of length `max_length` could be sampled from `graph` -- "
            "every entity reachable from a starting triple hit a dead end too soon."
        )

    model: torch.nn.Module = build_rsn_model(
        len(entity_to_id), num_relations, self.hidden_size, self.num_layers, self.keep_prob
    )
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=self.learning_rate)

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0
    best_state: dict[str, torch.Tensor] | None = None

    for epoch in range(self.num_epochs):
        train_epoch(model, optimizer, paths, rng, self.batch_size)

        if val_pairs and (epoch + 1) % eval_every == 0:
            model.eval()
            with torch.no_grad():
                current_embeds = _entity_embedding_weight(model).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0
                best_state = {k: v.detach().clone() for k, v in model.state_dict().items()}

    # Unlike this family's other linkers, RSN4EA's validation Hits@1
    # doesn't improve (or plateau) monotonically with more training --
    # diagnosed directly on the real EN-FR-15K-V1 benchmark: it peaks
    # within the first few `eval_every`-spaced checks, then steadily
    # degrades (the next-token training objective is only a proxy for
    # alignment quality, and keeps optimizing well past the point
    # where the two diverge). Restoring the best-seen checkpoint
    # (rather than using whichever epoch training happened to stop or
    # finish on, like every sibling linker here) is what actually
    # fixes this -- not a hyperparameter tweak.
    if best_state is not None:
        model.load_state_dict(best_state)

    model.eval()
    with torch.no_grad():
        final_embeds = _entity_embedding_weight(model).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/rsn4ea.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/rsn4ea.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

SEALinker

SEALinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    alpha_1: float = 2.5,
    alpha_2: float = 0.25,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via structural TransE plus a semi-supervised dual mapping.

Must be fit before link can be called.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings and both (square) mapping matrices. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each alternating one epoch of structural triple-loss training with one epoch of mapping-loss training. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for structural triple-loss training, split proportionally between the two KGs by triple count each epoch. The mapping loss reuses the resulting step count.	5000
learning_rate	float	Adam's learning rate, for both the structural and mapping optimizers -- OpenEA's published value is 0.01.	0.01
margin	float	Hinge margin for the structural triple loss. OpenEA's published value is 1.5.	1.5
alpha_1	float	Weight applied to the supervised mapping loss (seed-pair approximation error). OpenEA's published value is 2.5.	2.5
alpha_2	float	Weight applied to the semi-supervised (cycle-consistency) mapping loss over unlabeled entities. OpenEA's published value is 0.25.	0.25
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/sea.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    alpha_1: float = 2.5,
    alpha_2: float = 0.25,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.margin = margin
    self.alpha_1 = alpha_1
    self.alpha_2 = alpha_2
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._mapping: npt.NDArray[np.floating[Any]] | None = None
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> SEALinker

Train shared structural TransE embeddings plus a dual cross-KG mapping.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples for structural negative sampling and to build the semi-supervised unlabeled pool (see the module docstring).	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs used to train the supervised half of the mapping loss. Unlike KGELinker, these are not injected into graph as extra triples.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
SEALinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples -- nothing to train the mapping with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/sea.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> SEALinker:
    """Train shared structural TransE embeddings plus a dual cross-KG mapping.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples for structural negative
            sampling and to build the semi-supervised unlabeled pool
            (see the module docstring).
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs used to train the supervised half of the mapping
            loss. Unlike ``KGELinker``, these are *not* injected into
            ``graph`` as extra triples.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must
            already be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples -- nothing to train
            the mapping with -- or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("SEALinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(triples1_labels, mapped1)
    ctx2 = build_kg_context(triples2_labels, mapped2)

    seed_pairs = [(s, t) for s, t in ground_truth if s in entity_to_id and t in entity_to_id]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train SEA's mapping with."
        )
    seed_source = torch.tensor([entity_to_id[s] for s, _ in seed_pairs], device=device)
    seed_target = torch.tensor([entity_to_id[t] for _, t in seed_pairs], device=device)

    seed_sources = {s for s, _ in seed_pairs}
    seed_targets = {t for _, t in seed_pairs}
    unlabeled_pool1 = np.array(
        sorted({entity_to_id[e.id] for e in dataset1 if e.id in entity_to_id} - seed_sources),
        dtype=np.int64,
    )
    unlabeled_pool2 = np.array(
        sorted({entity_to_id[e.id] for e in dataset2 if e.id in entity_to_id} - seed_targets),
        dtype=np.int64,
    )

    entity_embeds, relation_embeds = self._init_structural_embeddings(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    mapping_shape = (self.embedding_dim, self.embedding_dim)
    mapping_mat_1 = torch.nn.Parameter(
        torch.nn.init.orthogonal_(torch.empty(*mapping_shape, device=device))
    )
    mapping_mat_2 = torch.nn.Parameter(
        torch.nn.init.orthogonal_(torch.empty(*mapping_shape, device=device))
    )

    triple_optimizer = torch.optim.Adam([entity_embeds, relation_embeds], lr=self.learning_rate)
    # `entity_embeds` is included here (not just the two mapping
    # matrices) to match OpenEA's `generate_optimizer(mapping_loss,
    # lr, var_list=None, opt='Adam')` -- see the module docstring.
    mapping_optimizer = torch.optim.Adam(
        [mapping_mat_1, mapping_mat_2, entity_embeds], lr=self.learning_rate
    )

    total_triples = len(ctx1.triples) + len(ctx2.triples)
    triple_steps = max(1, (total_triples + self.batch_size - 1) // self.batch_size)

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_structural_epoch(
            entity_embeds,
            relation_embeds,
            triple_optimizer,
            ctx1,
            ctx2,
            rng,
            self.batch_size,
            self.margin,
        )
        train_mapping_epoch(
            entity_embeds,
            mapping_mat_1,
            mapping_mat_2,
            mapping_optimizer,
            seed_source,
            seed_target,
            unlabeled_pool1,
            unlabeled_pool2,
            rng,
            triple_steps,
            self.alpha_1,
            self.alpha_2,
        )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
                current_mapping = mapping_mat_1.cpu().numpy()
            hits1 = validation_hits1(current_embeds, current_mapping, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
        final_mapping = mapping_mat_1.cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._mapping = final_mapping
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Projected via mapping_mat_1.

Source code in src/linkingtk/algorithms/ea/sea.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side.

    Projected via ``mapping_mat_1``.
    """
    return self._project(self._embedding(entity_id))

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side (unprojected).

Source code in src/linkingtk/algorithms/ea/sea.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side (unprojected)."""
    return self._embedding(entity_id)

SimpleHHEALinker

SimpleHHEALinker(
    name_model: str = "albert-base-v2",
    emb_size: int = 64,
    structure_size: int = 8,
    time_size: int = 8,
    use_structure: bool = True,
    walk_length: int = 80,
    num_walks: int = 10,
    structure_dim: int = 64,
    max_degree: int = 1000,
    structure_workers: int = 1,
    structure_epochs: int = 5,
    num_epochs: int = 1500,
    learning_rate: float = 0.01,
    weight_decay: float = 0.001,
    gamma: float = 1.0,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Fuses whitened name embeddings, node2vec structural embeddings, and (optionally) Time2Vec temporal embeddings for Entity Alignment.

Must be fit before link. Exposes source_embedding/target_embedding like KGELinker, so it's a drop-in for rank_exhaustive -- and, once #22 (ChatEA) lands, a candidate base_linker for it too.

Parameters:

Name	Type	Description	Default
name_model	str	Hugging Face model id for name embedding, loaded via AutoModel/AutoTokenizer.from_pretrained. Defaults to the reference's own albert-base-v2.	'albert-base-v2'
emb_size	int	Final fused embedding dimension.	64
structure_size	int	Post-projection size of the structure branch before concatenation into the final embedding.	8
time_size	int	Post-projection size of the time branch before concatenation into the final embedding.	8
use_structure	bool	Whether to compute and use node2vec structural embeddings at all (the reference's own --no_structure).	True
walk_length	int	node2vec walk hyperparameter -- see build_structure_embeddings.	80
num_walks	int	node2vec walk hyperparameter -- see build_structure_embeddings.	10
structure_dim	int	node2vec walk hyperparameter -- see build_structure_embeddings.	64
max_degree	int	node2vec walk hyperparameter -- see build_structure_embeddings.	1000
structure_workers	int	Worker threads for the structural Word2Vec fit -- see build_structure_embeddings. Walk simulation itself is always single-threaded, matching the reference's own default invocation.	1
structure_epochs	int	Word2Vec training epochs for the structural skip-gram model -- see build_structure_embeddings.	5
num_epochs	int	Training hyperparameter -- see train.	1500
learning_rate	float	Training hyperparameter -- see train.	0.01
weight_decay	float	Training hyperparameter -- see train.	0.001
gamma	float	Training hyperparameter -- see train.	1.0
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train/embed on, e.g. "cpu" (default) or "cuda".	'cpu'

Source code in src/linkingtk/algorithms/ea/simple_hhea.py

def __init__(
    self,
    name_model: str = "albert-base-v2",
    emb_size: int = 64,
    structure_size: int = 8,
    time_size: int = 8,
    use_structure: bool = True,
    walk_length: int = 80,
    num_walks: int = 10,
    structure_dim: int = 64,
    max_degree: int = 1000,
    structure_workers: int = 1,
    structure_epochs: int = 5,
    num_epochs: int = 1500,
    learning_rate: float = 0.01,
    weight_decay: float = 0.001,
    gamma: float = 1.0,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.name_model = name_model
    self.emb_size = emb_size
    self.structure_size = structure_size
    self.time_size = time_size
    self.use_structure = use_structure
    self.walk_length = walk_length
    self.num_walks = num_walks
    self.structure_dim = structure_dim
    self.max_degree = max_degree
    self.structure_workers = structure_workers
    self.structure_epochs = structure_epochs
    self.num_epochs = num_epochs
    self.learning_rate = learning_rate
    self.weight_decay = weight_decay
    self.gamma = gamma
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    temporal_triples: list[TemporalTriple] | None = None,
    random_state: int | None = None,
) -> SimpleHHEALinker

Trains the fusion model over dataset1/dataset2's entities.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities.	required
dataset2	list[Entity]	Target entities.	required
ground_truth	list[tuple[str, str]]	Known-correct (source_id, target_id) pairs. Doubles as both the margin-ranking training signal (like KGELinker.fit's ground_truth) and the node2vec node-merge seed (see build_structure_embeddings's train_pairs).	required
graph	Graph	Combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2)).	required
temporal_triples	list[TemporalTriple] | None	Optional (subject_id, relation_id, object_id, start_label, end_label) temporal facts (e.g. from IcewsWikiDataset.load_temporal_graphs), start_label/end_label as "YYYY-MM" strings or None. None (default) disables the time-feature branch entirely -- correct for non-ICEWS datasets, which have no temporal facts at all.	None
random_state	int | None	Seed for reproducible training.	None

Returns:

Type	Description
SimpleHHEALinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If no ground_truth pair's ids are actually present in dataset1/dataset2.

Source code in src/linkingtk/algorithms/ea/simple_hhea.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    temporal_triples: list[TemporalTriple] | None = None,
    random_state: int | None = None,
) -> SimpleHHEALinker:
    """Trains the fusion model over `dataset1`/`dataset2`'s entities.

    Args:
        dataset1: Source entities.
        dataset2: Target entities.
        ground_truth: Known-correct `(source_id, target_id)` pairs.
            Doubles as both the margin-ranking training signal (like
            [KGELinker.fit][linkingtk.algorithms.ea.kge.KGELinker.fit]'s
            `ground_truth`) and the node2vec node-merge seed (see
            [build_structure_embeddings][linkingtk.algorithms.ea._simple_hhea_structure.build_structure_embeddings]'s
            `train_pairs`).
        graph: Combined relational structure of both KGs (e.g.
            `to_triples(graph1) + to_triples(graph2)`).
        temporal_triples: Optional `(subject_id, relation_id,
            object_id, start_label, end_label)` temporal facts (e.g.
            from
            [IcewsWikiDataset.load_temporal_graphs][linkingtk.datasets.icews.IcewsWikiDataset.load_temporal_graphs]),
            `start_label`/`end_label` as `"YYYY-MM"` strings or
            `None`. `None` (default) disables the time-feature branch
            entirely -- correct for non-ICEWS datasets, which have no
            temporal facts at all.
        random_state: Seed for reproducible training.

    Returns:
        `self`, for chaining.

    Raises:
        LinkingTKError: If no `ground_truth` pair's ids are actually
            present in `dataset1`/`dataset2`.
    """
    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    train_pairs = [(s, t) for s, t in ground_truth if s in ids1 and t in ids2]
    if not train_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have ids present in `dataset1`/"
            "`dataset2`; fit() has no training pairs."
        )

    if random_state is not None:
        torch.manual_seed(random_state)

    entities = dataset1 + dataset2
    entity_ids = [entity.id for entity in entities]
    entity_to_index = {entity_id: index for index, entity_id in enumerate(entity_ids)}

    raw_name_emb = embed_names(entities, model_name=self.name_model, device=self.device)
    kernel, bias = compute_kernel_bias(raw_name_emb, self.emb_size)
    name_emb = whiten(raw_name_emb, kernel, bias)

    structure_emb = None
    if self.use_structure:
        triples = to_triples(graph)
        structure_by_id = build_structure_embeddings(
            entity_ids,
            triples,
            train_pairs,
            dimensions=self.structure_dim,
            walk_length=self.walk_length,
            num_walks=self.num_walks,
            max_degree=self.max_degree,
            workers=self.structure_workers,
            epochs=self.structure_epochs,
            random_state=random_state,
        )
        structure_emb = np.stack([structure_by_id[entity_id] for entity_id in entity_ids])

    time_emb = (
        build_time_histogram(entity_ids, temporal_triples)
        if temporal_triples is not None
        else None
    )

    resolved_device = resolve_device(self.device)
    model = SimpleHHEAModel(
        name_emb,
        time_emb,
        structure_emb,
        emb_size=self.emb_size,
        structure_size=self.structure_size,
        time_size=self.time_size,
        device=resolved_device,
    ).to(resolved_device)

    train_pairs_idx = np.array(
        [(entity_to_index[s], entity_to_index[t]) for s, t in train_pairs], dtype=np.int64
    )
    train_simple_hhea(
        model,
        train_pairs_idx,
        len(entity_ids),
        num_epochs=self.num_epochs,
        learning_rate=self.learning_rate,
        weight_decay=self.weight_decay,
        gamma=self.gamma,
        random_state=random_state,
    )

    model.eval()
    with torch.no_grad():
        embeddings = model().cpu().numpy()
    self._id_to_vector = {
        entity_id: embeddings[index] for index, entity_id in enumerate(entity_ids)
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/simple_hhea.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score `entity_id` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/simple_hhea.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score `entity_id` as a scored pair's target side."""
    return self._embedding(entity_id)

entmatcher

EntMatcher-style Entity Alignment linker.

References

EntMatcher: An Open-Source Library for Entity Alignment via Feature Matching in Knowledge Graphs. https://github.com/DexterZeng/EntMatcher (research repo, not a dependency of this project — see linkingtk.algorithms.feature_classifier for what's reused).

EntMatcherLinker

EntMatcherLinker()

Bases: FeatureClassifierLinker

EA linker combining hand-crafted similarity features with optimal one-to-one matching.

A preconfigured FeatureClassifierLinker (matching=OptimalMatcher()) — see that class for the feature set, training, and matching behavior. Kept as its own named class since it maps to a recognizable reference (EntMatcher's key insight: a globally optimal assignment can outperform independent per-source matching).

Source code in src/linkingtk/algorithms/ea/entmatcher.py

def __init__(self) -> None:
    super().__init__(matching=OptimalMatcher())

Knowledge Graph Embeddings

kge

Knowledge-graph-embedding linker for Entity Alignment.

Wraps pykeen (https://github.com/pykeen/pykeen) translational and bilinear embedding models -- TransE, DistMult, RotatE by default -- to score candidate pairs by embedding similarity rather than string/feature similarity. pykeen has no built-in Entity Alignment mode (it's a link-prediction library); the approach here is a simplified version of the MTransE/IPTransE family's core idea, not a reimplementation of either paper's full method (no separate per-KG embedding spaces plus a learned transform, no triple -swapping): both KGs' triples are combined into a single graph, the known training alignment pairs are added to it as extra (source_id, "__seed_alignment__", target_id) triples, and a single KGE model is trained jointly over all of it -- pulling aligned (and structurally analogous) entities' embeddings together via that one pseudo-relation.

KGELinker

KGELinker(
    model: str = "TransE",
    embedding_dim: int = 50,
    num_epochs: int = 50,
    batch_size: int | None = None,
    learning_rate: float = 0.01,
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

Scores candidate pairs by trained knowledge-graph-embedding similarity.

Must be fit before link can be called. Unlike FeatureClassifierLinker, training needs the KGs' relational structure, not just entity text -- see fit's graph argument.

Parameters:

Name	Type	Description	Default
model	str	A model name from pykeen's registry (pykeen.models.model_resolver, ~40 entries), e.g. "TransE", "DistMult", "RotatE". Not restricted to those three -- pykeen already provides its own named-model resolution, so there's nothing to add on top of it here -- but they're what this class is validated against; other models are best-effort, since not every pykeen model shares the same constructor keyword shape (e.g. some need more than embedding_dim).	'TransE'
embedding_dim	int	Dimensionality of the trained entity/relation embeddings.	50
num_epochs	int	Training epochs.	50
batch_size	int | None	Training batch size. Defaults to pykeen's own auto-selection (None).	None
learning_rate	float	Adam's learning rate. pykeen's own default (0.001, via torch.optim.Adam's default) converges too slowly for the seed-alignment triples to pull related entities together within a practical epoch budget at small scale; 0.01 converges reliably in testing.	0.01
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher. OptimalMatcher's dense Hungarian-algorithm cost matrix is explicitly not meant for KG-scale entity counts (see its docstring) -- avoid it here for anything but small graphs.	DEFAULT_MATCHER

Source code in src/linkingtk/algorithms/ea/kge.py

def __init__(
    self,
    model: str = "TransE",
    embedding_dim: int = 50,
    num_epochs: int = 50,
    batch_size: int | None = None,
    learning_rate: float = 0.01,
    matching: Matcher = DEFAULT_MATCHER,
) -> None:
    self.model = model
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.matching = matching
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
) -> KGELinker

Train KGE embeddings jointly over graph plus seed alignment triples.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities.	required
dataset2	list[Entity]	Target entities.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs, added to the training graph as seed alignment triples (see the module docstring).	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, pykeen trains non-deterministically.	None

Returns:

Type	Description
KGELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If no ground_truth pair's ids are actually present in dataset1/dataset2 -- nothing to seed training with.
OptionalDependencyError	If pykeen isn't installed.

Source code in src/linkingtk/algorithms/ea/kge.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
) -> KGELinker:
    """Train KGE embeddings jointly over ``graph`` plus seed alignment triples.

    Args:
        dataset1: Source entities.
        dataset2: Target entities.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs, added to the training graph as seed alignment
            triples (see the module docstring).
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must
            already be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left
            unspecified, pykeen trains non-deterministically.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If no ``ground_truth`` pair's ids are actually
            present in ``dataset1``/``dataset2`` -- nothing to seed
            training with.
        OptionalDependencyError: If pykeen isn't installed.
    """
    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    seed_pairs = [(s, t) for s, t in ground_truth if s in ids1 and t in ids2]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have ids present in `dataset1`/"
            "`dataset2`; fit() has no seed alignment triples to train with."
        )

    try:
        import torch
        from pykeen.models import model_resolver
        from pykeen.training import SLCWATrainingLoop
        from pykeen.triples import TriplesFactory
        from pykeen.utils import set_random_seed
        from torch.optim import Adam
    except ImportError as exc:
        raise OptionalDependencyError("KGELinker", "kge") from exc

    if random_state is not None:
        set_random_seed(random_state)

    triples = to_triples(graph) + [(s, _SEED_RELATION, t) for s, t in seed_pairs]
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped_triples = map_triples_to_ids(triples, entity_to_id, relation_to_id)
    triples_factory = TriplesFactory(
        mapped_triples=mapped_triples,
        entity_to_id=entity_to_id,
        relation_to_id=relation_to_id,
    )
    model: Any = model_resolver.make(
        self.model,
        triples_factory=triples_factory,
        embedding_dim=self.embedding_dim,
        random_seed=random_state,
    )
    training_loop = SLCWATrainingLoop(
        model=model,
        triples_factory=triples_factory,
        optimizer=Adam(params=model.get_grad_params(), lr=self.learning_rate),
    )
    training_loop.train(
        triples_factory=triples_factory,
        num_epochs=self.num_epochs,
        batch_size=self.batch_size,
        use_tqdm=False,
    )

    model.eval()
    with torch.no_grad():
        embeddings = model.entity_representations[0](indices=None)
    if embeddings.is_complex():
        embeddings = torch.view_as_real(embeddings).reshape(embeddings.shape[0], -1)
    embeddings_array = embeddings.detach().numpy()
    self._id_to_vector = {
        entity_id: embeddings_array[index]
        for entity_id, index in triples_factory.entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/kge.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/kge.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

MTransE

mtranse

MTransE knowledge-graph-embedding linker for Entity Alignment.

Chen, M., Tian, Y., Yang, M., & Zaniolo, C. (2017). Multilingual Knowledge Graph Embeddings for Cross-lingual Knowledge Alignment. IJCAI 2017. https://www.ijcai.org/proceedings/2017/0209.pdf

Unlike KGELinker's simplified single-shared-pseudo-relation trick, this is a faithful reimplementation of MTransE's actual training procedure, ported directly from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/mtranse.py and .../modules/base/mapping.py) rather than from a from-scratch reading of the paper, since OpenEA's actual algorithm differs from a naive paper reading in ways that matter for reproducing its published numbers:

• Entities and relations from both KGs share one embedding table (not two independently-trained per-KG spaces). Since each KG's own triples only ever reference that KG's own entities, this has the same practical effect as two separate spaces, but reuses this repo's existing combined-graph id-mapping utilities directly.
• The embedding step uses pure positive-triple loss (sum(||h+r-t||^2) over real triples only) with no negative sampling -- unlike KGELinker/pykeen's SLCWATrainingLoop. This is why training is hand-rolled in PyTorch here rather than routed through pykeen.
• Entities and relations are L2-normalized to the unit sphere on every forward pass (not just at initialization) -- this is what keeps the no-negative-sampling loss from collapsing to the trivial zero solution.
• Alignment signal comes from a learned square mapping matrix (orthogonally initialized, trained against a soft-orthogonality-regularized loss over seed pairs), not from injecting pseudo-triples into the embedding space itself. Critically, the mapping loss's optimizer updates both the mapping matrix and the shared entity embeddings themselves -- OpenEA's generate_optimizer(loss, lr, var_list=None, ...) leaves var_list unset, which makes TensorFlow's compute_gradients differentiate against every trainable variable connected to the loss, not just the mapping matrix. Restricting this port's mapping optimizer to only the mapping matrix (the more "obvious" reading of the algorithm) silently drops that second training signal and leaves the mapping loss converging roughly an order of magnitude slower.

MTransELinker

MTransELinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    alpha: float = 5.0,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a learned cross-KG mapping over shared TransE embeddings.

Must be fit before link can be called.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings and the (square) mapping matrix. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each alternating one epoch of triple-loss training with one epoch of mapping-loss training. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for triple-loss training; the per-epoch mapping-loss batch size is derived from it (matching OpenEA: the same number of mini-batches per epoch for both).	5000
learning_rate	float	Adagrad's learning rate, for both the embedding and mapping optimizers -- OpenEA's published value is 0.01.	0.01
alpha	float	Weight applied to the mapping loss (approximation error plus orthogonality regularization). OpenEA's published value is 5.0.	5.0
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/mtranse.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    alpha: float = 5.0,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.alpha = alpha
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._mapping: npt.NDArray[np.floating[Any]] | None = None
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> MTransELinker

Train shared TransE embeddings plus a learned cross-KG mapping.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities. Unused beyond mirroring KGELinker.fit's signature -- which entities get embeddings is determined entirely by graph.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs used to train the mapping matrix. Unlike KGELinker, these are not injected into graph as extra triples.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
MTransELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples -- nothing to train the mapping with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/mtranse.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> MTransELinker:
    """Train shared TransE embeddings plus a learned cross-KG mapping.

    Args:
        dataset1: Source entities. Unused beyond mirroring
            [KGELinker.fit][linkingtk.algorithms.ea.kge.KGELinker.fit]'s
            signature -- which entities get embeddings is determined
            entirely by ``graph``.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs used to train the mapping matrix. Unlike
            ``KGELinker``, these are *not* injected into ``graph`` as
            extra triples.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must
            already be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples -- nothing to train
            the mapping with -- or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("MTransELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)

    triples = to_triples(graph)
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped_triples = map_triples_to_ids(triples, entity_to_id, relation_to_id)
    heads = torch.from_numpy(mapped_triples[:, 0]).long().to(device)
    rels = torch.from_numpy(mapped_triples[:, 1]).long().to(device)
    tails = torch.from_numpy(mapped_triples[:, 2]).long().to(device)

    seed_pairs = [(s, t) for s, t in ground_truth if s in entity_to_id and t in entity_to_id]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train MTransE's mapping with."
        )
    seed_source = torch.tensor([entity_to_id[s] for s, _ in seed_pairs], device=device)
    seed_target = torch.tensor([entity_to_id[t] for _, t in seed_pairs], device=device)

    entity_embeds = torch.nn.Parameter(
        functional.normalize(
            torch.randn(len(entity_to_id), self.embedding_dim, device=device), dim=1
        )
    )
    relation_embeds = torch.nn.Parameter(
        functional.normalize(
            torch.randn(len(relation_to_id), self.embedding_dim, device=device), dim=1
        )
    )
    mapping_mat = torch.nn.Parameter(
        torch.nn.init.orthogonal_(
            torch.empty(self.embedding_dim, self.embedding_dim, device=device)
        )
    )
    eye = torch.eye(self.embedding_dim, device=device)

    triple_optimizer = torch.optim.Adagrad(
        [entity_embeds, relation_embeds], lr=self.learning_rate
    )
    # `entity_embeds` is included here (not just `mapping_mat`) to match
    # OpenEA's `generate_optimizer(mapping_loss, lr, var_list=None, ...)`,
    # which implicitly differentiates against every trainable variable
    # the loss touches -- see the module docstring.
    mapping_optimizer = torch.optim.Adagrad([mapping_mat, entity_embeds], lr=self.learning_rate)

    triple_steps = max(1, math.ceil(len(heads) / self.batch_size))
    seed_batch_size = max(1, len(seed_source) // triple_steps)

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        perm = torch.randperm(len(heads), device=device)
        for step in range(triple_steps):
            idx = perm[step * self.batch_size : (step + 1) * self.batch_size]
            h = functional.normalize(entity_embeds[heads[idx]], dim=1)
            r = functional.normalize(relation_embeds[rels[idx]], dim=1)
            t = functional.normalize(entity_embeds[tails[idx]], dim=1)
            triple_loss = ((h + r - t) ** 2).sum()
            triple_optimizer.zero_grad()
            triple_loss.backward()  # type: ignore[no-untyped-call]
            triple_optimizer.step()

        for _step in range(triple_steps):
            batch_idx = torch.randint(0, len(seed_source), (seed_batch_size,), device=device)
            s_vecs = functional.normalize(entity_embeds[seed_source[batch_idx]], dim=1)
            t_vecs = functional.normalize(entity_embeds[seed_target[batch_idx]], dim=1)
            map_loss = ((t_vecs - s_vecs @ mapping_mat) ** 2).sum()
            orth_loss = ((mapping_mat @ mapping_mat.T - eye) ** 2).sum()
            mapping_loss = self.alpha * (map_loss + orth_loss)
            mapping_optimizer.zero_grad()
            mapping_loss.backward()  # type: ignore[no-untyped-call]
            mapping_optimizer.step()

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
                current_mapping = mapping_mat.cpu().numpy()
            hits1 = _validation_hits1(current_embeds, current_mapping, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
        final_mapping = mapping_mat.cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._mapping = final_mapping
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side (projected).

Source code in src/linkingtk/algorithms/ea/mtranse.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side (projected)."""
    return self._project(self._embedding(entity_id))

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side (unprojected).

Source code in src/linkingtk/algorithms/ea/mtranse.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side (unprojected)."""
    return self._embedding(entity_id)

IPTransE

iptranse

IPTransE knowledge-graph-embedding linker for Entity Alignment.

Zhu, H., Xu, R., Yang, Y., Lin, X., & Cheng, X. (2017). Iterative Entity Alignment via Joint Knowledge Embeddings. IJCAI 2017. https://www.ijcai.org/proceedings/2017/0595.pdf

Like MTransELinker, this is a faithful reimplementation of IPTransE's actual training procedure, ported directly from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/iptranse.py and its modules/bootstrapping, modules/train/batch.py, modules/load/read.py support code) rather than from a from-scratch reading of the paper. Three things distinguish it from MTransE:

• Seed pairs share one embedding row, not two rows plus a learned mapping (OpenEA's alignment_module="sharing"): a seed pair's target entity is aliased to its source entity's id at training-tensor-build time. Alignment signal for held-out entities comes entirely from this shared space plus the two mechanisms below, not from any post-hoc projection.
• Training loss is margin-based (margin=1.5, one negative per positive, corrupting the head or tail from that triple's own KG), not MTransE's negative-sampling-free positive-only loss -- plus a relation-composition ("path") loss over 2-hop chains within each KG, weighted by path_parm, that encourages composed relations (r_x + r_y) to approximate a direct edge (r) when both exist.
• Iterative bootstrapping: every bootstrap_every epochs, an unsupervised self-training round computes embedding similarity between every not-yet-seeded entity on each side, keeps each source entity's best match if it clears sim_th, and turns each newly-found pair's real edges into weighted pseudo-triples trained for one extra epoch via a separate optimizer -- this is what "iterative" refers to in the method's name.

Ported into two private helper modules: _iptranse_training.py (plain numpy, independently testable without torch) and _iptranse_torch.py (the training-step functions that build/consume PyTorch tensors).

Deviations from OpenEA's own code, beyond the negative-sampling/path-generation simplifications documented in those two modules:

• OpenEA partitions relation triples into two KGs via its own KGs loader; this repo's EA linker family instead takes one combined graph in fit(). IPTransE genuinely needs the KG1/KG2 boundary (per-KG negative sampling, path generation, pseudo-triple dicts), so it's re-derived here from dataset1/dataset2's own id sets: a triple belongs to KG1 iff its head id is one of dataset1's ids. This relies on the same assumption MTransELinker already documents (ids are disjoint across the two KGs) for a new purpose.
• OpenEA's bootstrap reference pool is valid_entities + test_entities from its own loader (entities only, no labels needed). This repo's fit() only receives dataset1/dataset2 (full entity rosters) and ground_truth (the train split) -- so the reference pool here is every dataset1/dataset2 entity not already a ground_truth source/target, which is broader than OpenEA's val+test-only pool. At real dataset scale this is a dense similarity matrix over most of both KGs' entities; bootstrap_pool_size lets callers cap it (random subsample, fixed for the whole run) if that's a problem.

IPTransELinker

IPTransELinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    path_parm: float = 0.1,
    sim_th: float = 0.7,
    bootstrap_every: int = 100,
    bootstrap_pool_size: int | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a shared TransE space, path loss, and bootstrapping.

Must be fit before link can be called. See the module docstring for how this differs from MTransELinker.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each one pass over both KGs' triples (plus path-loss minibatches). OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for triple-loss training, split proportionally between the two KGs by triple count each epoch (matching OpenEA).	5000
learning_rate	float	Adagrad's learning rate, for both the main and bootstrap optimizers -- OpenEA's published value is 0.01.	0.01
margin	float	Hinge margin for the triple and path losses. OpenEA's published value is 1.5.	1.5
path_parm	float	Weight applied to the path (relation-composition) loss. OpenEA's published value is 0.1.	0.1
sim_th	float	Minimum embedding similarity for a bootstrap round to accept a newly found pair. OpenEA's published value is 0.7.	0.7
bootstrap_every	int	Run a bootstrapping round every this many epochs. OpenEA's published value (bp_freq) is 100.	100
bootstrap_pool_size	int | None	Cap on each side's bootstrap reference pool size (random subsample, fixed for the whole run) -- see the module docstring's deviation note. None (default) uses the full pool.	None
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/iptranse.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    path_parm: float = 0.1,
    sim_th: float = 0.7,
    bootstrap_every: int = 100,
    bootstrap_pool_size: int | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.margin = margin
    self.path_parm = path_parm
    self.sim_th = sim_th
    self.bootstrap_every = bootstrap_every
    self.bootstrap_pool_size = bootstrap_pool_size
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> IPTransELinker

Train a shared TransE space with margin+path loss and periodic bootstrapping.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (see the module docstring).	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs. Seeds the shared embedding table (each pair's target is aliased to its source's id) and is excluded from the bootstrap reference pool.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
IPTransELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in dataset1/dataset2 -- nothing to seed the shared embedding table with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/iptranse.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> IPTransELinker:
    """Train a shared TransE space with margin+path loss and periodic bootstrapping.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (see the module docstring).
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs. Seeds the shared embedding table (each pair's
            target is aliased to its source's id) and is excluded
            from the bootstrap reference pool.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must
            already be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``dataset1``/``dataset2`` -- nothing to
            seed the shared embedding table with -- or if ``device``
            is invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("IPTransELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    seed_pairs = [(s, t) for s, t in ground_truth if s in ids1 and t in ids2]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in "
            "`dataset1`/`dataset2`; fit() has no seed pairs to train "
            "IPTransE's shared embedding table with."
        )

    entity_to_id, relation_to_id = build_shared_id_mappings(
        triples1_labels + triples2_labels, seed_pairs
    )
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(triples1_labels, mapped1, relation_to_id)
    ctx2 = build_kg_context(triples2_labels, mapped2, relation_to_id)

    pool1_ids, pool2_ids = self._bootstrap_pools(
        dataset1, dataset2, entity_to_id, seed_pairs, rng
    )
    combined_entity_pool = np.union1d(ctx1.entity_pool, ctx2.entity_pool)

    entity_embeds, relation_embeds = self._init_embeddings(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    optimizer = torch.optim.Adagrad([entity_embeds, relation_embeds], lr=self.learning_rate)
    alignment_optimizer = torch.optim.Adagrad(
        [entity_embeds, relation_embeds], lr=self.learning_rate
    )

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_epoch(
            entity_embeds,
            relation_embeds,
            optimizer,
            ctx1,
            ctx2,
            rng,
            self.batch_size,
            self.margin,
            self.path_parm,
        )

        if (epoch + 1) % self.bootstrap_every == 0:
            bootstrap_round(
                entity_embeds,
                relation_embeds,
                alignment_optimizer,
                ctx1,
                ctx2,
                pool1_ids,
                pool2_ids,
                combined_entity_pool,
                rng,
                self.batch_size,
                self.margin,
                self.sim_th,
            )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/iptranse.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/iptranse.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

_iptranse_training

Private training helpers for IPTransELinker.

Pure, independently-testable functions ported from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA -- approaches/iptranse.py, modules/bootstrapping/alignment_finder.py, modules/train/batch.py, modules/load/read.py). Not part of the public API -- see iptranse.py for the orchestrating class and the deviations from OpenEA's own code documented in its docstring.

build_shared_id_mappings

build_shared_id_mappings(
    triples: list[Triple], seed_pairs: list[tuple[str, str]]
) -> tuple[dict[str, int], dict[str, int]]

Entity/relation id mappings with IPTransE's "sharing" alignment semantics.

Every seed pair's target entity is aliased to its source entity's id -- both get exactly one embedding row, not two independently-trained rows linked by a learned mapping (contrast MTransELinker). Relations are not merged across KGs (mirrors OpenEA's generate_sharing_id([], ...) fallthrough for relations -- the same plain sorted-union behavior build_id_mappings already implements, since real relation label strings never coincide across the two KGs in these datasets).

Parameters:

Name	Type	Description	Default
triples	list[Triple]	Combined triples from both KGs.	required
seed_pairs	list[tuple[str, str]]	(source_id, target_id) pairs to alias together. Pairs whose ids aren't present in triples are ignored.	required

Returns:

Type	Description
dict[str, int]	(entity_to_id, relation_to_id). The entity id range has
dict[str, int]	exactly len(entities) - len(usable seed pairs) distinct
tuple[dict[str, int], dict[str, int]]	values, matching OpenEA's kgs.entities_num after merging.

Source code in src/linkingtk/algorithms/ea/_iptranse_training.py

def build_shared_id_mappings(
    triples: list[Triple],
    seed_pairs: list[tuple[str, str]],
) -> tuple[dict[str, int], dict[str, int]]:
    """Entity/relation id mappings with IPTransE's "sharing" alignment semantics.

    Every seed pair's target entity is aliased to its source entity's id --
    both get exactly one embedding row, not two independently-trained rows
    linked by a learned mapping (contrast
    [MTransELinker][linkingtk.algorithms.ea.mtranse.MTransELinker]).
    Relations are *not* merged across KGs (mirrors OpenEA's
    ``generate_sharing_id([], ...)`` fallthrough for relations -- the same
    plain sorted-union behavior
    [build_id_mappings][linkingtk.utils.graph.build_id_mappings] already
    implements, since real relation label strings never coincide across
    the two KGs in these datasets).

    Args:
        triples: Combined triples from both KGs.
        seed_pairs: ``(source_id, target_id)`` pairs to alias together.
            Pairs whose ids aren't present in ``triples`` are ignored.

    Returns:
        ``(entity_to_id, relation_to_id)``. The entity id range has
        exactly ``len(entities) - len(usable seed pairs)`` distinct
        values, matching OpenEA's ``kgs.entities_num`` after merging.
    """
    entities = sorted({s for s, _, _ in triples} | {o for _, _, o in triples})
    relations = sorted({p for _, p, _ in triples})
    relation_to_id = {label: index for index, label in enumerate(relations)}

    entity_set = set(entities)
    alias = {
        target: source
        for source, target in seed_pairs
        if source in entity_set and target in entity_set
    }

    def canonical(entity: str) -> str:
        return alias.get(entity, entity)

    representatives = sorted({canonical(e) for e in entities})
    representative_to_id = {label: index for index, label in enumerate(representatives)}
    entity_to_id = {e: representative_to_id[canonical(e)] for e in entities}
    return entity_to_id, relation_to_id

find_mutual_pairs

find_mutual_pairs(
    sim_mat: NDArray[floating[Any]], sim_th: float = 0.0
) -> list[tuple[int, int, float]]

Reciprocal top-1 match finder for bootstrapping under saturated similarities.

Keeps (i, j) only if j is row i's best match and i is column j's best match (both argmaxes agree), unlike find_new_pairs's one-sided row-argmax. Added for KDCoE's description pathway (#29): once a similarity matrix saturates near 1.0 for nearly every entry (as the description encoder's cosine similarities do on real EN-FR-15K-V1 -- median ~0.9998), an absolute threshold like desc_sim_th can't discriminate anything -- almost every row clears it. Reciprocity stays useful under saturation because the argmax positions (which column is largest per row, which row is largest per column) still vary even when the values themselves are nearly indistinguishable, and requiring both sides to agree rejects most of the one-sided noise a lone threshold can't. sim_th is kept only as a coarse sanity floor (default 0.0, i.e. no real filtering) -- reciprocity is the operative filter here, not the threshold.

Parameters:

Name	Type	Description	Default
sim_mat	NDArray[floating[Any]]	(n1, n2) similarity matrix (e.g. dot products of unit-normalized embeddings).	required
sim_th	float	Minimum similarity for a reciprocal match to count. Defaults to 0.0 (no real floor) since reciprocity itself is the filter this function is for.	0.0

Returns:

Type	Description
list[tuple[int, int, float]]	(row_index, col_index, similarity) tuples for every reciprocal
list[tuple[int, int, float]]	best-match pair clearing sim_th. Empty if sim_mat is empty.

Source code in src/linkingtk/algorithms/ea/_iptranse_training.py

def find_mutual_pairs(
    sim_mat: npt.NDArray[np.floating[Any]],
    sim_th: float = 0.0,
) -> list[tuple[int, int, float]]:
    """Reciprocal top-1 match finder for bootstrapping under saturated similarities.

    Keeps ``(i, j)`` only if ``j`` is row ``i``'s best match **and** ``i`` is
    column ``j``'s best match (both argmaxes agree), unlike
    [find_new_pairs][linkingtk.algorithms.ea._iptranse_training.find_new_pairs]'s
    one-sided row-argmax. Added for KDCoE's description pathway (#29): once a
    similarity matrix saturates near ``1.0`` for nearly every entry (as the
    description encoder's cosine similarities do on real EN-FR-15K-V1 --
    median ~0.9998), an absolute threshold like ``desc_sim_th`` can't
    discriminate anything -- almost every row clears it. Reciprocity stays
    useful under saturation because the *argmax positions* (which column is
    largest per row, which row is largest per column) still vary even when
    the values themselves are nearly indistinguishable, and requiring both
    sides to agree rejects most of the one-sided noise a lone threshold
    can't. ``sim_th`` is kept only as a coarse sanity floor (default ``0.0``,
    i.e. no real filtering) -- reciprocity is the operative filter here, not
    the threshold.

    Args:
        sim_mat: ``(n1, n2)`` similarity matrix (e.g. dot products of
            unit-normalized embeddings).
        sim_th: Minimum similarity for a reciprocal match to count. Defaults
            to ``0.0`` (no real floor) since reciprocity itself is the
            filter this function is for.

    Returns:
        ``(row_index, col_index, similarity)`` tuples for every reciprocal
        best-match pair clearing ``sim_th``. Empty if ``sim_mat`` is empty.
    """
    if sim_mat.size == 0:
        return []
    best_col = np.argmax(sim_mat, axis=1)
    best_row = np.argmax(sim_mat, axis=0)
    rows = np.arange(sim_mat.shape[0])
    mutual = rows[best_row[best_col[rows]] == rows]
    mutual = mutual[sim_mat[mutual, best_col[mutual]] > sim_th]
    return [(int(i), int(best_col[i]), float(sim_mat[i, best_col[i]])) for i in mutual]

find_new_pairs

find_new_pairs(
    sim_mat: NDArray[floating[Any]], sim_th: float
) -> list[tuple[int, int, float]]

Row-wise top-1 + threshold match finder for bootstrapping.

For each row i, keeps (i, argmax_j sim[i, j]) iff sim[i, argmax_j] > sim_th. Reimplements OpenEA's find_potential_alignment_greedily (= find_alignment(sim_mat, sim_th, k=1)) directly with numpy -- with k=1, intersecting "threshold-filtered pairs" with "each row's top-1 neighbor" reduces exactly to this row-wise argmax+threshold check, so there's no need to port filter_sim_mat/ search_nearest_k's two-pass set-intersection, or the unused (by IPTransE) igraph/graph-tool-based find_potential_alignment_mwgm variant.

This is not mutual/bipartite matching -- two different rows can map to the same column.

Parameters:

Name	Type	Description	Default
sim_mat	NDArray[floating[Any]]	(n1, n2) similarity matrix (e.g. dot products of unit-normalized embeddings).	required
sim_th	float	Minimum similarity for a row's best match to count.	required

Returns:

Type	Description
list[tuple[int, int, float]]	(row_index, col_index, similarity) tuples for every row whose
list[tuple[int, int, float]]	best match clears sim_th. Empty if sim_mat is empty.

Source code in src/linkingtk/algorithms/ea/_iptranse_training.py

def find_new_pairs(
    sim_mat: npt.NDArray[np.floating[Any]],
    sim_th: float,
) -> list[tuple[int, int, float]]:
    """Row-wise top-1 + threshold match finder for bootstrapping.

    For each row ``i``, keeps ``(i, argmax_j sim[i, j])`` iff
    ``sim[i, argmax_j] > sim_th``. Reimplements OpenEA's
    ``find_potential_alignment_greedily`` (=
    ``find_alignment(sim_mat, sim_th, k=1)``) directly with numpy -- with
    ``k=1``, intersecting "threshold-filtered pairs" with "each row's
    top-1 neighbor" reduces exactly to this row-wise argmax+threshold
    check, so there's no need to port ``filter_sim_mat``/
    ``search_nearest_k``'s two-pass set-intersection, or the unused (by
    IPTransE) igraph/graph-tool-based ``find_potential_alignment_mwgm``
    variant.

    This is **not** mutual/bipartite matching -- two different rows can
    map to the same column.

    Args:
        sim_mat: ``(n1, n2)`` similarity matrix (e.g. dot products of
            unit-normalized embeddings).
        sim_th: Minimum similarity for a row's best match to count.

    Returns:
        ``(row_index, col_index, similarity)`` tuples for every row whose
        best match clears ``sim_th``. Empty if ``sim_mat`` is empty.
    """
    if sim_mat.size == 0:
        return []
    best_col = np.argmax(sim_mat, axis=1)
    best_val = sim_mat[np.arange(sim_mat.shape[0]), best_col]
    rows = np.nonzero(best_val > sim_th)[0]
    return [(int(i), int(best_col[i]), float(best_val[i])) for i in rows]

generate_two_step_paths

generate_two_step_paths(
    triples: list[Triple],
) -> list[tuple[str, str, str, int]]

2-hop relation-composition paths within one KG's own triples.

For every chain h_x --r_x--> m --r_y--> t_y where a direct edge h_x --r--> t_y also exists, emits (r_x, r_y, r, path_weight) where path_weight = fanout(h_x, r_x) * fanout(m, r_y) -- a proxy for how generic/common this 2-hop composition is. Chains with path_weight >= 101 are dropped (OpenEA's own hub-node cutoff).

Mirrors OpenEA's generate_2steps_path (approaches/iptranse.py), reimplemented with dict grouping instead of a pandas groupby/self-join/inner-join pipeline -- this repo has no pandas dependency and none should be added for this.

Parameters:

Name	Type	Description	Default
triples	list[Triple]	One KG's own triples (not combined across KGs -- paths are meaningless across disjoint entity spaces).	required

Returns:

Type	Description
list[tuple[str, str, str, int]]	(r_x, r_y, r, path_weight) tuples, relation labels still
list[tuple[str, str, str, int]]	unmapped to ids.

Source code in src/linkingtk/algorithms/ea/_iptranse_training.py

def generate_two_step_paths(triples: list[Triple]) -> list[tuple[str, str, str, int]]:
    """2-hop relation-composition paths within one KG's own triples.

    For every chain ``h_x --r_x--> m --r_y--> t_y`` where a *direct* edge
    ``h_x --r--> t_y`` also exists, emits ``(r_x, r_y, r, path_weight)``
    where ``path_weight = fanout(h_x, r_x) * fanout(m, r_y)`` -- a proxy
    for how generic/common this 2-hop composition is. Chains with
    ``path_weight >= 101`` are dropped (OpenEA's own hub-node cutoff).

    Mirrors OpenEA's ``generate_2steps_path``
    (``approaches/iptranse.py``), reimplemented with dict grouping instead
    of a pandas groupby/self-join/inner-join pipeline -- this repo has no
    pandas dependency and none should be added for this.

    Args:
        triples: One KG's own triples (not combined across KGs -- paths
            are meaningless across disjoint entity spaces).

    Returns:
        ``(r_x, r_y, r, path_weight)`` tuples, relation labels still
        unmapped to ids.
    """
    fanout: dict[tuple[str, str], int] = defaultdict(int)
    by_head_rel: dict[tuple[str, str], list[str]] = defaultdict(list)
    by_head: dict[str, list[tuple[str, str]]] = defaultdict(list)
    direct_edge: dict[tuple[str, str], list[str]] = defaultdict(list)

    for h, r, t in triples:
        fanout[(h, r)] += 1
        by_head_rel[(h, r)].append(t)
        by_head[h].append((r, t))
        direct_edge[(h, t)].append(r)

    paths: list[tuple[str, str, str, int]] = []
    for (h_x, r_x), mids in by_head_rel.items():
        size_x = fanout[(h_x, r_x)]
        for m in mids:
            for r_y, t_y in by_head.get(m, ()):
                weight = size_x * fanout[(m, r_y)]
                if weight >= 101:
                    continue
                paths.extend((r_x, r_y, r, weight) for r in direct_edge.get((h_x, t_y), ()))
    return paths

pseudo_triples_for_pair

pseudo_triples_for_pair(
    entity_a: int,
    entity_b: int,
    weight: float,
    by_head: dict[int, list[tuple[int, int]]],
    by_tail: dict[int, list[tuple[int, int]]],
) -> set[tuple[int, int, int, float]]

Weighted pseudo-triples generated from a newly bootstrapped pair.

If entity_a and entity_b are plausibly the same real-world entity, each of entity_a's edges plausibly also holds with entity_b substituted in. by_head/by_tail must be entity_a's own KG's edge dicts; call this twice per pair (once per direction, with by_head/by_tail swapped for the other KG's dicts and the arguments swapped) to cover both entities' edges, per OpenEA's generate_triples_of_latent_ents.

Parameters:

Name	Type	Description	Default
entity_a	int	The entity whose real edges are being copied over.	required
entity_b	int	The (bootstrapped) counterpart substituted into those edges.	required
weight	float	The pair's similarity score, carried through as the pseudo-triple's training weight.	required
by_head	dict[int, list[tuple[int, int]]]	entity_a's own KG's head -> [(relation, tail)] index.	required
by_tail	dict[int, list[tuple[int, int]]]	entity_a's own KG's tail -> [(head, relation)] index.	required

Returns:

Type	Description
set[tuple[int, int, int, float]]	(head, relation, tail, weight) pseudo-triples with
set[tuple[int, int, int, float]]	entity_a replaced by entity_b.

Source code in src/linkingtk/algorithms/ea/_iptranse_training.py

def pseudo_triples_for_pair(
    entity_a: int,
    entity_b: int,
    weight: float,
    by_head: dict[int, list[tuple[int, int]]],
    by_tail: dict[int, list[tuple[int, int]]],
) -> set[tuple[int, int, int, float]]:
    """Weighted pseudo-triples generated from a newly bootstrapped pair.

    If ``entity_a`` and ``entity_b`` are plausibly the same real-world
    entity, each of ``entity_a``'s edges plausibly also holds with
    ``entity_b`` substituted in. ``by_head``/``by_tail`` must be
    ``entity_a``'s *own* KG's edge dicts; call this twice per pair (once
    per direction, with ``by_head``/``by_tail`` swapped for the other
    KG's dicts and the arguments swapped) to cover both entities' edges,
    per OpenEA's ``generate_triples_of_latent_ents``.

    Args:
        entity_a: The entity whose real edges are being copied over.
        entity_b: The (bootstrapped) counterpart substituted into those
            edges.
        weight: The pair's similarity score, carried through as the
            pseudo-triple's training weight.
        by_head: ``entity_a``'s own KG's ``head -> [(relation, tail)]``
            index.
        by_tail: ``entity_a``'s own KG's ``tail -> [(head, relation)]``
            index.

    Returns:
        ``(head, relation, tail, weight)`` pseudo-triples with
        ``entity_a`` replaced by ``entity_b``.
    """
    newly: set[tuple[int, int, int, float]] = set()
    for r, t in by_head.get(entity_a, ()):
        newly.add((entity_b, r, t, weight))
    for h, r in by_tail.get(entity_a, ()):
        newly.add((h, r, entity_b, weight))
    return newly

sample_negative_path_relations

sample_negative_path_relations(
    positive_paths: NDArray[int64],
    relation_pool: NDArray[int64],
    rng: Generator,
) -> npt.NDArray[np.int64]

Corrupt a path batch's composed relation for the path loss's negatives.

Replaces column 2 (r, the composed/direct relation) with a uniformly random relation from relation_pool -- no real-triple exclusion check, mirroring OpenEA's generate_neg_paths, which has none either.

Parameters:

Name	Type	Description	Default
positive_paths	NDArray[int64]	(n, 4) int64 array of (r_x, r_y, r, weight) rows.	required
relation_pool	NDArray[int64]	That KG's own relation ids to sample replacements from.	required
rng	Generator	Random generator, for reproducibility.	required

Returns:

Type	Description
NDArray[int64]	(n, 4) int64 array, same shape as positive_paths, with
NDArray[int64]	column 2 replaced.

Source code in src/linkingtk/algorithms/ea/_iptranse_training.py

def sample_negative_path_relations(
    positive_paths: npt.NDArray[np.int64],
    relation_pool: npt.NDArray[np.int64],
    rng: np.random.Generator,
) -> npt.NDArray[np.int64]:
    """Corrupt a path batch's composed relation for the path loss's negatives.

    Replaces column 2 (``r``, the composed/direct relation) with a
    uniformly random relation from ``relation_pool`` -- no real-triple
    exclusion check, mirroring OpenEA's ``generate_neg_paths``, which has
    none either.

    Args:
        positive_paths: ``(n, 4)`` int64 array of ``(r_x, r_y, r, weight)``
            rows.
        relation_pool: That KG's own relation ids to sample replacements
            from.
        rng: Random generator, for reproducibility.

    Returns:
        ``(n, 4)`` int64 array, same shape as ``positive_paths``, with
        column 2 replaced.
    """
    negatives = positive_paths.copy()
    negatives[:, 2] = rng.choice(relation_pool, size=len(positive_paths))
    return negatives

sample_negative_triples

sample_negative_triples(
    positive_triples: NDArray[int64],
    entity_pool: NDArray[int64],
    real_triples: set[tuple[int, int, int]],
    rng: Generator,
    max_tries: int = 10,
) -> npt.NDArray[np.int64]

Corrupt one endpoint of each positive triple to build a negative.

For each row, a fair coin decides whether the head or tail is corrupted; the replacement is drawn uniformly from entity_pool (that triple's own KG's entities for the main triple loss, or the combined KG1+KG2 pool for the bootstrap/alignment loss -- see callers). Best-effort avoids reproducing a real triple: retries up to max_tries, then accepts the last draw regardless.

Deliberately simplified from OpenEA's vectorized generate_neg_triples_fast to a plain per-row retry loop, for readability -- same statistical behavior, same precedent as mtranse.py's _validation_hits1 simplification.

Parameters:

Name	Type	Description	Default
positive_triples	NDArray[int64]	(n, 3) int64 array of (head, rel, tail) ids.	required
entity_pool	NDArray[int64]	Candidate replacement entity ids.	required
real_triples	set[tuple[int, int, int]]	Set of (head, rel, tail) id tuples to avoid reproducing. Pass an empty set to skip avoidance entirely (matches OpenEA's alignment-loss negative sampling, which has no such check).	required
rng	Generator	Random generator, for reproducibility.	required
max_tries	int	Retries per row before giving up and accepting the last (possibly still-real) draw.	10

Returns:

Type	Description
NDArray[int64]	(n, 3) int64 array of corrupted triples, same shape as
NDArray[int64]	positive_triples.

Source code in src/linkingtk/algorithms/ea/_iptranse_training.py

def sample_negative_triples(
    positive_triples: npt.NDArray[np.int64],
    entity_pool: npt.NDArray[np.int64],
    real_triples: set[tuple[int, int, int]],
    rng: np.random.Generator,
    max_tries: int = 10,
) -> npt.NDArray[np.int64]:
    """Corrupt one endpoint of each positive triple to build a negative.

    For each row, a fair coin decides whether the head or tail is
    corrupted; the replacement is drawn uniformly from ``entity_pool``
    (that triple's own KG's entities for the main triple loss, or the
    combined KG1+KG2 pool for the bootstrap/alignment loss -- see
    callers). Best-effort avoids reproducing a real triple: retries up to
    ``max_tries``, then accepts the last draw regardless.

    Deliberately simplified from OpenEA's vectorized
    ``generate_neg_triples_fast`` to a plain per-row retry loop, for
    readability -- same statistical behavior, same precedent as
    ``mtranse.py``'s ``_validation_hits1`` simplification.

    Args:
        positive_triples: ``(n, 3)`` int64 array of ``(head, rel, tail)``
            ids.
        entity_pool: Candidate replacement entity ids.
        real_triples: Set of ``(head, rel, tail)`` id tuples to avoid
            reproducing. Pass an empty set to skip avoidance entirely
            (matches OpenEA's alignment-loss negative sampling, which has
            no such check).
        rng: Random generator, for reproducibility.
        max_tries: Retries per row before giving up and accepting the
            last (possibly still-real) draw.

    Returns:
        ``(n, 3)`` int64 array of corrupted triples, same shape as
        ``positive_triples``.
    """
    negatives = positive_triples.copy()
    corrupt_head = rng.random(len(positive_triples)) < 0.5
    for i in range(len(positive_triples)):
        h, r, t = (int(x) for x in positive_triples[i])
        for _ in range(max_tries):
            if corrupt_head[i]:
                h = int(rng.choice(entity_pool))
            else:
                t = int(rng.choice(entity_pool))
            if (h, r, t) not in real_triples:
                break
        negatives[i] = (h, r, t)
    return negatives

JAPE

jape

JAPE knowledge-graph-embedding linker for Entity Alignment.

Sun, Z., Hu, W., & Li, C. (2017). Cross-Lingual Entity Alignment via Joint Attribute-Preserving Embedding. ISWC 2017. https://link.springer.com/chapter/10.1007/978-3-319-68288-4_37

Like MTransELinker and IPTransELinker, this is a faithful reimplementation of JAPE's actual training procedure, ported directly from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/jape.py and .../attr2vec.py) rather than from a from-scratch reading of the paper. JAPE combines two signals:

• Structural embedding: the same shared-id ("sharing") mechanism as IPTransELinker -- reuses _iptranse_training.py's build_shared_id_mappings and sample_negative_triples directly -- trained via sum(||h+r-t||^2) - neg_alpha * sum(||h'+r-t'||^2). Unlike MTransE/IPTransE, this has no margin/hinge -- it's OpenEA's actual (unbounded-below) formula for JAPE specifically, not a simplification.
• Attribute-correlation embedding: a skip-gram-style model (ported from OpenEA's Attr2Vec) trained over which attribute predicates co-occur on the same entity -- entities that are a known seed pair have their attribute vocabularies merged before training, which is what makes this cross-lingual-aware. The trained attribute embeddings produce a per-entity attribute vector (mean-pooled, L2-normalized) and a similarity matrix between the two KGs' not-yet-seeded entities, thresholded to keep only confident correlations -- exposed post-fit() as attr_sim_mat_/attr_pool1_ids_/attr_pool2_ids_ (see below for why this is no longer used to perturb the structural embeddings).

Found while diagnosing #28's exhaustive-evaluation gap: OpenEA's own reference jape.py builds a sim_optimizer (an apply_gradients op meant to regularize ent_embeds toward the attribute-similarity blend each epoch) but its launch_sim_1epo only ever fetches self.sim_loss via session.run -- self.sim_optimizer is never passed to session.run anywhere in the file. In TensorFlow 1.x an unfetched op simply never executes, so OpenEA's own published JAPE numbers were produced by pure structural-TransE training -- the attribute pipeline runs and computes a loss value (used only for a print statement) that never actually touches the entity embeddings. Confirmed empirically: running OpenEA's real reference implementation reproduces its published Hits@1 (0.261 vs. target 0.263); a version of this port that does apply the attribute-similarity gradient (either OpenEA's literal dense/thresholded matrix, or this port's previous row-argmax-sparsified variant) scores far below both that number and structural-only training under exhaustive ranking. This port now matches OpenEA's real (not its apparent) behavior: fit() still trains Attr2Vec and builds the attribute-similarity matrix -- exposed for inspection -- but never uses it to update the structural embeddings.

Deviations from OpenEA's own code:

• fit() gains two new optional parameters, attribute_triples1/ attribute_triples2 -- the paper's method genuinely doesn't fit the EA linker family's existing signature otherwise, since attribute triples are JAPE's whole point. Left as None (default), fit() skips the attribute pipeline entirely (a no-op difference now that its output never affects training anyway). Most of this repo's EA dataset loaders have no attribute triples at all; see linkingtk.datasets.openea_native's module docstring for the one that does and why it's a separate loader family from linkingtk.datasets.openea's.
• Attr2Vec's NCE loss uses uniform negative sampling, not OpenEA's tf.nn.nce_loss default (a log-uniform sampler assuming frequency- sorted class ids, which our attribute-predicate ids aren't, and which PyTorch has no drop-in equivalent for) -- see _jape_torch.py's train_attr2vec docstring.
• The attribute-similarity reference pool is "every dataset1/ dataset2 entity not already a ground_truth source/target" rather than OpenEA's own valid_entities + test_entities -- the same deviation, for the same reason, already documented on IPTransELinker's bootstrap pool. Only affects the exposed attr_sim_mat_, not training.

JAPELinker

JAPELinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    neg_alpha: float = 0.1,
    top_attr_threshold: float = 0.9,
    attr_sim_mat_threshold: float = 0.95,
    attr_max_epoch: int = 200,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a shared TransE space, plus an inspectable attribute similarity.

Must be fit before link can be called. See the module docstring for how this differs from MTransELinker/ IPTransELinker, and in particular why the attribute-similarity matrix built from attribute_triples1/attribute_triples2 no longer perturbs the trained structural embeddings (it's exposed post-fit() as attr_sim_mat_ instead).

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained structural and attribute embeddings. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Structural training epochs. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for both the structural triple loss and (if attribute triples are given) Attr2Vec's training.	5000
learning_rate	float	Adagrad's learning rate, for every optimizer used here -- OpenEA's published value is 0.01.	0.01
neg_alpha	float	Weight on the negative-triple term in the structural loss (pos_loss - neg_alpha * neg_loss). OpenEA's published value is 0.1.	0.1
top_attr_threshold	float	Fraction of each KG's distinct attribute predicates to keep, by descending frequency. OpenEA's published value is 0.9.	0.9
attr_sim_mat_threshold	float	Attribute-similarity values below this are zeroed out in the exposed attr_sim_mat_. OpenEA's published value is 0.95.	0.95
attr_max_epoch	int	Attr2Vec training epochs. OpenEA's published value is 200.	200
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/jape.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    neg_alpha: float = 0.1,
    top_attr_threshold: float = 0.9,
    attr_sim_mat_threshold: float = 0.95,
    attr_max_epoch: int = 200,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.neg_alpha = neg_alpha
    self.top_attr_threshold = top_attr_threshold
    self.attr_sim_mat_threshold = attr_sim_mat_threshold
    self.attr_max_epoch = attr_max_epoch
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False
    self.attr_sim_mat_: npt.NDArray[np.floating[Any]] | None = None
    """Thresholded attribute-similarity matrix from the last `fit()` call, or ``None`` if no
    attribute triples were given. Rows/columns correspond to
    `attr_pool1_ids_`/`attr_pool2_ids_`. Diagnostic only -- see the
    module docstring for why this doesn't affect the trained
    embeddings."""
    self.attr_pool1_ids_: list[str] = []
    """Entity ids indexing `attr_sim_mat_`'s rows."""
    self.attr_pool2_ids_: list[str] = []
    """Entity ids indexing `attr_sim_mat_`'s columns."""

attr_pool1_ids_ instance-attribute

attr_pool1_ids_: list[str] = []

Entity ids indexing attr_sim_mat_'s rows.

attr_pool2_ids_ instance-attribute

attr_pool2_ids_: list[str] = []

Entity ids indexing attr_sim_mat_'s columns.

attr_sim_mat_ instance-attribute

attr_sim_mat_: NDArray[floating[Any]] | None = None

Thresholded attribute-similarity matrix from the last fit() call, or None if no attribute triples were given. Rows/columns correspond to attr_pool1_ids_/attr_pool2_ids_. Diagnostic only -- see the module docstring for why this doesn't affect the trained embeddings.

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> JAPELinker

Train a shared TransE space; optionally build an inspectable attribute-similarity matrix.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (same approach as IPTransELinker.fit, for the same reason: negative sampling needs each KG's own entity pool).	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs. Seeds the shared embedding table (each pair's target is aliased to its source's id) and, if attribute triples are given, merges each pair's attribute vocabularies for Attr2Vec training.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2)) -- entity ids on both sides must already be disjoint.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10
attribute_triples1	list[Triple] | None	KG1's (entity_id, predicate, value) attribute triples, e.g. from a GraphDatasetLoader that supports load_attribute_triples() (see linkingtk.datasets.openea_native). If None (default) or empty, together with attribute_triples2, the attribute pipeline is skipped entirely -- doesn't affect the trained structural embeddings either way (see the module docstring's discovered-upstream-bug note), but skips training Attr2Vec and populating attr_sim_mat_.	None
attribute_triples2	list[Triple] | None	KG2's attribute triples. See attribute_triples1.	None

Returns:

Type	Description
JAPELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in dataset1/dataset2 -- nothing to seed the shared embedding table with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/jape.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> JAPELinker:
    """Train a shared TransE space; optionally build an inspectable attribute-similarity matrix.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (same approach as
            `IPTransELinker.fit`, for the same reason: negative
            sampling needs each KG's own entity pool).
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs. Seeds the shared embedding table (each pair's
            target is aliased to its source's id) and, if attribute
            triples are given, merges each pair's attribute
            vocabularies for Attr2Vec training.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)``) -- entity ids
            on both sides must already be disjoint.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples, e.g. from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]
            that supports
            ``load_attribute_triples()``
            (see [linkingtk.datasets.openea_native][]). If ``None``
            (default) or empty, together with ``attribute_triples2``,
            the attribute pipeline is skipped entirely -- doesn't
            affect the trained structural embeddings either way (see
            the module docstring's discovered-upstream-bug note), but
            skips training Attr2Vec and populating ``attr_sim_mat_``.
        attribute_triples2: KG2's attribute triples. See
            ``attribute_triples1``.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``dataset1``/``dataset2`` -- nothing to
            seed the shared embedding table with -- or if ``device``
            is invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("JAPELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    seed_pairs = [(s, t) for s, t in ground_truth if s in ids1 and t in ids2]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in "
            "`dataset1`/`dataset2`; fit() has no seed pairs to train "
            "JAPE's shared embedding table with."
        )

    entity_to_id, relation_to_id = build_shared_id_mappings(
        triples1_labels + triples2_labels, seed_pairs
    )
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(mapped1)
    ctx2 = build_kg_context(mapped2)

    entity_embeds, relation_embeds = self._init_structural_embeddings(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    optimizer = torch.optim.Adagrad([entity_embeds, relation_embeds], lr=self.learning_rate)

    self.attr_sim_mat_, self.attr_pool1_ids_, self.attr_pool2_ids_ = (
        self._setup_attribute_pipeline(
            dataset1,
            dataset2,
            entity_to_id,
            seed_pairs,
            attribute_triples1,
            attribute_triples2,
            rng,
            device,
        )
    )

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_triple_epoch(
            entity_embeds,
            relation_embeds,
            optimizer,
            ctx1,
            ctx2,
            rng,
            self.batch_size,
            self.neg_alpha,
        )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/jape.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/jape.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

KDCoE

kdcoe

KDCoE knowledge-graph-embedding linker for Entity Alignment.

Chen, M., Tian, Y., Chang, K.-W., Skiena, S., & Zaniolo, C. (2018). Co-training Embeddings of Knowledge Graphs and Entity Descriptions for Cross-lingual Entity Alignment. IJCAI 2018. https://www.ijcai.org/proceedings/2018/0556.pdf

Like MTransELinker and IPTransELinker, this is a faithful reimplementation of KDCoE's actual training procedure, ported directly from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/kdcoe.py and its modules/base/losses.py, modules/bootstrapping/alignment_finder.py support code) rather than from a from-scratch reading of the paper. KDCoE co-trains two signals with a bootstrapping loop between them:

• A structural embedding: margin-based TransE (margin=1.5, one negative per positive, corrupting within each triple's own KG -- reuses _iptranse_training.py's sample_negative_triples) bridged across KGs by a learned mapping matrix, the same "mapping" alignment module MTransE uses (not IPTransE's/JAPE's shared-id "sharing" trick) -- entities keep separate per-KG embedding rows in one combined table.
• A cross-lingual entity-description encoder: each entity's description text (or, when it has none, its own label -- see _kdcoe_text.extract_description_text) is embedded via pretrained word vectors, then run through a GRU -> Conv1D -> attention-pool -> GRU -> attention-pool -> dense pipeline (_kdcoe_torch.build_description_encoder), trained with an in-batch contrastive loss that pulls each known-aligned pair's description vectors together and pushes apart other pairs in the same batch.

Each outer co-training iteration trains the description encoder to convergence, uses it to find new high-confidence pairs among entities outside the seed set (row-wise best-match above desc_sim_th, reusing _iptranse_training.py's find_new_pairs -- confirmed to be exactly OpenEA's own find_potential_alignment_greedily/ find_alignment logic), trains the structural+mapping model to convergence (now also pulled toward those newly found pairs via a second, new_param-weighted mapping loss), and uses it to find new pairs the same way -- stopping when a round re-discovers nothing beyond what's already known (max_co_training_iters caps the worst case).

Key finding shared with MTransELinker: both mapping-loss optimizers (the seed-pair one and the newly-bootstrapped-pair one) must update the shared entity embeddings as well as the mapping matrix, matching OpenEA's generate_optimizer(loss, lr, var_list=None, ...) -- leaving var_list unset makes TensorFlow differentiate against every trainable variable the loss touches, not just the mapping matrix. Restricting either optimizer here to just the mapping matrix converges roughly an order of magnitude slower.

Key finding from reading the reference source, not obvious from the paper: _get_desc_input doesn't require every entity to have a real description triple -- entities without one fall back to their own local name as pseudo-description text. Since this repo's dataset loaders already populate Entity.labels with exactly that fallback text (see linkingtk.datasets.openea_native), the description pathway here always has some text to work with. Unlike JAPELinker, where omitting attribute_triples1/attribute_triples2 disables the attribute pathway entirely, here it only reduces the description signal to label-only text -- worth noting since it's a real interface difference within this linker family.

Word embeddings: pretrained vectors are required (the description encoder's whole point is leveraging generic-language priors that aren't learnable from a 15K-entity alignment task alone). This fetches the same file OpenEA's own published config uses -- wiki-news-300d-1M.vec.zip from fastText's stable, Facebook-hosted URL -- via linkingtk.datasets._util's fetch_cached, the same mechanism every dataset loader in this repo already uses. It's a much larger one-time download (~681MB) than any other asset this repo fetches, so _kdcoe_text.py's load_fasttext_vectors streams the archived .vec member line-by-line and keeps only rows whose word is in this run's own description vocabulary, rather than materializing the full ~1M-word file.

Deviations beyond what _kdcoe_text.py/_kdcoe_torch.py already document:

• Reference pool for bootstrapping: OpenEA restricts the "not yet aligned" search pool to its own dataset loader's valid_entities + test_entities specifically. This repo's fit() only receives dataset1/dataset2 and ground_truth (the train split), so the pool here is every entity not already a ground_truth source/target -- broader than OpenEA's pool, same already-established deviation as IPTransELinker's bootstrap pool and JAPELinker's attribute reference pool (reuses _jape_training.py's reference_pools directly).
• Final scoring uses the structural+mapping embedding only: the description encoder drives bootstrapping during fit() but doesn't contribute to link()'s own similarity score -- link()/ _embedding()/_project() are near-identical to MTransE's.
• Mapping-loss training samples with replacement each step rather than OpenEA's random.sample (already-established MTransE deviation, same rationale).
• Early stopping uses this repo's own plain patience-counter Hits@1 check (twice -- once per phase, each reset every outer co-training iteration, matching OpenEA's own per-phase flag1/flag2 reset) rather than OpenEA's exact two-flag early_stop().
• Description-found bootstrap pairs use reciprocal best-match, not desc_sim_th. Diagnosed while benchmarking on real EN-FR-15K-V1 data (#29): the description encoder's cosine similarities saturate near 1.0 for nearly every reference-pool pair regardless of correctness (median ~0.9998 in testing), so an absolute threshold like desc_sim_th=0.95 accepts ~99.9% of the pool each round and can't discriminate anything -- find_new_pairs's row-argmax+threshold (OpenEA's own find_potential_alignment_greedily) isn't a usable filter once similarities are this uniform. This dataset's real .../description triple coverage is also sparse (10.5% of KG1 entities, 0.5% of KG2), so most descriptions are single-word label fallbacks -- but 60.7% of EN-FR-15K-V1's ground-truth pairs share a literally identical local entity name across languages (OpenEA's own documented v1-dataset "name bias"; their README recommends v2.0 to remove it), which the label-fallback pathway can exploit even without rich text: the description encoder's own validation Hits@1 is a real 44-46% across co-training iterations, well above chance. The earlier fix here (before this reciprocal-match change) fully separated description-found pairs from structural training to stop the ~99.9% false-positive rate from collapsing structural Hits@1 from a clean ~0.40 to near-zero -- but that also meant the description encoder's genuine 44-46% signal could never improve the final structural+mapping embedding at all. Switching the description pathway's bootstrap rounds to _iptranse_training.find_mutual_pairs (reciprocal row/column argmax agreement, which stays discriminative under saturation since the argmax positions still vary even when the values don't) raises precision enough that description-found pairs are now merged back into the same pool that feeds the structural mapping-new loss, matching OpenEA's own unified self.new_alignment/self.new_alignment_index design rather than permanently working around it.

KDCoELinker

KDCoELinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    alpha: float = 5.0,
    new_param: float = 0.1,
    wv_dim: int = 300,
    default_desc_length: int = 4,
    desc_batch_size: int = 512,
    desc_sim_th: float = 0.95,
    sim_th: float = 0.8,
    max_co_training_iters: int = 5,
    word_embed_url: str = _FASTTEXT_URL,
    cache_dir: Path | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Co-trains a structural TransE+mapping space with a description-text encoder.

Must be fit before link can be called. See the module docstring for how this differs from MTransELinker and IPTransELinker.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings and the (square) mapping matrix. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Per-phase training-epoch cap (structural or description), each co-training iteration. OpenEA's config allows up to 2000 per phase with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for structural triple-loss training; the mapping-loss batch size is derived from it (same number of minibatches per epoch for both, matching OpenEA).	5000
learning_rate	float	Adagrad's learning rate, for every optimizer here (structural, mapping, mapping-new, description) -- OpenEA's published value is 0.01.	0.01
margin	float	Hinge margin for the structural triple loss. OpenEA's published value is 1.5.	1.5
alpha	float	Weight applied to the seed-pairs mapping loss. OpenEA's published value is 5.0 (matches MTransELinker's own default).	5.0
new_param	float	Weight applied to the bootstrapped-pairs mapping loss. OpenEA's published value is 0.1.	0.1
wv_dim	int	Pretrained word-vector dimensionality. Must match word_embed_url's file -- OpenEA's published fastText file is 300-dimensional.	300
default_desc_length	int	Fixed word-sequence length each entity's description is padded/truncated to. OpenEA's published value is 4.	4
desc_batch_size	int	Mini-batch size for description-encoder training. OpenEA's published value is 512.	512
desc_sim_th	float	Kept for interface parity with OpenEA's published value (0.95), but not the operative filter for the description pathway's bootstrap rounds -- see the module docstring's "reciprocal best-match" deviation. Description similarities saturate near 1.0 on real data, making an absolute threshold meaningless; bootstrap rounds instead use _iptranse_training.find_mutual_pairs's reciprocal-agreement filter.	0.95
sim_th	float	Minimum (mapped) structural-embedding similarity for a bootstrap round to accept a newly found pair. OpenEA's published value is 0.8.	0.8
max_co_training_iters	int	Cap on outer co-training iterations (each alternating a description-encoder phase and a structural phase). OpenEA's published value (max_iter) is 5.	5
word_embed_url	str	URL of a fastText-format .vec.zip (see _kdcoe_text.py's load_fasttext_vectors). Defaults to the real, published pretrained file; tests override this with a file:// url pointing at a tiny local fixture.	_FASTTEXT_URL
cache_dir	Path | None	Passed through to fetch_cached for both the word-vector download and (indirectly, via callers' attribute_triples) any dataset assets. None uses fetch_cached's own default cache location.	None
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/kdcoe.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    alpha: float = 5.0,
    new_param: float = 0.1,
    wv_dim: int = 300,
    default_desc_length: int = 4,
    desc_batch_size: int = 512,
    desc_sim_th: float = 0.95,
    sim_th: float = 0.8,
    max_co_training_iters: int = 5,
    word_embed_url: str = _FASTTEXT_URL,
    cache_dir: Path | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.margin = margin
    self.alpha = alpha
    self.new_param = new_param
    self.wv_dim = wv_dim
    self.default_desc_length = default_desc_length
    self.desc_batch_size = desc_batch_size
    self.desc_sim_th = desc_sim_th
    self.sim_th = sim_th
    self.max_co_training_iters = max_co_training_iters
    self.word_embed_url = word_embed_url
    self.cache_dir = cache_dir
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._mapping: npt.NDArray[np.floating[Any]] | None = None
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> KDCoELinker

Co-train the structural+mapping space and the description encoder.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (same approach IPTransELinker uses) and as the description pathway's label-fallback source.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs, used as both phases' seed training pairs.	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping in both phases -- every eval_every epochs, Hits@1 is checked, and that phase stops after patience checks with no improvement. If None (default), each phase trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping a phase early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10
attribute_triples1	list[Triple] | None	KG1's (entity_id, predicate, value) attribute triples, if available -- entities with a .../description-like predicate triple use its value as description text; others fall back to their own label (see the module docstring). May be omitted (or partial) -- the description pathway still runs, just with less real description text.	None
attribute_triples2	list[Triple] | None	KG2's attribute triples. See attribute_triples1.	None

Returns:

Type	Description
KDCoELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples -- nothing to seed either phase's training with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/kdcoe.py

def fit(  # noqa: PLR0915 -- co-training loop genuinely needs this much sequential setup
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> KDCoELinker:
    """Co-train the structural+mapping space and the description encoder.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (same approach
            [IPTransELinker][linkingtk.algorithms.ea.iptranse.IPTransELinker]
            uses) and as the description pathway's label-fallback source.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs, used as both phases' seed training pairs.
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping in *both* phases -- every ``eval_every`` epochs,
            Hits@1 is checked, and that phase stops after ``patience``
            checks with no improvement. If ``None`` (default), each
            phase trains the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping a phase early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples, if available -- entities with a
            ``.../description``-like predicate triple use its value as
            description text; others fall back to their own label
            (see the module docstring). May be omitted (or partial) --
            the description pathway still runs, just with less real
            description text.
        attribute_triples2: KG2's attribute triples. See
            ``attribute_triples1``.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples -- nothing to seed
            either phase's training with -- or if ``device`` is
            invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("KDCoELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    entity_to_id, relation_to_id = build_id_mappings(triples1_labels + triples2_labels)
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(mapped1)
    ctx2 = build_kg_context(mapped2)

    seed_pairs = [(s, t) for s, t in ground_truth if s in entity_to_id and t in entity_to_id]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train KDCoE's structural "
            "and description phases with."
        )
    seed_source_ids = np.array([entity_to_id[s] for s, _ in seed_pairs], dtype=np.int64)
    seed_target_ids = np.array([entity_to_id[t] for _, t in seed_pairs], dtype=np.int64)

    word_embeds, word_ids = self._build_description_inputs(
        dataset1, dataset2, attribute_triples1 or [], attribute_triples2 or [], device
    )

    entity_embeds, relation_embeds, mapping_mat, eye = self._init_structural(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    structural_optimizer = torch.optim.Adagrad(
        [entity_embeds, relation_embeds], lr=self.learning_rate
    )
    mapping_optimizer = torch.optim.Adagrad([mapping_mat, entity_embeds], lr=self.learning_rate)
    mapping_optimizer_new = torch.optim.Adagrad(
        [mapping_mat, entity_embeds], lr=self.learning_rate
    )
    desc_encoder = build_description_encoder(self.wv_dim, self.default_desc_length).to(device)
    desc_optimizer = torch.optim.Adagrad(desc_encoder.parameters(), lr=self.learning_rate)

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    pool1, pool2 = reference_pools(sorted(ids1), sorted(ids2), seed_pairs)

    desc_bootstrapped_pairs: set[tuple[str, str]] = set()
    rel_bootstrapped_pairs: set[tuple[str, str]] = set()
    triple_steps = max(1, math.ceil((len(ctx1.triples) + len(ctx2.triples)) / self.batch_size))

    for _iteration in range(self.max_co_training_iters):
        self._train_desc_phase(
            torch,
            desc_encoder,
            desc_optimizer,
            word_embeds,
            word_ids,
            seed_pairs + list(desc_bootstrapped_pairs),
            val_pairs,
            rng,
            patience,
            eval_every,
        )
        desc_found = self._find_new_pairs_via_desc(
            torch, desc_encoder, word_embeds, word_ids, pool1, pool2
        )
        if desc_found and desc_found <= desc_bootstrapped_pairs:
            break
        desc_bootstrapped_pairs |= desc_found

        mapping_new_pairs = rel_bootstrapped_pairs | desc_bootstrapped_pairs
        bootstrap_source_ids = np.array(
            [entity_to_id[s] for s, _ in mapping_new_pairs], dtype=np.int64
        )
        bootstrap_target_ids = np.array(
            [entity_to_id[t] for _, t in mapping_new_pairs], dtype=np.int64
        )
        self._train_structural_phase(
            torch,
            functional,
            entity_embeds,
            relation_embeds,
            mapping_mat,
            eye,
            structural_optimizer,
            mapping_optimizer,
            mapping_optimizer_new,
            ctx1,
            ctx2,
            seed_source_ids,
            seed_target_ids,
            bootstrap_source_ids,
            bootstrap_target_ids,
            entity_to_id,
            val_pairs,
            rng,
            triple_steps,
            patience,
            eval_every,
        )
        rel_found = self._find_new_pairs_via_structural(
            torch, functional, entity_embeds, mapping_mat, entity_to_id, pool1, pool2
        )
        if rel_found and rel_found <= rel_bootstrapped_pairs:
            break
        rel_bootstrapped_pairs |= rel_found

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
        final_mapping = mapping_mat.cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._mapping = final_mapping
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side (projected).

Source code in src/linkingtk/algorithms/ea/kdcoe.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side (projected)."""
    return self._project(self._embedding(entity_id))

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side (unprojected).

Source code in src/linkingtk/algorithms/ea/kdcoe.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side (unprojected)."""
    return self._embedding(entity_id)

_kdcoe_text

Private, plain-numpy (torch-free) helpers for KDCoELinker.

Ported from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA -- approaches/kdcoe.py's _get_desc_input/_get_local_name_by_name_triple). Independently testable without torch installed -- see _kdcoe_torch.py for the description-encoder network and training-step functions.

build_word_ids

build_word_ids(
    tokens_by_entity: dict[str, list[str]],
    word_to_index: dict[str, int],
    default_desc_length: int,
    oov_id: int,
) -> dict[str, npt.NDArray[np.int64]]

Each entity's tokens -> a fixed-length array of word ids.

Ports lookup_and_padding: truncates to default_desc_length words, padding with oov_id if shorter. A token absent from word_to_index also maps to oov_id -- the same handling OpenEA gives words unseen in the pretrained vectors (see load_fasttext_vectors).

Parameters:

Name	Type	Description	Default
tokens_by_entity	dict[str, list[str]]	Per-entity token lists, e.g. from tokenizing extract_description_text's output.	required
word_to_index	dict[str, int]	Vocabulary word -> row index (e.g. into a pretrained embedding table).	required
default_desc_length	int	Fixed output length. OpenEA's published value is 4.	required
oov_id	int	Id used for padding and out-of-vocabulary words -- conventionally one past the real vocabulary's last index.	required

Returns:

Type	Description
dict[str, NDArray[int64]]	entity_id -> (default_desc_length,) int64 array.

Source code in src/linkingtk/algorithms/ea/_kdcoe_text.py

def build_word_ids(
    tokens_by_entity: dict[str, list[str]],
    word_to_index: dict[str, int],
    default_desc_length: int,
    oov_id: int,
) -> dict[str, npt.NDArray[np.int64]]:
    """Each entity's tokens -> a fixed-length array of word ids.

    Ports ``lookup_and_padding``: truncates to ``default_desc_length``
    words, padding with ``oov_id`` if shorter. A token absent from
    ``word_to_index`` also maps to ``oov_id`` -- the same handling OpenEA
    gives words unseen in the pretrained vectors (see
    [load_fasttext_vectors][linkingtk.algorithms.ea._kdcoe_text.load_fasttext_vectors]).

    Args:
        tokens_by_entity: Per-entity token lists, e.g. from tokenizing
            [extract_description_text][linkingtk.algorithms.ea._kdcoe_text.extract_description_text]'s
            output.
        word_to_index: Vocabulary word -> row index (e.g. into a
            pretrained embedding table).
        default_desc_length: Fixed output length. OpenEA's published
            value is ``4``.
        oov_id: Id used for padding and out-of-vocabulary words --
            conventionally one past the real vocabulary's last index.

    Returns:
        ``entity_id -> (default_desc_length,) int64 array``.
    """
    result: dict[str, npt.NDArray[np.int64]] = {}
    for entity_id, tokens in tokens_by_entity.items():
        ids = [word_to_index.get(token, oov_id) for token in tokens[:default_desc_length]]
        ids += [oov_id] * (default_desc_length - len(ids))
        result[entity_id] = np.array(ids, dtype=np.int64)
    return result

extract_description_text

extract_description_text(
    entities: list[Entity], attribute_triples: list[Triple]
) -> dict[str, str]

Each entity's description text, falling back to its own label.

Ports _get_desc_input's entity-description selection: an entity's text is the first attribute-triple value whose predicate URI contains "escription" (catching description/Description, e.g. dc:description, dbo:description) -- the same substring check OpenEA's own code uses. Entities with no such triple fall back to their own local name (_get_local_name_by_name_triple's fallback -- this repo's dataset loaders already populate Entity.labels with exactly that text, see linkingtk.datasets.openea_native, so no separate name-triple lookup is needed here).

This is why the description pathway never strictly needs attribute_triples to produce some text for every entity -- see KDCoELinker's docstring.

Parameters:

Name	Type	Description	Default
entities	list[Entity]	Every entity needing description text (both KGs, concatenated by the caller).	required
attribute_triples	list[Triple]	(entity_id, predicate, value) triples (both KGs, concatenated by the caller) -- values are plain text (already RDF-literal-stripped by the dataset loader).	required

Returns:

Type	Description
dict[str, str]	entity_id -> description text (untokenized -- see
dict[str, str]	tokenize_description).

Source code in src/linkingtk/algorithms/ea/_kdcoe_text.py

def extract_description_text(
    entities: list[Entity],
    attribute_triples: list[Triple],
) -> dict[str, str]:
    """Each entity's description text, falling back to its own label.

    Ports ``_get_desc_input``'s entity-description selection: an entity's
    text is the first attribute-triple value whose predicate URI contains
    ``"escription"`` (catching ``description``/``Description``, e.g.
    ``dc:description``, ``dbo:description``) -- the same substring check
    OpenEA's own code uses. Entities with no such triple fall back to
    their own local name (``_get_local_name_by_name_triple``'s fallback --
    this repo's dataset loaders already populate ``Entity.labels`` with
    exactly that text, see ``linkingtk.datasets.openea_native``, so no
    separate name-triple lookup is needed here).

    This is why the description pathway never strictly needs
    ``attribute_triples`` to produce *some* text for every entity -- see
    [KDCoELinker][linkingtk.algorithms.ea.kdcoe.KDCoELinker]'s docstring.

    Args:
        entities: Every entity needing description text (both KGs,
            concatenated by the caller).
        attribute_triples: ``(entity_id, predicate, value)`` triples (both
            KGs, concatenated by the caller) -- values are plain text
            (already RDF-literal-stripped by the dataset loader).

    Returns:
        ``entity_id -> description text`` (untokenized -- see
        [tokenize_description][linkingtk.algorithms.ea._kdcoe_text.tokenize_description]).
    """
    descriptions: dict[str, str] = {}
    for entity_id, predicate, value in attribute_triples:
        if entity_id not in descriptions and _DESCRIPTION_PREDICATE_RE.search(predicate):
            descriptions[entity_id] = value

    result: dict[str, str] = {}
    for entity in entities:
        if entity.id in descriptions:
            result[entity.id] = descriptions[entity.id]
        else:
            texts = label_texts(entity)
            result[entity.id] = texts[0] if texts else ""
    return result

load_fasttext_vectors

load_fasttext_vectors(
    zip_url: str,
    vocabulary: set[str],
    cache_dir: Path | None = None,
) -> dict[str, npt.NDArray[np.floating[Any]]]

Pretrained word vectors for vocabulary's words only, from a fastText .vec.zip.

Fetches zip_url via fetch_cached (linkingtk.datasets._util -- cached on disk after the first call; file:// urls, used by tests to point at a tiny local fixture, are never cached, matching fetch_cached's existing behavior) and streams the archived .vec member line-by-line (zipfile.ZipFile(...).open(member)) rather than decoding the whole decompressed text at once -- for the real fastText file (wiki-news-300d-1M.vec, ~1M words, ~2.3GB decompressed) this keeps peak new memory bounded by len(vocabulary) rather than the file's full vocabulary.

Parameters:

Name	Type	Description	Default
zip_url	str	URL of a zip containing exactly one .vec member in fastText's plain-text format (first line "<count> <dim>", then one "word f1 f2 ... fN" line per word).	required
vocabulary	set[str]	Words to keep; every other row is skipped without being parsed into an array.	required
cache_dir	Path | None	Passed through to fetch_cached.	None

Returns:

Type	Description
dict[str, NDArray[floating[Any]]]	word -> (dim,) float32 array, restricted to words that are
dict[str, NDArray[floating[Any]]]	both in vocabulary and the file. Words in vocabulary but
dict[str, NDArray[floating[Any]]]	absent from the file are simply omitted -- callers treat a
dict[str, NDArray[floating[Any]]]	missing word as out-of-vocabulary, same as
dict[str, NDArray[floating[Any]]]	build_word_ids's
dict[str, NDArray[floating[Any]]]	oov_id fallback.

Source code in src/linkingtk/algorithms/ea/_kdcoe_text.py

def load_fasttext_vectors(
    zip_url: str,
    vocabulary: set[str],
    cache_dir: Path | None = None,
) -> dict[str, npt.NDArray[np.floating[Any]]]:
    """Pretrained word vectors for ``vocabulary``'s words only, from a fastText ``.vec.zip``.

    Fetches ``zip_url`` via ``fetch_cached`` (``linkingtk.datasets._util``
    -- cached on disk after the first call; ``file://`` urls, used by
    tests to point at a tiny local fixture, are never cached, matching
    ``fetch_cached``'s existing behavior) and streams the archived
    ``.vec`` member line-by-line (``zipfile.ZipFile(...).open(member)``)
    rather than decoding the whole decompressed text at once -- for the
    real fastText file (``wiki-news-300d-1M.vec``, ~1M words, ~2.3GB
    decompressed) this keeps peak new memory bounded by ``len(vocabulary)``
    rather than the file's full vocabulary.

    Args:
        zip_url: URL of a zip containing exactly one ``.vec`` member in
            fastText's plain-text format (first line ``"<count> <dim>"``,
            then one ``"word f1 f2 ... fN"`` line per word).
        vocabulary: Words to keep; every other row is skipped without
            being parsed into an array.
        cache_dir: Passed through to ``fetch_cached``.

    Returns:
        ``word -> (dim,) float32 array``, restricted to words that are
        both in ``vocabulary`` and the file. Words in ``vocabulary`` but
        absent from the file are simply omitted -- callers treat a
        missing word as out-of-vocabulary, same as
        [build_word_ids][linkingtk.algorithms.ea._kdcoe_text.build_word_ids]'s
        ``oov_id`` fallback.
    """
    archive_bytes = fetch_cached(zip_url, cache_dir)
    vectors: dict[str, npt.NDArray[np.floating[Any]]] = {}
    with zipfile.ZipFile(BytesIO(archive_bytes)) as archive:
        member_name = next(name for name in archive.namelist() if name.endswith(".vec"))
        with archive.open(member_name) as member:
            for line_number, raw_line in enumerate(member):
                if line_number == 0:
                    continue  # header: "<count> <dim>"
                parts = raw_line.decode("utf-8").split()
                if not parts:
                    continue
                word = parts[0]
                if word in vocabulary:
                    vectors[word] = np.array(parts[1:], dtype=np.float32)
    return vectors

tokenize_description

tokenize_description(text: str) -> list[str]

Lowercased, punctuation-stripped whitespace tokens.

Ports the string.punctuation regex replace + whitespace split OpenEA applies to both description and local-name text before word-id lookup.

Source code in src/linkingtk/algorithms/ea/_kdcoe_text.py

def tokenize_description(text: str) -> list[str]:
    """Lowercased, punctuation-stripped whitespace tokens.

    Ports the ``string.punctuation`` regex replace + whitespace split
    OpenEA applies to both description and local-name text before word-id
    lookup.
    """
    stripped = _PUNCTUATION_RE.sub("", text).lower()
    return [token for token in stripped.split() if token]

AttrE

attre

AttrE knowledge-graph-embedding linker for Entity Alignment.

Trisedya, B. D., Qi, J., & Zhang, R. (2019). Entity Alignment between Knowledge Graphs Using Attribute Embeddings. AAAI 2019. https://people.eng.unimelb.edu.au/jianzhongq/papers/AAAI2019_EntityAlignment.pdf

Like IPTransELinker and JAPELinker, this is a faithful reimplementation of AttrE's actual training procedure, ported directly from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/attre.py and its modules/train/batch.py support code) rather than from a from-scratch reading of the paper.

Important finding from reading the reference source, not obvious from a one-line method summary: despite AttrE's paper framing it as aligning entities "using attribute triples only," OpenEA's actual benchmarked implementation trains a full structural TransE component over relation triples too, alongside a separate character/attribute-embedding (CE) space -- the two are cross-regularized by a joint cosine-similarity loss. This is not an optional add-on; it's what OpenEA's own published number comes from, so this port includes both halves rather than a structure-free implementation. Two embedding spaces, trained jointly:

• Structural (SE): ent_embeds/rel_embeds, plain margin-based TransE + negative sampling over relation triples -- reuses _kdcoe_torch.py's KGContext, build_kg_context, and train_structural_epoch directly (already fully generic, no path loss, no mapping).
• Character/attribute (CE): ent_embeds_ce (a second, independent entity embedding table, sharing the same ids as ent_embeds), attr_embeds (one row per attribute predicate), char_embeds (one row per selected character). Attribute triples are treated as TransE-style triples -- entity_ce + attribute ≈ compose(value's characters) -- with the value's embedding built compositionally from its characters (see _attre_torch.py's compose_value_embeddings). Negative sampling corrupts only the entity endpoint (see _attre_text.py's sample_negative_attribute_triples).

Alignment module is "sharing" (attre_args_15K.json), identical to IPTransE's/JAPE's mechanism -- seed pairs get one shared embedding row (in both the SE and CE tables), no mapping matrix. This is the only cross-KG alignment signal; there is no separate seed-pair loss anywhere in AttrE's own code. Reuses _iptranse_training.py's build_shared_id_mappings directly.

Joint loss: sum(1 - cos_sim(ent_embeds[e], ent_embeds_ce[e])) over every entity in both KGs (not just seed pairs) -- this is what actually ties the two spaces together, pulling both toward mutual agreement rather than fixing one side to a projection of the other (contrast MTransELinker's mapping loss, which only updates its mapping matrix). See _attre_torch.py's train_joint_epoch for a note on a curiosity ported literally from OpenEA's own code (it re-runs this loss over the entire entity list multiple times per epoch, not sub-batched).

No bootstrapping, no co-training loop -- simpler control flow than KDCoELinker. Per epoch: one SE epoch, one CE epoch, one joint-loss pass, then periodic validation.

optimizer="SGD" in OpenEA's published config -- the first method in this family to use plain SGD rather than Adagrad; ported literally (torch.optim.SGD).

Deviations beyond what _attre_text.py/_attre_torch.py already document:

• attribute_triples1/attribute_triples2 are required, not optional -- unlike JAPELinker (structural-only fallback) and KDCoELinker (label-fallback always gives some description text), AttrE has no meaningful degraded mode without attribute triples -- the entire CE half and joint loss depend on them, and there's no fallback text source. fit() raises LinkingTKError if both are empty.
• Final scoring uses the structural (SE) embedding only: same precedent as every other linker in this family that has more than one trained representation -- ent_embeds_ce regularizes ent_embeds during training via the joint loss but isn't itself used by link().
• Early stopping uses this repo's own plain patience-counter Hits@1 check rather than OpenEA's own two-flag early_stop(), and doesn't gate on OpenEA's start_valid (checks are simply skipped until val_ground_truth is given and eval_every divides the epoch, same convention as every other linker in this family).

AttrELinker

AttrELinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    literal_len: int = 5,
    char_freq_threshold: float = 0.0001,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a shared TransE space regularized by character embeddings.

Must be fit before link can be called. See the module docstring for how this differs from IPTransELinker and JAPELinker.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of every trained embedding table (entities, relations, character-space entities, attributes, characters). OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each one pass of structural training, one pass of attribute/character training, and one (repeated) pass of joint-loss training. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for both the structural and attribute-triple training (each derives its own step count from its own triple count), and the divisor for the joint loss's repeat count.	5000
learning_rate	float	SGD's learning rate, for every optimizer here. OpenEA's published value is 0.01.	0.01
margin	float	Hinge margin for both the structural and attribute-triple losses. OpenEA's published value is 1.5.	1.5
literal_len	int	Fixed character-sequence length each attribute value is truncated/padded to. OpenEA's published value is 5.	5
char_freq_threshold	float	Minimum share of total (distinct-value) character occurrences required to keep a character in the vocabulary; rarer characters map to the padding id. OpenEA's published value is 0.0001.	0.0001
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/attre.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    literal_len: int = 5,
    char_freq_threshold: float = 0.0001,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.margin = margin
    self.literal_len = literal_len
    self.char_freq_threshold = char_freq_threshold
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> AttrELinker

Train a shared TransE space jointly with a character/attribute embedding space.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (same approach IPTransELinker uses).	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs. Seeds the shared embedding tables (each pair's target is aliased to its source's id, in both the SE and CE spaces).	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint.	required
attribute_triples1	list[Triple]	KG1's (entity_id, predicate, value) attribute triples. Required -- see the module docstring's deviation note.	required
attribute_triples2	list[Triple]	KG2's attribute triples. Required.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
AttrELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If attribute_triples1 and attribute_triples2 are both empty, or if none of ground_truth's pairs have both ids present in dataset1/dataset2, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/attre.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> AttrELinker:
    """Train a shared TransE space jointly with a character/attribute embedding space.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (same approach
            [IPTransELinker][linkingtk.algorithms.ea.iptranse.IPTransELinker]
            uses).
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs. Seeds the shared embedding tables (each pair's
            target is aliased to its source's id, in both the SE and
            CE spaces).
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples. **Required** -- see the module
            docstring's deviation note.
        attribute_triples2: KG2's attribute triples. **Required**.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If ``attribute_triples1`` and
            ``attribute_triples2`` are both empty, or if none of
            ``ground_truth``'s pairs have both ids present in
            ``dataset1``/``dataset2``, or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    if not attribute_triples1 and not attribute_triples2:
        raise LinkingTKError(
            "AttrELinker.fit() needs attribute triples -- `attribute_triples1` "
            "and `attribute_triples2` are both empty, leaving nothing to train "
            "the character/attribute embedding half with."
        )

    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("AttrELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    seed_pairs = [(s, t) for s, t in ground_truth if s in ids1 and t in ids2]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in "
            "`dataset1`/`dataset2`; fit() has no seed pairs to train "
            "AttrE's shared embedding tables with."
        )

    entity_to_id, relation_to_id = build_shared_id_mappings(
        triples1_labels + triples2_labels, seed_pairs
    )
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)
    se_ctx1 = build_kg_context(mapped1)
    se_ctx2 = build_kg_context(mapped2)

    (
        ce_ctx1,
        ce_ctx2,
        attribute_to_id,
        char_to_id,
        value_char_ids_tensor,
    ) = self._build_attribute_context(
        torch, entity_to_id, attribute_triples1, attribute_triples2, device
    )

    entity_embeds = self._init_embedding(torch, functional, len(entity_to_id), device)
    relation_embeds = self._init_embedding(torch, functional, len(relation_to_id), device)
    entity_embeds_ce = self._init_embedding(torch, functional, len(entity_to_id), device)
    attr_embeds = self._init_embedding(torch, functional, max(1, len(attribute_to_id)), device)
    char_embeds = self._init_embedding(torch, functional, len(char_to_id) + 1, device)

    se_optimizer = torch.optim.SGD([entity_embeds, relation_embeds], lr=self.learning_rate)
    ce_optimizer = torch.optim.SGD(
        [entity_embeds_ce, attr_embeds, char_embeds], lr=self.learning_rate
    )
    joint_optimizer = torch.optim.SGD([entity_embeds, entity_embeds_ce], lr=self.learning_rate)

    joint_entity_ids = np.array(sorted(entity_to_id.values()), dtype=np.int64)
    joint_steps = max(1, math.ceil(len(joint_entity_ids) / self.batch_size))

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_structural_epoch(
            entity_embeds,
            relation_embeds,
            se_optimizer,
            se_ctx1,
            se_ctx2,
            rng,
            self.batch_size,
            self.margin,
        )
        train_attr_epoch(
            entity_embeds_ce,
            attr_embeds,
            char_embeds,
            value_char_ids_tensor,
            ce_optimizer,
            ce_ctx1,
            ce_ctx2,
            rng,
            self.batch_size,
            self.margin,
        )
        train_joint_epoch(
            entity_embeds, entity_embeds_ce, joint_optimizer, joint_entity_ids, joint_steps
        )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/attre.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/attre.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

IMUSE

imuse

IMUSE knowledge-graph-embedding linker for Entity Alignment.

He, F., Li, Z., Qiao, Q., Liu, A., & Zhao, L. (2019). Unsupervised Entity Alignment Using Attribute Triples and Relation Triples. DASFAA 2019. https://link.springer.com/content/pdf/10.1007%2F978-3-030-18576-3_22.pdf

Like KDCoELinker and AttrELinker, this is a faithful reimplementation of IMUSE's actual training procedure, ported directly from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/imuse.py, its modules/base/initializers.py/modules/base/losses.py, and models/basic_model.py) rather than from a from-scratch reading of the paper.

IMUSE is unsupervised -- unlike every other linker in this family, fit takes no ground_truth parameter at all. Reading imuse.py confirms there's no seed-pair argument anywhere in OpenEA's own IMUSE class:

• Bootstrap alignment from attribute string similarity, before any embedding training starts (ports interactive_model): entity attribute predicates are paired across KGs by Levenshtein-ratio similarity of their URIs' local names, kept to the 10 pairs with the most combined triple support; entities are then paired by the average Levenshtein ratio of their values under those aligned predicates. This fixed set is IMUSE's only source of cross-KG signal.
• Likely bug found while reading the reference source, not obvious from a paper summary: OpenEA's own run_one_ea/run_one_ae accept a candidate match inside the inner per-candidate scan loop, on every improving match, rather than after scanning all candidates -- so an earlier, since-superseded match for the same entity stays in the result set and permanently claims its target, even after a strictly better match is found later in the same scan. This is inconsistent with the textbook-correct version of the identical greedy-accept pattern used a few dozen lines away in the same file (get_aligned_attr_pair_by_name_similarity). Treated as an unintentional implementation bug rather than a deliberate design choice -- unlike this family's usual "port literal quirks faithfully" posture for behavior that's actually load-bearing for a published number (e.g. this linker's own full-list epoch re-runs, below) -- this port accepts only the single best match found after a full scan. See _imuse_text.py's align_entities_by_attribute_values docstring.
• Tractability: OpenEA brute-forces every candidate entity pair (~225M comparisons at this dataset's scale), parallelized across 8 processes. A single-process port instead indexes attribute values per KG and only scores entity pairs that share at least one value (after casefolding) under an aligned predicate -- entities with no shared aligned-predicate value can only ever score 0 similarity anyway, so this loses no correct matches from OpenEA's own similarity formula, only ones whose only signal is a near-but-not-exact value match. See _imuse_text.py's _value_index docstring.
• A real, diagnosed benchmark finding: at OpenEA's own hardcoded 0.6 name-similarity threshold, live EN-FR-15K-V1 data lets a spurious attribute-predicate pair through -- .../ontology/games/foaf/0.1/name score 0.667 on local-name Levenshtein ratio (genuinely above 0.6, not a bug), and having more combined triple support than the correct self-match foaf/0.1/name/foaf/0.1/name, it wins that predicate's top-10 slot instead. Verified directly (not assumed): this one wrong predicate pair drags the entity-bootstrap step's precision against this dataset's own known links down to 54.7%; raising name_sim_threshold to 0.9 removes it and raises precision to 82.1%, taking live Hits@1 from 0.470 to 0.571 with no other change. This is very likely a rehost-specific attribute-predicate-vocabulary difference from OpenEA's own original dataset files (not present in whatever they benchmarked on), not a bug in this port -- see examples/imuse_benchmark.py, which uses name_sim_threshold=0.9 for exactly this reason, while the class default here stays at OpenEA's own literal 0.6.
• Structural (SE) half: once bootstrapping is done, IMUSE trains a margin-based TransE embedding over relation triples -- the same shape every linker in this family uses, reusing _kdcoe_torch.py's KGContext/build_kg_context/train_structural_epoch directly. Entity/relation ids come from one combined id space across both KGs (like MTransELinker/ KDCoELinker), not the shared-id "sharing" merge IPTransELinker/ JAPELinker/ AttrELinker use -- IMUSE never merges rows for aligned pairs; alignment is a separate loss term (next point).
• Alignment loss -- the genuinely new part: unlike every other linker here, IMUSE has no mapping matrix and no shared-id merge. Instead, a direct sum(||ent_embeds[e1] - ent_embeds[e2]||^2) loss pulls the two (still separate-row) embeddings of each bootstrapped pair together. OpenEA's own launch_align_training_1epo re-runs this over the full bootstrapped-pair list steps times per epoch rather than sub-batching it -- the same full-list-re-run quirk AttrELinker's train_joint_epoch already documents; ported the same way here (see _imuse_torch.py's train_align_epoch).
• optimizer="SGD" -- like AttrE, not Adagrad (this family's more common default).
• Even though training is unsupervised, early stopping is not: val_ground_truth is still accepted and used exactly like every other linker here (OpenEA's own run() checks Hits@1 against kgs.valid_links every eval_freq epochs) -- the "unsupervised" claim is specifically about training signal, not about early stopping.

New dependency: rapidfuzz (added to the kge optional-dependency group alongside pykeen/torch) for Levenshtein-ratio string similarity -- OpenEA's own bootstrap relies on python-Levenshtein's .ratio(); rapidfuzz.distance.Indel.normalized_similarity computes the identical formula (rapidfuzz.distance.Levenshtein alone uses a different substitution cost/normalization -- see _imuse_text.py's levenshtein_ratio docstring). Imported lazily (inside functions, not at module top) in _imuse_text.py, matching this package's established torch-avoidance convention, so importing linkingtk.algorithms.ea doesn't require the kge extra just to reach an unrelated linker.

Same dataset note as JAPE/KDCoE/AttrE: this uses EnFr15KAttrDataset, not EnFr15KDataset, since IMUSE's entire mechanism depends on attribute triples and that loader has none.

IMUSELinker

IMUSELinker(
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    name_sim_threshold: float = 0.6,
    entity_sim_threshold: float = 0.6,
    attr_sim_threshold: float = 0.6,
    top_k_attribute_pairs: int = 10,
    bootstrap_iterations: int = 1,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Unsupervised EA: bootstraps its own seed pairs from attribute-value similarity.

Must be fit before link can be called. Unlike every other linker in this family, fit takes no ground_truth -- see the module docstring.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each pairing one epoch of structural triple-loss training with one epoch of alignment-loss training. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	2000
batch_size	int	Mini-batch size for structural triple-loss training; the alignment-loss step count is derived from it (matching OpenEA).	5000
learning_rate	float	SGD's learning rate, for both optimizers -- OpenEA's published value is 0.01.	0.01
margin	float	Hinge margin for the structural triple loss. OpenEA's published value is 1.5.	1.5
name_sim_threshold	float	Minimum Levenshtein ratio for two attribute predicates' local names to be considered aligned. OpenEA hardcodes 0.6 at this call site.	0.6
entity_sim_threshold	float	Minimum average Levenshtein ratio (across shared aligned-predicate values) for two entities to be considered aligned. OpenEA's published value (sim_thresholds_ent) is 0.6.	0.6
attr_sim_threshold	float	Same as entity_sim_threshold but for re-aligning attribute predicates from aligned entities (sim_thresholds_attr, only used if ``bootstrap_iterations 1). OpenEA's published value is0.6``.	0.6
top_k_attribute_pairs	int	Maximum aligned attribute-predicate pairs kept after name-similarity matching, ranked by combined triple count. OpenEA hardcodes 10.	10
bootstrap_iterations	int	Outer bootstrap rounds (entity alignment, then optionally attribute-predicate re-alignment from those entities). OpenEA's published default (interactive_model_iter_num) is 1, at which the attribute-predicate-realignment step never runs -- see the module docstring.	1
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/imuse.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    name_sim_threshold: float = 0.6,
    entity_sim_threshold: float = 0.6,
    attr_sim_threshold: float = 0.6,
    top_k_attribute_pairs: int = 10,
    bootstrap_iterations: int = 1,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.margin = margin
    self.name_sim_threshold = name_sim_threshold
    self.entity_sim_threshold = entity_sim_threshold
    self.attr_sim_threshold = attr_sim_threshold
    self.top_k_attribute_pairs = top_k_attribute_pairs
    self.bootstrap_iterations = bootstrap_iterations
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> IMUSELinker

Bootstrap an alignment from attribute similarity, then train jointly.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- used to partition graph into each KG's own triples (same approach KDCoELinker uses).	required
dataset2	list[Entity]	Target entities. See dataset1.	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint.	required
attribute_triples1	list[Triple]	KG1's (entity_id, predicate, value) attribute triples -- required, not optional (unlike JAPELinker/ KDCoELinker): IMUSE has no ground truth at all, so bootstrapping from attributes is its only source of cross-KG signal, with no fallback.	required
attribute_triples2	list[Triple]	KG2's attribute triples. See attribute_triples1.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. Unlike ground_truth, this is accepted here -- see the module docstring's note on what "unsupervised" does and doesn't cover. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
IMUSELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If both attribute_triples1 and attribute_triples2 are empty (nothing to bootstrap an alignment with), or if bootstrapping finds no confident aligned entity pairs at all (nothing to train the alignment loss with), or if device is invalid or unavailable.
OptionalDependencyError	If torch or rapidfuzz isn't installed.

Source code in src/linkingtk/algorithms/ea/imuse.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> IMUSELinker:
    """Bootstrap an alignment from attribute similarity, then train jointly.

    Args:
        dataset1: Source entities -- used to partition ``graph`` into
            each KG's own triples (same approach
            [KDCoELinker][linkingtk.algorithms.ea.kdcoe.KDCoELinker]
            uses).
        dataset2: Target entities. See ``dataset1``.
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples -- **required**, not optional (unlike
            [JAPELinker][linkingtk.algorithms.ea.jape.JAPELinker]/
            [KDCoELinker][linkingtk.algorithms.ea.kdcoe.KDCoELinker]):
            IMUSE has no ground truth at all, so bootstrapping from
            attributes is its only source of cross-KG signal, with no
            fallback.
        attribute_triples2: KG2's attribute triples. See
            ``attribute_triples1``.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. Unlike ``ground_truth``, this
            *is* accepted here -- see the module docstring's note on
            what "unsupervised" does and doesn't cover. If ``None``
            (default), trains the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If both ``attribute_triples1`` and
            ``attribute_triples2`` are empty (nothing to bootstrap an
            alignment with), or if bootstrapping finds no confident
            aligned entity pairs at all (nothing to train the
            alignment loss with), or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch or rapidfuzz isn't installed.
    """
    try:
        import rapidfuzz  # noqa: F401 -- import-only check, used inside _imuse_text.py
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("IMUSELinker", "kge") from exc

    if not attribute_triples1 and not attribute_triples2:
        raise LinkingTKError(
            "Both `attribute_triples1` and `attribute_triples2` are empty; "
            "IMUSELinker.fit() has no attribute signal to bootstrap an "
            "alignment with (it has no ground_truth to fall back on)."
        )

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1 = [t for t in triples if t[0] in ids1]
    triples2 = [t for t in triples if t[0] in ids2]

    entity_to_id, relation_to_id = build_id_mappings(triples1 + triples2)
    mapped1 = map_triples_to_ids(triples1, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(mapped1)
    ctx2 = build_kg_context(mapped2)

    bootstrapped_pairs = bootstrap_alignment(
        attribute_triples1,
        attribute_triples2,
        self.name_sim_threshold,
        self.entity_sim_threshold,
        self.attr_sim_threshold,
        self.top_k_attribute_pairs,
        self.bootstrap_iterations,
    )
    seed_pairs = [
        (s, t) for s, t in bootstrapped_pairs if s in entity_to_id and t in entity_to_id
    ]
    if not seed_pairs:
        raise LinkingTKError(
            "Bootstrapping from attribute-value similarity found no confident "
            "aligned entity pairs (or none with both ids present in `graph`'s "
            "own triples); fit() has no seed pairs to train IMUSE's alignment "
            "loss with."
        )
    seed_source_ids = np.array([entity_to_id[s] for s, _ in seed_pairs], dtype=np.int64)
    seed_target_ids = np.array([entity_to_id[t] for _, t in seed_pairs], dtype=np.int64)

    entity_embeds, relation_embeds = self._init_embeddings(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    triple_optimizer = torch.optim.SGD([entity_embeds, relation_embeds], lr=self.learning_rate)
    align_optimizer = torch.optim.SGD([entity_embeds], lr=self.learning_rate)
    align_steps = max(1, math.ceil(len(seed_pairs) / self.batch_size))

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_structural_epoch(
            entity_embeds,
            relation_embeds,
            triple_optimizer,
            ctx1,
            ctx2,
            rng,
            self.batch_size,
            self.margin,
        )
        train_align_epoch(
            entity_embeds, align_optimizer, seed_source_ids, seed_target_ids, align_steps
        )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/imuse.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/imuse.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

MultiKE

multike

MultiKE knowledge-graph-embedding linker for Entity Alignment.

Zhang, Q., Sun, Z., Hu, W., Chen, M., Guo, L., & Qu, Y. (2019). Multi-view Knowledge Graph Embedding for Entity Alignment. IJCAI 2019. https://www.ijcai.org/proceedings/2019/0754.pdf

By far the most complex method in this family (see MTransELinker, IPTransELinker, JAPELinker, KDCoELinker, AttrELinker, IMUSELinker) -- a faithful reimplementation of MultiKE's actual training procedure, ported directly from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/multi_ke.py, 931 lines, read in full, plus literal_encoder.py, predicate_alignmnet.py, modules/base/losses.py, modules/load/read.py, models/basic_model.py) rather than a from-scratch paper reading.

MultiKE is supervised (unlike IMUSE): seed pairs are its only real supervision channel, via literal entity-id substitution into each KG's own triples (see _multike_text.generate_cross_kg_relation_triples/ generate_cross_kg_attribute_triples, ported from read.py's generate_sup_relation_triples/generate_sup_attribute_triples).

Architecture: three views unified into one shared table

1. Name view (name_embeds): a literal-text embedding of each entity's local name, computed once and frozen for the rest of training. See "The literal encoder" below for the one place this port genuinely departs from a faithful reimplementation.
2. Relation view (rv_ent_embeds/rel_embeds): logistic-loss TransE over relation triples (softplus(||h+r-t||^2) + softplus(-||h'+r-t'||^2)) -- not the margin-based loss every other method in this family uses; MultiKE hardcodes this, ignoring the reference's own configurable args.loss.
3. Attribute view (av_ent_embeds/attr_embeds): a small CNN (_multike_torch.build_attribute_scorer) predicts a "head" representation from (attribute, value) pairs, scored via negative squared distance against the real head.
4. Shared space (ent_embeds, used for every eval call -- confirmed in basic_model.py's valid()/test(), both always read self.ent_embeds regardless of which views exist): trained by alignment terms baked directly into the relation/attribute view losses and a separate common-space-learning step (_multike_torch.train_common_space_epoch) pulling ent_embeds toward all three views for every entity, at its own learning rate (common_space_learning_rate, distinct from the main learning_rate).
5. Cross-KG entity inference (the actual supervision channel): for each seed pair, every KG1 triple touching the source entity is relabeled with the target's id and trained back into the shared tables -- see _multike_text.py's substitution functions.
6. Predicate soft-alignment: relation and attribute predicates are aligned across KGs by mutual-best-match Levenshtein ratio on local names (_multike_text.align_predicates_by_name), feeding a per-triple weight in the attribute view's main loss and two more cross-KG inference loss terms, gated to start at start_predicate_soft_alignment.

Three deliberate scope cuts, documented plainly

• Space mapping (nv_mapping/rv_mapping/av_mapping orthogonal matrices, train_shared_space_mapping_1epo): defined in the reference's init() but the training loop that calls it is commented out in its own run() (multi_ke.py:923-926) -- genuinely dead code in the reference's own published run, confirmed by reading run() itself. Not ported at all.
• Periodic predicate-alignment re-estimation (update_predicate_alignment/find_predicate_alignment_by_embedding, every 10 epochs): refines the Levenshtein-based predicate pairs using the training embeddings' own cosine similarity. This port computes the Levenshtein-based alignment once at fit() start and keeps it fixed for the whole run (mirrors IMUSE's own fixed-bootstrap precedent) -- a refinement pass over an already-working initial signal, not the primary mechanism.
• Truncated/nearest-neighbor negative sampling: not ported, matching every other linker in this family's established uniform-negative-sampling precedent (not a new deviation specific to MultiKE).

Also confirmed by reading every loss graph in full (not obvious from a method-name skim, see _multike_torch.py's module docstring for the full breakdown): only the main relation-view loss ever uses negative sampling -- every other loss (attribute view, all four cross-KG inference terms) is purely positive/regression-style. And PredicateAlignModel.relation_triples_w_weights is computed by the reference but never actually consumed anywhere in its own training methods -- not ported, since it's unused in the reference itself.

The literal encoder

OpenEA's own name_embeds/literal_embeds come from a pretrained English-only word2vec/fastText file plus a custom stacked autoencoder trained via reconstruction loss to compress bag-of-word-vectors down to embedding_dim. This port instead uses a small pretrained multilingual transformer (transformers -- already a hard dependency of this repo, zero new install) via _multike_literal.encode_literals: mean-pooled over the attention mask, projected down to embedding_dim via a fixed random orthogonal projection (no training loop needed -- the transformer already produces well-formed semantic structure). Real multilingual semantic signal for cross-lingual name matching (e.g. "United States"/"États-Unis"), which an English-only fastText file (_kdcoe_text.py's existing load_fasttext_vectors infrastructure was considered and is a real precedent in this repo, but doesn't fit a cross-lingual signal) or a from-scratch in-corpus word2vec couldn't provide. Default model: distilbert-base-multilingual-cased. First run downloads it from the Hub (~540MB), cached afterward.

Other implementation notes

• Combined, not shared-id, entity space: like MTransELinker/ KDCoELinker/ IMUSELinker -- MultiKE never merges seed pairs into one row; alignment is entirely loss-driven.
• Which embedding tables get L2-normalized on every read isn't uniform -- rv_ent_embeds/rel_embeds/av_ent_embeds/ ent_embeds always are; attr_embeds never is (ported from the reference's own xavier_init(..., is_l2_norm=False) for that one table); see _multike_torch.py's module docstring for the full detail.
• optimizer="SGD", matching AttrE/IMUSE's precedent in this family.
• Final scoring uses the shared ent_embeds directly, cosine similarity, no mapping -- matches basic_model.py's valid()/ test() always reading self.ent_embeds.
• attribute_triples1/attribute_triples2 are required (raises LinkingTKError if both empty) -- the attribute view, the literal encoder's value vocabulary, and the predicate soft-alignment mechanism all depend on them, same posture as AttrE.
• Early stopping uses this repo's own plain patience-counter Hits@1 check rather than OpenEA's own two-flag early_stop(), and doesn't gate on OpenEA's start_valid, matching every other linker in this family.

MultiKELinker

MultiKELinker(
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    batch_size: int = 5000,
    entity_batch_size: int = 5000,
    attribute_batch_size: int = 5000,
    learning_rate: float = 0.001,
    common_space_learning_rate: float = 0.004,
    neg_triple_num: int = 10,
    cv_weight: float = 1.0,
    predicate_init_sim: float = 0.9,
    predicate_soft_sim: float = 0.8,
    start_predicate_soft_alignment: int = 10,
    min_predicate_triple_count: int = 10,
    literal_encoder_model: str = "distilbert-base-multilingual-cased",
    literal_max_length: int = 16,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a shared embedding table unifying name/relation/attribute views.

Must be fit before link can be called. See the module docstring for the full architecture and the three documented scope cuts.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of every trained embedding table. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each one pass of relation-view, common-space, cross-KG relation-entity (and, once warmed up, cross-KG relation-predicate) training, then the same four steps for the attribute view. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	2000
batch_size	int	Mini-batch size for the relation view's main loss and its two cross-KG inference terms.	5000
entity_batch_size	int	Mini-batch size for common-space learning.	5000
attribute_batch_size	int	Mini-batch size for the attribute view's main loss and its two cross-KG inference terms.	5000
learning_rate	float	SGD's learning rate for every optimizer except common-space learning. OpenEA's published value is 0.001.	0.001
common_space_learning_rate	float	SGD's learning rate for common-space learning specifically (OpenEA's own ITC_learning_rate, deliberately distinct from learning_rate). Published value is 0.004.	0.004
neg_triple_num	int	Negatives sampled per positive relation triple -- only the main relation-view loss uses negative sampling at all (see the module docstring). Published value is 10.	10
cv_weight	float	Weight applied to the common-space-learning loss. Published value is 1.0.	1.0
predicate_init_sim	float	Levenshtein-ratio threshold (of humanized local names) for the one-time mutual-best-match predicate alignment. Published value is 0.9.	0.9
predicate_soft_sim	float	Similarity floor used to rescale an aligned predicate pair's weight (_multike_text.zoom_weight). Published value is 0.80.	0.8
start_predicate_soft_alignment	int	Epoch (0-indexed) at which the two predicate-alignment cross-KG loss terms start contributing -- a warm-up gate. Published value is 10.	10
min_predicate_triple_count	int	Attribute predicates with fewer than this many combined (both-KG) triples are dropped before anything else. Published value is 10.	10
literal_encoder_model	str	Hugging Face model id for the literal encoder (see the module docstring's "The literal encoder" section). Defaults to distilbert-base-multilingual-cased.	'distilbert-base-multilingual-cased'
literal_max_length	int	Token-truncation length passed to the literal encoder for every name/value string.	16
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Also used for the literal encoder's forward pass. Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/multike.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    batch_size: int = 5000,
    entity_batch_size: int = 5000,
    attribute_batch_size: int = 5000,
    learning_rate: float = 0.001,
    common_space_learning_rate: float = 0.004,
    neg_triple_num: int = 10,
    cv_weight: float = 1.0,
    predicate_init_sim: float = 0.9,
    predicate_soft_sim: float = 0.80,
    start_predicate_soft_alignment: int = 10,
    min_predicate_triple_count: int = 10,
    literal_encoder_model: str = "distilbert-base-multilingual-cased",
    literal_max_length: int = 16,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.entity_batch_size = entity_batch_size
    self.attribute_batch_size = attribute_batch_size
    self.learning_rate = learning_rate
    self.common_space_learning_rate = common_space_learning_rate
    self.neg_triple_num = neg_triple_num
    self.cv_weight = cv_weight
    self.predicate_init_sim = predicate_init_sim
    self.predicate_soft_sim = predicate_soft_sim
    self.start_predicate_soft_alignment = start_predicate_soft_alignment
    self.min_predicate_triple_count = min_predicate_triple_count
    self.literal_encoder_model = literal_encoder_model
    self.literal_max_length = literal_max_length
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> MultiKELinker

Train the three-view embedding and a shared space unifying them.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs -- MultiKE's only real supervision channel (see the module docstring's "cross-KG entity inference").	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint.	required
attribute_triples1	list[Triple]	KG1's (entity_id, predicate, value) attribute triples. Required -- see the module docstring's deviation note.	required
attribute_triples2	list[Triple]	KG2's attribute triples. Required.	required
random_state	int | None	Seed for reproducible training (also seeds the literal encoder's random projection). Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
MultiKELinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If attribute_triples1 and attribute_triples2 are both empty, or if none of ground_truth's pairs have both ids present in graph's own triples, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/multike.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> MultiKELinker:
    """Train the three-view embedding and a shared space unifying them.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs -- MultiKE's only real supervision channel (see the
            module docstring's "cross-KG entity inference").
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples. **Required** -- see the module
            docstring's deviation note.
        attribute_triples2: KG2's attribute triples. **Required**.
        random_state: Seed for reproducible training (also seeds the
            literal encoder's random projection). Left unspecified,
            training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If ``attribute_triples1`` and
            ``attribute_triples2`` are both empty, or if none of
            ``ground_truth``'s pairs have both ids present in
            ``graph``'s own triples, or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    if not attribute_triples1 and not attribute_triples2:
        raise LinkingTKError(
            "MultiKELinker.fit() needs attribute triples -- `attribute_triples1` "
            "and `attribute_triples2` are both empty, leaving nothing to train "
            "the attribute view or literal encoder with."
        )

    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("MultiKELinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1 = [t for t in triples if t[0] in ids1]
    triples2 = [t for t in triples if t[0] in ids2]

    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped1 = map_triples_to_ids(triples1, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(mapped1)
    ctx2 = build_kg_context(mapped2)

    seed_pairs = [(s, t) for s, t in ground_truth if s in entity_to_id and t in entity_to_id]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train MultiKE's cross-KG "
            "entity inference with -- its only real supervision channel."
        )

    attr1_in = [t for t in attribute_triples1 if t[0] in entity_to_id]
    attr2_in = [t for t in attribute_triples2 if t[0] in entity_to_id]
    attr1_filtered, attr2_filtered = filter_frequent_predicates(
        attr1_in, attr2_in, self.min_predicate_triple_count
    )
    attr1_clean = self._clean_triples(attr1_filtered)
    attr2_clean = self._clean_triples(attr2_filtered)

    attribute_to_id = {
        predicate: index
        for index, predicate in enumerate(
            sorted({p for _, p, _ in attr1_clean} | {p for _, p, _ in attr2_clean})
        )
    }
    values = sorted({v for _, _, v in attr1_clean} | {v for _, _, v in attr2_clean})
    value_to_id = {value: index for index, value in enumerate(values)}

    def map_attribute_triples(mapped_triples: list[Triple]) -> npt.NDArray[np.int64]:
        rows = [
            (entity_to_id[e], attribute_to_id[a], value_to_id[v]) for e, a, v in mapped_triples
        ]
        return np.array(rows, dtype=np.int64) if rows else np.empty((0, 3), dtype=np.int64)

    attr_mapped1 = map_attribute_triples(attr1_clean)
    attr_mapped2 = map_attribute_triples(attr2_clean)

    aligned_relations = align_predicates_by_name(
        {p for _, p, _ in triples1}, {p for _, p, _ in triples2}, self.predicate_init_sim
    )
    aligned_attributes = align_predicates_by_name(
        {p for _, p, _ in attr1_clean}, {p for _, p, _ in attr2_clean}, self.predicate_init_sim
    )

    weighted1 = weight_attribute_triples(
        attr1_clean, aligned_attributes, self.predicate_soft_sim, is_kg1=True
    )
    weighted2 = weight_attribute_triples(
        attr2_clean, aligned_attributes, self.predicate_soft_sim, is_kg1=False
    )
    actx1 = build_attribute_context(
        attr_mapped1, np.array([w for *_, w in weighted1], dtype=np.float32)
    )
    actx2 = build_attribute_context(
        attr_mapped2, np.array([w for *_, w in weighted2], dtype=np.float32)
    )

    cross_rel1, cross_rel2 = generate_cross_kg_relation_triples(seed_pairs, triples1, triples2)
    cross_rel_ids = np.concatenate(
        [
            map_triples_to_ids(cross_rel1, entity_to_id, relation_to_id),
            map_triples_to_ids(cross_rel2, entity_to_id, relation_to_id),
        ],
        axis=0,
    )
    cross_attr1, cross_attr2 = generate_cross_kg_attribute_triples(
        seed_pairs, attr1_clean, attr2_clean
    )
    cross_attr_ids = np.concatenate(
        [map_attribute_triples(cross_attr1), map_attribute_triples(cross_attr2)], axis=0
    )

    pred_rel1_ids, pred_rel1_w = self._relation_ids_and_weights(
        substitute_predicate_triples(triples1, aligned_relations, is_kg1=True),
        entity_to_id,
        relation_to_id,
    )
    pred_rel2_ids, pred_rel2_w = self._relation_ids_and_weights(
        substitute_predicate_triples(triples2, aligned_relations, is_kg1=False),
        entity_to_id,
        relation_to_id,
    )
    pred_rel_ids = np.concatenate([pred_rel1_ids, pred_rel2_ids], axis=0)
    pred_rel_w = np.concatenate([pred_rel1_w, pred_rel2_w], axis=0)

    pred_attr1_ids, pred_attr1_w = self._attribute_ids_and_weights(
        substitute_predicate_triples(attr1_clean, aligned_attributes, is_kg1=True),
        entity_to_id,
        attribute_to_id,
        value_to_id,
    )
    pred_attr2_ids, pred_attr2_w = self._attribute_ids_and_weights(
        substitute_predicate_triples(attr2_clean, aligned_attributes, is_kg1=False),
        entity_to_id,
        attribute_to_id,
        value_to_id,
    )
    pred_attr_ids = np.concatenate([pred_attr1_ids, pred_attr2_ids], axis=0)
    pred_attr_w = np.concatenate([pred_attr1_w, pred_attr2_w], axis=0)

    name_embeds, literal_embeds = self._encode_literals(
        dataset1, dataset2, entity_to_id, values, random_state, device
    )
    name_embeds_t = torch.from_numpy(name_embeds).to(device)
    literal_embeds_t = torch.from_numpy(literal_embeds).to(device)

    rv_ent_embeds = self._init_embedding(torch, functional, len(entity_to_id), device)
    rel_embeds = self._init_embedding(torch, functional, max(1, len(relation_to_id)), device)
    av_ent_embeds = self._init_embedding(torch, functional, len(entity_to_id), device)
    attr_embeds = self._init_embedding_raw(torch, max(1, len(attribute_to_id)), device)
    ent_embeds = self._init_embedding(torch, functional, len(entity_to_id), device)
    scorer = build_attribute_scorer(self.embedding_dim).to(device)

    relation_optimizer = torch.optim.SGD(
        [rv_ent_embeds, rel_embeds, ent_embeds], lr=self.learning_rate
    )
    attribute_optimizer = torch.optim.SGD(
        [av_ent_embeds, attr_embeds, ent_embeds, *scorer.parameters()], lr=self.learning_rate
    )
    common_space_optimizer = torch.optim.SGD(
        [ent_embeds, rv_ent_embeds, av_ent_embeds], lr=self.common_space_learning_rate
    )
    cross_kg_relation_entity_optimizer = torch.optim.SGD(
        [rv_ent_embeds, rel_embeds], lr=self.learning_rate
    )
    cross_kg_attribute_entity_optimizer = torch.optim.SGD(
        [av_ent_embeds, attr_embeds, *scorer.parameters()], lr=self.learning_rate
    )
    cross_kg_relation_predicate_optimizer = torch.optim.SGD(
        [rv_ent_embeds, rel_embeds], lr=self.learning_rate
    )
    cross_kg_attribute_predicate_optimizer = torch.optim.SGD(
        [av_ent_embeds, attr_embeds, *scorer.parameters()], lr=self.learning_rate
    )

    entity_ids = np.arange(len(entity_to_id), dtype=np.int64)
    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_relation_view_epoch(
            rv_ent_embeds,
            rel_embeds,
            ent_embeds,
            name_embeds_t,
            relation_optimizer,
            ctx1,
            ctx2,
            rng,
            self.batch_size,
            self.neg_triple_num,
        )
        train_common_space_epoch(
            ent_embeds,
            rv_ent_embeds,
            av_ent_embeds,
            name_embeds_t,
            common_space_optimizer,
            entity_ids,
            rng,
            self.entity_batch_size,
            self.cv_weight,
        )
        train_cross_kg_relation_entity_epoch(
            rv_ent_embeds,
            rel_embeds,
            cross_kg_relation_entity_optimizer,
            cross_rel_ids,
            rng,
            self.batch_size,
        )
        if epoch >= self.start_predicate_soft_alignment:
            train_cross_kg_relation_predicate_epoch(
                rv_ent_embeds,
                rel_embeds,
                cross_kg_relation_predicate_optimizer,
                pred_rel_ids,
                pred_rel_w,
                rng,
                self.batch_size,
            )

        train_attribute_view_epoch(
            av_ent_embeds,
            attr_embeds,
            literal_embeds_t,
            ent_embeds,
            name_embeds_t,
            scorer,
            attribute_optimizer,
            actx1,
            actx2,
            rng,
            self.attribute_batch_size,
        )
        train_common_space_epoch(
            ent_embeds,
            rv_ent_embeds,
            av_ent_embeds,
            name_embeds_t,
            common_space_optimizer,
            entity_ids,
            rng,
            self.entity_batch_size,
            self.cv_weight,
        )
        train_cross_kg_attribute_entity_epoch(
            av_ent_embeds,
            attr_embeds,
            literal_embeds_t,
            scorer,
            cross_kg_attribute_entity_optimizer,
            cross_attr_ids,
            rng,
            self.attribute_batch_size,
        )
        if epoch >= self.start_predicate_soft_alignment:
            train_cross_kg_attribute_predicate_epoch(
                av_ent_embeds,
                attr_embeds,
                literal_embeds_t,
                scorer,
                cross_kg_attribute_predicate_optimizer,
                pred_attr_ids,
                pred_attr_w,
                rng,
                self.attribute_batch_size,
            )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(ent_embeds, dim=1).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(ent_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/multike.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/multike.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

RSN4EA

rsn4ea

RSN4EA knowledge-graph-embedding linker for Entity Alignment.

Guo, L., Sun, Z., & Hu, W. (2019). Learning to Exploit Long-term Relational Dependencies in Knowledge Graphs. ICML 2019. http://proceedings.mlr.press/v97/guo19c/guo19c.pdf

Structurally the most different method in this family: instead of a triple-local translational embedding (TransE-style, like every other linker in linkingtk.algorithms.ea), RSN4EA trains a Recurrent Skipping Network -- an LSTM run over sampled random-walk paths through the combined graph, with a next-token prediction loss (entity/relation language-modeling, not a distance loss), and a residual "skip" connection letting each relation-step's prediction shortcut directly from the preceding entity's embedding. Cross-KG alignment comes from a data-construction trick, not a mapping matrix or a shared-id merge: seed pairs make their two entity ids mutually substitutable throughout the walkable graph, so a random walk that reaches one can carry on through the other KG's own edges from that point.

Like MTransELinker, this is a faithful reimplementation of RSN4EA's actual training procedure, ported directly from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/rsn4ea.py) rather than from a from-scratch reading of the paper. Several of OpenEA's published config values turn out to be dead -- never actually consulted by rsn4ea.py's own code -- and this port follows what's actually executed:

• ent_l2_norm/rel_l2_norm are true in OpenEA's config, but _define_variables builds embeddings via plain tf.get_variable(..., xavier_initializer), never through the init_embeddings() utility that applies l2_normalize -- unlike every other linker in this family, RSN4EA's embeddings are never L2-normalized, at init or per-forward-pass.
• "optimizer": "Adagrad" in the config, but _define_variables hardcodes tf.train.AdamOptimizer(options.learning_rate) directly, bypassing generate_optimizer(..., opt=self.args.optimizer) -- Adam is what actually runs, so this port uses Adam unconditionally (not a constructor option).
• "dim": 100 is never read anywhere in rsn4ea.py -- embedding size is driven entirely by hidden_size.
• add_weight's w_h/w_r/w_t alias-tracking columns are computed but never consumed by sample_paths/cal_loss -- only the structural effect of alias substitution (which triples exist) is load-bearing.
• add_weight's relation-alias variants are always empty (RSN4EA's own rel_mapping is always an empty table -- no relation-alignment signal exists for this task), so only the 4 live head/tail-alias variants are ported (see _rsn4ea_training.build_augmented_kb).
• alignment_module: "mapping" has zero effect on RSN4EA's own training (run() fully overrides BasicModel.run(), never calling launch_training_1epo's mapping-loss branch) -- it only matters at the id-space level, which is exactly build_id_mappings's disjoint-per-KG shape (matching MTransELinker/ SEALinker), used here too.

Ported into two private helper modules: _rsn4ea_training.py (plain numpy -- graph augmentation and path sampling, independently testable without torch) and _rsn4ea_torch.py (the model and training-step functions that build/consume PyTorch tensors).

Deliberate deviations from OpenEA's own code, beyond what those two modules already document:

• Full softmax cross-entropy instead of tf.nn.nce_loss. OpenEA's NCE loss approximates a full softmax via log-uniform negative sampling -- a 2019-TF1-era efficiency hack for large vocabularies. At this dataset's real scale (~15K entities, low hundreds of relations doubled), OpenEA's own num_samples is already capped to vocab_size // 3 for entities -- nearly a third of the whole vocabulary sampled as noise every step regardless. A full softmax is exact (no noise-distribution correction term needed) and tractable at this vocab size on current GPU hardware -- the same "swap a reference-specific efficiency mechanism for a modern equivalent once its own justification no longer applies at this scale" precedent as BootEALinker's scipy.optimize.linear_sum_assignment substitution for graph_tool/igraph. num_samples is therefore not exposed as a constructor parameter.
• No disk-cached path table. OpenEA caches sampled paths to {data_path}/paths_{alpha}_{beta} and reloads if present. This repo's fit() takes in-memory data with no persistent working directory, so paths are always freshly sampled each call.
• max_paths cap. repeat_times walks per (already-expanded) row of the augmented KB can produce a large path table at real dataset scale. An optional max_paths cap (random subsample, fixed for the run) is exposed as a pragmatic safety valve -- same precedent as IPTransELinker's bootstrap_pool_size. Default None (unbounded, matching OpenEA).

RSN4EALinker

RSN4EALinker(
    hidden_size: int = 100,
    num_layers: int = 2,
    keep_prob: float = 0.6,
    max_length: int = 15,
    alpha: float = 0.7,
    beta: float = 0.7,
    repeat_times: int = 2,
    max_paths: int | None = None,
    num_epochs: int = 30,
    batch_size: int = 512,
    learning_rate: float = 0.0005,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a Recurrent Skipping Network over sampled walk paths.

Must be fit before link can be called. Unlike MTransELinker/ SEALinker, scoring uses raw (unprojected) embeddings on both sides -- RSN4EA never builds a mapping matrix; see the module docstring.

Parameters:

Name	Type	Description	Default
hidden_size	int	Entity/relation embedding dimensionality and LSTM hidden size. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_layers	int	LSTM depth. OpenEA's published value is 2.	2
keep_prob	float	Per-LSTM-layer output dropout keep-probability. OpenEA's published value is 0.6.	0.6
max_length	int	Sampled path length -- must be odd and >= 3 (entities at even positions, relations at odd positions). OpenEA's published value is 15.	15
alpha	float	Path-sampling depth-bias weight (down-weights continuations already directly reachable from a walk's own start entity). OpenEA's published value is 0.7.	0.7
beta	float	Path-sampling cross-KG bias weight (up-weights continuations whose tail is a seed-pair entity). OpenEA's published value is 0.7.	0.7
repeat_times	int	Walks sampled per row of the augmented knowledge base. OpenEA's published value is 2.	2
max_paths	int | None	Optional cap on the number of sampled paths kept for training (random subsample, fixed for the run) -- see the module docstring. None (default) keeps every successfully sampled path.	None
num_epochs	int	Training epochs. OpenEA's published value is 30 (RSN4EA converges in far fewer epochs than this family's TransE-based methods).	30
batch_size	int	Mini-batch size for path training. OpenEA's published value is 512.	512
learning_rate	float	Adam's learning rate (see the module docstring for why Adam, not the config's stated Adagrad). OpenEA's published value is 0.0005.	0.0005
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/rsn4ea.py

def __init__(
    self,
    hidden_size: int = 100,
    num_layers: int = 2,
    keep_prob: float = 0.6,
    max_length: int = 15,
    alpha: float = 0.7,
    beta: float = 0.7,
    repeat_times: int = 2,
    max_paths: int | None = None,
    num_epochs: int = 30,
    batch_size: int = 512,
    learning_rate: float = 0.0005,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    if max_length < 3 or max_length % 2 == 0:
        raise LinkingTKError(f"max_length must be odd and >= 3, got {max_length!r}.")
    self.hidden_size = hidden_size
    self.num_layers = num_layers
    self.keep_prob = keep_prob
    self.max_length = max_length
    self.alpha = alpha
    self.beta = beta
    self.repeat_times = repeat_times
    self.max_paths = max_paths
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> RSN4EALinker

Sample walk paths over an alias-expanded graph and train the RSN.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities. Unused beyond mirroring sibling linkers' signatures -- which entities get embeddings is determined entirely by graph.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs. Used to build the alias-substituted graph (_rsn4ea_training.build_augmented_kb) that lets sampled walks cross between KGs -- RSN4EA's only alignment-training mechanism, no separate loss channel.	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint, as they are from a GraphDatasetLoader's load_graphs().	required
random_state	int | None	Seed for reproducible path sampling and training. Left unspecified, both are non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
RSN4EALinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples, if no walk path could be sampled at all, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/rsn4ea.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> RSN4EALinker:
    """Sample walk paths over an alias-expanded graph and train the RSN.

    Args:
        dataset1: Source entities. Unused beyond mirroring sibling
            linkers' signatures -- which entities get embeddings is
            determined entirely by ``graph``.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs. Used to build the alias-substituted graph
            (`_rsn4ea_training.build_augmented_kb`)
            that lets sampled walks cross between KGs -- RSN4EA's only
            alignment-training mechanism, no separate loss channel.
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint, as they are
            from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``.
        random_state: Seed for reproducible path sampling and training.
            Left unspecified, both are non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples, if no walk path
            could be sampled at all, or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
    except ImportError as exc:
        raise OptionalDependencyError("RSN4EALinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    triples = to_triples(graph)
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped = map_triples_to_ids(triples, entity_to_id, relation_to_id)

    seed_pairs = [
        (entity_to_id[s], entity_to_id[t])
        for s, t in ground_truth
        if s in entity_to_id and t in entity_to_id
    ]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to build RSN4EA's alias-expanded "
            "graph with."
        )

    augmented_kb, num_relations = build_augmented_kb(
        mapped, len(entity_to_id), len(relation_to_id), seed_pairs
    )
    paths = sample_paths(
        augmented_kb,
        seed_pairs,
        rng,
        self.max_length,
        self.alpha,
        self.beta,
        self.repeat_times,
        self.max_paths,
    )
    if len(paths) == 0:
        raise LinkingTKError(
            "No walk paths of length `max_length` could be sampled from `graph` -- "
            "every entity reachable from a starting triple hit a dead end too soon."
        )

    model: torch.nn.Module = build_rsn_model(
        len(entity_to_id), num_relations, self.hidden_size, self.num_layers, self.keep_prob
    )
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=self.learning_rate)

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0
    best_state: dict[str, torch.Tensor] | None = None

    for epoch in range(self.num_epochs):
        train_epoch(model, optimizer, paths, rng, self.batch_size)

        if val_pairs and (epoch + 1) % eval_every == 0:
            model.eval()
            with torch.no_grad():
                current_embeds = _entity_embedding_weight(model).cpu().numpy()
            hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0
                best_state = {k: v.detach().clone() for k, v in model.state_dict().items()}

    # Unlike this family's other linkers, RSN4EA's validation Hits@1
    # doesn't improve (or plateau) monotonically with more training --
    # diagnosed directly on the real EN-FR-15K-V1 benchmark: it peaks
    # within the first few `eval_every`-spaced checks, then steadily
    # degrades (the next-token training objective is only a proxy for
    # alignment quality, and keeps optimizing well past the point
    # where the two diverge). Restoring the best-seen checkpoint
    # (rather than using whichever epoch training happened to stop or
    # finish on, like every sibling linker here) is what actually
    # fixes this -- not a hyperparameter tweak.
    if best_state is not None:
        model.load_state_dict(best_state)

    model.eval()
    with torch.no_grad():
        final_embeds = _entity_embedding_weight(model).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/rsn4ea.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/rsn4ea.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

SEA

sea

SEA knowledge-graph-embedding linker for Entity Alignment.

Pei, S., Yu, L., Hoehndorf, R., & Zhang, X. (2019). Semi-Supervised Entity Alignment via Knowledge Graph Embedding with Awareness of Degree Difference. WWW 2019. https://dl.acm.org/citation.cfm?id=3313646

Like MTransELinker, this is a faithful reimplementation of SEA's actual training procedure, ported directly from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/sea.py, plus its models/basic_model.py and modules/train/batch.py support code) rather than from a from-scratch reading of the paper.

SEA trains a single shared entity/relation embedding table with disjoint per-KG id ranges (OpenEA's alignment_module="mapping", the same shape MTransELinker uses -- unlike IPTransELinker's id-sharing or BootEALinker's pseudo-triple-injecting "swapping"), via standard margin-based TransE with uniform per-KG negative sampling -- no path loss, no bootstrapping, the simplest structural half in this family after MTransE. What makes SEA distinct is its mapping half:

• Two orthogonally-initialized square mapping matrices, mat_1 (KG1 -> KG2) and mat_2 (KG2 -> KG1) -- not MTransE's single matrix.
• A supervised term over seed pairs: mat_1-mapped source entities are pulled toward their real target (and symmetrically mat_2-mapped targets toward their real source).
• A semi-supervised (cycle-consistency) term over unlabeled entities: a KG1 entity mapped through mat_1 then back through mat_2 should return to itself (and symmetrically for KG2) -- this needs no ground-truth pairing at all, just entity ids from both sides, which is what "semi-supervised" refers to in the method's name.

Ported into a private helper module, _sea_torch.py (the training-step functions that build/consume PyTorch tensors; negative sampling itself reuses _iptranse_training.py's sample_negative_triples directly, since it's already fully generic).

Deliberately not ported: OpenEA's sea.py defines eye_mat_1/ eye_mat_2 (the same orthogonality-regularizer constants MTransELinker's mapping loss uses) but never references them in SEA's actual mapping loss -- no orthogonality regularization is applied here, unlike MTransE. Same "port actual, not apparent, reference behavior" situation as BootEALinker's skipped likelihood() and JAPELinker's inert sim_optimizer (#28) -- skipped here entirely.

Deviations from OpenEA's own code:

• var_list=None gotcha, checked and applicable: like MTransELinker/ KDCoELinker (see feedback_openea_var_list_none), OpenEA's generate_optimizer(mapping_loss, lr, opt='Adam') leaves var_list unset, so TensorFlow differentiates the mapping loss against every trainable variable it touches -- both mapping matrices and the shared entity embeddings (every term reads from self.ent_embeds). This port's mapping optimizer is scoped to [mapping_mat_1, mapping_mat_2, entity_embeds], not just "the two obvious" mapping matrices.
• Optimizer is Adam, not Adagrad -- unlike MTransE/IPTransE/BootEA. OpenEA's own SEA.init() asserts self.args.optimizer == 'Adam' and its published sea_args_15K.json confirms it.
• Unlabeled/semi-supervised pool: OpenEA samples from its own loader's test_links + valid_links (using only entity ids, never the gold pairing). This repo's fit() only receives dataset1/dataset2 (full entity rosters) and ground_truth (the train split), so the pool here is every entity not already a ground_truth source/ target, sampled independently per side -- the same broader-pool deviation already established for IPTransELinker's and BootEALinker's bootstrap pools, and in practice covers the same entities OpenEA's own pool would (val+test).
• KG1/KG2 triple partition: like IPTransELinker, this repo takes one combined graph in fit() rather than OpenEA's own separately-loaded per-KG KGs; a triple belongs to KG1 iff its head id is one of dataset1's ids, relying on the same disjoint-id assumption MTransELinker already documents.

SEALinker

SEALinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    alpha_1: float = 2.5,
    alpha_2: float = 0.25,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via structural TransE plus a semi-supervised dual mapping.

Must be fit before link can be called.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings and both (square) mapping matrices. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs, each alternating one epoch of structural triple-loss training with one epoch of mapping-loss training. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	500
batch_size	int	Mini-batch size for structural triple-loss training, split proportionally between the two KGs by triple count each epoch. The mapping loss reuses the resulting step count.	5000
learning_rate	float	Adam's learning rate, for both the structural and mapping optimizers -- OpenEA's published value is 0.01.	0.01
margin	float	Hinge margin for the structural triple loss. OpenEA's published value is 1.5.	1.5
alpha_1	float	Weight applied to the supervised mapping loss (seed-pair approximation error). OpenEA's published value is 2.5.	2.5
alpha_2	float	Weight applied to the semi-supervised (cycle-consistency) mapping loss over unlabeled entities. OpenEA's published value is 0.25.	0.25
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/sea.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    margin: float = 1.5,
    alpha_1: float = 2.5,
    alpha_2: float = 0.25,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.margin = margin
    self.alpha_1 = alpha_1
    self.alpha_2 = alpha_2
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._mapping: npt.NDArray[np.floating[Any]] | None = None
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> SEALinker

Train shared structural TransE embeddings plus a dual cross-KG mapping.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples for structural negative sampling and to build the semi-supervised unlabeled pool (see the module docstring).	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs used to train the supervised half of the mapping loss. Unlike KGELinker, these are not injected into graph as extra triples.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
SEALinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples -- nothing to train the mapping with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/sea.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> SEALinker:
    """Train shared structural TransE embeddings plus a dual cross-KG mapping.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples for structural negative
            sampling and to build the semi-supervised unlabeled pool
            (see the module docstring).
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs used to train the supervised half of the mapping
            loss. Unlike ``KGELinker``, these are *not* injected into
            ``graph`` as extra triples.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must
            already be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples -- nothing to train
            the mapping with -- or if ``device`` is invalid or
            unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("SEALinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)
    ctx1 = build_kg_context(triples1_labels, mapped1)
    ctx2 = build_kg_context(triples2_labels, mapped2)

    seed_pairs = [(s, t) for s, t in ground_truth if s in entity_to_id and t in entity_to_id]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train SEA's mapping with."
        )
    seed_source = torch.tensor([entity_to_id[s] for s, _ in seed_pairs], device=device)
    seed_target = torch.tensor([entity_to_id[t] for _, t in seed_pairs], device=device)

    seed_sources = {s for s, _ in seed_pairs}
    seed_targets = {t for _, t in seed_pairs}
    unlabeled_pool1 = np.array(
        sorted({entity_to_id[e.id] for e in dataset1 if e.id in entity_to_id} - seed_sources),
        dtype=np.int64,
    )
    unlabeled_pool2 = np.array(
        sorted({entity_to_id[e.id] for e in dataset2 if e.id in entity_to_id} - seed_targets),
        dtype=np.int64,
    )

    entity_embeds, relation_embeds = self._init_structural_embeddings(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    mapping_shape = (self.embedding_dim, self.embedding_dim)
    mapping_mat_1 = torch.nn.Parameter(
        torch.nn.init.orthogonal_(torch.empty(*mapping_shape, device=device))
    )
    mapping_mat_2 = torch.nn.Parameter(
        torch.nn.init.orthogonal_(torch.empty(*mapping_shape, device=device))
    )

    triple_optimizer = torch.optim.Adam([entity_embeds, relation_embeds], lr=self.learning_rate)
    # `entity_embeds` is included here (not just the two mapping
    # matrices) to match OpenEA's `generate_optimizer(mapping_loss,
    # lr, var_list=None, opt='Adam')` -- see the module docstring.
    mapping_optimizer = torch.optim.Adam(
        [mapping_mat_1, mapping_mat_2, entity_embeds], lr=self.learning_rate
    )

    total_triples = len(ctx1.triples) + len(ctx2.triples)
    triple_steps = max(1, (total_triples + self.batch_size - 1) // self.batch_size)

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        train_structural_epoch(
            entity_embeds,
            relation_embeds,
            triple_optimizer,
            ctx1,
            ctx2,
            rng,
            self.batch_size,
            self.margin,
        )
        train_mapping_epoch(
            entity_embeds,
            mapping_mat_1,
            mapping_mat_2,
            mapping_optimizer,
            seed_source,
            seed_target,
            unlabeled_pool1,
            unlabeled_pool2,
            rng,
            triple_steps,
            self.alpha_1,
            self.alpha_2,
        )

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
                current_mapping = mapping_mat_1.cpu().numpy()
            hits1 = validation_hits1(current_embeds, current_mapping, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
        final_mapping = mapping_mat_1.cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._mapping = final_mapping
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Projected via mapping_mat_1.

Source code in src/linkingtk/algorithms/ea/sea.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side.

    Projected via ``mapping_mat_1``.
    """
    return self._project(self._embedding(entity_id))

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side (unprojected).

Source code in src/linkingtk/algorithms/ea/sea.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side (unprojected)."""
    return self._embedding(entity_id)

BootEA

bootea

BootEA knowledge-graph-embedding linker for Entity Alignment.

Sun, Z., Hu, W., Zhang, Q., & Qu, Y. (2018). Bootstrapping Entity Alignment with Knowledge Graph Embedding. IJCAI 2018. https://www.ijcai.org/proceedings/2018/0611.pdf

Like MTransELinker and its siblings, this is a faithful reimplementation of BootEA's actual training procedure, ported directly from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/bootea.py, which extends approaches/aligne.py, plus its modules/bootstrapping/alignment_finder.py, modules/train/batch.py, modules/load/kgs.py support code) rather than from a from-scratch reading of the paper.

BootEA trains a single shared entity/relation embedding table (disjoint per-KG id ranges, no id merging -- OpenEA's alignment_module="swapping", unlike IPTransELinker's "sharing" id-merge and unlike MTransELinker's learned mapping matrix) with the "limited" loss -- a double-margin hinge (pos_loss = relu(pos_score - pos_margin), neg_loss = relu(neg_margin - neg_score)), not the single-hinge margin loss MTransE/IPTransE use. Three mechanisms distinguish it from a plain structural baseline:

1. Seed pairs are baked into the main structural triple set from the start, not fed through a separate loss channel: for each seed pair (a, b), every real KG1 triple touching a becomes an extra pseudo-triple with a replaced by b (and symmetrically for KG2), folded directly into each side's own triple set (OpenEA's generate_sup_relation_triples, called once at KGs-construction time) -- so the main loss pulls a's and b's embeddings together implicitly via shared structural neighbors, no mapping matrix needed.
2. Truncated (hard) negative sampling: each triple's corrupted head/tail is drawn from that entity's current K-nearest-neighbor set by embedding similarity (only the top 1 - truncated_epsilon fraction of same-KG entities are eligible), not uniformly -- refreshed every outer iteration from the live embeddings.
3. Bootstrapping with two-sided editing + max-weight bipartite matching, run once per outer iteration (every sub_epoch epochs): candidate pairs are filtered by similarity threshold and row-wise top-k, then resolved into a true 1-to-1 assignment via maximum- weight bipartite matching (see _bootea_training for the scipy-based substitution for OpenEA's graph_tool/igraph), then merged into an accumulated labeled-alignment set with editing (a later, higher-confidence match can overwrite an earlier one). The entire current labeled-alignment set -- not just this round's deltas -- is turned into pseudo-triples the same way seed pairs are, and trained for exactly 1 epoch through a separate alignment loss/optimizer (a positive-only sigmoid loss, no negative sampling -- a different shape than the main limited loss).

Deliberately not ported: OpenEA's bootea.py also defines a likelihood() EM-style step (_define_likelihood_graph), but run()'s call to it is commented out (# self.likelihood(labeled_align)) -- the reference never actually runs it. Same "port the actual, not apparent, reference behavior" situation as JAPELinker's sim_optimizer (#28) -- skipped here entirely.

Ported into two private helper modules: _bootea_training.py (plain numpy, independently testable without torch) and _bootea_torch.py (the training-step functions that build/consume PyTorch tensors).

Deviations from OpenEA's own code, beyond what those two modules already document:

• Bootstrap reference pool: OpenEA restricts the "not yet aligned" search pool to its own dataset loader's valid_entities + test_entities. This repo's fit() only receives dataset1/dataset2 and ground_truth (the train split), so the pool here is every entity not already a ground_truth source/target -- broader than OpenEA's pool, same already-established deviation as IPTransELinker's bootstrap pool and JAPELinker's attribute reference pool (reuses _jape_training.py's reference_pools directly).
• Max-weight bipartite matching via scipy.optimize.linear_sum_assignment instead of OpenEA's graph_tool/igraph -- see _bootea_training.py's find_mwgm_pairs docstring.
• var_list=None gotcha checked, not applicable: unlike MTransE/KDCoE, BootEA has no separate mapping-matrix parameter to under-scope an optimizer against -- ent_embeds/rel_embeds are the only trainable variables in scope for every optimizer this method defines, so leaving an optimizer's variable list unset already resolves to exactly the intended target.

BootEALinker

BootEALinker(
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    sub_epoch: int = 10,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    pos_margin: float = 0.01,
    neg_margin: float = 2.0,
    neg_margin_balance: float = 0.2,
    neg_triple_num: int = 10,
    truncated_epsilon: float = 0.9,
    sim_th: float = 0.7,
    k: int = 10,
    bootstrap_pool_size: int | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a shared TransE space, grown by editing-based bootstrapping.

Must be fit before link can be called. See the module docstring for how this differs from MTransELinker and IPTransELinker.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the trained entity/relation embeddings. OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Total structural-training-epoch cap, split across outer bootstrapping iterations of sub_epoch epochs each. OpenEA's config allows up to 2000 with early stopping (see val_ground_truth on fit).	2000
sub_epoch	int	Epochs per outer iteration, between bootstrapping rounds. OpenEA's published value is 10.	10
batch_size	int	Mini-batch size for both the structural and alignment losses. OpenEA's published value is 5000.	5000
learning_rate	float	Adagrad's learning rate, for both optimizers. OpenEA's published value is 0.01.	0.01
pos_margin	float	Positive-triple margin in the "limited" loss -- positive scores aren't penalized until they exceed this. OpenEA's published value is 0.01.	0.01
neg_margin	float	Negative-triple margin -- negative scores aren't penalized until they drop below this. OpenEA's published value is 2.0.	2.0
neg_margin_balance	float	Weight applied to the negative-triple loss term relative to the positive term. OpenEA's published value is 0.2.	0.2
neg_triple_num	int	Negative triples sampled per positive. OpenEA's published value is 10.	10
truncated_epsilon	float	Fraction of each KG's entities excluded from an entity's truncated negative-sampling candidate set (only the top 1 - truncated_epsilon most similar entities are eligible). OpenEA's published value is 0.9.	0.9
sim_th	float	Minimum similarity for a bootstrap round's candidate pair. OpenEA's published value is 0.7.	0.7
k	int	Max candidate matches kept per reference-pool row before max-weight bipartite matching. OpenEA's published value is 10.	10
bootstrap_pool_size	int | None	Cap on each side's bootstrap reference pool size (random subsample, fixed for the whole run) -- same deviation-mitigation as IPTransELinker's own parameter: this repo's broader-than-OpenEA reference pool (see the module docstring) means find_mwgm_pairs's scipy.optimize.linear_sum_assignment step runs over a denser candidate set than OpenEA's own val+test-only pool would produce. None (default) uses the full pool.	None
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/bootea.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 2000,
    sub_epoch: int = 10,
    batch_size: int = 5000,
    learning_rate: float = 0.01,
    pos_margin: float = 0.01,
    neg_margin: float = 2.0,
    neg_margin_balance: float = 0.2,
    neg_triple_num: int = 10,
    truncated_epsilon: float = 0.9,
    sim_th: float = 0.7,
    k: int = 10,
    bootstrap_pool_size: int | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.sub_epoch = sub_epoch
    self.batch_size = batch_size
    self.learning_rate = learning_rate
    self.pos_margin = pos_margin
    self.neg_margin = neg_margin
    self.neg_margin_balance = neg_margin_balance
    self.neg_triple_num = neg_triple_num
    self.truncated_epsilon = truncated_epsilon
    self.sim_th = sim_th
    self.k = k
    self.bootstrap_pool_size = bootstrap_pool_size
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> BootEALinker

Train a shared TransE space with seed pseudo-triples and iterative bootstrapping.

See the module docstring for the three mechanisms that distinguish this from a plain structural baseline.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (same approach IPTransELinker uses) and as the bootstrap reference pool's source.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs. Seeds the main structural triple set (via pseudo- triples) and is excluded from the bootstrap reference pool.	required
graph	Graph	The combined relational structure of both KGs -- entity ids on both sides must already be disjoint.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked, and training stops after patience checks with no improvement (no further bootstrapping rounds run once stopped, matching OpenEA's own early-stop-before- bootstrap ordering).	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
BootEALinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples -- nothing to seed the structural triple set with -- or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/bootea.py

def fit(  # noqa: PLR0915 -- the bootstrapping loop genuinely needs this much sequential setup
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> BootEALinker:
    """Train a shared TransE space with seed pseudo-triples and iterative bootstrapping.

    See the module docstring for the three mechanisms that
    distinguish this from a plain structural baseline.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (same approach
            [IPTransELinker][linkingtk.algorithms.ea.iptranse.IPTransELinker]
            uses) and as the bootstrap reference pool's source.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs. Seeds the main structural triple set (via pseudo-
            triples) and is excluded from the bootstrap reference pool.
        graph: The combined relational structure of both KGs -- entity
            ids on both sides must already be disjoint.
        random_state: Seed for reproducible training. Left
            unspecified, training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked,
            and training stops after ``patience`` checks with no
            improvement (no further bootstrapping rounds run once
            stopped, matching OpenEA's own early-stop-before-
            bootstrap ordering).
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples -- nothing to seed
            the structural triple set with -- or if ``device`` is
            invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
        import torch.nn.functional as functional
    except ImportError as exc:
        raise OptionalDependencyError("BootEALinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    triples1_labels = [t for t in triples if t[0] in ids1]
    triples2_labels = [t for t in triples if t[0] in ids2]

    entity_to_id, relation_to_id = build_id_mappings(triples1_labels + triples2_labels)
    mapped1 = map_triples_to_ids(triples1_labels, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2_labels, entity_to_id, relation_to_id)

    seed_pairs = [(s, t) for s, t in ground_truth if s in entity_to_id and t in entity_to_id]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train BootEA's structural "
            "triple set with."
        )
    seed_pair_ids = [(entity_to_id[s], entity_to_id[t]) for s, t in seed_pairs]

    real_ctx1 = build_kg_context(mapped1)
    real_ctx2 = build_kg_context(mapped2)
    seed_triples1, seed_triples2 = pseudo_triples_for_pairs(
        seed_pair_ids,
        real_ctx1.by_head,
        real_ctx1.by_tail,
        real_ctx2.by_head,
        real_ctx2.by_tail,
    )
    combined1 = self._augment_triples(mapped1, seed_triples1)
    combined2 = self._augment_triples(mapped2, seed_triples2)
    ctx1 = build_kg_context(combined1)
    ctx2 = build_kg_context(combined2)

    kg1_entity_ids = np.array(
        sorted(entity_to_id[e] for e in ids1 if e in entity_to_id), dtype=np.int64
    )
    kg2_entity_ids = np.array(
        sorted(entity_to_id[e] for e in ids2 if e in entity_to_id), dtype=np.int64
    )
    neighbor_k1 = max(1, int(round((1 - self.truncated_epsilon) * len(kg1_entity_ids))))
    neighbor_k2 = max(1, int(round((1 - self.truncated_epsilon) * len(kg2_entity_ids))))

    entity_embeds, relation_embeds = self._init_embeddings(
        torch, functional, len(entity_to_id), len(relation_to_id), device
    )
    structural_optimizer = torch.optim.Adagrad(
        [entity_embeds, relation_embeds], lr=self.learning_rate
    )
    alignment_optimizer = torch.optim.Adagrad(
        [entity_embeds, relation_embeds], lr=self.learning_rate
    )

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]

    pool1_labels, pool2_labels = reference_pools(sorted(ids1), sorted(ids2), seed_pairs)
    pool1_ids = np.array([entity_to_id[e] for e in pool1_labels], dtype=np.int64)
    pool2_ids = np.array([entity_to_id[e] for e in pool2_labels], dtype=np.int64)
    if self.bootstrap_pool_size is not None:
        if len(pool1_ids) > self.bootstrap_pool_size:
            pool1_ids = rng.choice(pool1_ids, size=self.bootstrap_pool_size, replace=False)
        if len(pool2_ids) > self.bootstrap_pool_size:
            pool2_ids = rng.choice(pool2_ids, size=self.bootstrap_pool_size, replace=False)

    labeled_alignment: dict[int, int] = {}
    neighbors1: dict[int, npt.NDArray[np.int64]] = {}
    neighbors2: dict[int, npt.NDArray[np.int64]] = {}

    num_outer_iters = max(1, self.num_epochs // self.sub_epoch)
    best_hits1 = -1.0
    epochs_without_improvement = 0
    epoch = 0
    stopped_early = False

    for _outer in range(num_outer_iters):
        for _ in range(self.sub_epoch):
            train_structural_epoch(
                entity_embeds,
                relation_embeds,
                structural_optimizer,
                ctx1,
                ctx2,
                neighbors1,
                neighbors2,
                rng,
                self.batch_size,
                self.pos_margin,
                self.neg_margin,
                self.neg_margin_balance,
                self.neg_triple_num,
            )
            epoch += 1

            if val_pairs and epoch % eval_every == 0:
                with torch.no_grad():
                    current_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
                hits1 = validation_hits1(current_embeds, entity_to_id, val_pairs)
                if hits1 <= best_hits1:
                    epochs_without_improvement += 1
                    if epochs_without_improvement >= patience:
                        stopped_early = True
                        break
                else:
                    best_hits1 = hits1
                    epochs_without_improvement = 0

        if stopped_early:
            break

        if len(pool1_ids) and len(pool2_ids):
            with torch.no_grad():
                current = functional.normalize(entity_embeds, dim=1).cpu().numpy()
            sim_mat = current[pool1_ids] @ current[pool2_ids].T
            curr_matches = find_mwgm_pairs(sim_mat, self.sim_th, self.k)
            labeled_alignment = edit_labeled_alignment(
                labeled_alignment, {(i, j) for i, j, _ in curr_matches}, sim_mat
            )

            if labeled_alignment:
                aligned_pairs = [
                    (int(pool1_ids[i]), int(pool2_ids[j])) for i, j in labeled_alignment.items()
                ]
                align_triples1, align_triples2 = pseudo_triples_for_pairs(
                    aligned_pairs, ctx1.by_head, ctx1.by_tail, ctx2.by_head, ctx2.by_tail
                )
                train_alignment_epoch(
                    entity_embeds,
                    relation_embeds,
                    alignment_optimizer,
                    np.array(align_triples1, dtype=np.int64).reshape(-1, 3),
                    np.array(align_triples2, dtype=np.int64).reshape(-1, 3),
                    self.batch_size,
                )

        with torch.no_grad():
            current = functional.normalize(entity_embeds, dim=1).cpu().numpy()
        neighbors1 = compute_truncated_neighbors(current, kg1_entity_ids, neighbor_k1)
        neighbors2 = compute_truncated_neighbors(current, kg2_entity_ids, neighbor_k2)

    with torch.no_grad():
        final_embeds = functional.normalize(entity_embeds, dim=1).cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/bootea.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/bootea.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

_bootea_training

Private, plain-numpy (torch-free) helpers for BootEALinker.

Ported from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA -- approaches/bootea.py, modules/bootstrapping/alignment_finder.py). Independently testable without torch installed -- see _bootea_torch.py for the loss/training-step functions.

compute_truncated_neighbors

compute_truncated_neighbors(
    entity_embeds: NDArray[floating[Any]],
    entity_ids: NDArray[int64],
    k: int,
) -> dict[int, npt.NDArray[np.int64]]

Each of entity_ids' own top-k nearest neighbors, restricted to entity_ids itself.

Ports find_neighbours: similarity is a plain dot product (callers pass already L2-normalized embeddings), restricted to one KG's own entities (entity_ids) so negative-sampling replacements never cross the KG boundary. Matches the reference in not excluding an entity from its own neighbor set (self-similarity is always highest, so an entity typically appears as its own top neighbor) -- harmless downstream, since sample_truncated_negative_triples's real-triple check just retries if a self-corruption reproduces the positive triple.

Parameters:

Name	Type	Description	Default
entity_embeds	NDArray[floating[Any]]	Full (num_entities, dim) embedding table.	required
entity_ids	NDArray[int64]	This KG's own entity ids (a subset of entity_embeds's row indices).	required
k	int	Number of neighbors to keep per entity (OpenEA's truncated_epsilon-derived count).	required

Returns:

Type	Description
dict[int, NDArray[int64]]	entity_id -> (k,) int64 array of neighbor candidate ids, one
dict[int, NDArray[int64]]	entry per id in entity_ids. Empty dict if entity_ids is
dict[int, NDArray[int64]]	empty.

Source code in src/linkingtk/algorithms/ea/_bootea_training.py

def compute_truncated_neighbors(
    entity_embeds: npt.NDArray[np.floating[Any]],
    entity_ids: npt.NDArray[np.int64],
    k: int,
) -> dict[int, npt.NDArray[np.int64]]:
    """Each of ``entity_ids``' own top-``k`` nearest neighbors, restricted to ``entity_ids`` itself.

    Ports ``find_neighbours``: similarity is a plain dot product (callers
    pass already L2-normalized embeddings), restricted to one KG's own
    entities (``entity_ids``) so negative-sampling replacements never
    cross the KG boundary. Matches the reference in not excluding an
    entity from its own neighbor set (self-similarity is always highest,
    so an entity typically appears as its own top neighbor) --
    harmless downstream, since
    [sample_truncated_negative_triples][linkingtk.algorithms.ea._bootea_training.sample_truncated_negative_triples]'s
    real-triple check just retries if a self-corruption reproduces the
    positive triple.

    Args:
        entity_embeds: Full ``(num_entities, dim)`` embedding table.
        entity_ids: This KG's own entity ids (a subset of
            ``entity_embeds``'s row indices).
        k: Number of neighbors to keep per entity (OpenEA's
            ``truncated_epsilon``-derived count).

    Returns:
        ``entity_id -> (k,) int64 array`` of neighbor candidate ids, one
        entry per id in ``entity_ids``. Empty dict if ``entity_ids`` is
        empty.
    """
    if len(entity_ids) == 0:
        return {}
    k = min(k, len(entity_ids))
    sub = entity_embeds[entity_ids]
    sim = sub @ sub.T
    # Vectorized top-k per row (unordered within the top-k, which is fine
    # -- callers only need a candidate set) via one `argpartition` call
    # over the whole matrix, rather than a per-entity `argsort` loop --
    # the loop measurably bottlenecked real-dataset-scale training
    # (~15K entities/side), called every outer bootstrapping iteration.
    top = np.argpartition(-sim, k - 1, axis=1)[:, :k]
    neighbor_ids = entity_ids[top]
    return {int(entity_id): neighbor_ids[row] for row, entity_id in enumerate(entity_ids)}

edit_labeled_alignment

edit_labeled_alignment(
    pre_alignment: dict[int, int],
    curr_alignment: set[tuple[int, int]],
    sim_mat: NDArray[floating[Any]],
) -> dict[int, int]

Merge this round's matched pairs into the accumulated labeled alignment, with editing.

Ports update_labeled_alignment_x (composed with) update_labeled_alignment_y:

1. For each (i, j) freshly matched this round: if i already has an earlier label, replace it only if the new pair's similarity is >= the old label's similarity (walks back an earlier mistaken label when a later round finds better evidence); otherwise it's a brand-new label.
2. Resolve multiple sources now claiming the same target: keep only the claimant with the highest similarity to that target.

Parameters:

Name	Type	Description	Default
pre_alignment	dict[int, int]	The accumulated source_index -> target_index labeled alignment from previous rounds ({} on the first round).	required
curr_alignment	set[tuple[int, int]]	This round's freshly matched (source_index, target_index) pairs, e.g. from find_mwgm_pairs.	required
sim_mat	NDArray[floating[Any]]	The same similarity matrix curr_alignment was matched against -- also used to compare candidate labels' confidence during editing. Indices in both pre_alignment and curr_alignment must index into this matrix.	required

Returns:

Type	Description
dict[int, int]	The updated, edited source_index -> target_index alignment.

Source code in src/linkingtk/algorithms/ea/_bootea_training.py

def edit_labeled_alignment(
    pre_alignment: dict[int, int],
    curr_alignment: set[tuple[int, int]],
    sim_mat: npt.NDArray[np.floating[Any]],
) -> dict[int, int]:
    """Merge this round's matched pairs into the accumulated labeled alignment, with editing.

    Ports ``update_labeled_alignment_x`` (composed with)
    ``update_labeled_alignment_y``:

    1. For each ``(i, j)`` freshly matched this round: if ``i`` already has
       an earlier label, replace it only if the new pair's similarity is
       ``>=`` the old label's similarity (walks back an earlier mistaken
       label when a later round finds better evidence); otherwise it's a
       brand-new label.
    2. Resolve multiple sources now claiming the same target: keep only
       the claimant with the highest similarity to that target.

    Args:
        pre_alignment: The accumulated ``source_index -> target_index``
            labeled alignment from previous rounds (``{}`` on the first
            round).
        curr_alignment: This round's freshly matched ``(source_index,
            target_index)`` pairs, e.g. from
            [find_mwgm_pairs][linkingtk.algorithms.ea._bootea_training.find_mwgm_pairs].
        sim_mat: The same similarity matrix ``curr_alignment`` was matched
            against -- also used to compare candidate labels' confidence
            during editing. Indices in both ``pre_alignment`` and
            ``curr_alignment`` must index into this matrix.

    Returns:
        The updated, edited ``source_index -> target_index`` alignment.
    """
    updated = dict(pre_alignment)
    for i, j in curr_alignment:
        if i in updated:
            if sim_mat[i, j] >= sim_mat[i, updated[i]]:
                updated[i] = j
        else:
            updated[i] = j

    by_target: dict[int, list[int]] = defaultdict(list)
    for i, j in updated.items():
        by_target[j].append(i)

    result: dict[int, int] = {}
    for j, sources in by_target.items():
        best_i = max(sources, key=lambda i: sim_mat[i, j])
        result[best_i] = j
    return result

encode_triples

encode_triples(
    triples: NDArray[int64],
) -> npt.NDArray[np.int64]

Pack (head, relation, tail) id rows into single int64 keys for fast set membership.

(h * _ID_ENCODE_SHIFT + r) * _ID_ENCODE_SHIFT + t -- safe for id spaces up to a million entities/relations each (int64 easily covers the ~10^18 result), and lets triple-membership checks use np.isin (vectorized, sort-based) instead of a Python set of tuples, which the real-triple-avoidance check in sample_truncated_negative_triples calls often enough (BootEA's neg_triple_num=10, an order of magnitude more than every other linker) for the pure-Python tuple-construction-plus-set-lookup cost to dominate real-dataset-scale training wall-clock time (confirmed via profiling).

Parameters:

Name	Type	Description	Default
triples	NDArray[int64]	(n, 3) int64 array of (head, relation, tail) ids.	required

Returns:

Type	Description
NDArray[int64]	(n,) int64 array of encoded keys, same row order.

Source code in src/linkingtk/algorithms/ea/_bootea_training.py

def encode_triples(triples: npt.NDArray[np.int64]) -> npt.NDArray[np.int64]:
    """Pack ``(head, relation, tail)`` id rows into single ``int64`` keys for fast set membership.

    ``(h * _ID_ENCODE_SHIFT + r) * _ID_ENCODE_SHIFT + t`` -- safe for id
    spaces up to a million entities/relations each (``int64`` easily
    covers the ``~10^18`` result), and lets triple-membership checks use
    ``np.isin`` (vectorized, sort-based) instead of a Python ``set`` of
    tuples, which the real-triple-avoidance check in
    [sample_truncated_negative_triples][linkingtk.algorithms.ea._bootea_training.sample_truncated_negative_triples]
    calls often enough (BootEA's ``neg_triple_num=10``, an order of
    magnitude more than every other linker) for the pure-Python
    tuple-construction-plus-set-lookup cost to dominate real-dataset-scale
    training wall-clock time (confirmed via profiling).

    Args:
        triples: ``(n, 3)`` int64 array of ``(head, relation, tail)`` ids.

    Returns:
        ``(n,)`` int64 array of encoded keys, same row order.
    """
    if len(triples) == 0:
        return np.empty(0, dtype=np.int64)
    h, r, t = triples[:, 0], triples[:, 1], triples[:, 2]
    return (h * _ID_ENCODE_SHIFT + r) * _ID_ENCODE_SHIFT + t

find_mwgm_pairs

find_mwgm_pairs(
    sim_mat: NDArray[floating[Any]], sim_th: float, k: int
) -> list[tuple[int, int, float]]

Candidate filtering (threshold ∩ row-wise top-k) then maximum-weight bipartite matching.

Ports find_potential_alignment_mwgm/find_alignment: a pair (i, j) is a candidate iff sim[i, j] > sim_th and j is among row i's k largest entries (unlike find_new_pairs's plain row-argmax, this keeps up to k candidates per row, not just the single best). Candidates then go through maximum-weight bipartite matching, via scipy.optimize.linear_sum_assignment (already a dependency) rather than OpenEA's graph_tool/igraph -- a deliberate, documented dependency substitution (same rationale as IMUSE #32's rapidfuzz): build a dense cost submatrix over just the distinct row/col indices appearing in the candidate set (non-candidate cells get a large sentinel cost so they're only chosen if the assignment is forced to), solve it, then drop any returned pair that wasn't an actual candidate -- this gives the same "true 1-to-1 assignment, greedy pairs discarded if a better global assignment exists" property without the extra dependency.

Uses a plain per-row argsort for the top-k step rather than OpenEA's argpartition, for readability and to sidestep argpartition's "kth must be less than array size" edge case on small matrices -- same simplification precedent as sample_negative_triples.

Parameters:

Name	Type	Description	Default
sim_mat	NDArray[floating[Any]]	(n1, n2) similarity matrix.	required
sim_th	float	Minimum similarity for a pair to be a candidate at all.	required
k	int	Max candidates kept per row (OpenEA's published value is 10).	required

Returns:

Type	Description
list[tuple[int, int, float]]	(row_index, col_index, similarity) tuples for the matched
list[tuple[int, int, float]]	pairs. Empty if sim_mat is empty or no candidates clear
list[tuple[int, int, float]]	sim_th.

Source code in src/linkingtk/algorithms/ea/_bootea_training.py

def find_mwgm_pairs(
    sim_mat: npt.NDArray[np.floating[Any]],
    sim_th: float,
    k: int,
) -> list[tuple[int, int, float]]:
    """Candidate filtering (threshold ∩ row-wise top-k) then maximum-weight bipartite matching.

    Ports ``find_potential_alignment_mwgm``/``find_alignment``: a pair
    ``(i, j)`` is a *candidate* iff ``sim[i, j] > sim_th`` and ``j`` is
    among row ``i``'s ``k`` largest entries (unlike
    [find_new_pairs][linkingtk.algorithms.ea._iptranse_training.find_new_pairs]'s
    plain row-argmax, this keeps up to ``k`` candidates per row, not just
    the single best). Candidates then go through **maximum-weight
    bipartite matching**, via ``scipy.optimize.linear_sum_assignment``
    (already a dependency) rather than OpenEA's ``graph_tool``/``igraph``
    -- a deliberate, documented dependency substitution (same rationale as
    IMUSE #32's ``rapidfuzz``): build a dense cost submatrix over just the
    distinct row/col indices appearing in the candidate set (non-candidate
    cells get a large sentinel cost so they're only chosen if the
    assignment is forced to), solve it, then drop any returned pair that
    wasn't an actual candidate -- this gives the same "true 1-to-1
    assignment, greedy pairs discarded if a better global assignment
    exists" property without the extra dependency.

    Uses a plain per-row ``argsort`` for the top-k step rather than
    OpenEA's ``argpartition``, for readability and to sidestep
    ``argpartition``'s "kth must be less than array size" edge case on
    small matrices -- same simplification precedent as
    ``sample_negative_triples``.

    Args:
        sim_mat: ``(n1, n2)`` similarity matrix.
        sim_th: Minimum similarity for a pair to be a candidate at all.
        k: Max candidates kept per row (OpenEA's published value is
            ``10``).

    Returns:
        ``(row_index, col_index, similarity)`` tuples for the matched
        pairs. Empty if ``sim_mat`` is empty or no candidates clear
        ``sim_th``.
    """
    if sim_mat.size == 0:
        return []
    candidates = _threshold_and_topk_candidates(sim_mat, sim_th, k)
    if not candidates:
        return []

    rows = sorted({i for i, _ in candidates})
    cols = sorted({j for _, j in candidates})
    row_index = {i: idx for idx, i in enumerate(rows)}
    col_index = {j: idx for idx, j in enumerate(cols)}

    sentinel = 1e6
    cost = np.full((len(rows), len(cols)), sentinel, dtype=np.float64)
    for i, j in candidates:
        cost[row_index[i], col_index[j]] = -float(sim_mat[i, j])

    row_assign, col_assign = linear_sum_assignment(cost)
    matched: list[tuple[int, int, float]] = []
    for r, c in zip(row_assign, col_assign, strict=True):
        i, j = rows[r], cols[c]
        if (i, j) in candidates:
            matched.append((i, j, float(sim_mat[i, j])))
    return matched

pseudo_triples_for_pairs

pseudo_triples_for_pairs(
    pairs: list[tuple[int, int]],
    by_head1: dict[int, list[tuple[int, int]]],
    by_tail1: dict[int, list[tuple[int, int]]],
    by_head2: dict[int, list[tuple[int, int]]],
    by_tail2: dict[int, list[tuple[int, int]]],
) -> tuple[
    list[tuple[int, int, int]], list[tuple[int, int, int]]
]

Unweighted pseudo-triples for a batch of aligned pairs, split by originating KG.

Ports generate_supervised_triples/generate_newly_triples: for each (source, target) pair, every real KG1 edge touching source becomes an extra KG1-side triple with source replaced by target (and symmetrically, every real KG2 edge touching target becomes an extra KG2-side triple with target replaced by source). Reuses pseudo_triples_for_pair for the actual substitution (called with weight=1.0 and the weight discarded) -- unlike IPTransE's own bootstrap pseudo-triples, BootEA's reference treats every pseudo-triple equally, with no similarity-derived weight.

This is used both for seed pairs (baked into the main structural triple set once, at fit() start) and for the bootstrap-found labeled-alignment set (retrained fresh via the separate alignment loss every outer iteration) -- see BootEALinker's module docstring.

Parameters:

Name	Type	Description	Default
pairs	list[tuple[int, int]]	(source_id, target_id) id pairs, already mapped to this run's combined entity id space.	required
by_head1	dict[int, list[tuple[int, int]]]	KG1's own head -> [(relation, tail)] index (from real triples only).	required
by_tail1	dict[int, list[tuple[int, int]]]	KG1's own tail -> [(head, relation)] index.	required
by_head2	dict[int, list[tuple[int, int]]]	KG2's own head -> [(relation, tail)] index.	required
by_tail2	dict[int, list[tuple[int, int]]]	KG2's own tail -> [(head, relation)] index.	required

Returns:

Type	Description
list[tuple[int, int, int]]	(kg1_side_triples, kg2_side_triples) -- deduplicated, sorted
list[tuple[int, int, int]]	(head, relation, tail) id tuples. Note these can mix KG1/KG2
tuple[list[tuple[int, int, int]], list[tuple[int, int, int]]]	ids in one triple (the substituted endpoint comes from the other
tuple[list[tuple[int, int, int]], list[tuple[int, int, int]]]	side) -- this is deliberate, matching OpenEA's own behavior of
tuple[list[tuple[int, int, int]], list[tuple[int, int, int]]]	folding them into each side's own triple set regardless.

Source code in src/linkingtk/algorithms/ea/_bootea_training.py

def pseudo_triples_for_pairs(
    pairs: list[tuple[int, int]],
    by_head1: dict[int, list[tuple[int, int]]],
    by_tail1: dict[int, list[tuple[int, int]]],
    by_head2: dict[int, list[tuple[int, int]]],
    by_tail2: dict[int, list[tuple[int, int]]],
) -> tuple[list[tuple[int, int, int]], list[tuple[int, int, int]]]:
    """Unweighted pseudo-triples for a batch of aligned pairs, split by originating KG.

    Ports ``generate_supervised_triples``/``generate_newly_triples``: for
    each ``(source, target)`` pair, every real KG1 edge touching ``source``
    becomes an extra KG1-side triple with ``source`` replaced by
    ``target`` (and symmetrically, every real KG2 edge touching ``target``
    becomes an extra KG2-side triple with ``target`` replaced by
    ``source``). Reuses
    [pseudo_triples_for_pair][linkingtk.algorithms.ea._iptranse_training.pseudo_triples_for_pair]
    for the actual substitution (called with ``weight=1.0`` and the weight
    discarded) -- unlike IPTransE's own bootstrap pseudo-triples, BootEA's
    reference treats every pseudo-triple equally, with no
    similarity-derived weight.

    This is used both for seed pairs (baked into the main structural
    triple set once, at ``fit()`` start) and for the bootstrap-found
    labeled-alignment set (retrained fresh via the separate alignment loss
    every outer iteration) -- see
    [BootEALinker][linkingtk.algorithms.ea.bootea.BootEALinker]'s module
    docstring.

    Args:
        pairs: ``(source_id, target_id)`` id pairs, already mapped to this
            run's combined entity id space.
        by_head1: KG1's own ``head -> [(relation, tail)]`` index (from real
            triples only).
        by_tail1: KG1's own ``tail -> [(head, relation)]`` index.
        by_head2: KG2's own ``head -> [(relation, tail)]`` index.
        by_tail2: KG2's own ``tail -> [(head, relation)]`` index.

    Returns:
        ``(kg1_side_triples, kg2_side_triples)`` -- deduplicated, sorted
        ``(head, relation, tail)`` id tuples. Note these can mix KG1/KG2
        ids in one triple (the substituted endpoint comes from the other
        side) -- this is deliberate, matching OpenEA's own behavior of
        folding them into each side's own triple set regardless.
    """
    triples1: set[tuple[int, int, int]] = set()
    triples2: set[tuple[int, int, int]] = set()
    for source, target in pairs:
        for h, r, t, _ in pseudo_triples_for_pair(source, target, 1.0, by_head1, by_tail1):
            triples1.add((h, r, t))
        for h, r, t, _ in pseudo_triples_for_pair(target, source, 1.0, by_head2, by_tail2):
            triples2.add((h, r, t))
    return sorted(triples1), sorted(triples2)

sample_truncated_negative_triples

sample_truncated_negative_triples(
    positive_triples: NDArray[int64],
    neighbor_candidates: dict[int, NDArray[int64]],
    fallback_pool: NDArray[int64],
    real_triples: NDArray[int64],
    rng: Generator,
    max_tries: int = 10,
) -> npt.NDArray[np.int64]

Like sample_negative_triples but corrupted from each entity's own truncated neighbor set.

See sample_negative_triples.

Ports generate_neg_triples_fast's neighbor-restricted candidate behavior: the replacement head/tail is drawn from neighbor_candidates.get(entity, fallback_pool) -- typically each entity's current K-nearest-neighbor set by embedding similarity (see compute_truncated_neighbors), which produces harder negatives than uniform sampling. Falls back to fallback_pool for an entity with no computed neighbor set yet (e.g. before the first neighbor refresh).

Vectorized rather than a plain per-row retry loop (unlike sample_negative_triples's established simplification): BootEA's own published neg_triple_num is 10 (vs. every other linker's 1), so this gets called an order of magnitude more per training step -- a scalar-Python retry loop measurably bottlenecks real-dataset-scale training (confirmed via profiling on EN-FR-15K-V1: 0% GPU utilization, all wall-clock time spent in this function). Draws are batched per unique corrupted entity via numpy fancy-indexing; the real-triple-avoidance check uses encode_triples + np.isin (vectorized) rather than a Python set of tuples, and each retry pass only re-draws rows that actually collided (typically few, and shrinking), not the whole array every pass.

Parameters:

Name	Type	Description	Default
positive_triples	NDArray[int64]	(n, 3) int64 array of (head, rel, tail) ids.	required
neighbor_candidates	dict[int, NDArray[int64]]	entity_id -> candidate replacement ids (all arrays the same length, as compute_truncated_neighbors produces).	required
fallback_pool	NDArray[int64]	Candidate pool used when an entity has no entry in neighbor_candidates.	required
real_triples	NDArray[int64]	Sorted encode_triples-encoded keys to avoid reproducing (e.g. KGContext.real_triples).	required
rng	Generator	Random generator, for reproducibility.	required
max_tries	int	Retry passes over still-colliding rows before accepting the last draw regardless.	10

Returns:

Type	Description
NDArray[int64]	(n, 3) int64 array of corrupted triples, same shape as
NDArray[int64]	positive_triples.

Source code in src/linkingtk/algorithms/ea/_bootea_training.py

def sample_truncated_negative_triples(
    positive_triples: npt.NDArray[np.int64],
    neighbor_candidates: dict[int, npt.NDArray[np.int64]],
    fallback_pool: npt.NDArray[np.int64],
    real_triples: npt.NDArray[np.int64],
    rng: np.random.Generator,
    max_tries: int = 10,
) -> npt.NDArray[np.int64]:
    """Like ``sample_negative_triples`` but corrupted from each entity's own truncated neighbor set.

    See
    [sample_negative_triples][linkingtk.algorithms.ea._iptranse_training.sample_negative_triples].

    Ports ``generate_neg_triples_fast``'s ``neighbor``-restricted
    candidate behavior: the replacement head/tail is drawn from
    ``neighbor_candidates.get(entity, fallback_pool)`` -- typically each
    entity's current K-nearest-neighbor set by embedding similarity (see
    [compute_truncated_neighbors][linkingtk.algorithms.ea._bootea_training.compute_truncated_neighbors]),
    which produces harder negatives than uniform sampling. Falls back to
    ``fallback_pool`` for an entity with no computed neighbor set yet
    (e.g. before the first neighbor refresh).

    Vectorized rather than a plain per-row retry loop (unlike
    [sample_negative_triples][linkingtk.algorithms.ea._iptranse_training.sample_negative_triples]'s
    established simplification): BootEA's own published ``neg_triple_num``
    is ``10`` (vs. every other linker's ``1``), so this gets called an
    order of magnitude more per training step -- a scalar-Python retry
    loop measurably bottlenecks real-dataset-scale training (confirmed via
    profiling on EN-FR-15K-V1: 0% GPU utilization, all wall-clock time
    spent in this function). Draws are batched per unique corrupted
    entity via numpy fancy-indexing; the real-triple-avoidance check uses
    [encode_triples][linkingtk.algorithms.ea._bootea_training.encode_triples]
    + ``np.isin`` (vectorized) rather than a Python ``set`` of tuples, and
    each retry pass only re-draws rows that actually collided (typically
    few, and shrinking), not the whole array every pass.

    Args:
        positive_triples: ``(n, 3)`` int64 array of ``(head, rel, tail)``
            ids.
        neighbor_candidates: ``entity_id -> candidate replacement ids``
            (all arrays the same length, as
            [compute_truncated_neighbors][linkingtk.algorithms.ea._bootea_training.compute_truncated_neighbors]
            produces).
        fallback_pool: Candidate pool used when an entity has no entry in
            ``neighbor_candidates``.
        real_triples: Sorted
            [encode_triples][linkingtk.algorithms.ea._bootea_training.encode_triples]-encoded
            keys to avoid reproducing (e.g. ``KGContext.real_triples``).
        rng: Random generator, for reproducibility.
        max_tries: Retry passes over still-colliding rows before accepting
            the last draw regardless.

    Returns:
        ``(n, 3)`` int64 array of corrupted triples, same shape as
        ``positive_triples``.
    """
    n = len(positive_triples)
    negatives = positive_triples.copy()
    if n == 0:
        return negatives

    corrupt_head = rng.random(n) < 0.5
    corrupt_col = np.where(corrupt_head, 0, 2)
    corrupt_entities = positive_triples[np.arange(n), corrupt_col]

    if neighbor_candidates:
        candidate_width = next(iter(neighbor_candidates.values())).shape[0]
        default = (
            fallback_pool
            if len(fallback_pool) == candidate_width
            else np.resize(fallback_pool, candidate_width)
        )
        unique_entities, inverse = np.unique(corrupt_entities, return_inverse=True)
        candidate_rows = np.stack(
            [neighbor_candidates.get(int(e), default) for e in unique_entities]
        )

        def draw(rows: npt.NDArray[np.int64]) -> npt.NDArray[np.int64]:
            cols = rng.integers(0, candidate_width, size=len(rows))
            result: npt.NDArray[np.int64] = candidate_rows[inverse[rows], cols]
            return result
    else:

        def draw(rows: npt.NDArray[np.int64]) -> npt.NDArray[np.int64]:
            return rng.choice(fallback_pool, size=len(rows))

    all_rows = np.arange(n)
    negatives[all_rows, corrupt_col] = draw(all_rows)
    if len(real_triples) == 0:
        return negatives

    # Only re-check rows still colliding after the previous pass, not the
    # whole array every time -- collisions are rare, so this set of rows
    # shrinks fast. Rescanning everything each pass (an earlier version of
    # this function did) turns `max_tries` full-array checks into the
    # actual bottleneck at real-dataset scale.
    check_rows = all_rows
    for _ in range(max_tries):
        if len(check_rows) == 0:
            break
        # `assume_unique=True` is safe here (query-side duplicates don't
        # affect per-element correctness -- verified) and matters a lot:
        # without it, `np.isin` re-runs `np.unique` over the *whole*
        # `real_triples` array (tens of thousands of rows) on every call,
        # which was the actual remaining bottleneck after switching from a
        # Python set.
        colliding = np.isin(encode_triples(negatives[check_rows]), real_triples, assume_unique=True)
        check_rows = check_rows[colliding]
        if len(check_rows) == 0:
            break
        negatives[check_rows, corrupt_col[check_rows]] = draw(check_rows)
    return negatives

GCN-Align

gcn_align

GCN-Align knowledge-graph-embedding linker for Entity Alignment.

Wang, Z., Lv, Q., Lan, X., & Zhang, Y. (2018). Cross-lingual Knowledge Graph Alignment via Graph Convolutional Networks. EMNLP 2018. https://aclanthology.org/D18-1032/

The first GNN-based EA linker in this package (see #18's parent issue) -- unlike every linker in linkingtk.algorithms.ea so far, entity representations come from propagating embeddings through a fixed graph convolution over the combined relational structure, not from a triple-local translational/distance loss over embeddings directly. This is a faithful port of OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/gcn_align.py), not a from-scratch reading of the paper, following this repo's established convention (e.g. MTransELinker, RSN4EALinker).

Both of OpenEA's branches are implemented: the structural (se) branch propagates a directly-learnable per-entity embedding table; the attribute (ae) branch propagates a learned transform of each entity's one-hot attribute-predicate-presence vector (see build_attribute_features for why predicates, not values). Both are 2-layer GCNs over the same structural adjacency, trained independently (separate optimizers, never a combined loss) with the same shared negatives each step, matching OpenEA's own GCN_Align.train_embeddings. At evaluation time (and for early-stopping/final embeddings), scores use the concatenation [se * beta, ae * (1 - beta)] -- OpenEA's own test_method: "sa", the config that produced the published EN-FR-15K-V1 numbers. use_attributes=False (default) skips the ae branch entirely, needing nothing beyond EnFr15KDataset's plain relational triples -- pass use_attributes=True plus EnFr15KAttrDataset's attribute triples to fit() for numbers directly comparable to OpenEA's published Hits@1=0.338, Hits@10=0.680, MRR=0.451 (docs/detailed_results_current_approaches_15K.csv).

Deliberate deviations, confirmed against OpenEA's own reference config (run/args/gcnalign_args_15K.json) and source:

• Plain SGD, not Adam/Adagrad. OpenEA's own GCN_Align_Unit uses tf.train.GradientDescentOptimizer and asserts args.learning_rate >= 0.01 -- its own published value is an unusually high learning_rate=8, which is intentional fidelity to the reference, not a typo. Every other hand-rolled linker in this package uses Adam/Adagrad; this one deliberately doesn't, to match what the reference actually runs.
• No mapping matrix. OpenEA's generic pipeline config labels this "alignment_module": "mapping", but gcn_align.py's own model code never builds or trains a projection matrix -- entities from both KGs share one embedding table (disjoint ids, same pattern as RSN4EALinker), and source_embedding/target_embedding are identical passthroughs.
• Negative sampling refresh cadence is fixed at every 10 epochs (matching OpenEA's if i % 10 == 1), independent of eval_every.
• Candidate scoring uses Manhattan (L1) distance, not cosine similarity. OpenEA's own config for this method is eval_metric: "manhattan", eval_norm: false (confirmed by reading run/args/gcnalign_args_15K.json directly) -- and the margin loss both branches train under (margin_ranking_loss_l1) is itself L1-based, so embeddings are optimized to be close in L1 distance, not high-cosine. Scoring with a metric that matches the training geometry, rather than an arbitrary different one, is a real correctness fix, not just fidelity to OpenEA's own choice -- confirmed empirically on the real EN-FR-15K-V1 benchmark (structural-only): Hits@1 0.130 (cosine) -> 0.147 (manhattan, no CSLS). CSLS (Cross-domain Similarity Local Scaling, csls_k=10 in OpenEA's own config) is a further, separate refinement -- see rank_exhaustive's csls_k parameter, used by examples/gcn_align_benchmark.py but not by link() itself (CSLS needs the full candidate pool's neighbor structure, which doesn't fit link()'s post-blocking, per-source candidate-list shape).

GCNAlignLinker

GCNAlignLinker(
    embedding_dim: int = 100,
    num_epochs: int = 500,
    learning_rate: float = 8.0,
    neg_triple_num: int = 5,
    gamma: float = 3.0,
    min_weight: float = 0.3,
    use_attributes: bool = False,
    attr_dim: int = 100,
    attr_top_fraction: float = 0.7,
    beta: float = 0.9,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via a 2-layer GCN's propagated embeddings.

Must be fit before link can be called. See the module docstring for what's ported vs. deliberately deviated from OpenEA, and for the structural-only-vs.-sa (structural+attribute) tradeoff.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Dimensionality of the propagated structural embeddings (OpenEA's se_dim). OpenEA's published EN-FR-15K-V1 config uses 100.	100
num_epochs	int	Training epochs. OpenEA's published value allows up to 2000 with early stopping.	500
learning_rate	float	Plain SGD's learning rate -- see the module docstring for why this is unusually high by this package's conventions. OpenEA's published value is 8.	8.0
neg_triple_num	int	Negatives sampled per seed pair per side (k), refreshed every 10 epochs. OpenEA's published value is 5.	5
gamma	float	Margin for margin_ranking_loss_l1. OpenEA's published value is 3.	3.0
min_weight	float	Floor applied to every adjacency edge weight -- see build_weighted_adjacency. OpenEA's own value (hardcoded, not configurable in the reference) is 0.3.	0.3
use_attributes	bool	Whether to also train the attribute (ae) branch -- needs fit()'s attribute_triples1/attribute_triples2 to be given. Defaults to False (structural-only, needs nothing beyond plain relational triples).	False
attr_dim	int	Dimensionality of the propagated attribute embeddings (OpenEA's ae_dim). Only used if use_attributes. OpenEA's published value is 100.	100
attr_top_fraction	float	Fraction of the combined attribute-predicate vocabulary kept as feature columns -- see build_attribute_features. Only used if use_attributes. OpenEA's own value (hardcoded) is 0.7.	0.7
beta	float	Weight on the structural half of the concatenated [se * beta, ae * (1 - beta)] scoring embedding. Only used if use_attributes. OpenEA's published value is 0.9.	0.9
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/gcn_align.py

def __init__(
    self,
    embedding_dim: int = 100,
    num_epochs: int = 500,
    learning_rate: float = 8.0,
    neg_triple_num: int = 5,
    gamma: float = 3.0,
    min_weight: float = 0.3,
    use_attributes: bool = False,
    attr_dim: int = 100,
    attr_top_fraction: float = 0.7,
    beta: float = 0.9,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.learning_rate = learning_rate
    self.neg_triple_num = neg_triple_num
    self.gamma = gamma
    self.min_weight = min_weight
    self.use_attributes = use_attributes
    self.attr_dim = attr_dim
    self.attr_top_fraction = attr_top_fraction
    self.beta = beta
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> GCNAlignLinker

Propagate structural (and, if enabled, attribute) embeddings through the graph.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities. Unused beyond mirroring sibling linkers' signatures -- which entities get embeddings is determined entirely by graph.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs used as positive pairs in the margin loss.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader. Also drives the structural adjacency both branches propagate over, even when use_attributes is set.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10
attribute_triples1	list[Triple] | None	KG1's (entity_id, predicate, value) attribute triples, e.g. from EnFr15KAttrDataset.load_attribute_triples. Required if use_attributes (set in __init__); ignored otherwise.	None
attribute_triples2	list[Triple] | None	KG2's own attribute triples. See attribute_triples1.	None

Returns:

Type	Description
GCNAlignLinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples, if use_attributes is set but no attribute triples (or no attribute predicate clears attr_top_fraction's cutoff) are given, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/gcn_align.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
    attribute_triples1: list[Triple] | None = None,
    attribute_triples2: list[Triple] | None = None,
) -> GCNAlignLinker:
    """Propagate structural (and, if enabled, attribute) embeddings through the graph.

    Args:
        dataset1: Source entities. Unused beyond mirroring sibling
            linkers' signatures -- which entities get embeddings is
            determined entirely by ``graph``.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs used as positive pairs in the margin loss.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must already
            be disjoint, as they are from that loader. Also drives the
            structural adjacency both branches propagate over, even
            when ``use_attributes`` is set.
        random_state: Seed for reproducible training. Left unspecified,
            training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            attribute triples, e.g. from
            [EnFr15KAttrDataset.load_attribute_triples][linkingtk.datasets.openea_native._OpenEANativeDataset.load_attribute_triples].
            Required if ``use_attributes`` (set in ``__init__``);
            ignored otherwise.
        attribute_triples2: KG2's own attribute triples. See
            ``attribute_triples1``.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples, if ``use_attributes``
            is set but no attribute triples (or no attribute predicate
            clears ``attr_top_fraction``'s cutoff) are given, or if
            ``device`` is invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
    except ImportError as exc:
        raise OptionalDependencyError("GCNAlignLinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    triples = to_triples(graph)
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped = map_triples_to_ids(triples, entity_to_id, relation_to_id)
    num_entities = len(entity_to_id)

    seed_pairs = [
        (entity_to_id[s], entity_to_id[t])
        for s, t in ground_truth
        if s in entity_to_id and t in entity_to_id
    ]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train GCN-Align's margin loss with."
        )

    r2f, r2if = compute_relation_functionality(mapped)
    raw_indices, raw_values = build_weighted_adjacency(mapped, r2f, r2if, self.min_weight)
    norm_indices, norm_values = normalize_adjacency_coo(raw_indices, raw_values, num_entities)
    adjacency = coo_to_torch_sparse(
        norm_indices, norm_values, (num_entities, num_entities), device
    )

    model_se = build_gcn_align_branch(num_entities, self.embedding_dim).to(device)
    optimizer_se = torch.optim.SGD(model_se.parameters(), lr=self.learning_rate)

    model_ae: torch.nn.Module | None = None
    optimizer_ae: torch.optim.Optimizer | None = None
    attr_features: torch.Tensor | None = None
    if self.use_attributes:
        if not attribute_triples1 and not attribute_triples2:
            raise LinkingTKError(
                "GCNAlignLinker(use_attributes=True) needs `attribute_triples1`/"
                "`attribute_triples2` -- pass entities and attribute triples from "
                "EnFr15KAttrDataset (or similar), not EnFr15KDataset."
            )
        attr_indices, attr_values, num_attrs = build_attribute_features(
            attribute_triples1 or [],
            attribute_triples2 or [],
            entity_to_id,
            self.attr_top_fraction,
        )
        if num_attrs == 0:
            raise LinkingTKError(
                "GCNAlignLinker(use_attributes=True): no attribute predicate cleared "
                "`attr_top_fraction`'s cutoff -- nothing to train the attribute branch with."
            )
        attr_features = coo_to_torch_sparse(
            attr_indices, attr_values, (num_entities, num_attrs), device
        )
        model_ae = build_gcn_align_attr_branch(num_attrs, self.attr_dim).to(device)
        optimizer_ae = torch.optim.SGD(model_ae.parameters(), lr=self.learning_rate)

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    pos_left = torch.tensor([s for s, _ in seed_pairs], device=device)
    pos_right = torch.tensor([t for _, t in seed_pairs], device=device)

    def _resample_negatives() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        neg_left_np, neg_right_np, neg2_left_np, neg2_right_np = sample_negatives(
            seed_pairs, num_entities, self.neg_triple_num, rng
        )
        return (
            torch.from_numpy(neg_left_np).long().to(device),
            torch.from_numpy(neg_right_np).long().to(device),
            torch.from_numpy(neg2_left_np).long().to(device),
            torch.from_numpy(neg2_right_np).long().to(device),
        )

    def _combined_embeddings(
        se: npt.NDArray[Any], ae: npt.NDArray[Any] | None
    ) -> npt.NDArray[Any]:
        if ae is None:
            return se
        return np.concatenate([se * self.beta, ae * (1.0 - self.beta)], axis=1)

    neg_left, neg_right, neg2_left, neg2_right = _resample_negatives()

    for epoch in range(self.num_epochs):
        if epoch % _NEG_REFRESH_INTERVAL == 0 and epoch > 0:
            neg_left, neg_right, neg2_left, neg2_right = _resample_negatives()

        embeddings_se = model_se(adjacency)
        loss_se = margin_ranking_loss_l1(
            embeddings_se,
            pos_left,
            pos_right,
            neg_left,
            neg_right,
            neg2_left,
            neg2_right,
            self.gamma,
        )
        optimizer_se.zero_grad()
        loss_se.backward()  # type: ignore[no-untyped-call]
        optimizer_se.step()

        if model_ae is not None and optimizer_ae is not None and attr_features is not None:
            embeddings_ae = model_ae(adjacency, attr_features)
            loss_ae = margin_ranking_loss_l1(
                embeddings_ae,
                pos_left,
                pos_right,
                neg_left,
                neg_right,
                neg2_left,
                neg2_right,
                self.gamma,
            )
            optimizer_ae.zero_grad()
            loss_ae.backward()  # type: ignore[no-untyped-call]
            optimizer_ae.step()

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                se_np = model_se(adjacency).cpu().numpy()
                ae_np = (
                    model_ae(adjacency, attr_features).cpu().numpy()
                    if model_ae is not None and attr_features is not None
                    else None
                )
            current_embeds = _combined_embeddings(se_np, ae_np)
            hits1 = _validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_se = model_se(adjacency).cpu().numpy()
        final_ae = (
            model_ae(adjacency, attr_features).cpu().numpy()
            if model_ae is not None and attr_features is not None
            else None
        )
    final_embeds = _combined_embeddings(final_se, final_ae)
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/gcn_align.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/gcn_align.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

_gcn_align_training

Private, plain-numpy (torch-free) helpers for GCNAlignLinker.

Ported from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA -- approaches/gcn_align.py's GCN_Utils.func/ifunc/get_weighted_adj, GCN_Align.train_embeddings's uniform negative-sampling arrays, and load_attr for the attribute-presence feature matrix). Independently testable without torch installed -- see _gcn_align_torch.py for the model/training-step functions.

build_attribute_features

build_attribute_features(
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    entity_to_id: dict[str, int],
    top_fraction: float = _DEFAULT_TOP_FRACTION,
) -> tuple[
    npt.NDArray[np.int64], npt.NDArray[np.float64], int
]

One-hot attribute-predicate-presence features, per OpenEA's load_attr.

Ports load_attr exactly: the feature vocabulary is the top top_fraction (by how many distinct entities carry each predicate, combined across both KGs) attribute predicates -- not values, since load_attr reads its input from entity_attributes_dict (entity -> set of predicate URIs that entity has an attribute triple for), discarding literal values entirely. An entity gets feature 1.0 at a predicate's column iff it has at least one attribute triple with that predicate; the specific value(s) never matter. Predicates outside the top top_fraction are dropped (contribute no feature column at all, not folded into an "other" bucket), matching OpenEA's own if v in attr2id guard.

Parameters:

Name	Type	Description	Default
attribute_triples1	list[Triple]	KG1's (entity_id, predicate, value) triples, e.g. from _OpenEANativeDataset.load_attribute_triples.	required
attribute_triples2	list[Triple]	KG2's own attribute triples.	required
entity_to_id	dict[str, int]	Combined entity label -> id mapping, e.g. from build_id_mappings. Triples whose entity id isn't a key here are dropped.	required
top_fraction	float	Fraction of the combined, frequency-ranked predicate vocabulary to keep as feature columns. OpenEA's own value (hardcoded, not configurable in the reference) is 0.7.	_DEFAULT_TOP_FRACTION

Returns:

Type	Description
NDArray[int64]	(indices, values, num_attributes): indices is (2, nnz)
NDArray[float64]	int64 (row 0 = entity id, row 1 = attribute-column id -- ready for
int	normalize_adjacency_coo-free
tuple[NDArray[int64], NDArray[float64], int]	use directly via
tuple[NDArray[int64], NDArray[float64], int]	coo_to_torch_sparse
tuple[NDArray[int64], NDArray[float64], int]	with size=(num_entities, num_attributes)), values is
tuple[NDArray[int64], NDArray[float64], int]	(nnz,) float64 (always 1.0), num_attributes is the
tuple[NDArray[int64], NDArray[float64], int]	feature vocabulary size (0 if there are no attribute triples
tuple[NDArray[int64], NDArray[float64], int]	at all).

Source code in src/linkingtk/algorithms/ea/_gcn_align_training.py

def build_attribute_features(
    attribute_triples1: list[Triple],
    attribute_triples2: list[Triple],
    entity_to_id: dict[str, int],
    top_fraction: float = _DEFAULT_TOP_FRACTION,
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float64], int]:
    """One-hot attribute-*predicate*-presence features, per OpenEA's ``load_attr``.

    Ports ``load_attr`` exactly: the feature vocabulary is the top
    ``top_fraction`` (by how many distinct entities carry each predicate,
    combined across both KGs) attribute *predicates* -- not values, since
    ``load_attr`` reads its input from ``entity_attributes_dict`` (entity
    -> set of predicate URIs that entity has an attribute triple for),
    discarding literal values entirely. An entity gets feature ``1.0`` at
    a predicate's column iff it has at least one attribute triple with
    that predicate; the specific value(s) never matter. Predicates outside
    the top ``top_fraction`` are dropped (contribute no feature column at
    all, not folded into an "other" bucket), matching OpenEA's own
    ``if v in attr2id`` guard.

    Args:
        attribute_triples1: KG1's ``(entity_id, predicate, value)``
            triples, e.g. from
            [_OpenEANativeDataset.load_attribute_triples][linkingtk.datasets.openea_native._OpenEANativeDataset.load_attribute_triples].
        attribute_triples2: KG2's own attribute triples.
        entity_to_id: Combined entity label -> id mapping, e.g. from
            [build_id_mappings][linkingtk.utils.graph.build_id_mappings].
            Triples whose entity id isn't a key here are dropped.
        top_fraction: Fraction of the combined, frequency-ranked predicate
            vocabulary to keep as feature columns. OpenEA's own value
            (hardcoded, not configurable in the reference) is ``0.7``.

    Returns:
        ``(indices, values, num_attributes)``: ``indices`` is ``(2, nnz)``
        int64 (row 0 = entity id, row 1 = attribute-column id -- ready for
        [normalize_adjacency_coo][linkingtk.utils.sparse_gcn.normalize_adjacency_coo]-free
        use directly via
        [coo_to_torch_sparse][linkingtk.utils.sparse_gcn.coo_to_torch_sparse]
        with ``size=(num_entities, num_attributes)``), ``values`` is
        ``(nnz,)`` float64 (always ``1.0``), ``num_attributes`` is the
        feature vocabulary size (``0`` if there are no attribute triples
        at all).
    """
    predicates_by_entity: dict[str, set[str]] = defaultdict(set)
    for entity_id, predicate, _value in attribute_triples1 + attribute_triples2:
        predicates_by_entity[entity_id].add(predicate)

    counts: Counter[str] = Counter()
    for predicates in predicates_by_entity.values():
        counts.update(predicates)

    ranked_predicates = [predicate for predicate, _count in counts.most_common()]
    num_attributes = int(top_fraction * len(ranked_predicates))
    predicate_to_id = {
        predicate: i for i, predicate in enumerate(ranked_predicates[:num_attributes])
    }

    rows: list[int] = []
    cols: list[int] = []
    for entity_id, predicates in predicates_by_entity.items():
        row = entity_to_id.get(entity_id)
        if row is None:
            continue
        for predicate in predicates:
            col = predicate_to_id.get(predicate)
            if col is not None:
                rows.append(row)
                cols.append(col)

    if not rows:
        return np.empty((2, 0), dtype=np.int64), np.empty(0, dtype=np.float64), num_attributes

    indices = np.array([rows, cols], dtype=np.int64)
    values = np.ones(len(rows), dtype=np.float64)
    return indices, values, num_attributes

build_weighted_adjacency

build_weighted_adjacency(
    triples: NDArray[int64],
    r2f: dict[int, float],
    r2if: dict[int, float],
    min_weight: float = _DEFAULT_MIN_WEIGHT,
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]]

Directed, relation-functionality-weighted adjacency, per OpenEA's get_weighted_adj.

For every triple (h, r, t) with h != t: edge h -> t gets weight max(r2if[r], min_weight) and edge t -> h gets weight max(r2f[r], min_weight) -- edges through more "selective"/functional relations carry more weight. Weights from multiple triples sharing the same (h, t) pair are summed. Self-loop triples (h == t) are skipped entirely (no edge contributed either direction), matching OpenEA's own if tri[0] == tri[2]: continue.

Parameters:

Name	Type	Description	Default
triples	NDArray[int64]	(n, 3) int64 array of (head, relation, tail) ids.	required
r2f	dict[int, float]	Per-relation functionality, from compute_relation_functionality.	required
r2if	dict[int, float]	Per-relation inverse functionality, same source.	required
min_weight	float	Floor applied to every edge weight.	_DEFAULT_MIN_WEIGHT

Returns:

Type	Description
NDArray[int64]	(indices, values): indices is (2, nnz) int64 (row 0 =
NDArray[float64]	source/row id, row 1 = target/column id -- ready for
tuple[NDArray[int64], NDArray[float64]]	normalize_adjacency_coo),
tuple[NDArray[int64], NDArray[float64]]	values is (nnz,) float64. Empty (shape (2, 0)/(0,))
tuple[NDArray[int64], NDArray[float64]]	if every triple is a self-loop or triples is empty.

Source code in src/linkingtk/algorithms/ea/_gcn_align_training.py

def build_weighted_adjacency(
    triples: npt.NDArray[np.int64],
    r2f: dict[int, float],
    r2if: dict[int, float],
    min_weight: float = _DEFAULT_MIN_WEIGHT,
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]]:
    """Directed, relation-functionality-weighted adjacency, per OpenEA's ``get_weighted_adj``.

    For every triple ``(h, r, t)`` with ``h != t``: edge ``h -> t`` gets
    weight ``max(r2if[r], min_weight)`` and edge ``t -> h`` gets weight
    ``max(r2f[r], min_weight)`` -- edges through more "selective"/functional
    relations carry more weight. Weights from multiple triples sharing the
    same ``(h, t)`` pair are summed. Self-loop triples (``h == t``) are
    skipped entirely (no edge contributed either direction), matching
    OpenEA's own ``if tri[0] == tri[2]: continue``.

    Args:
        triples: ``(n, 3)`` int64 array of ``(head, relation, tail)`` ids.
        r2f: Per-relation functionality, from
            [compute_relation_functionality][linkingtk.algorithms.ea._gcn_align_training.compute_relation_functionality].
        r2if: Per-relation inverse functionality, same source.
        min_weight: Floor applied to every edge weight.

    Returns:
        ``(indices, values)``: ``indices`` is ``(2, nnz)`` int64 (row 0 =
        source/row id, row 1 = target/column id -- ready for
        [normalize_adjacency_coo][linkingtk.utils.sparse_gcn.normalize_adjacency_coo]),
        ``values`` is ``(nnz,)`` float64. Empty (shape ``(2, 0)``/``(0,)``)
        if every triple is a self-loop or ``triples`` is empty.
    """
    weights: dict[tuple[int, int], float] = defaultdict(float)
    for head, relation, tail in triples:
        head, relation, tail = int(head), int(relation), int(tail)
        if head == tail:
            continue
        weights[(head, tail)] += max(r2if[relation], min_weight)
        weights[(tail, head)] += max(r2f[relation], min_weight)

    if not weights:
        return np.empty((2, 0), dtype=np.int64), np.empty(0, dtype=np.float64)

    pairs = list(weights.keys())
    rows = np.array([pair[0] for pair in pairs], dtype=np.int64)
    cols = np.array([pair[1] for pair in pairs], dtype=np.int64)
    values = np.array([weights[pair] for pair in pairs], dtype=np.float64)
    return np.stack([rows, cols]), values

compute_relation_functionality

compute_relation_functionality(
    triples: NDArray[int64],
) -> tuple[dict[int, float], dict[int, float]]

Per-relation functionality/inverse-functionality, per OpenEA's func/ifunc.

r2f[r] = (# distinct head entities of r) / (# triples with r) -- how "selective" r is as a function of its head. r2if[r] is the same but over tail entities. Both are used by build_weighted_adjacency to weight each triple's structural-influence edges.

Parameters:

Name	Type	Description	Default
triples	NDArray[int64]	(n, 3) int64 array of (head, relation, tail) ids.	required

Returns:

Type	Description
dict[int, float]	(r2f, r2if), both {relation_id: float}. Empty dicts if
dict[int, float]	triples is empty.

Source code in src/linkingtk/algorithms/ea/_gcn_align_training.py

def compute_relation_functionality(
    triples: npt.NDArray[np.int64],
) -> tuple[dict[int, float], dict[int, float]]:
    """Per-relation functionality/inverse-functionality, per OpenEA's ``func``/``ifunc``.

    ``r2f[r]`` = (# distinct head entities of ``r``) / (# triples with
    ``r``) -- how "selective" ``r`` is as a function of its head.
    ``r2if[r]`` is the same but over tail entities. Both are used by
    [build_weighted_adjacency][linkingtk.algorithms.ea._gcn_align_training.build_weighted_adjacency]
    to weight each triple's structural-influence edges.

    Args:
        triples: ``(n, 3)`` int64 array of ``(head, relation, tail)`` ids.

    Returns:
        ``(r2f, r2if)``, both ``{relation_id: float}``. Empty dicts if
        ``triples`` is empty.
    """
    heads_by_relation: dict[int, set[int]] = defaultdict(set)
    tails_by_relation: dict[int, set[int]] = defaultdict(set)
    counts: dict[int, int] = defaultdict(int)
    for head, relation, tail in triples:
        head, relation, tail = int(head), int(relation), int(tail)
        heads_by_relation[relation].add(head)
        tails_by_relation[relation].add(tail)
        counts[relation] += 1
    r2f = {relation: len(heads_by_relation[relation]) / count for relation, count in counts.items()}
    r2if = {
        relation: len(tails_by_relation[relation]) / count for relation, count in counts.items()
    }
    return r2f, r2if

sample_negatives

sample_negatives(
    seed_pairs: list[tuple[int, int]],
    num_entities: int,
    neg_triple_num: int,
    rng: Generator,
) -> tuple[
    npt.NDArray[np.int64],
    npt.NDArray[np.int64],
    npt.NDArray[np.int64],
    npt.NDArray[np.int64],
]

Two-sided uniform negative sampling, per OpenEA's train_embeddings setup.

For t seed pairs and k = neg_triple_num negatives per pair: right-corruption keeps each pair's own left id fixed (repeated k times) and draws k random right ids; left-corruption is the mirror image. Both id sides are drawn from the same combined [0, num_entities) id space -- GCN-Align trains one shared embedding table over both KGs' entities, not two independently-indexed spaces.

Parameters:

Name	Type	Description	Default
seed_pairs	list[tuple[int, int]]	(source_id, target_id) id pairs (already mapped via entity_to_id, not raw entity labels).	required
num_entities	int	Total entity count (both KGs combined) to sample negatives from.	required
neg_triple_num	int	Negatives per pair per side (k).	required
rng	Generator	Source of randomness.	required

Returns:

Type	Description
NDArray[int64]	(neg_left, neg_right, neg2_left, neg2_right), each (t * k,)
NDArray[int64]	int64 -- pass directly to
NDArray[int64]	margin_ranking_loss_l1.

Source code in src/linkingtk/algorithms/ea/_gcn_align_training.py

def sample_negatives(
    seed_pairs: list[tuple[int, int]],
    num_entities: int,
    neg_triple_num: int,
    rng: np.random.Generator,
) -> tuple[
    npt.NDArray[np.int64], npt.NDArray[np.int64], npt.NDArray[np.int64], npt.NDArray[np.int64]
]:
    """Two-sided uniform negative sampling, per OpenEA's ``train_embeddings`` setup.

    For ``t`` seed pairs and ``k = neg_triple_num`` negatives per pair:
    right-corruption keeps each pair's own left id fixed (repeated ``k``
    times) and draws ``k`` random right ids; left-corruption is the mirror
    image. Both id sides are drawn from the same combined ``[0,
    num_entities)`` id space -- GCN-Align trains one shared embedding table
    over both KGs' entities, not two independently-indexed spaces.

    Args:
        seed_pairs: ``(source_id, target_id)`` id pairs (already mapped via
            `entity_to_id`, not raw entity labels).
        num_entities: Total entity count (both KGs combined) to sample
            negatives from.
        neg_triple_num: Negatives per pair per side (``k``).
        rng: Source of randomness.

    Returns:
        ``(neg_left, neg_right, neg2_left, neg2_right)``, each ``(t * k,)``
        int64 -- pass directly to
        [margin_ranking_loss_l1][linkingtk.algorithms.ea._ea_losses.margin_ranking_loss_l1].
    """
    left = np.array([source for source, _ in seed_pairs], dtype=np.int64)
    right = np.array([target for _, target in seed_pairs], dtype=np.int64)
    k = neg_triple_num
    neg_left = np.repeat(left, k)
    neg_right = rng.integers(0, num_entities, size=len(left) * k, dtype=np.int64)
    neg2_left = rng.integers(0, num_entities, size=len(left) * k, dtype=np.int64)
    neg2_right = np.repeat(right, k)
    return neg_left, neg_right, neg2_left, neg2_right

_gcn_align_torch

PyTorch model and training-step functions for GCNAlignLinker.

Ports OpenEA's GCN_Align_Unit/GraphConvolution (approaches/gcn_align.py) both branches:

• The structural (se) branch: a 2-layer, purely-propagation GCN over a directly-learnable per-entity embedding table -- layer 1 is "featureless" (its only weight is the learnable embedding table itself, transformed by nothing but the adjacency), layer 2 has no weight at all (transform=False, input_dim == output_dim), so the whole branch is A_norm @ relu(A_norm @ E) for a learnable E.
• The attribute (ae) branch: the same 2-layer shape, but layer 1 has a real learnable weight matrix transforming the sparse attribute-presence input (from build_attribute_features) before propagating -- A_norm @ relu(A_norm @ (attrs @ W)).

Uses plain torch.sparse.mm throughout (see coo_to_torch_sparse) rather than a GNN library -- see the module docstring on gcn_align for why.

Callers must already have confirmed torch is importable (gcn_align.py's fit() does this via OptionalDependencyError before any of these run) -- same precedent as _rsn4ea_torch.py/ _kdcoe_torch.py.

build_gcn_align_attr_branch

build_gcn_align_attr_branch(
    num_attrs: int, dim: int
) -> torch.nn.Module

OpenEA's GCN_Align_Unit attribute branch: a real layer-1 weight over sparse features.

Same factory-function precedent as build_gcn_align_branch. Layer 1's weight uses the same trunc_normal init (truncated normal, stddev 1/sqrt(input_dim), then L2-row-normalized) OpenEA's own GCN_Align_Unit._build applies regardless of branch (init=trunc_normal is passed for both the se and ae branches' layer 1, confirmed by reading _build directly -- only featureless/sparse_inputs differ between the two branches' constructor args).

Parameters:

Name	Type	Description	Default
num_attrs	int	Attribute feature-vocabulary size (layer 1's input dimension), e.g. from build_attribute_features.	required
dim	int	Output dimensionality (OpenEA's ae_dim).	required

Returns:

Type	Description
Module	A module whose forward(adjacency, attr_features) takes a
Module	coalesced sparse [num_entities, num_entities] normalized
Module	adjacency and a sparse [num_entities, num_attrs] attribute
Module	feature matrix, returning (num_entities, dim) entity
Module	embeddings.

Source code in src/linkingtk/algorithms/ea/_gcn_align_torch.py

def build_gcn_align_attr_branch(num_attrs: int, dim: int) -> torch.nn.Module:
    """OpenEA's ``GCN_Align_Unit`` attribute branch: a real layer-1 weight over sparse features.

    Same factory-function precedent as
    [build_gcn_align_branch][linkingtk.algorithms.ea._gcn_align_torch.build_gcn_align_branch].
    Layer 1's weight uses the same ``trunc_normal`` init (truncated normal,
    stddev ``1/sqrt(input_dim)``, then L2-row-normalized) OpenEA's own
    ``GCN_Align_Unit._build`` applies regardless of branch (``init=trunc_normal``
    is passed for both the ``se`` and ``ae`` branches' layer 1, confirmed by
    reading ``_build`` directly -- only ``featureless``/``sparse_inputs``
    differ between the two branches' constructor args).

    Args:
        num_attrs: Attribute feature-vocabulary size (layer 1's input
            dimension), e.g. from
            [build_attribute_features][linkingtk.algorithms.ea._gcn_align_training.build_attribute_features].
        dim: Output dimensionality (OpenEA's ``ae_dim``).

    Returns:
        A module whose ``forward(adjacency, attr_features)`` takes a
        coalesced sparse ``[num_entities, num_entities]`` normalized
        adjacency and a sparse ``[num_entities, num_attrs]`` attribute
        feature matrix, returning ``(num_entities, dim)`` entity
        embeddings.
    """
    import torch
    import torch.nn.functional as functional
    from torch import nn

    class _GCNAlignAttrBranch(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            weight = torch.empty(num_attrs, dim)
            nn.init.trunc_normal_(weight, std=1.0 / (num_attrs**0.5))
            self.weight = nn.Parameter(functional.normalize(weight, dim=1))

        def forward(self, adjacency: torch.Tensor, attr_features: torch.Tensor) -> torch.Tensor:
            pre_sup = torch.sparse.mm(attr_features, self.weight)
            hidden = torch.relu(torch.sparse.mm(adjacency, pre_sup))
            return torch.sparse.mm(adjacency, hidden)  # type: ignore[no-any-return]

    return _GCNAlignAttrBranch()

build_gcn_align_branch

build_gcn_align_branch(
    num_entities: int, dim: int
) -> torch.nn.Module

OpenEA's GCN_Align_Unit structural branch: two featureless-first-layer GCN steps.

A factory function rather than a module-level torch.nn.Module subclass, since torch.nn.Module can't be named at module scope without torch installed -- same precedent as _kdcoe_torch.build_description_encoder/_rsn4ea_torch.build_rsn_model.

Parameters:

Name	Type	Description	Default
num_entities	int	Total entity count (both KGs combined) -- both the embedding table's row count and (per OpenEA's trunc_normal init) the scale of its initialization stddev.	required
dim	int	Embedding/output dimensionality (OpenEA's se_dim).	required

Returns:

Type	Description
Module	A module whose forward(adjacency) takes a coalesced sparse
Module	[num_entities, num_entities] normalized adjacency (e.g. from
Module	coo_to_torch_sparse)
Module	and returns (num_entities, dim) entity embeddings.

Source code in src/linkingtk/algorithms/ea/_gcn_align_torch.py

def build_gcn_align_branch(num_entities: int, dim: int) -> torch.nn.Module:
    """OpenEA's ``GCN_Align_Unit`` structural branch: two featureless-first-layer GCN steps.

    A factory function rather than a module-level ``torch.nn.Module``
    subclass, since ``torch.nn.Module`` can't be named at module scope
    without ``torch`` installed -- same precedent as
    ``_kdcoe_torch.build_description_encoder``/``_rsn4ea_torch.build_rsn_model``.

    Args:
        num_entities: Total entity count (both KGs combined) -- both the
            embedding table's row count and (per OpenEA's ``trunc_normal``
            init) the scale of its initialization stddev.
        dim: Embedding/output dimensionality (OpenEA's ``se_dim``).

    Returns:
        A module whose ``forward(adjacency)`` takes a coalesced sparse
        ``[num_entities, num_entities]`` normalized adjacency (e.g. from
        [coo_to_torch_sparse][linkingtk.utils.sparse_gcn.coo_to_torch_sparse])
        and returns ``(num_entities, dim)`` entity embeddings.
    """
    import torch
    import torch.nn.functional as functional
    from torch import nn

    class _GCNAlignBranch(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            table = torch.empty(num_entities, dim)
            nn.init.trunc_normal_(table, std=1.0 / (num_entities**0.5))
            self.embedding = nn.Parameter(functional.normalize(table, dim=1))

        def forward(self, adjacency: torch.Tensor) -> torch.Tensor:
            hidden = torch.relu(torch.sparse.mm(adjacency, self.embedding))
            return torch.sparse.mm(adjacency, hidden)  # type: ignore[no-any-return]

    return _GCNAlignBranch()

_ea_losses

Shared margin-ranking loss functions for the GNN-based EA linkers.

Both losses below are ported once from OpenEA's reference implementation and reused across linkers, rather than each one reimplementing the same math in its own private _<name>_torch.py:

• margin_ranking_loss_l1: the L1-distance, double-sided-negative margin loss GCN-Align's align_loss (approaches/gcn_align.py) and RDGCN's get_loss (approaches/rdgcn.py) both use, verbatim -- confirmed by reading both directly, the two functions are identical modulo variable names.
• margin_ranking_loss_l2_squared: AliNet's squared-Euclidean single-sided variant (AliNet.compute_loss, approaches/alinet.py), with its own neg_margin_balance weight on the negative term.

Both take an already-built embedding tensor plus index tensors into it (rather than pre-gathered vectors) so callers can pass whichever entity representation is currently live (e.g. AliNet's per-round concatenated multi-layer output) without an extra gather step at the call site.

Callers must already have confirmed torch is importable (each linker's fit() does this via OptionalDependencyError before any of these run) -- same precedent as _rsn4ea_torch.py/_kdcoe_torch.py.

margin_ranking_loss_l1

margin_ranking_loss_l1(
    embeddings: Tensor,
    pos_left: Tensor,
    pos_right: Tensor,
    neg_left: Tensor,
    neg_right: Tensor,
    neg2_left: Tensor,
    neg2_right: Tensor,
    gamma: float,
) -> torch.Tensor

L1-distance margin loss over t positive pairs and k negatives per side.

Ports OpenEA's align_loss/get_loss exactly: for each positive pair, both a left-corrupted (neg_left/neg_right) and a right-corrupted (neg2_left/neg2_right) set of k negatives contribute a relu(margin + pos_dist - neg_dist) term, averaged over 2 * k * t.

Parameters:

Name	Type	Description	Default
embeddings	Tensor	(num_entities, dim) embedding table to index into.	required
pos_left	Tensor	(t,) ids, the positive pairs' left/source side.	required
pos_right	Tensor	(t,) ids, the positive pairs' right/target side.	required
neg_left	Tensor	(t * k,) ids -- left-corrupted negatives' left side (each positive pair's own left id, repeated k times; see each linker's negative-sampling helper).	required
neg_right	Tensor	(t * k,) ids -- left-corrupted negatives' right side (k random/hard-mined ids per positive pair).	required
neg2_left	Tensor	(t * k,) ids -- right-corrupted negatives' left side.	required
neg2_right	Tensor	(t * k,) ids -- right-corrupted negatives' right side (each positive pair's own right id, repeated k times).	required
gamma	float	Margin.	required

Returns:

Type	Description
Tensor	A scalar loss tensor.

Source code in src/linkingtk/algorithms/ea/_ea_losses.py

def margin_ranking_loss_l1(
    embeddings: torch.Tensor,
    pos_left: torch.Tensor,
    pos_right: torch.Tensor,
    neg_left: torch.Tensor,
    neg_right: torch.Tensor,
    neg2_left: torch.Tensor,
    neg2_right: torch.Tensor,
    gamma: float,
) -> torch.Tensor:
    """L1-distance margin loss over ``t`` positive pairs and ``k`` negatives per side.

    Ports OpenEA's ``align_loss``/``get_loss`` exactly: for each positive
    pair, both a left-corrupted (``neg_left``/``neg_right``) and a
    right-corrupted (``neg2_left``/``neg2_right``) set of ``k`` negatives
    contribute a ``relu(margin + pos_dist - neg_dist)`` term, averaged over
    ``2 * k * t``.

    Args:
        embeddings: ``(num_entities, dim)`` embedding table to index into.
        pos_left: ``(t,)`` ids, the positive pairs' left/source side.
        pos_right: ``(t,)`` ids, the positive pairs' right/target side.
        neg_left: ``(t * k,)`` ids -- left-corrupted negatives' left side
            (each positive pair's own left id, repeated ``k`` times; see
            each linker's negative-sampling helper).
        neg_right: ``(t * k,)`` ids -- left-corrupted negatives' right side
            (``k`` random/hard-mined ids per positive pair).
        neg2_left: ``(t * k,)`` ids -- right-corrupted negatives' left side.
        neg2_right: ``(t * k,)`` ids -- right-corrupted negatives' right
            side (each positive pair's own right id, repeated ``k`` times).
        gamma: Margin.

    Returns:
        A scalar loss tensor.
    """
    import torch

    t = pos_left.shape[0]
    k = neg_left.shape[0] // t
    pos_dist = (embeddings[pos_left] - embeddings[pos_right]).abs().sum(dim=1)
    margin = (pos_dist + gamma).reshape(t, 1)

    def _side_loss(left_ids: torch.Tensor, right_ids: torch.Tensor) -> torch.Tensor:
        neg_dist = (embeddings[left_ids] - embeddings[right_ids]).abs().sum(dim=1)
        return torch.relu(margin - neg_dist.reshape(t, k)).sum()

    loss = _side_loss(neg_left, neg_right) + _side_loss(neg2_left, neg2_right)
    return loss / (2.0 * k * t)

margin_ranking_loss_l2_squared

margin_ranking_loss_l2_squared(
    embeddings: Tensor,
    pos_left: Tensor,
    pos_right: Tensor,
    neg_left: Tensor,
    neg_right: Tensor,
    margin: float,
    neg_margin_balance: float,
) -> torch.Tensor

Squared-Euclidean margin loss: pulls positives together, pushes negatives apart.

Ports AliNet's compute_loss exactly: sum(||e1 - e2||^2) over positive pairs, plus neg_margin_balance * sum(relu(margin - ||e1 - e2||^2)) over negative pairs (a single corruption side, unlike margin_ranking_loss_l1's two).

Parameters:

Name	Type	Description	Default
embeddings	Tensor	(num_entities, dim) embedding table to index into.	required
pos_left	Tensor	(t,) ids, positive pairs' left side.	required
pos_right	Tensor	(t,) ids, positive pairs' right side.	required
neg_left	Tensor	(n,) ids, negative pairs' left side.	required
neg_right	Tensor	(n,) ids, negative pairs' right side.	required
margin	float	Margin subtracted from each negative pair's squared distance.	required
neg_margin_balance	float	Weight applied to the summed negative-pair loss.	required

Returns:

Type	Description
Tensor	A scalar loss tensor.

Source code in src/linkingtk/algorithms/ea/_ea_losses.py

def margin_ranking_loss_l2_squared(
    embeddings: torch.Tensor,
    pos_left: torch.Tensor,
    pos_right: torch.Tensor,
    neg_left: torch.Tensor,
    neg_right: torch.Tensor,
    margin: float,
    neg_margin_balance: float,
) -> torch.Tensor:
    """Squared-Euclidean margin loss: pulls positives together, pushes negatives apart.

    Ports AliNet's ``compute_loss`` exactly: ``sum(||e1 - e2||^2)`` over
    positive pairs, plus ``neg_margin_balance * sum(relu(margin -
    ||e1 - e2||^2))`` over negative pairs (a single corruption side, unlike
    `margin_ranking_loss_l1`'s two).

    Args:
        embeddings: ``(num_entities, dim)`` embedding table to index into.
        pos_left: ``(t,)`` ids, positive pairs' left side.
        pos_right: ``(t,)`` ids, positive pairs' right side.
        neg_left: ``(n,)`` ids, negative pairs' left side.
        neg_right: ``(n,)`` ids, negative pairs' right side.
        margin: Margin subtracted from each negative pair's squared distance.
        neg_margin_balance: Weight applied to the summed negative-pair loss.

    Returns:
        A scalar loss tensor.
    """
    import torch

    pos_dist = ((embeddings[pos_left] - embeddings[pos_right]) ** 2).sum(dim=1)
    pos_loss = pos_dist.sum()
    neg_dist = ((embeddings[neg_left] - embeddings[neg_right]) ** 2).sum(dim=1)
    neg_loss = torch.relu(margin - neg_dist).sum()
    return pos_loss + neg_margin_balance * neg_loss

RDGCN

rdgcn

RDGCN knowledge-graph-embedding linker for Entity Alignment.

Wu, Y., Liu, X., Feng, Y., Wang, Z., Yan, R., & Zhao, D. (2019). Relation-Aware Entity Alignment for Heterogeneous Knowledge Graphs. IJCAI 2019. https://www.ijcai.org/proceedings/2019/0733.pdf

The second GNN-based EA linker in this package (see #18's parent issue and

42's GCNAlignLinker).

RDGCN builds entity representations from two graphs jointly: the "primal" graph over entities (structural, like GCN-Align's), and a "dual" graph over relations (two relations are dual-graph-adjacent if their head-/tail-entity sets overlap a lot), attending back and forth between them for two rounds before a final 2-layer diagonal-weight GCN with highway gates. This is a faithful port of OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/rdgcn.py), not a from-scratch reading of the paper -- see _rdgcn_training and _rdgcn_torch for the algorithm's own docstrings, and this module's deviations below (confirmed against OpenEA's own reference config, run/args/rdgcn_args_15K.json, and source):

• Requires EnFr15KAttrDataset, not EnFr15KDataset. RDGCN's entity representations are seeded from pretrained word-vector embeddings of each entity's local name -- EnFr15KDataset's numeric-id, URI-free format has no name text to embed (same reason KDCoELinker (#29) needs it). fit() raises LinkingTKError early if no entity in dataset1/ dataset2 has a non-empty label, as a cheap guard.
• Local names come from Entity.labels, not OpenEA's own URI-parsing. OpenEA's _get_local_name_by_name_triple has a dataset-specific name-attribute matching step (special-cased predicate lists for D_Y/ D_W) that falls through to a generic else: name_attribute_list = {} for every other dataset, including EN-FR-15K -- meaning OpenEA's own code already just uses the entity URI's tail segment for this dataset, the same text this repo's loader already populates as the label. Confirmed a no-op deviation for this benchmark, not a fidelity-losing simplification.
• Pretrained word vectors via _kdcoe_text.load_fasttext_vectors, reusing exactly the file OpenEA's own reference hardcodes (wiki-news-300d-1M.vec.zip) and the same streaming-fetch mechanism KDCoELinker already established, rather than a separate implementation.
• Hard-negative mining excludes each anchor's own id from its own candidate pool -- OpenEA's own get_neg doesn't, which (since an entity's distance to itself is always the global minimum) makes an anchor frequently rank as its own "hardest negative," a training signal that's structurally meaningless. See mine_hard_negatives's docstring.
• Duplicate-relation primal edges are coalesce-summed, not kept as separate softmax competitors -- a torch.sparse.softmax-driven simplification vs. TensorFlow's uncoalesced-sparse-tensor behavior; see _SparseAttentionLayer's docstring (inside build_rdgcn_model).
• Candidate scoring uses Manhattan (L1) distance, not cosine similarity. OpenEA's own config for this method is eval_metric: "manhattan", eval_norm: false (confirmed by reading run/args/rdgcn_args_15K.json directly) -- and margin_ranking_loss_l1 itself trains embeddings to be close in L1 distance, not high-cosine. This is a real correctness fix matching the training geometry, not just fidelity to OpenEA's own choice -- confirmed empirically: on the real EN-FR-15K-V1 benchmark, switching the evaluation metric alone (same trained embeddings, via rank_exhaustive's metric/ csls_k) moved Hits@1 from 0.666 (cosine) to 0.735 (manhattan+CSLS k=10, OpenEA's own config) against a published 0.755 -- closing the gap from ~88% to ~97% relative. CSLS (Cross-domain Similarity Local Scaling) is a further, separate refinement used by examples/rdgcn_benchmark.py but not by link() itself (CSLS needs the full candidate pool's neighbor structure, which doesn't fit link()'s post-blocking, per-source candidate-list shape).

RDGCNLinker

RDGCNLinker(
    embedding_dim: int = 300,
    num_epochs: int = 500,
    learning_rate: float = 0.002,
    neg_triple_num: int = 125,
    gamma: float = 1.0,
    alpha: float = 0.1,
    beta: float = 0.3,
    default_name_length: int = 4,
    word_embed_url: str = _FASTTEXT_URL,
    cache_dir: Path | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via RDGCN's dual/primal-attention-plus-GCN embeddings.

Must be fit before link can be called. See the module docstring for what's ported vs. deliberately deviated from OpenEA, and why dataset1/dataset2 must come from EnFr15KAttrDataset rather than EnFr15KDataset.

Parameters:

Name	Type	Description	Default
embedding_dim	int	Entity-embedding dimensionality -- must equal the pretrained word vectors' own dimensionality (no projection layer sits between them). OpenEA's published EN-FR-15K-V1 config uses 300 (fastText's native size).	300
num_epochs	int	Training epochs. OpenEA's config allows up to 2000 with early stopping.	500
learning_rate	float	Adam's learning rate. OpenEA's published value is 0.002.	0.002
neg_triple_num	int	Hard negatives mined per seed pair per side (k), refreshed every 10 epochs. OpenEA's published value is 125.	125
gamma	float	Margin for margin_ranking_loss_l1. OpenEA's published value is 1.0.	1.0
alpha	float	Round-1 primal-embedding mixing weight. OpenEA's published value is 0.1.	0.1
beta	float	Round-2 primal-embedding mixing weight. OpenEA's published value is 0.3.	0.3
default_name_length	int	Max local-name tokens summed per entity for the initial embedding. OpenEA's published value is 4.	4
word_embed_url	str	URL of a fastText-format .vec.zip (see _kdcoe_text.py's load_fasttext_vectors). Defaults to the real, published pretrained file; tests override this with a file:// url pointing at a tiny local fixture.	_FASTTEXT_URL
cache_dir	Path | None	Passed through to fetch_cached for the word-vector download. None uses fetch_cached's own default cache location.	None
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/rdgcn.py

def __init__(
    self,
    embedding_dim: int = 300,
    num_epochs: int = 500,
    learning_rate: float = 0.002,
    neg_triple_num: int = 125,
    gamma: float = 1.0,
    alpha: float = 0.1,
    beta: float = 0.3,
    default_name_length: int = 4,
    word_embed_url: str = _FASTTEXT_URL,
    cache_dir: Path | None = None,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.embedding_dim = embedding_dim
    self.num_epochs = num_epochs
    self.learning_rate = learning_rate
    self.neg_triple_num = neg_triple_num
    self.gamma = gamma
    self.alpha = alpha
    self.beta = beta
    self.default_name_length = default_name_length
    self.word_embed_url = word_embed_url
    self.cache_dir = cache_dir
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> RDGCNLinker

Train the two-round dual/primal-attention GCN over the combined graph.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- must carry non-empty labels (e.g. from EnFr15KAttrDataset), used to seed each entity's initial embedding from pretrained word vectors of its local name.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs used as positive pairs in the margin loss.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. If None (default), trains the full num_epochs unconditionally.	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth. Only used if val_ground_truth is given.	10

Returns:

Type	Description
RDGCNLinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If no entity in dataset1/dataset2 has a non-empty label, if none of ground_truth's pairs have both ids present in graph's own triples, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/rdgcn.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> RDGCNLinker:
    """Train the two-round dual/primal-attention GCN over the combined graph.

    Args:
        dataset1: Source entities -- must carry non-empty `labels`
            (e.g. from
            [EnFr15KAttrDataset][linkingtk.datasets.EnFr15KAttrDataset]),
            used to seed each entity's initial embedding from
            pretrained word vectors of its local name.
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs used as positive pairs in the margin loss.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must already
            be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left unspecified,
            training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally.
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``.
            Only used if ``val_ground_truth`` is given.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If no entity in ``dataset1``/``dataset2`` has a
            non-empty label, if none of ``ground_truth``'s pairs have
            both ids present in ``graph``'s own triples, or if
            ``device`` is invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    entities = dataset1 + dataset2
    if not any(label_texts(entity) for entity in entities):
        raise LinkingTKError(
            "RDGCNLinker.fit() needs entities with non-empty `labels` to seed name "
            "embeddings from -- pass entities from EnFr15KAttrDataset (or similar), "
            "not EnFr15KDataset. See the module docstring."
        )

    try:
        import torch
    except ImportError as exc:
        raise OptionalDependencyError("RDGCNLinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)

    triples = to_triples(graph)
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped = map_triples_to_ids(triples, entity_to_id, relation_to_id)
    num_entities = len(entity_to_id)
    num_relations = len(relation_to_id)

    seed_pairs = [
        (entity_to_id[s], entity_to_id[t])
        for s, t in ground_truth
        if s in entity_to_id and t in entity_to_id
    ]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train RDGCN's margin loss with."
        )

    entity_by_id = {entity.id: entity for entity in entities}
    vocabulary = {
        token
        for entity in entities
        for token in tokenize_description(label_texts(entity)[0] if label_texts(entity) else "")
    }
    word_vectors = load_fasttext_vectors(self.word_embed_url, vocabulary, self.cache_dir)
    ordered_entities = [
        entity_by_id.get(entity_id, Entity(id=entity_id, labels=[]))
        for entity_id, _index in sorted(entity_to_id.items(), key=lambda item: item[1])
    ]
    name_embeddings = init_name_embeddings(
        ordered_entities, word_vectors, self.embedding_dim, self.default_name_length
    )
    initial_embeds_np = np.stack(
        [name_embeddings[entity.id] for entity in ordered_entities]
    ).astype(np.float32)
    initial_embeds = torch.from_numpy(initial_embeds_np).to(device)

    primal_indices, primal_values = build_primal_adjacency(mapped, num_entities)
    primal_indices, primal_values = normalize_adjacency_coo(
        primal_indices, primal_values, num_entities
    )
    primal_adjacency = coo_to_torch_sparse(
        primal_indices, primal_values, (num_entities, num_entities), device
    )

    dual_adjacency_np = build_dual_graph(mapped, num_relations, num_entities)
    dual_adjacency = torch.from_numpy(dual_adjacency_np).float().to(device)

    (head_indices, head_values), (tail_indices, tail_values) = build_relation_masks(
        mapped, num_relations
    )
    head_mask = coo_to_torch_sparse(
        head_indices, head_values, (num_relations, num_entities), device
    )
    tail_mask = coo_to_torch_sparse(
        tail_indices, tail_values, (num_relations, num_entities), device
    )

    edge_heads_np, edge_tails_np, edge_relations_np = build_edge_relations(mapped)
    edge_heads = torch.from_numpy(edge_heads_np).long().to(device)
    edge_tails = torch.from_numpy(edge_tails_np).long().to(device)
    edge_relations = torch.from_numpy(edge_relations_np).long().to(device)

    model = build_rdgcn_model(initial_embeds, self.embedding_dim, self.alpha, self.beta).to(
        device
    )
    optimizer = torch.optim.Adam(model.parameters(), lr=self.learning_rate)

    pos_left = torch.tensor([s for s, _ in seed_pairs], device=device)
    pos_right = torch.tensor([t for _, t in seed_pairs], device=device)

    def _forward() -> torch.Tensor:
        return model(  # type: ignore[no-any-return]
            primal_adjacency,
            dual_adjacency,
            head_mask,
            tail_mask,
            edge_heads,
            edge_tails,
            edge_relations,
        )

    with torch.no_grad():
        neg_left, neg_right, neg2_left, neg2_right = mine_hard_negatives(
            _forward(), pos_left, pos_right, self.neg_triple_num
        )

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        if epoch % _NEG_REFRESH_INTERVAL == 0 and epoch > 0:
            with torch.no_grad():
                neg_left, neg_right, neg2_left, neg2_right = mine_hard_negatives(
                    _forward(), pos_left, pos_right, self.neg_triple_num
                )

        embeddings = _forward()
        loss = margin_ranking_loss_l1(
            embeddings,
            pos_left,
            pos_right,
            neg_left,
            neg_right,
            neg2_left,
            neg2_right,
            self.gamma,
        )
        optimizer.zero_grad()
        loss.backward()  # type: ignore[no-untyped-call]
        optimizer.step()

        if val_pairs and (epoch + 1) % eval_every == 0:
            with torch.no_grad():
                current_embeds = _forward().cpu().numpy()
            hits1 = _validation_hits1(current_embeds, entity_to_id, val_pairs)
            if hits1 <= best_hits1:
                epochs_without_improvement += 1
                if epochs_without_improvement >= patience:
                    break
            else:
                best_hits1 = hits1
                epochs_without_improvement = 0

    with torch.no_grad():
        final_embeds = _forward().cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/rdgcn.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/rdgcn.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

_rdgcn_training

Private, plain-numpy (torch-free) helpers for RDGCNLinker.

Ported from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA -- approaches/rdgcn.py's get_mat/get_sparse_tensor (primal adjacency), rfunc/compute_r/ get_dual_input (relation "dual graph" and per-relation features), and _get_desc_input/_get_local_name_by_name_triple (name-embedding init)). Independently testable without torch installed -- see _rdgcn_torch.py for the attention/GCN model and training-step functions.

build_dual_graph

build_dual_graph(
    triples: NDArray[int64],
    num_relations: int,
    num_entities: int,
) -> npt.NDArray[np.float64]

Dense relation-to-relation "dual graph", per OpenEA's get_dual_input.

Two relations are dual-graph-adjacent with weight jaccard(head-entity sets) + jaccard(tail-entity sets) -- including the diagonal (i == j, weight 2.0, both Jaccard similarities of a set with itself). Kept dense (not sparse COO) since the relation count is small (low hundreds for OpenEA-scale datasets) and every pair genuinely participates in the downstream dense-attention softmax's masking, unlike the primal/entity-scale adjacencies elsewhere in this package. Vectorized via head/tail membership matrices rather than OpenEA's own O(R^2) nested-loop-over-Python-sets -- same result, just computed as one matrix multiply per side.

Parameters:

Name	Type	Description	Default
triples	NDArray[int64]	(n, 3) int64 array of (head, relation, tail) ids.	required
num_relations	int	Total relation count.	required
num_entities	int	Total entity count (both KGs combined).	required

Returns:

Type	Description
NDArray[float64]	(num_relations, num_relations) float64 dense array.

Source code in src/linkingtk/algorithms/ea/_rdgcn_training.py

def build_dual_graph(
    triples: npt.NDArray[np.int64], num_relations: int, num_entities: int
) -> npt.NDArray[np.float64]:
    """Dense relation-to-relation "dual graph", per OpenEA's ``get_dual_input``.

    Two relations are dual-graph-adjacent with weight ``jaccard(head-entity
    sets) + jaccard(tail-entity sets)`` -- including the diagonal (``i ==
    j``, weight ``2.0``, both Jaccard similarities of a set with itself).
    Kept dense (not sparse COO) since the relation count is small (low
    hundreds for OpenEA-scale datasets) and every pair genuinely
    participates in the downstream dense-attention softmax's masking, unlike
    the primal/entity-scale adjacencies elsewhere in this package.
    Vectorized via head/tail membership matrices rather than OpenEA's own
    ``O(R^2)`` nested-loop-over-Python-sets -- same result, just computed as
    one matrix multiply per side.

    Args:
        triples: ``(n, 3)`` int64 array of ``(head, relation, tail)`` ids.
        num_relations: Total relation count.
        num_entities: Total entity count (both KGs combined).

    Returns:
        ``(num_relations, num_relations)`` float64 dense array.
    """
    head_matrix = np.zeros((num_relations, num_entities), dtype=np.float64)
    tail_matrix = np.zeros((num_relations, num_entities), dtype=np.float64)
    for head, relation, tail in triples:
        head_matrix[int(relation), int(head)] = 1.0
        tail_matrix[int(relation), int(tail)] = 1.0

    def _jaccard(membership: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        counts = membership.sum(axis=1)
        intersection = membership @ membership.T
        union = counts[:, None] + counts[None, :] - intersection
        with np.errstate(divide="ignore", invalid="ignore"):
            result = np.where(union > 0, intersection / union, 0.0)
        return result

    return _jaccard(head_matrix) + _jaccard(tail_matrix)

build_edge_relations

build_edge_relations(
    triples: NDArray[int64],
) -> tuple[
    npt.NDArray[np.int64],
    npt.NDArray[np.int64],
    npt.NDArray[np.int64],
]

Parallel per-triple (head, tail, relation) id arrays, per OpenEA's rfunc's r_mat.

Unlike build_primal_adjacency, this is not deduplicated or symmetrized -- one row per triple, directed (head -> tail only), exactly as given. Used by the sparse attention layer to look up, for every primal edge, which relation connects it.

Parameters:

Name	Type	Description	Default
triples	NDArray[int64]	(n, 3) int64 array of (head, relation, tail) ids.	required

Returns:

Type	Description
tuple[NDArray[int64], NDArray[int64], NDArray[int64]]	(edge_heads, edge_tails, edge_relations), each (n,) int64.

Source code in src/linkingtk/algorithms/ea/_rdgcn_training.py

def build_edge_relations(
    triples: npt.NDArray[np.int64],
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64], npt.NDArray[np.int64]]:
    """Parallel per-triple ``(head, tail, relation)`` id arrays, per OpenEA's ``rfunc``'s ``r_mat``.

    Unlike [build_primal_adjacency][linkingtk.algorithms.ea._rdgcn_training.build_primal_adjacency],
    this is **not** deduplicated or symmetrized -- one row per triple,
    directed (head -> tail only), exactly as given. Used by the sparse
    attention layer to look up, for every primal edge, which relation
    connects it.

    Args:
        triples: ``(n, 3)`` int64 array of ``(head, relation, tail)`` ids.

    Returns:
        ``(edge_heads, edge_tails, edge_relations)``, each ``(n,)`` int64.
    """
    if len(triples) == 0:
        empty = np.empty(0, dtype=np.int64)
        return empty, empty, empty
    return (
        triples[:, 0].astype(np.int64),
        triples[:, 2].astype(np.int64),
        triples[:, 1].astype(np.int64),
    )

build_primal_adjacency

build_primal_adjacency(
    triples: NDArray[int64], num_entities: int
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]]

Unweighted, symmetric primal (entity) adjacency, per OpenEA's get_mat.

For every triple (h, r, t) with h != t, both h -> t and t -> h get weight 1.0 -- a plain undirected adjacency, deduplicated (unlike build_weighted_adjacency, repeated triples between the same entity pair do not accumulate weight, matching OpenEA's own if (h, t) not in pos: pos[(h, t)] = 1 dedup-to-1 behavior). Self-loop triples (h == t) are skipped, same as build_weighted_adjacency. Pass the result to normalize_adjacency_coo (add_self_loops=True) -- symmetric here, so its row-sum-vs-column-sum distinction is a no-op.

Parameters:

Name	Type	Description	Default
triples	NDArray[int64]	(n, 3) int64 array of (head, relation, tail) ids.	required
num_entities	int	Total entity count (both KGs combined).	required

Returns:

Type	Description
tuple[NDArray[int64], NDArray[float64]]	(indices, values) -- (2, nnz) int64 / (nnz,) float64.

Source code in src/linkingtk/algorithms/ea/_rdgcn_training.py

def build_primal_adjacency(
    triples: npt.NDArray[np.int64], num_entities: int
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]]:
    """Unweighted, symmetric primal (entity) adjacency, per OpenEA's ``get_mat``.

    For every triple ``(h, r, t)`` with ``h != t``, both ``h -> t`` and
    ``t -> h`` get weight ``1.0`` -- a plain undirected adjacency,
    deduplicated (unlike
    [build_weighted_adjacency][linkingtk.algorithms.ea._gcn_align_training.build_weighted_adjacency],
    repeated triples between the same entity pair do **not** accumulate
    weight, matching OpenEA's own ``if (h, t) not in pos: pos[(h, t)] = 1``
    dedup-to-1 behavior). Self-loop triples (``h == t``) are skipped, same
    as
    [build_weighted_adjacency][linkingtk.algorithms.ea._gcn_align_training.build_weighted_adjacency].
    Pass the result to
    [normalize_adjacency_coo][linkingtk.utils.sparse_gcn.normalize_adjacency_coo]
    (``add_self_loops=True``) -- symmetric here, so its row-sum-vs-column-sum
    distinction is a no-op.

    Args:
        triples: ``(n, 3)`` int64 array of ``(head, relation, tail)`` ids.
        num_entities: Total entity count (both KGs combined).

    Returns:
        ``(indices, values)`` -- ``(2, nnz)`` int64 / ``(nnz,)`` float64.
    """
    pairs: set[tuple[int, int]] = set()
    for head, _relation, tail in triples:
        head, tail = int(head), int(tail)
        if head == tail:
            continue
        pairs.add((head, tail))
        pairs.add((tail, head))

    if not pairs:
        return np.empty((2, 0), dtype=np.int64), np.empty(0, dtype=np.float64)

    rows = np.array([pair[0] for pair in pairs], dtype=np.int64)
    cols = np.array([pair[1] for pair in pairs], dtype=np.int64)
    values = np.ones(len(pairs), dtype=np.float64)
    return np.stack([rows, cols]), values

build_relation_masks

build_relation_masks(
    triples: NDArray[int64], num_relations: int
) -> tuple[
    tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]],
    tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]],
]

Row-normalized head/tail entity-membership matrices, per OpenEA's rfunc/compute_r.

head_mask[r, e] = 1/|heads(r)| if e is ever a head of relation r (0 otherwise); tail_mask is the mirror image over tails. Membership is binary per (relation, entity) pair before normalizing -- an entity appearing as relation r's head in multiple triples still only contributes once, matching OpenEA's own head_r[h][r] = 1 (assignment, not accumulation). Pre-normalizing (dividing by each relation's head/tail count) here means the torch side's per-relation feature computation is a single sparse matmul (head_mask @ primal_embeddings) rather than a matmul-then-divide.

Parameters:

Name	Type	Description	Default
triples	NDArray[int64]	(n, 3) int64 array of (head, relation, tail) ids.	required
num_relations	int	Total relation count.	required

Returns:

Type	Description
tuple[NDArray[int64], NDArray[float64]]	((head_indices, head_values), (tail_indices, tail_values)),
tuple[NDArray[int64], NDArray[float64]]	each indices shaped (2, nnz) (row 0 = relation id, row 1 =
tuple[tuple[NDArray[int64], NDArray[float64]], tuple[NDArray[int64], NDArray[float64]]]	entity id) and values shaped (nnz,) float64 -- ready for
tuple[tuple[NDArray[int64], NDArray[float64]], tuple[NDArray[int64], NDArray[float64]]]	coo_to_torch_sparse
tuple[tuple[NDArray[int64], NDArray[float64]], tuple[NDArray[int64], NDArray[float64]]]	with size=(num_relations, num_entities).

Source code in src/linkingtk/algorithms/ea/_rdgcn_training.py

def build_relation_masks(
    triples: npt.NDArray[np.int64], num_relations: int
) -> tuple[
    tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]],
    tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]],
]:
    """Row-normalized head/tail entity-membership matrices, per OpenEA's ``rfunc``/``compute_r``.

    ``head_mask[r, e] = 1/|heads(r)|`` if ``e`` is ever a head of relation
    ``r`` (``0`` otherwise); ``tail_mask`` is the mirror image over tails.
    Membership is binary per (relation, entity) pair before normalizing --
    an entity appearing as relation ``r``'s head in multiple triples still
    only contributes once, matching OpenEA's own ``head_r[h][r] = 1``
    (assignment, not accumulation). Pre-normalizing (dividing by each
    relation's head/tail count) here means the torch side's per-relation
    feature computation is a single sparse matmul (``head_mask @
    primal_embeddings``) rather than a matmul-then-divide.

    Args:
        triples: ``(n, 3)`` int64 array of ``(head, relation, tail)`` ids.
        num_relations: Total relation count.

    Returns:
        ``((head_indices, head_values), (tail_indices, tail_values))``,
        each ``indices`` shaped ``(2, nnz)`` (row 0 = relation id, row 1 =
        entity id) and ``values`` shaped ``(nnz,)`` float64 -- ready for
        [coo_to_torch_sparse][linkingtk.utils.sparse_gcn.coo_to_torch_sparse]
        with ``size=(num_relations, num_entities)``.
    """
    heads_by_relation: dict[int, set[int]] = {}
    tails_by_relation: dict[int, set[int]] = {}
    for head, relation, tail in triples:
        relation = int(relation)
        heads_by_relation.setdefault(relation, set()).add(int(head))
        tails_by_relation.setdefault(relation, set()).add(int(tail))

    def _mask(
        by_relation: dict[int, set[int]],
    ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]]:
        rows, cols, values = [], [], []
        for relation, entities in by_relation.items():
            weight = 1.0 / len(entities)
            for entity in entities:
                rows.append(relation)
                cols.append(entity)
                values.append(weight)
        if not rows:
            return np.empty((2, 0), dtype=np.int64), np.empty(0, dtype=np.float64)
        indices = np.array([rows, cols], dtype=np.int64)
        return indices, np.array(values, dtype=np.float64)

    return _mask(heads_by_relation), _mask(tails_by_relation)

init_name_embeddings

init_name_embeddings(
    entities: list[Entity],
    word_vectors: dict[str, NDArray[floating[Any]]],
    dim: int,
    default_length: int = 4,
) -> dict[str, npt.NDArray[np.float32]]

Each entity's initial embedding: summed word vectors of its first label's tokens.

Ports OpenEA's _get_desc_input's name-embedding init (word_em[e_desc_input].sum(axis=1)): tokenize the entity's local name, look up up to default_length tokens' pretrained word vectors, sum them. A token with no vector (out of word_vectors, or padding past the name's real token count) contributes a zero vector -- not a random placeholder, matching OpenEA's own np.zeros row for out-of-vocabulary words, so a name entirely absent from word_vectors just gets a zero-vector init rather than a fabricated one.

Uses Entity.labels (already populated by this repo's dataset loaders, see EnFr15KAttrDataset) as the "local name" source rather than OpenEA's own _get_local_name_by_name_triple URI-parsing -- confirmed a no-op deviation for EN-FR-15K specifically: that function's dataset-specific name-attribute matching (D_Y/D_W special-cased predicate lists) never triggers for this dataset (falls to its else: name_attribute_list = {} branch), so OpenEA's own code already falls back to the entity URI's tail segment here -- the same text this repo's loader already populates as the label.

Parameters:

Name	Type	Description	Default
entities	list[Entity]	Entities needing an initial embedding (both KGs, concatenated by the caller).	required
word_vectors	dict[str, NDArray[floating[Any]]]	Pretrained word vectors, e.g. from load_fasttext_vectors.	required
dim	int	Embedding dimensionality -- must match word_vectors' own vector dimensionality (no projection is applied).	required
default_length	int	Max tokens summed per entity. OpenEA's published value is 4.	4

Returns:

Type	Description
dict[str, NDArray[float32]]	entity_id -> (dim,) float32 array.

Source code in src/linkingtk/algorithms/ea/_rdgcn_training.py

def init_name_embeddings(
    entities: list[Entity],
    word_vectors: dict[str, npt.NDArray[np.floating[Any]]],
    dim: int,
    default_length: int = 4,
) -> dict[str, npt.NDArray[np.float32]]:
    """Each entity's initial embedding: summed word vectors of its first label's tokens.

    Ports OpenEA's ``_get_desc_input``'s name-embedding init
    (``word_em[e_desc_input].sum(axis=1)``): tokenize the entity's local
    name, look up up to ``default_length``
    tokens' pretrained word vectors, sum them. A token with no vector (out
    of ``word_vectors``, or padding past the name's real token count)
    contributes a zero vector -- **not** a random placeholder, matching
    OpenEA's own ``np.zeros`` row for out-of-vocabulary words, so a
    name entirely absent from ``word_vectors`` just gets a zero-vector
    init rather than a fabricated one.

    Uses `Entity.labels` (already populated by this repo's dataset
    loaders, see
    [EnFr15KAttrDataset][linkingtk.datasets.EnFr15KAttrDataset]) as the
    "local name" source rather than OpenEA's own
    ``_get_local_name_by_name_triple`` URI-parsing -- confirmed a no-op
    deviation for EN-FR-15K specifically: that function's dataset-specific
    name-attribute matching (``D_Y``/``D_W`` special-cased predicate
    lists) never triggers for this dataset (falls to its
    ``else: name_attribute_list = {}`` branch), so OpenEA's own code
    already falls back to the entity URI's tail segment here -- the same
    text this repo's loader already populates as the label.

    Args:
        entities: Entities needing an initial embedding (both KGs,
            concatenated by the caller).
        word_vectors: Pretrained word vectors, e.g. from
            [load_fasttext_vectors][linkingtk.algorithms.ea._kdcoe_text.load_fasttext_vectors].
        dim: Embedding dimensionality -- must match ``word_vectors``' own
            vector dimensionality (no projection is applied).
        default_length: Max tokens summed per entity. OpenEA's published
            value is ``4``.

    Returns:
        ``entity_id -> (dim,) float32 array``.
    """
    result: dict[str, npt.NDArray[np.float32]] = {}
    for entity in entities:
        texts = label_texts(entity)
        text = texts[0] if texts else ""
        tokens = tokenize_description(text)[:default_length]
        vectors = [word_vectors[token] for token in tokens if token in word_vectors]
        if vectors:
            result[entity.id] = np.sum(vectors, axis=0).astype(np.float32)
        else:
            result[entity.id] = np.zeros(dim, dtype=np.float32)
    return result

_rdgcn_torch

PyTorch model and training-step functions for RDGCNLinker.

Ports OpenEA's Layer/RDGCN (approaches/rdgcn.py): two rounds of dual-graph (relation) attention projected onto the primal (entity) graph via sparse attention, each mixed back into the running entity embedding by a learned scalar (alpha/beta), followed by two diagonal-weight GCN layers each wrapped in a highway gate. Uses plain torch.sparse throughout (including torch.sparse.softmax, PyTorch's own built-in sparse-row-softmax primitive) -- no torch_geometric/dgl dependency, same posture as _gcn_align_torch.py.

Callers must already have confirmed torch is importable (rdgcn.py's fit() does this via OptionalDependencyError before any of these run) -- same precedent as _gcn_align_torch.py/ _kdcoe_torch.py.

Dimension convention (confirmed by reading rdgcn.py's Layer.build() directly): entity/primal embeddings have dimensionality D (the constructor's embedding_dim, which must equal the pretrained word vectors' own dimensionality -- no projection layer exists between them). Per-relation "dual graph" features are 2 * D (concat([mean head-entity embeds, mean tail-entity embeds])) -- every dense-attention layer's query projection takes 2 * D in and produces D out; the sparse attention layer's relation-score projection takes 2 * D in and produces a single scalar per edge.

build_rdgcn_model

build_rdgcn_model(
    initial_primal_embeddings: Tensor,
    embedding_dim: int,
    alpha: float,
    beta: float,
) -> torch.nn.Module

Ports Layer.build()'s computation graph as a live PyTorch model.

A factory function rather than a module-level torch.nn.Module subclass, since torch.nn.Module can't be named at module scope without torch installed -- same precedent as _gcn_align_torch.build_gcn_align_branch/ _kdcoe_torch.build_description_encoder.

Unlike build_gcn_align_branch, the returned model's forward() takes the fixed graph structures (adjacencies/masks/edge lists) as arguments rather than storing them internally -- they never change during training (RDGCN has no AliNet-style bootstrapping adjacency mutation), so this just keeps device-placement simple: build every fixed tensor directly on the target device before constructing the model, same as gcn_align.py's own fit() does for its adjacency.

Confirmed from Layer.build() directly: OpenEA's own get_pretrained_input wraps the name-embedding init in a plain tf.Variable (trainable by default, no trainable=False) -- primal_x0 is a trained parameter seeded from the pretrained name embeddings, not a frozen input. This is why the per-relation "dual graph" features (which depend on the live primal embeddings) must be recomputed every forward pass, not cached once.

Parameters:

Name	Type	Description	Default
initial_primal_embeddings	Tensor	(num_entities, embedding_dim) -- the entity-embedding parameter's initial value, e.g. from init_name_embeddings, already placed on the target device.	required
embedding_dim	int	Entity-embedding dimensionality (D). Must equal initial_primal_embeddings's own last dimension.	required
alpha	float	Round-1 mixing weight (primal_X_1 = primal_X_0 + alpha * primal_H_1). OpenEA's published value is 0.1.	required
beta	float	Round-2 mixing weight. OpenEA's published value is 0.3.	required

Returns:

Type	Description
Module	A module whose ``forward(primal_adjacency, dual_adjacency,
Module	head_mask, tail_mask, edge_heads, edge_tails, edge_relations)``
Module	returns (num_entities, embedding_dim) entity embeddings --
Module	see each argument's shape in
Module	RDGCNLinker.fit's
Module	call site.

Source code in src/linkingtk/algorithms/ea/_rdgcn_torch.py

def build_rdgcn_model(
    initial_primal_embeddings: torch.Tensor,
    embedding_dim: int,
    alpha: float,
    beta: float,
) -> torch.nn.Module:
    """Ports ``Layer.build()``'s computation graph as a live PyTorch model.

    A factory function rather than a module-level ``torch.nn.Module``
    subclass, since ``torch.nn.Module`` can't be named at module scope
    without ``torch`` installed -- same precedent as
    ``_gcn_align_torch.build_gcn_align_branch``/
    ``_kdcoe_torch.build_description_encoder``.

    Unlike
    [build_gcn_align_branch][linkingtk.algorithms.ea._gcn_align_torch.build_gcn_align_branch],
    the returned model's ``forward()`` takes the fixed graph structures
    (adjacencies/masks/edge lists) as *arguments* rather than storing them
    internally -- they never change during training (RDGCN has no
    AliNet-style bootstrapping adjacency mutation), so this just keeps
    device-placement simple: build every fixed tensor directly on the
    target device before constructing the model, same as
    ``gcn_align.py``'s own ``fit()`` does for its adjacency.

    Confirmed from ``Layer.build()`` directly: OpenEA's own
    ``get_pretrained_input`` wraps the name-embedding init in a plain
    ``tf.Variable`` (trainable by default, no ``trainable=False``) --
    ``primal_x0`` is a **trained** parameter seeded from the pretrained
    name embeddings, not a frozen input. This is why the per-relation
    "dual graph" features (which depend on the live primal embeddings)
    must be recomputed every forward pass, not cached once.

    Args:
        initial_primal_embeddings: ``(num_entities, embedding_dim)`` --
            the entity-embedding parameter's initial value, e.g. from
            [init_name_embeddings][linkingtk.algorithms.ea._rdgcn_training.init_name_embeddings],
            already placed on the target device.
        embedding_dim: Entity-embedding dimensionality (``D``). Must equal
            ``initial_primal_embeddings``'s own last dimension.
        alpha: Round-1 mixing weight (``primal_X_1 = primal_X_0 + alpha *
            primal_H_1``). OpenEA's published value is ``0.1``.
        beta: Round-2 mixing weight. OpenEA's published value is ``0.3``.

    Returns:
        A module whose ``forward(primal_adjacency, dual_adjacency,
        head_mask, tail_mask, edge_heads, edge_tails, edge_relations)``
        returns ``(num_entities, embedding_dim)`` entity embeddings --
        see each argument's shape in
        [RDGCNLinker.fit][linkingtk.algorithms.ea.rdgcn.RDGCNLinker.fit]'s
        call site.
    """
    import torch
    import torch.nn.functional as functional
    from torch import nn

    dual_dim = 2 * embedding_dim

    class _DenseGraphAttentionLayer(nn.Module):
        """Ports ``add_self_att_layer``/``add_dual_att_layer`` (identical math,
        differing only in whether the query and value inputs coincide, and
        whether the query projection has a bias)."""

        def __init__(self, use_bias: bool) -> None:
            super().__init__()
            self.proj = nn.Linear(dual_dim, embedding_dim, bias=use_bias)
            self.score1 = nn.Linear(embedding_dim, 1)
            self.score2 = nn.Linear(embedding_dim, 1)

        def forward(
            self, query_input: torch.Tensor, value_input: torch.Tensor, adjacency: torch.Tensor
        ) -> torch.Tensor:
            feats = self.proj(query_input)
            f1 = self.score1(feats)
            f2 = self.score2(feats)
            logits = f1 + f2.transpose(0, 1)
            logits = adjacency * logits
            bias_mat = -1e9 * (1.0 - (adjacency > 0).float())
            coefs = torch.softmax(functional.leaky_relu(logits) + bias_mat, dim=1)
            return torch.relu(coefs @ value_input)

    class _SparseAttentionLayer(nn.Module):
        """Ports ``add_sparse_att_layer``: projects the dual-graph output to a
        per-relation scalar, gathers it onto every primal edge via
        ``edge_relations``, softmax-normalizes per source entity (row), and
        aggregates the primal layer through the result.

        Duplicate ``(head, tail)`` edges connected by *different* relations
        (the same entity pair linked by more than one relation) are summed
        by ``.coalesce()`` before the softmax, rather than kept as
        separate competing entries the way TensorFlow's uncoalesced sparse
        tensor lets OpenEA's own code do -- a deliberate simplification
        (`torch.sparse.softmax` needs a coalesced input; replicating
        per-edge-not-per-cell softmax without `torch_scatter` isn't worth
        it for what's a rare case in practice, one entity pair connected by
        multiple distinct relations simultaneously).
        """

        def __init__(self) -> None:
            super().__init__()
            self.transform = nn.Linear(dual_dim, 1)

        def forward(
            self,
            dual_layer: torch.Tensor,
            primal_layer: torch.Tensor,
            edge_heads: torch.Tensor,
            edge_tails: torch.Tensor,
            edge_relations: torch.Tensor,
            num_entities: int,
        ) -> torch.Tensor:
            relation_scores = self.transform(dual_layer).squeeze(-1)
            edge_logits = functional.leaky_relu(relation_scores[edge_relations])
            indices = torch.stack([edge_heads, edge_tails])
            sparse_logits = torch.sparse_coo_tensor(
                indices, edge_logits, size=(num_entities, num_entities)
            ).coalesce()
            coefs = torch.sparse.softmax(sparse_logits, dim=1)
            return torch.relu(torch.sparse.mm(coefs, primal_layer))

    class _DiagonalGCNLayer(nn.Module):
        """Ports ``add_diag_layer``: an elementwise (not full-matrix) learnable
        weight, propagated through the primal adjacency."""

        def __init__(self) -> None:
            super().__init__()
            self.scale = nn.Parameter(torch.ones(1, embedding_dim))

        def forward(self, inputs: torch.Tensor, adjacency: torch.Tensor) -> torch.Tensor:
            return torch.relu(torch.sparse.mm(adjacency, inputs * self.scale))

    class _HighwayGate(nn.Module):
        """Ports ``highway``."""

        def __init__(self) -> None:
            super().__init__()
            self.gate = nn.Linear(embedding_dim, embedding_dim)

        def forward(self, layer1: torch.Tensor, layer2: torch.Tensor) -> torch.Tensor:
            transform_gate = torch.sigmoid(self.gate(layer1))
            carry_gate = 1.0 - transform_gate
            return transform_gate * layer2 + carry_gate * layer1

    class _RDGCNModel(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.primal_x0 = nn.Parameter(initial_primal_embeddings.clone())
            self.round1_att = _DenseGraphAttentionLayer(use_bias=False)
            self.round1_sparse_att = _SparseAttentionLayer()
            self.round2_att = _DenseGraphAttentionLayer(use_bias=True)
            self.round2_sparse_att = _SparseAttentionLayer()
            self.diag1 = _DiagonalGCNLayer()
            self.diag2 = _DiagonalGCNLayer()
            self.highway1 = _HighwayGate()
            self.highway2 = _HighwayGate()
            self.alpha = alpha
            self.beta = beta

        def _relation_features(
            self, primal_embeds: torch.Tensor, head_mask: torch.Tensor, tail_mask: torch.Tensor
        ) -> torch.Tensor:
            head_means = torch.sparse.mm(head_mask, primal_embeds)
            tail_means = torch.sparse.mm(tail_mask, primal_embeds)
            return torch.cat([head_means, tail_means], dim=-1)

        def forward(
            self,
            primal_adjacency: torch.Tensor,
            dual_adjacency: torch.Tensor,
            head_mask: torch.Tensor,
            tail_mask: torch.Tensor,
            edge_heads: torch.Tensor,
            edge_tails: torch.Tensor,
            edge_relations: torch.Tensor,
        ) -> torch.Tensor:
            num_entities = self.primal_x0.shape[0]
            primal_x0 = self.primal_x0

            dual_x1 = self._relation_features(primal_x0, head_mask, tail_mask)
            dual_h1 = self.round1_att(dual_x1, dual_x1, dual_adjacency)
            primal_h1 = self.round1_sparse_att(
                dual_h1, primal_x0, edge_heads, edge_tails, edge_relations, num_entities
            )
            primal_x1 = primal_x0 + self.alpha * primal_h1

            dual_x2 = self._relation_features(primal_x1, head_mask, tail_mask)
            dual_h2 = self.round2_att(dual_x2, dual_h1, dual_adjacency)
            primal_h2 = self.round2_sparse_att(
                dual_h2, primal_x1, edge_heads, edge_tails, edge_relations, num_entities
            )
            primal_x2 = primal_x0 + self.beta * primal_h2

            gcn1 = self.diag1(primal_x2, primal_adjacency)
            gcn1 = self.highway1(primal_x2, gcn1)
            gcn2 = self.diag2(gcn1, primal_adjacency)
            output: torch.Tensor = self.highway2(gcn1, gcn2)
            return output

    return _RDGCNModel()

mine_hard_negatives

mine_hard_negatives(
    embeddings: Tensor,
    pos_left: Tensor,
    pos_right: Tensor,
    k: int,
) -> tuple[
    torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
]

Manhattan-distance hard-negative mining, per OpenEA's get_neg.

For each seed pair's left/right anchor entity, the k nearest (L1/cityblock distance) other entities in the current embedding space become that side's negatives -- unlike GCN-Align's uniform random sampling. Uses torch.cdist (GPU-friendly) rather than scipy.spatial.distance.cdist -- with a real seed-pair count in the thousands and an entity count in the tens of thousands, computing this on CPU every refresh would be a genuine bottleneck (same "vectorize/ accelerate before trusting a smaller-scale precedent" lesson as _bootea_training.encode_triples's profiling note).

Deliberate deviation: OpenEA's own get_neg never excludes the anchor entity itself from its own candidate pool -- since the anchor's distance to itself is always the global minimum (0), it's very often ranked as its own "hardest negative," which is meaningless as a training signal (the corresponding loss term becomes relu(margin + pos_dist), effectively a constant that can never be reduced by training). This function excludes each anchor's own id from its ranking before taking the top k -- a deliberate fix, not reproduced as OpenEA's own apparent oversight would have it.

Parameters:

Name	Type	Description	Default
embeddings	Tensor	(num_entities, dim) current entity embeddings.	required
pos_left	Tensor	(t,) ids, seed pairs' left/source side.	required
pos_right	Tensor	(t,) ids, seed pairs' right/target side.	required
k	int	Negatives mined per pair per side.	required

Returns:

Type	Description
Tensor	(neg_left, neg_right, neg2_left, neg2_right), each (t * k,)
Tensor	-- pass directly to
Tensor	margin_ranking_loss_l1.
Tensor	neg_left/neg2_right are each anchor's own id repeated
tuple[Tensor, Tensor, Tensor, Tensor]	k times; neg_right/neg2_left are the hard-mined
tuple[Tensor, Tensor, Tensor, Tensor]	neighbor ids.

Source code in src/linkingtk/algorithms/ea/_rdgcn_torch.py

def mine_hard_negatives(
    embeddings: torch.Tensor,
    pos_left: torch.Tensor,
    pos_right: torch.Tensor,
    k: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Manhattan-distance hard-negative mining, per OpenEA's ``get_neg``.

    For each seed pair's left/right anchor entity, the ``k`` nearest
    (L1/cityblock distance) *other* entities in the current embedding
    space become that side's negatives -- unlike GCN-Align's uniform
    random sampling. Uses `torch.cdist` (GPU-friendly) rather than
    `scipy.spatial.distance.cdist` -- with a real seed-pair count in the
    thousands and an entity count in the tens of thousands, computing this
    on CPU every refresh would be a genuine bottleneck (same "vectorize/
    accelerate before trusting a smaller-scale precedent" lesson as
    `_bootea_training.encode_triples`'s profiling note).

    **Deliberate deviation**: OpenEA's own ``get_neg`` never excludes the
    anchor entity itself from its own candidate pool -- since the anchor's
    distance to itself is always the global minimum (``0``), it's very
    often ranked as its own "hardest negative," which is meaningless as a
    training signal (the corresponding loss term becomes ``relu(margin +
    pos_dist)``, effectively a constant that can never be reduced by
    training). This function excludes each anchor's own id from its
    ranking before taking the top ``k`` -- a deliberate fix, not
    reproduced as OpenEA's own apparent oversight would have it.

    Args:
        embeddings: ``(num_entities, dim)`` current entity embeddings.
        pos_left: ``(t,)`` ids, seed pairs' left/source side.
        pos_right: ``(t,)`` ids, seed pairs' right/target side.
        k: Negatives mined per pair per side.

    Returns:
        ``(neg_left, neg_right, neg2_left, neg2_right)``, each ``(t * k,)``
        -- pass directly to
        [margin_ranking_loss_l1][linkingtk.algorithms.ea._ea_losses.margin_ranking_loss_l1].
        ``neg_left``/``neg2_right`` are each anchor's own id repeated
        ``k`` times; ``neg_right``/``neg2_left`` are the hard-mined
        neighbor ids.
    """
    import torch

    def _hard_neighbors(anchor_ids: torch.Tensor) -> torch.Tensor:
        anchor_vecs = embeddings[anchor_ids]
        distances = torch.cdist(anchor_vecs, embeddings, p=1)
        distances.scatter_(1, anchor_ids.unsqueeze(1), float("inf"))
        ranked = torch.topk(distances, k, dim=1, largest=False).indices
        return ranked.reshape(-1)

    neg_right = _hard_neighbors(pos_left)
    neg2_left = _hard_neighbors(pos_right)
    neg_left = pos_left.repeat_interleave(k)
    neg2_right = pos_right.repeat_interleave(k)
    return neg_left, neg_right, neg2_left, neg2_right

AliNet

alinet

AliNet knowledge-graph-embedding linker for Entity Alignment.

Sun, Z., Wang, C., Hu, W., Chen, M., Dai, J., Zhang, W., & Qu, Y. (2020). Knowledge Graph Alignment Network with Gated Multi-hop Neighborhood Aggregation. AAAI 2020. https://ojs.aaai.org/index.php/AAAI/article/view/5354

The third GNN-based EA linker in this package (see #18's parent issue,

42's GCNAlignLinker,

43's RDGCNLinker). AliNet

combines a 1-hop GraphConvolution branch with a 2-hop gated-attention branch at every layer but the last, fusing them with a highway gate, and represents each entity by concatenating every layer's (plus the initial embedding's) output -- letting non-isomorphic local neighborhoods between the two KGs still align through a longer-range, multi-hop signal. This is a faithful port of OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA/blob/master/src/openea/approaches/alinet.py), not a from-scratch reading of the paper -- see _alinet_training and _alinet_torch for the algorithm's own docstrings, and this module's deviations below (confirmed against OpenEA's own reference config, run/args/alinet_args_15K.json, and source):

• Bootstrapping is implemented but off by default, matching what OpenEA's own published config actually runs -- not the milestone planning assumption that it's "core to published performance." Reading AliNet.run() directly: bootstrapping (augment_neighborhood()) is only called if self.args.sim_th > 0.0, and the real published config's own value is "sim_th": 0.0 -- 0.0 > 0.0 is False, so OpenEA's own published EN-FR-15K-V1 numbers never invoke bootstrapping at all. This port replicates that exactly: sim_th defaults to 0.0 (bootstrapping off); pass sim_th > 0.0 to opt in.
• No published EN-FR-15K-V1 number exists for AliNet to benchmark against in docs/detailed_results_current_approaches_15K.csv -- added to OpenEA after that CSV's last update; its numbers live only in an external Google Sheet linked from OpenEA's README, not reliably fetchable. examples/alinet_benchmark.py exists for manual sanity-checking only, with no target to compare against.
• The relation-consistency auxiliary loss (compute_rel_loss, rel_param) is not implemented. OpenEA's own published config uses a small non-zero rel_param: 0.01 (i.e. it is nominally active, correcting an earlier assumption that it's off by default) -- skipped here since no numeric target exists to validate it against anyway, and it needs its own rel_win_size-batched relation-window sampling machinery for what's a secondary loss term.
• Uniform negative sampling only. OpenEA's own published config uses neg_sampling: "truncated" (nearest-neighbor-restricted negatives after the first validation checkpoint) -- not implemented; see sample_uniform_cross_kg_negatives's docstring.
• Dropout and L2 weight regularization are not implemented -- both faithful no-ops (OpenEA's own config uses dropout: 0.0, and its training graph never sums the declared weight regularizers into the actual loss). See _alinet_torch's module docstring.
• Candidate scoring uses cosine similarity, unlike GCNAlignLinker/ RDGCNLinker (both switched to Manhattan distance -- see their own module docstrings for why). Checked, not overlooked: OpenEA's own config for this method is eval_metric: "inner" (raw dot product) over embeddings its own _eval_test_embeddings/save L2-normalize immediately before scoring -- and AliNetModel's own forward() already returns L2-normalized output (the final functional.normalize(concatenated, dim=1) call), so raw inner product over these embeddings is mathematically identical to cosine similarity (a . b == cosine(a, b) when |a| = |b| = 1). No change needed here.

AliNetLinker

AliNetLinker(
    layer_dims: list[int] | None = None,
    num_epochs: int = 500,
    learning_rate: float = 0.001,
    batch_size: int = 3000,
    neg_triple_num: int = 10,
    neg_margin: float = 1.5,
    neg_margin_balance: float = 0.1,
    sim_th: float = 0.0,
    bootstrap_start_epoch: int = 20,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Scores candidate pairs via AliNet's gated multi-hop GCN/attention embeddings.

Must be fit before link can be called. See the module docstring for what's ported vs. deliberately deviated from OpenEA.

Parameters:

Name	Type	Description	Default
layer_dims	list[int] | None	Layer width sequence -- layer_dims[0] is the initial embedding's dimensionality, each subsequent entry one more GCN+attention+highway layer's output width. OpenEA's published EN-FR-15K-V1 config uses [500, 400, 300].	None
num_epochs	int	Training epochs. OpenEA's config allows up to 2000 with early stopping.	500
learning_rate	float	Adam's learning rate. OpenEA's published value is 0.001.	0.001
batch_size	int	Positive pairs sampled per training step (with replacement, matching OpenEA). OpenEA's published value is 3000.	3000
neg_triple_num	int	Negatives drawn per positive pair (k, both sides independently). OpenEA's published value is 10.	10
neg_margin	float	Margin for margin_ranking_loss_l2_squared. OpenEA's published value is 1.5.	1.5
neg_margin_balance	float	Weight applied to the negative-pair loss term. OpenEA's published value is 0.1.	0.1
sim_th	float	Minimum similarity for a bootstrapped pair to be accepted -- also the bootstrapping on/off switch: OpenEA's own run() only calls its bootstrapping step if sim_th > 0.0, and its own published config uses 0.0 (off). See the module docstring.	0.0
bootstrap_start_epoch	int	First epoch a bootstrapping round may run, if sim_th > 0. OpenEA's published value (start_augment * eval_freq) is 2 * 10 = 20.	20
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda"/"cuda:0". Trained embeddings are always stored as CPU numpy arrays regardless of this setting.	'cpu'

Source code in src/linkingtk/algorithms/ea/alinet.py

def __init__(
    self,
    layer_dims: list[int] | None = None,
    num_epochs: int = 500,
    learning_rate: float = 0.001,
    batch_size: int = 3000,
    neg_triple_num: int = 10,
    neg_margin: float = 1.5,
    neg_margin_balance: float = 0.1,
    sim_th: float = 0.0,
    bootstrap_start_epoch: int = 20,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.layer_dims = layer_dims if layer_dims is not None else [500, 400, 300]
    self.num_epochs = num_epochs
    self.learning_rate = learning_rate
    self.batch_size = batch_size
    self.neg_triple_num = neg_triple_num
    self.neg_margin = neg_margin
    self.neg_margin_balance = neg_margin_balance
    self.sim_th = sim_th
    self.bootstrap_start_epoch = bootstrap_start_epoch
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> AliNetLinker

Train the gated multi-hop GCN/attention stack, optionally bootstrapping new pairs.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities -- also used to partition graph into each KG's own triples (same approach IPTransELinker uses), needed for enhance_triples/ generate_2hop_pairs.	required
dataset2	list[Entity]	Target entities. See dataset1.	required
ground_truth	list[tuple[str, str]]	List of (source_id, target_id) known-correct pairs used as positive pairs in the margin loss, and as the initial bootstrapping seed set.	required
graph	Graph	The combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2) from a GraphDatasetLoader's load_graphs()) -- entity ids on both sides must already be disjoint, as they are from that loader.	required
random_state	int | None	Seed for reproducible training. Left unspecified, training is non-deterministic.	None
val_ground_truth	list[tuple[str, str]] | None	Optional held-out pairs used for early stopping -- every eval_every epochs, Hits@1 is checked against this set, and training stops after patience checks with no improvement. Also the cadence bootstrapping rounds run on, when enabled. If None (default), trains the full num_epochs unconditionally (bootstrapping still runs on the eval_every cadence if enabled).	None
patience	int	Number of non-improving eval_every-spaced checks to tolerate before stopping early. Only used if val_ground_truth is given.	5
eval_every	int	How often (in epochs) to check val_ground_truth and (if enabled) run a bootstrapping round.	10

Returns:

Type	Description
AliNetLinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If none of ground_truth's pairs have both ids present in graph's own triples, or if device is invalid or unavailable.
OptionalDependencyError	If torch isn't installed.

Source code in src/linkingtk/algorithms/ea/alinet.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    random_state: int | None = None,
    val_ground_truth: list[tuple[str, str]] | None = None,
    patience: int = 5,
    eval_every: int = 10,
) -> AliNetLinker:
    """Train the gated multi-hop GCN/attention stack, optionally bootstrapping new pairs.

    Args:
        dataset1: Source entities -- also used to partition ``graph``
            into each KG's own triples (same approach
            [IPTransELinker][linkingtk.algorithms.ea.iptranse.IPTransELinker]
            uses), needed for
            [enhance_triples][linkingtk.algorithms.ea._alinet_training.enhance_triples]/
            [generate_2hop_pairs][linkingtk.algorithms.ea._alinet_training.generate_2hop_pairs].
        dataset2: Target entities. See ``dataset1``.
        ground_truth: List of ``(source_id, target_id)`` known-correct
            pairs used as positive pairs in the margin loss, and as
            the initial bootstrapping seed set.
        graph: The combined relational structure of both KGs (e.g.
            ``to_triples(graph1) + to_triples(graph2)`` from a
            [GraphDatasetLoader][linkingtk.datasets.base.GraphDatasetLoader]'s
            ``load_graphs()``) -- entity ids on both sides must already
            be disjoint, as they are from that loader.
        random_state: Seed for reproducible training. Left unspecified,
            training is non-deterministic.
        val_ground_truth: Optional held-out pairs used for early
            stopping -- every ``eval_every`` epochs, Hits@1 is checked
            against this set, and training stops after ``patience``
            checks with no improvement. Also the cadence bootstrapping
            rounds run on, when enabled. If ``None`` (default), trains
            the full ``num_epochs`` unconditionally (bootstrapping
            still runs on the ``eval_every`` cadence if enabled).
        patience: Number of non-improving ``eval_every``-spaced checks
            to tolerate before stopping early. Only used if
            ``val_ground_truth`` is given.
        eval_every: How often (in epochs) to check ``val_ground_truth``
            and (if enabled) run a bootstrapping round.

    Returns:
        ``self``, for chaining.

    Raises:
        LinkingTKError: If none of ``ground_truth``'s pairs have both
            ids present in ``graph``'s own triples, or if ``device`` is
            invalid or unavailable.
        OptionalDependencyError: If torch isn't installed.
    """
    try:
        import torch
    except ImportError as exc:
        raise OptionalDependencyError("AliNetLinker", "kge") from exc

    device = resolve_device(self.device)
    if random_state is not None:
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)
    rng = np.random.default_rng(random_state)

    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    triples = to_triples(graph)
    entity_to_id, relation_to_id = build_id_mappings(triples)
    mapped = map_triples_to_ids(triples, entity_to_id, relation_to_id)
    num_entities = len(entity_to_id)

    triples1 = [t for t in triples if t[0] in ids1]
    triples2 = [t for t in triples if t[0] in ids2]
    mapped1 = map_triples_to_ids(triples1, entity_to_id, relation_to_id)
    mapped2 = map_triples_to_ids(triples2, entity_to_id, relation_to_id)

    seed_pairs = [
        (entity_to_id[s], entity_to_id[t])
        for s, t in ground_truth
        if s in entity_to_id and t in entity_to_id
    ]
    if not seed_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have both ids present in `graph`'s "
            "own triples; fit() has no seed pairs to train AliNet's margin loss with."
        )

    entities1_ids = [entity_to_id[eid] for eid in ids1 if eid in entity_to_id]
    entities2_ids = [entity_to_id[eid] for eid in ids2 if eid in entity_to_id]

    two_hop_pairs = generate_2hop_pairs(mapped1) | generate_2hop_pairs(mapped2)
    two_indices, two_values = pairs_to_symmetric_adjacency(two_hop_pairs, num_entities)
    two_indices, two_values = normalize_adjacency_coo(two_indices, two_values, num_entities)
    two_hop_adjacency = coo_to_torch_sparse(
        two_indices, two_values, (num_entities, num_entities), device
    )

    def _build_one_hop_adjacency(pairs: list[tuple[int, int]]) -> torch.Tensor:
        enhanced1, enhanced2 = enhance_triples(mapped1, mapped2, pairs)
        all_triples = np.array(
            mapped.tolist() + list(enhanced1) + list(enhanced2), dtype=np.int64
        )
        indices, values = build_primal_adjacency(all_triples, num_entities)
        indices, values = normalize_adjacency_coo(indices, values, num_entities)
        return coo_to_torch_sparse(indices, values, (num_entities, num_entities), device)

    model = build_alinet_model(num_entities, self.layer_dims).to(device)
    _set_adjacency(model, _build_one_hop_adjacency(seed_pairs), two_hop_adjacency)
    optimizer = torch.optim.Adam(model.parameters(), lr=self.learning_rate)

    aligned1 = {source for source, _target in seed_pairs}
    aligned2 = {target for _source, target in seed_pairs}
    unaligned1 = [entity for entity in entities1_ids if entity not in aligned1]
    unaligned2 = [entity for entity in entities2_ids if entity not in aligned2]

    val_pairs = [
        (s, t) for s, t in (val_ground_truth or []) if s in entity_to_id and t in entity_to_id
    ]
    best_hits1 = -1.0
    epochs_without_improvement = 0

    for epoch in range(self.num_epochs):
        model.train()
        batch_size = min(self.batch_size, len(seed_pairs))
        batch_idx = rng.integers(0, len(seed_pairs), size=batch_size)
        pos_left = torch.tensor([seed_pairs[i][0] for i in batch_idx], device=device)
        pos_right = torch.tensor([seed_pairs[i][1] for i in batch_idx], device=device)

        neg_left_np, neg_right_np = sample_uniform_cross_kg_negatives(
            entities1_ids,
            entities2_ids,
            batch_size,
            self.neg_triple_num,
            rng,
            exclude_pairs=set(seed_pairs),
        )
        neg_left = torch.from_numpy(neg_left_np).long().to(device)
        neg_right = torch.from_numpy(neg_right_np).long().to(device)

        embeddings = model()
        loss = margin_ranking_loss_l2_squared(
            embeddings,
            pos_left,
            pos_right,
            neg_left,
            neg_right,
            self.neg_margin,
            self.neg_margin_balance,
        )
        optimizer.zero_grad()
        loss.backward()  # type: ignore[no-untyped-call]
        optimizer.step()

        if (epoch + 1) % eval_every == 0:
            model.eval()
            with torch.no_grad():
                current_embeds = model().cpu().numpy()
            model.train()

            if val_pairs:
                hits1 = _validation_hits1(current_embeds, entity_to_id, val_pairs)
                if hits1 <= best_hits1:
                    epochs_without_improvement += 1
                    if epochs_without_improvement >= patience:
                        break
                else:
                    best_hits1 = hits1
                    epochs_without_improvement = 0

            if self.sim_th > 0.0 and epoch + 1 >= self.bootstrap_start_epoch:
                new_seed_pairs, unaligned1, unaligned2 = run_bootstrapping_round(
                    current_embeds, seed_pairs, unaligned1, unaligned2, self.sim_th
                )
                if len(new_seed_pairs) > len(seed_pairs):
                    seed_pairs = new_seed_pairs
                    _set_adjacency(
                        model, _build_one_hop_adjacency(seed_pairs), two_hop_adjacency
                    )

    model.eval()
    with torch.no_grad():
        final_embeds = model().cpu().numpy()
    self._id_to_vector = {
        entity_id: final_embeds[index] for entity_id, index in entity_to_id.items()
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/alinet.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/alinet.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score ``entity_id`` as a scored pair's target side."""
    return self._embedding(entity_id)

_alinet_training

Private, plain-numpy (torch-free) helpers for AliNetLinker.

Ported from OpenEA's reference implementation (https://github.com/nju-websoft/OpenEA -- approaches/alinet.py's generate_2hop_triples/enhance_triples/no_weighted_adj, and AliNet.augment/augment_neighborhood's bootstrapping loop). Independently testable without torch installed -- see _alinet_torch.py for the GCN/attention/highway model and training-step functions.

Reuses two functions already ported for sibling linkers rather than reimplementing them: build_primal_adjacency (RDGCN's unweighted-symmetric-adjacency builder is exactly OpenEA's own no_weighted_adj's edge-set construction -- used here for the 1-hop adjacency, given a triple list; the 2-hop adjacency's own pairs_to_symmetric_adjacency below is the same construction over a raw pair set instead) and find_new_pairs (IPTransE's row-argmax-plus-threshold match finder is exactly OpenEA's own find_alignment(sim_mat, sim_th, k=1), which is what AliNet's own augment calls too).

enhance_triples

enhance_triples(
    triples1: NDArray[int64],
    triples2: NDArray[int64],
    seed_pairs: list[tuple[int, int]],
) -> tuple[
    set[tuple[int, int, int]], set[tuple[int, int, int]]
]

Cross-KG-inferred triples, per OpenEA's enhance_triples.

For each KG1 triple (h1, r, t1) whose both entities have a seed counterpart (h2 = links[h1], t2 = links[t1]), infers a matching KG2 triple (h2, r, t2) -- unless KG2 already has a direct h2 -> t2 edge (avoiding redundant edges). Mirrored for KG2 -> KG1. This is what lets a single seed pair "bridge" one KG's own relational structure into the other's, making the eventual 1-hop adjacency cross-KG-connected rather than two disjoint components.

Parameters:

Name	Type	Description	Default
triples1	NDArray[int64]	KG1's own (head, relation, tail) id triples (already mapped into the combined id space).	required
triples2	NDArray[int64]	KG2's own triples, same id space.	required
seed_pairs	list[tuple[int, int]]	(kg1_id, kg2_id) known-correct pairs.	required

Returns:

Type	Description
set[tuple[int, int, int]]	(enhanced_triples1, enhanced_triples2) -- inferred triples to
set[tuple[int, int, int]]	add to KG1's/KG2's own triple sets before adjacency construction.

Source code in src/linkingtk/algorithms/ea/_alinet_training.py

def enhance_triples(
    triples1: npt.NDArray[np.int64],
    triples2: npt.NDArray[np.int64],
    seed_pairs: list[tuple[int, int]],
) -> tuple[set[tuple[int, int, int]], set[tuple[int, int, int]]]:
    """Cross-KG-inferred triples, per OpenEA's ``enhance_triples``.

    For each KG1 triple ``(h1, r, t1)`` whose *both* entities have a seed
    counterpart (``h2 = links[h1]``, ``t2 = links[t1]``), infers a matching
    KG2 triple ``(h2, r, t2)`` -- unless KG2 already has a direct ``h2 ->
    t2`` edge (avoiding redundant edges). Mirrored for KG2 -> KG1. This is
    what lets a single seed pair "bridge" one KG's own relational
    structure into the other's, making the eventual 1-hop adjacency
    cross-KG-connected rather than two disjoint components.

    Args:
        triples1: KG1's own ``(head, relation, tail)`` id triples (already
            mapped into the combined id space).
        triples2: KG2's own triples, same id space.
        seed_pairs: ``(kg1_id, kg2_id)`` known-correct pairs.

    Returns:
        ``(enhanced_triples1, enhanced_triples2)`` -- inferred triples to
        add to KG1's/KG2's own triple sets before adjacency construction.
    """
    links1 = dict(seed_pairs)
    links2 = {target: source for source, target in seed_pairs}

    out_edges1: dict[int, set[int]] = defaultdict(set)
    for head, _relation, tail in triples1:
        out_edges1[int(head)].add(int(tail))
    out_edges2: dict[int, set[int]] = defaultdict(set)
    for head, _relation, tail in triples2:
        out_edges2[int(head)].add(int(tail))

    enhanced2: set[tuple[int, int, int]] = set()
    for head, relation, tail in triples1:
        head2 = links1.get(int(head))
        tail2 = links1.get(int(tail))
        if head2 is not None and tail2 is not None and tail2 not in out_edges2.get(head2, set()):
            enhanced2.add((head2, int(relation), tail2))

    enhanced1: set[tuple[int, int, int]] = set()
    for head, relation, tail in triples2:
        head1 = links2.get(int(head))
        tail1 = links2.get(int(tail))
        if head1 is not None and tail1 is not None and tail1 not in out_edges1.get(head1, set()):
            enhanced1.add((head1, int(relation), tail1))

    return enhanced1, enhanced2

generate_2hop_pairs

generate_2hop_pairs(
    triples: NDArray[int64],
    skip_top_k: int = _SKIP_TOP_K_PATTERNS,
) -> set[tuple[int, int]]

Entity pairs reachable by a 2-hop path, filtered by relation-pattern frequency.

Ports OpenEA's generate_2hop_triples: finds every (h, m, t) path (h -r1-> m -r2-> t) whose (h, t) pair is not already a direct 1-hop edge (either direction), groups them by (r1, r2) relation-composition pattern, and keeps only paths whose pattern is not among the skip_top_k most frequent -- confirmed from reading generate_2hop_triples directly: it computes (but never uses) a p=0.05-proportional cutoff (num = int(p * len(relation_patterns))) and its commented-out original loop (# for i in range(20, num)), but the actual active loop is for i in range(5, len(relation_patterns)) -- i.e. only the top-5 most frequent patterns are ever excluded, num/p are dead code. Ported as what's actually executed, not the apparent original intent.

The relation ids and the synthetic self-loop/composed-relation triples OpenEA's own version builds around this selection are dropped here -- downstream (pairs_to_symmetric_adjacency) only needs (h, t) connectivity, not relation ids, and the synthetic self-loop triple (head, 0, head) OpenEA's own code adds is already redundant with normalize_adjacency_coo's own add_self_loops.

Parameters:

Name	Type	Description	Default
triples	NDArray[int64]	One KG's own (head, relation, tail) id triples (called separately per KG, matching OpenEA's own per-kg usage -- 2-hop paths are never combined across KGs here).	required
skip_top_k	int	Number of most-frequent (r1, r2) patterns to exclude.	_SKIP_TOP_K_PATTERNS

Returns:

Type	Description
set[tuple[int, int]]	{(head, tail)} pairs reachable by a selected 2-hop path.

Source code in src/linkingtk/algorithms/ea/_alinet_training.py

def generate_2hop_pairs(
    triples: npt.NDArray[np.int64], skip_top_k: int = _SKIP_TOP_K_PATTERNS
) -> set[tuple[int, int]]:
    """Entity pairs reachable by a 2-hop path, filtered by relation-pattern frequency.

    Ports OpenEA's ``generate_2hop_triples``: finds every ``(h, m, t)``
    path (``h -r1-> m -r2-> t``) whose ``(h, t)`` pair is **not** already a
    direct 1-hop edge (either direction), groups them by ``(r1, r2)``
    relation-composition pattern, and keeps only paths whose pattern is
    **not** among the ``skip_top_k`` most frequent -- confirmed from
    reading ``generate_2hop_triples`` directly: it computes (but never
    uses) a ``p=0.05``-proportional cutoff (``num = int(p *
    len(relation_patterns))``) and its commented-out original loop
    (``# for i in range(20, num)``), but the *actual* active loop is
    ``for i in range(5, len(relation_patterns))`` -- i.e. only the top-5
    most frequent patterns are ever excluded, `num`/`p` are dead code.
    Ported as what's actually executed, not the apparent original intent.

    The relation ids and the synthetic self-loop/composed-relation triples
    OpenEA's own version builds around this selection are dropped here --
    downstream
    ([pairs_to_symmetric_adjacency][linkingtk.algorithms.ea._alinet_training.pairs_to_symmetric_adjacency])
    only needs ``(h, t)`` connectivity, not relation ids, and the
    synthetic self-loop triple `(head, 0, head)` OpenEA's own code adds is
    already redundant with
    [normalize_adjacency_coo][linkingtk.utils.sparse_gcn.normalize_adjacency_coo]'s
    own ``add_self_loops``.

    Args:
        triples: One KG's own ``(head, relation, tail)`` id triples
            (called separately per KG, matching OpenEA's own per-``kg``
            usage -- 2-hop paths are never combined across KGs here).
        skip_top_k: Number of most-frequent ``(r1, r2)`` patterns to
            exclude.

    Returns:
        ``{(head, tail)}`` pairs reachable by a selected 2-hop path.
    """
    by_head: dict[int, list[tuple[int, int]]] = defaultdict(list)
    out_neighbors: dict[int, set[int]] = defaultdict(set)
    in_neighbors: dict[int, set[int]] = defaultdict(set)
    for head, relation, tail in triples:
        head, relation, tail = int(head), int(relation), int(tail)
        by_head[head].append((relation, tail))
        out_neighbors[head].add(tail)
        in_neighbors[tail].add(head)

    pattern_counts: Counter[tuple[int, int]] = Counter()
    quadruples: list[tuple[int, int, int, int]] = []
    for head, relation1, mid in triples:
        head, relation1, mid = int(head), int(relation1), int(mid)
        for relation2, tail in by_head.get(mid, []):
            if tail in out_neighbors.get(head, set()) or head in in_neighbors.get(tail, set()):
                continue
            pattern_counts[(relation1, relation2)] += 1
            quadruples.append((head, relation1, relation2, tail))

    ranked_patterns = [pattern for pattern, _count in pattern_counts.most_common()]
    selected_patterns = set(ranked_patterns[skip_top_k:])

    return {
        (head, tail)
        for head, relation1, relation2, tail in quadruples
        if (relation1, relation2) in selected_patterns
    }

pairs_to_symmetric_adjacency

pairs_to_symmetric_adjacency(
    pairs: set[tuple[int, int]], num_entities: int
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]]

Unweighted symmetric adjacency from a raw {(head, tail)} pair set.

Parameters:

Name	Type	Description	Default
pairs	set[tuple[int, int]]	Entity-id pairs, e.g. from generate_2hop_pairs.	required
num_entities	int	Total entity count (both KGs combined).	required

Returns:

Type	Description
NDArray[int64]	(indices, values) -- (2, nnz) int64 / (nnz,) float64,
NDArray[float64]	both directions present, ready for
tuple[NDArray[int64], NDArray[float64]]	normalize_adjacency_coo.

Source code in src/linkingtk/algorithms/ea/_alinet_training.py

def pairs_to_symmetric_adjacency(
    pairs: set[tuple[int, int]], num_entities: int
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]]:
    """Unweighted symmetric adjacency from a raw ``{(head, tail)}`` pair set.

    Args:
        pairs: Entity-id pairs, e.g. from
            [generate_2hop_pairs][linkingtk.algorithms.ea._alinet_training.generate_2hop_pairs].
        num_entities: Total entity count (both KGs combined).

    Returns:
        ``(indices, values)`` -- ``(2, nnz)`` int64 / ``(nnz,)`` float64,
        both directions present, ready for
        [normalize_adjacency_coo][linkingtk.utils.sparse_gcn.normalize_adjacency_coo].
    """
    symmetric = {(head, tail) for head, tail in pairs} | {(tail, head) for head, tail in pairs}
    if not symmetric:
        return np.empty((2, 0), dtype=np.int64), np.empty(0, dtype=np.float64)
    rows = np.array([pair[0] for pair in symmetric], dtype=np.int64)
    cols = np.array([pair[1] for pair in symmetric], dtype=np.int64)
    values = np.ones(len(symmetric), dtype=np.float64)
    return np.stack([rows, cols]), values

run_bootstrapping_round

run_bootstrapping_round(
    embeddings: NDArray[floating[Any]],
    seed_pairs: list[tuple[int, int]],
    unaligned1: list[int],
    unaligned2: list[int],
    sim_th: float = 0.0,
) -> tuple[list[tuple[int, int]], list[int], list[int]]

One bootstrapping round: find new high-confidence pairs, expand the seed set.

Ports OpenEA's augment/augment_neighborhood: cosine similarity between every still-unaligned KG1 entity and every still-unaligned KG2 entity, find_new_pairs (= OpenEA's own find_alignment(sim_mat, sim_th, k=1), row-wise top-1 plus threshold) for one-sided candidate matches, then a column-collision resolution step (keep only the highest-similarity match per KG2 target) standing in for OpenEA's separate, stateful update_labeled_alignment_x/_y editing passes -- once a pair is accepted here it's added to the seed set permanently, rather than remaining revisable in later rounds the way OpenEA's own labeled- alignment dict does. A simpler, one-directional approximation of the same "avoid many-to-one collisions" property, not a full port.

Parameters:

Name	Type	Description	Default
embeddings	NDArray[floating[Any]]	(num_entities, dim) current entity embeddings (any consistent representation -- L2-normalized here before comparing).	required
seed_pairs	list[tuple[int, int]]	Current seed set (both original and previously bootstrapped pairs).	required
unaligned1	list[int]	KG1 entity ids not yet in seed_pairs.	required
unaligned2	list[int]	KG2 entity ids not yet in seed_pairs.	required
sim_th	float	Minimum cosine similarity for a match to count. OpenEA's published value is 0.0 (no real floor -- reciprocal/ collision resolution is the operative filter, same posture as find_mutual_pairs's docstring).	0.0

Returns:

Type	Description
list[tuple[int, int]]	(new_seed_pairs, remaining_unaligned1, remaining_unaligned2).
list[int]	new_seed_pairs is seed_pairs plus every pair accepted this
list[int]	round; the two remaining lists have those entities removed.

Source code in src/linkingtk/algorithms/ea/_alinet_training.py

def run_bootstrapping_round(
    embeddings: npt.NDArray[np.floating[Any]],
    seed_pairs: list[tuple[int, int]],
    unaligned1: list[int],
    unaligned2: list[int],
    sim_th: float = 0.0,
) -> tuple[list[tuple[int, int]], list[int], list[int]]:
    """One bootstrapping round: find new high-confidence pairs, expand the seed set.

    Ports OpenEA's ``augment``/``augment_neighborhood``: cosine similarity
    between every still-unaligned KG1 entity and every still-unaligned KG2
    entity, [find_new_pairs][linkingtk.algorithms.ea._iptranse_training.find_new_pairs]
    (= OpenEA's own ``find_alignment(sim_mat, sim_th, k=1)``, row-wise
    top-1 plus threshold) for one-sided candidate matches, then a
    column-collision resolution step (keep only the highest-similarity
    match per KG2 target) standing in for OpenEA's separate, stateful
    ``update_labeled_alignment_x``/``_y`` editing passes -- once a pair is
    accepted here it's added to the seed set permanently, rather than
    remaining revisable in later rounds the way OpenEA's own labeled-
    alignment dict does. A simpler, one-directional approximation of the
    same "avoid many-to-one collisions" property, not a full port.

    Args:
        embeddings: ``(num_entities, dim)`` current entity embeddings
            (any consistent representation -- L2-normalized here before
            comparing).
        seed_pairs: Current seed set (both original and previously
            bootstrapped pairs).
        unaligned1: KG1 entity ids not yet in ``seed_pairs``.
        unaligned2: KG2 entity ids not yet in ``seed_pairs``.
        sim_th: Minimum cosine similarity for a match to count. OpenEA's
            published value is ``0.0`` (no real floor -- reciprocal/
            collision resolution is the operative filter, same posture as
            [find_mutual_pairs][linkingtk.algorithms.ea._iptranse_training.find_mutual_pairs]'s
            docstring).

    Returns:
        ``(new_seed_pairs, remaining_unaligned1, remaining_unaligned2)``.
        ``new_seed_pairs`` is ``seed_pairs`` plus every pair accepted this
        round; the two remaining lists have those entities removed.
    """
    if not unaligned1 or not unaligned2:
        return seed_pairs, unaligned1, unaligned2

    embeds1 = embeddings[unaligned1]
    embeds2 = embeddings[unaligned2]
    embeds1 = embeds1 / np.clip(np.linalg.norm(embeds1, axis=1, keepdims=True), 1e-12, None)
    embeds2 = embeds2 / np.clip(np.linalg.norm(embeds2, axis=1, keepdims=True), 1e-12, None)
    sim_mat = embeds1 @ embeds2.T

    matches = find_new_pairs(sim_mat, sim_th)
    best_by_col: dict[int, tuple[int, float]] = {}
    for row, col, score in matches:
        if col not in best_by_col or score > best_by_col[col][1]:
            best_by_col[col] = (row, score)

    new_pairs = [(unaligned1[row], unaligned2[col]) for col, (row, _score) in best_by_col.items()]
    new_ids1 = {pair[0] for pair in new_pairs}
    new_ids2 = {pair[1] for pair in new_pairs}
    remaining1 = [entity for entity in unaligned1 if entity not in new_ids1]
    remaining2 = [entity for entity in unaligned2 if entity not in new_ids2]
    return seed_pairs + new_pairs, remaining1, remaining2

sample_uniform_cross_kg_negatives

sample_uniform_cross_kg_negatives(
    entities1: list[int],
    entities2: list[int],
    batch_size: int,
    neg_triple_num: int,
    rng: Generator,
    exclude_pairs: set[tuple[int, int]],
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]

Fully-random cross-KG negative pairs, per OpenEA's generate_input_batch's uniform branch.

Unlike GCN-Align's/RDGCN's negatives (each anchored to a real positive pair's own entity), AliNet's uniform negatives are independently drawn from each KG's own entity pool -- neg_ent1[i]/neg_ent2[i] are unrelated random draws, not a corruption of a specific positive pair. Pairs that happen to collide with a known positive (exclude_pairs -- the current seed set, so a "negative" is never accidentally a real alignment) are dropped.

Deliberate scope reduction: OpenEA's own published config uses neg_sampling: "truncated", switching to nearest-neighbor-restricted negatives (via find_neighbors) after the first validation checkpoint -- not implemented here; every epoch uses this uniform scheme. See the module docstring on alinet for why.

Parameters:

Name	Type	Description	Default
entities1	list[int]	Candidate KG1-side entity ids to draw from (OpenEA's own pool is sup_ent1 + ref_ent1 -- every entity, not just currently-seeded ones).	required
entities2	list[int]	Candidate KG2-side entity ids.	required
batch_size	int	Positive pairs in this batch.	required
neg_triple_num	int	Negatives drawn per positive pair (k).	required
rng	Generator	Source of randomness.	required
exclude_pairs	set[tuple[int, int]]	Pairs to drop if drawn (the current seed set).	required

Returns:

Type	Description
NDArray[int64]	(neg_left, neg_right), each up to (batch_size * k,) int64
NDArray[int64]	(shorter if some draws were excluded).

Source code in src/linkingtk/algorithms/ea/_alinet_training.py

def sample_uniform_cross_kg_negatives(
    entities1: list[int],
    entities2: list[int],
    batch_size: int,
    neg_triple_num: int,
    rng: np.random.Generator,
    exclude_pairs: set[tuple[int, int]],
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
    """Fully-random cross-KG negative pairs, per OpenEA's ``generate_input_batch``'s uniform branch.

    Unlike GCN-Align's/RDGCN's negatives (each anchored to a real positive
    pair's own entity), AliNet's uniform negatives are **independently**
    drawn from each KG's own entity pool -- ``neg_ent1[i]``/``neg_ent2[i]``
    are unrelated random draws, not a corruption of a specific positive
    pair. Pairs that happen to collide with a known positive
    (``exclude_pairs`` -- the current seed set, so a "negative" is never
    accidentally a real alignment) are dropped.

    **Deliberate scope reduction**: OpenEA's own published config uses
    ``neg_sampling: "truncated"``, switching to nearest-neighbor-restricted
    negatives (via ``find_neighbors``) after the first validation
    checkpoint -- not implemented here; every epoch uses this uniform
    scheme. See the module docstring on
    [alinet][linkingtk.algorithms.ea.alinet] for why.

    Args:
        entities1: Candidate KG1-side entity ids to draw from (OpenEA's
            own pool is ``sup_ent1 + ref_ent1`` -- every entity, not just
            currently-seeded ones).
        entities2: Candidate KG2-side entity ids.
        batch_size: Positive pairs in this batch.
        neg_triple_num: Negatives drawn per positive pair (``k``).
        rng: Source of randomness.
        exclude_pairs: Pairs to drop if drawn (the current seed set).

    Returns:
        ``(neg_left, neg_right)``, each up to ``(batch_size * k,)`` int64
        (shorter if some draws were excluded).
    """
    size = batch_size * neg_triple_num
    left = rng.choice(entities1, size=size)
    right = rng.choice(entities2, size=size)
    keep = [i for i in range(size) if (int(left[i]), int(right[i])) not in exclude_pairs]
    return left[keep].astype(np.int64), right[keep].astype(np.int64)

_alinet_torch

PyTorch model and training-step functions for AliNetLinker.

Ports OpenEA's GraphConvolution/AliNetGraphAttentionLayer/ HighwayLayer/AliNet._define_model (approaches/alinet.py): a stack of layers where every layer but the last combines a 1-hop GraphConvolution branch and a 2-hop AliNetGraphAttentionLayer branch, fused by a HighwayLayer gate; the last layer is GraphConvolution only. The final representation concatenates every layer's (plus the initial embedding's) L2-normalized output -- a jumping-knowledge-style multi-scale representation. Uses plain torch.sparse throughout (including torch.sparse.softmax, same as _rdgcn_torch.py) -- no torch_geometric/dgl dependency.

Callers must already have confirmed torch is importable (alinet.py's fit() does this via OptionalDependencyError before any of these run) -- same precedent as _gcn_align_torch.py/ _rdgcn_torch.py.

Dropout and L2 weight regularization are not implemented. OpenEA's own published config uses dropout: 0.0 (never actually triggers the reference's own dropout branches), and while every learnable weight in the reference is declared with an l2_regularizer(scale=0.01), nothing in AliNet's own training graph construction (_generate_graph/ _generate_rel_graph) ever retrieves and sums TensorFlow's regularization-losses collection into the actual optimized loss -- confirmed by reading AliNet's class body directly. Both are faithful no-ops to skip, not simplifications with real fidelity cost.

This is the only linker in this package whose adjacency mutates during training -- AliNetModel.set_adjacency is called once before training and again after every bootstrapping round (see AliNetLinker.fit), updating every layer's stored sparse adjacency in place rather than rebuilding the model.

build_alinet_model

build_alinet_model(
    num_entities: int, layer_dims: list[int]
) -> torch.nn.Module

Ports AliNet._define_model's layer stack as a live PyTorch model.

A factory function rather than a module-level torch.nn.Module subclass, since torch.nn.Module can't be named at module scope without torch installed -- same precedent as _gcn_align_torch.build_gcn_align_branch/ _rdgcn_torch.build_rdgcn_model.

Parameters:

Name	Type	Description	Default
num_entities	int	Total entity count (both KGs combined).	required
layer_dims	list[int]	Layer width sequence, e.g. [500, 400, 300] for 2 layers (layer_dims[0] is the initial embedding's dimensionality). OpenEA's published EN-FR-15K-V1 config uses [500, 400, 300].	required

Returns:

Type	Description
Module	A module exposing ``set_adjacency(one_hop_adjacency,
Module	two_hop_adjacency)`` (must be called before the first
Module	forward() and again after every bootstrapping round; both
Module	are coalesced sparse [num_entities, num_entities] tensors,
Module	e.g. from
Module	coo_to_torch_sparse)
Module	and forward() -> torch.Tensor (returns ``(num_entities,
Module	sum(layer_dims) + layer_dims[0])`` -- the concatenated,
Module	L2-normalized multi-layer representation, every layer's output
Module	plus the initial embedding, matching OpenEA's own
Module	[input_embeds] + output_embeds_list concat order).

Source code in src/linkingtk/algorithms/ea/_alinet_torch.py

def build_alinet_model(num_entities: int, layer_dims: list[int]) -> torch.nn.Module:
    """Ports ``AliNet._define_model``'s layer stack as a live PyTorch model.

    A factory function rather than a module-level ``torch.nn.Module``
    subclass, since ``torch.nn.Module`` can't be named at module scope
    without ``torch`` installed -- same precedent as
    ``_gcn_align_torch.build_gcn_align_branch``/
    ``_rdgcn_torch.build_rdgcn_model``.

    Args:
        num_entities: Total entity count (both KGs combined).
        layer_dims: Layer width sequence, e.g. ``[500, 400, 300]`` for 2
            layers (``layer_dims[0]`` is the initial embedding's
            dimensionality). OpenEA's published EN-FR-15K-V1 config uses
            ``[500, 400, 300]``.

    Returns:
        A module exposing ``set_adjacency(one_hop_adjacency,
        two_hop_adjacency)`` (must be called before the first
        ``forward()`` and again after every bootstrapping round; both
        are coalesced sparse ``[num_entities, num_entities]`` tensors,
        e.g. from
        [coo_to_torch_sparse][linkingtk.utils.sparse_gcn.coo_to_torch_sparse])
        and ``forward() -> torch.Tensor`` (returns ``(num_entities,
        sum(layer_dims) + layer_dims[0])`` -- the concatenated,
        L2-normalized multi-layer representation, every layer's output
        plus the initial embedding, matching OpenEA's own
        ``[input_embeds] + output_embeds_list`` concat order).
    """
    import torch
    import torch.nn.functional as functional
    from torch import nn

    class _GraphConvolution(nn.Module):
        """Ports ``GraphConvolution`` (single-adjacency case, the only one
        AliNet's own code ever uses)."""

        def __init__(self, in_dim: int, out_dim: int) -> None:
            super().__init__()
            weight = torch.empty(in_dim, out_dim)
            nn.init.xavier_uniform_(weight)
            self.weight = nn.Parameter(weight)
            self.bias = nn.Parameter(torch.zeros(out_dim))
            self.batch_norm = nn.BatchNorm1d(in_dim)
            self.adjacency: torch.Tensor | None = None

        def forward(self, inputs: torch.Tensor) -> torch.Tensor:
            assert self.adjacency is not None  # noqa: S101 -- set via set_adjacency before use
            hidden = self.batch_norm(inputs)
            pre_sup = hidden @ self.weight
            out = torch.sparse.mm(self.adjacency, pre_sup) + self.bias
            return torch.tanh(out)

    class _AliNetGraphAttentionLayer(nn.Module):
        """Ports ``AliNetGraphAttentionLayer``: self+neighbor attention over
        the (typically 2-hop) adjacency."""

        def __init__(self, in_dim: int, out_dim: int) -> None:
            super().__init__()
            kernel = torch.empty(in_dim, out_dim)
            nn.init.xavier_uniform_(kernel)
            self.kernel = nn.Parameter(kernel)
            kernel1 = torch.empty(in_dim, in_dim)
            nn.init.xavier_uniform_(kernel1)
            self.kernel1 = nn.Parameter(kernel1)
            kernel2 = torch.empty(in_dim, in_dim)
            nn.init.xavier_uniform_(kernel2)
            self.kernel2 = nn.Parameter(kernel2)
            self.batch_norm = nn.BatchNorm1d(in_dim)
            self.adjacency: torch.Tensor | None = None

        def forward(self, inputs: torch.Tensor) -> torch.Tensor:
            assert self.adjacency is not None  # noqa: S101 -- set via set_adjacency before use
            adjacency = self.adjacency.coalesce()
            hidden = self.batch_norm(inputs)
            mapped = hidden @ self.kernel
            attention1 = hidden @ self.kernel1
            attention2 = hidden @ self.kernel2
            self_score = (attention1 * hidden).sum(dim=1, keepdim=True)
            neighbor_score = (attention2 * hidden).sum(dim=1, keepdim=True)
            self_score = torch.tanh(self_score)
            neighbor_score = torch.tanh(neighbor_score)

            indices = adjacency.indices()
            row, col = indices[0], indices[1]
            edge_logits = self_score[row].squeeze(-1) + neighbor_score[col].squeeze(-1)
            edge_logits = functional.leaky_relu(edge_logits)
            size = adjacency.shape
            sparse_logits = torch.sparse_coo_tensor(indices, edge_logits, size=size).coalesce()
            attention_weights = torch.sparse.softmax(sparse_logits, dim=1)
            value = torch.sparse.mm(attention_weights, mapped)
            return torch.tanh(value)

    class _HighwayLayer(nn.Module):
        """Ports ``HighwayLayer`` -- including its one shared `BatchNormalization`
        instance applied to *both* inputs sequentially (confirmed from
        reading ``HighwayLayer.__init__``/``call`` directly: one
        `self.batch_normal`, called twice), and its ``tanh`` -> ``relu``
        gate (not the more common sigmoid gate)."""

        def __init__(self, dim: int) -> None:
            super().__init__()
            weight = torch.empty(dim, dim)
            nn.init.xavier_uniform_(weight)
            self.weight = nn.Parameter(weight)
            self.batch_norm = nn.BatchNorm1d(dim)

        def forward(self, layer1: torch.Tensor, layer2: torch.Tensor) -> torch.Tensor:
            layer1 = self.batch_norm(layer1)
            layer2 = self.batch_norm(layer2)
            gate = torch.relu(torch.tanh(layer1 @ self.weight))
            output = layer2 * (1.0 - gate) + layer1 * gate
            return torch.tanh(output)

    class _AliNetModel(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            init_table = torch.empty(num_entities, layer_dims[0])
            nn.init.xavier_uniform_(init_table)
            self.init_embedding = nn.Parameter(init_table)

            self.num_layers = len(layer_dims) - 1
            self.gcn_layers = nn.ModuleList(
                [
                    _GraphConvolution(layer_dims[i], layer_dims[i + 1])
                    for i in range(self.num_layers)
                ]
            )
            self.gat_layers = nn.ModuleList(
                [
                    _AliNetGraphAttentionLayer(layer_dims[i], layer_dims[i + 1])
                    for i in range(self.num_layers - 1)
                ]
            )
            self.highway_layers = nn.ModuleList(
                [_HighwayLayer(layer_dims[i + 1]) for i in range(self.num_layers - 1)]
            )

        def set_adjacency(
            self, one_hop_adjacency: torch.Tensor, two_hop_adjacency: torch.Tensor
        ) -> None:
            for gcn in self.gcn_layers:
                gcn.adjacency = one_hop_adjacency
            for gat in self.gat_layers:
                gat.adjacency = two_hop_adjacency

        def forward(self) -> torch.Tensor:
            output = self.init_embedding
            layer_outputs = []
            for i in range(self.num_layers):
                gcn_out = self.gcn_layers[i](output)
                if i < self.num_layers - 1:
                    gat_out = self.gat_layers[i](output)
                    output = self.highway_layers[i](gat_out, gcn_out)
                else:
                    output = gcn_out
                layer_outputs.append(output)

            all_layers = [self.init_embedding, *layer_outputs]
            normalized = [functional.normalize(layer, dim=1) for layer in all_layers]
            concatenated = torch.cat(normalized, dim=1)
            result: torch.Tensor = functional.normalize(concatenated, dim=1)
            return result

    return _AliNetModel()

Simple-HHEA

simple_hhea

Simple-HHEA linker for Entity Alignment.

Ports Simple-HHEA (Jiang et al. 2024, "Toward Practical Entity Alignment Method Design: Insights from New Highly Heterogeneous Knowledge Graph Datasets", https://dl.acm.org/doi/pdf/10.1145/3589334.3645720) from its reference implementation at https://github.com/DataArcTech/Simple-HHEA -- a base embedding method for highly heterogeneous KG pairs (e.g. an event KG aligned to Wikipedia), fusing ALBERT name embeddings (whitened), node2vec structural embeddings, and (optionally) Time2Vec temporal embeddings through a small MLP trained with margin ranking loss.

Filed as #62, a prerequisite for #22 (ChatEA): the ChatEA paper's published llama2-13b numbers are ChatEA layered on top of Simple-HHEA specifically, not a generic KGE method, so SimpleHHEALinker is what ChatEALinker will use as its base_linker for a genuine comparison.

See linkingtk.algorithms.ea._simple_hhea_name/_simple_hhea_structure/ _simple_hhea_torch/_simple_hhea_training for the individual feature- group ports and their documented fidelity decisions vs. the reference.

SimpleHHEALinker

SimpleHHEALinker(
    name_model: str = "albert-base-v2",
    emb_size: int = 64,
    structure_size: int = 8,
    time_size: int = 8,
    use_structure: bool = True,
    walk_length: int = 80,
    num_walks: int = 10,
    structure_dim: int = 64,
    max_degree: int = 1000,
    structure_workers: int = 1,
    structure_epochs: int = 5,
    num_epochs: int = 1500,
    learning_rate: float = 0.01,
    weight_decay: float = 0.001,
    gamma: float = 1.0,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
)

Bases: BaseLinker

Fuses whitened name embeddings, node2vec structural embeddings, and (optionally) Time2Vec temporal embeddings for Entity Alignment.

Must be fit before link. Exposes source_embedding/target_embedding like KGELinker, so it's a drop-in for rank_exhaustive -- and, once #22 (ChatEA) lands, a candidate base_linker for it too.

Parameters:

Name	Type	Description	Default
name_model	str	Hugging Face model id for name embedding, loaded via AutoModel/AutoTokenizer.from_pretrained. Defaults to the reference's own albert-base-v2.	'albert-base-v2'
emb_size	int	Final fused embedding dimension.	64
structure_size	int	Post-projection size of the structure branch before concatenation into the final embedding.	8
time_size	int	Post-projection size of the time branch before concatenation into the final embedding.	8
use_structure	bool	Whether to compute and use node2vec structural embeddings at all (the reference's own --no_structure).	True
walk_length	int	node2vec walk hyperparameter -- see build_structure_embeddings.	80
num_walks	int	node2vec walk hyperparameter -- see build_structure_embeddings.	10
structure_dim	int	node2vec walk hyperparameter -- see build_structure_embeddings.	64
max_degree	int	node2vec walk hyperparameter -- see build_structure_embeddings.	1000
structure_workers	int	Worker threads for the structural Word2Vec fit -- see build_structure_embeddings. Walk simulation itself is always single-threaded, matching the reference's own default invocation.	1
structure_epochs	int	Word2Vec training epochs for the structural skip-gram model -- see build_structure_embeddings.	5
num_epochs	int	Training hyperparameter -- see train.	1500
learning_rate	float	Training hyperparameter -- see train.	0.01
weight_decay	float	Training hyperparameter -- see train.	0.001
gamma	float	Training hyperparameter -- see train.	1.0
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
device	str	Torch device to train/embed on, e.g. "cpu" (default) or "cuda".	'cpu'

Source code in src/linkingtk/algorithms/ea/simple_hhea.py

def __init__(
    self,
    name_model: str = "albert-base-v2",
    emb_size: int = 64,
    structure_size: int = 8,
    time_size: int = 8,
    use_structure: bool = True,
    walk_length: int = 80,
    num_walks: int = 10,
    structure_dim: int = 64,
    max_degree: int = 1000,
    structure_workers: int = 1,
    structure_epochs: int = 5,
    num_epochs: int = 1500,
    learning_rate: float = 0.01,
    weight_decay: float = 0.001,
    gamma: float = 1.0,
    matching: Matcher = DEFAULT_MATCHER,
    device: str = "cpu",
) -> None:
    self.name_model = name_model
    self.emb_size = emb_size
    self.structure_size = structure_size
    self.time_size = time_size
    self.use_structure = use_structure
    self.walk_length = walk_length
    self.num_walks = num_walks
    self.structure_dim = structure_dim
    self.max_degree = max_degree
    self.structure_workers = structure_workers
    self.structure_epochs = structure_epochs
    self.num_epochs = num_epochs
    self.learning_rate = learning_rate
    self.weight_decay = weight_decay
    self.gamma = gamma
    self.matching = matching
    self.device = device
    self._id_to_vector: dict[str, npt.NDArray[np.floating[Any]]] = {}
    self._fitted = False

fit

fit(
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    temporal_triples: list[TemporalTriple] | None = None,
    random_state: int | None = None,
) -> SimpleHHEALinker

Trains the fusion model over dataset1/dataset2's entities.

Parameters:

Name	Type	Description	Default
dataset1	list[Entity]	Source entities.	required
dataset2	list[Entity]	Target entities.	required
ground_truth	list[tuple[str, str]]	Known-correct (source_id, target_id) pairs. Doubles as both the margin-ranking training signal (like KGELinker.fit's ground_truth) and the node2vec node-merge seed (see build_structure_embeddings's train_pairs).	required
graph	Graph	Combined relational structure of both KGs (e.g. to_triples(graph1) + to_triples(graph2)).	required
temporal_triples	list[TemporalTriple] | None	Optional (subject_id, relation_id, object_id, start_label, end_label) temporal facts (e.g. from IcewsWikiDataset.load_temporal_graphs), start_label/end_label as "YYYY-MM" strings or None. None (default) disables the time-feature branch entirely -- correct for non-ICEWS datasets, which have no temporal facts at all.	None
random_state	int | None	Seed for reproducible training.	None

Returns:

Type	Description
SimpleHHEALinker	self, for chaining.

Raises:

Type	Description
LinkingTKError	If no ground_truth pair's ids are actually present in dataset1/dataset2.

Source code in src/linkingtk/algorithms/ea/simple_hhea.py

def fit(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity],
    ground_truth: list[tuple[str, str]],
    graph: Graph,
    temporal_triples: list[TemporalTriple] | None = None,
    random_state: int | None = None,
) -> SimpleHHEALinker:
    """Trains the fusion model over `dataset1`/`dataset2`'s entities.

    Args:
        dataset1: Source entities.
        dataset2: Target entities.
        ground_truth: Known-correct `(source_id, target_id)` pairs.
            Doubles as both the margin-ranking training signal (like
            [KGELinker.fit][linkingtk.algorithms.ea.kge.KGELinker.fit]'s
            `ground_truth`) and the node2vec node-merge seed (see
            [build_structure_embeddings][linkingtk.algorithms.ea._simple_hhea_structure.build_structure_embeddings]'s
            `train_pairs`).
        graph: Combined relational structure of both KGs (e.g.
            `to_triples(graph1) + to_triples(graph2)`).
        temporal_triples: Optional `(subject_id, relation_id,
            object_id, start_label, end_label)` temporal facts (e.g.
            from
            [IcewsWikiDataset.load_temporal_graphs][linkingtk.datasets.icews.IcewsWikiDataset.load_temporal_graphs]),
            `start_label`/`end_label` as `"YYYY-MM"` strings or
            `None`. `None` (default) disables the time-feature branch
            entirely -- correct for non-ICEWS datasets, which have no
            temporal facts at all.
        random_state: Seed for reproducible training.

    Returns:
        `self`, for chaining.

    Raises:
        LinkingTKError: If no `ground_truth` pair's ids are actually
            present in `dataset1`/`dataset2`.
    """
    ids1 = {entity.id for entity in dataset1}
    ids2 = {entity.id for entity in dataset2}
    train_pairs = [(s, t) for s, t in ground_truth if s in ids1 and t in ids2]
    if not train_pairs:
        raise LinkingTKError(
            "None of `ground_truth`'s pairs have ids present in `dataset1`/"
            "`dataset2`; fit() has no training pairs."
        )

    if random_state is not None:
        torch.manual_seed(random_state)

    entities = dataset1 + dataset2
    entity_ids = [entity.id for entity in entities]
    entity_to_index = {entity_id: index for index, entity_id in enumerate(entity_ids)}

    raw_name_emb = embed_names(entities, model_name=self.name_model, device=self.device)
    kernel, bias = compute_kernel_bias(raw_name_emb, self.emb_size)
    name_emb = whiten(raw_name_emb, kernel, bias)

    structure_emb = None
    if self.use_structure:
        triples = to_triples(graph)
        structure_by_id = build_structure_embeddings(
            entity_ids,
            triples,
            train_pairs,
            dimensions=self.structure_dim,
            walk_length=self.walk_length,
            num_walks=self.num_walks,
            max_degree=self.max_degree,
            workers=self.structure_workers,
            epochs=self.structure_epochs,
            random_state=random_state,
        )
        structure_emb = np.stack([structure_by_id[entity_id] for entity_id in entity_ids])

    time_emb = (
        build_time_histogram(entity_ids, temporal_triples)
        if temporal_triples is not None
        else None
    )

    resolved_device = resolve_device(self.device)
    model = SimpleHHEAModel(
        name_emb,
        time_emb,
        structure_emb,
        emb_size=self.emb_size,
        structure_size=self.structure_size,
        time_size=self.time_size,
        device=resolved_device,
    ).to(resolved_device)

    train_pairs_idx = np.array(
        [(entity_to_index[s], entity_to_index[t]) for s, t in train_pairs], dtype=np.int64
    )
    train_simple_hhea(
        model,
        train_pairs_idx,
        len(entity_ids),
        num_epochs=self.num_epochs,
        learning_rate=self.learning_rate,
        weight_decay=self.weight_decay,
        gamma=self.gamma,
        random_state=random_state,
    )

    model.eval()
    with torch.no_grad():
        embeddings = model().cpu().numpy()
    self._id_to_vector = {
        entity_id: embeddings[index] for index, entity_id in enumerate(entity_ids)
    }
    self._fitted = True
    return self

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/simple_hhea.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score `entity_id` as a scored pair's source side."""
    return self._embedding(entity_id)

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/simple_hhea.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score `entity_id` as a scored pair's target side."""
    return self._embedding(entity_id)

_simple_hhea_name

ALBERT name-embedding + whitening for SimpleHHEALinker.

Ports process_name_embedding.py from https://github.com/DataArcTech/Simple-HHEA: mean-pooled ALBERT embeddings of each entity's label text, followed by whitening (Su et al. 2021, "Whitening Sentence Representations for Better Semantics and Faster Retrieval") -- a linear transform derived from the SVD of the combined covariance of both KGs' raw embeddings that reduces dimensionality and improves cosine-similarity behavior of BERT-family sentence embeddings. transformers/torch are core dependencies already, so this needs no OptionalDependencyError guard (unlike the structural-embedding stage's new node2vec dependency).

compute_kernel_bias

compute_kernel_bias(
    vectors: NDArray[floating[Any]], n_components: int
) -> tuple[
    npt.NDArray[np.floating[Any]],
    npt.NDArray[np.floating[Any]],
]

Whitening kernel/bias from vectors' covariance.

Ports process_name_embedding.py's compute_kernel_bias -- the final transform is (x + bias) @ kernel.

Source code in src/linkingtk/algorithms/ea/_simple_hhea_name.py

def compute_kernel_bias(
    vectors: npt.NDArray[np.floating[Any]], n_components: int
) -> tuple[npt.NDArray[np.floating[Any]], npt.NDArray[np.floating[Any]]]:
    """Whitening kernel/bias from ``vectors``' covariance.

    Ports `process_name_embedding.py`'s ``compute_kernel_bias`` -- the
    final transform is ``(x + bias) @ kernel``.
    """
    mean = vectors.mean(axis=0, keepdims=True)
    covariance = np.cov(vectors.T)
    u, singular_values, _ = np.linalg.svd(covariance)
    kernel = u @ np.diag(1 / np.sqrt(singular_values))
    return kernel[:, :n_components], -mean

embed_names

embed_names(
    entities: list[Entity],
    model_name: str = "albert-base-v2",
    device: str = "cpu",
    batch_size: int = 64,
) -> npt.NDArray[np.floating[Any]]

Mean-pooled ALBERT embedding for each entity's cleaned label text.

Returns a raw (len(entities), hidden_size) array -- not whitened. compute_kernel_bias/ whiten must be fit jointly over both KGs' raw embeddings (matching the reference's own combined-covariance whitening), so aren't folded in here.

Source code in src/linkingtk/algorithms/ea/_simple_hhea_name.py

def embed_names(
    entities: list[Entity],
    model_name: str = "albert-base-v2",
    device: str = "cpu",
    batch_size: int = 64,
) -> npt.NDArray[np.floating[Any]]:
    """Mean-pooled ALBERT embedding for each entity's cleaned label text.

    Returns a raw ``(len(entities), hidden_size)`` array -- not whitened.
    [compute_kernel_bias][linkingtk.algorithms.ea._simple_hhea_name.compute_kernel_bias]/
    [whiten][linkingtk.algorithms.ea._simple_hhea_name.whiten] must be fit
    jointly over *both* KGs' raw embeddings (matching the reference's own
    combined-covariance whitening), so aren't folded in here.
    """
    import torch
    from transformers import AutoModel, AutoTokenizer

    resolved = resolve_device(device)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name).to(resolved)
    model.eval()

    texts = [entity_name_text(entity) for entity in entities]
    chunks: list[npt.NDArray[np.floating[Any]]] = []
    with torch.no_grad():
        for start in range(0, len(texts), batch_size):
            encoded = tokenizer(
                texts[start : start + batch_size],
                padding=True,
                truncation=True,
                return_tensors="pt",
            ).to(resolved)
            hidden_states = model(**encoded).last_hidden_state
            mask = encoded["attention_mask"].unsqueeze(-1).float()
            pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
            chunks.append(pooled.cpu().numpy())
    if not chunks:
        hidden_size: int = model.config.hidden_size
        return np.empty((0, hidden_size), dtype=np.float32)
    return np.concatenate(chunks, axis=0)

entity_name_text

entity_name_text(entity: Entity) -> str

Cleaned label text for name-embedding, per the reference's own name.split("/")[-1].replace("_", " ") (strips a URI down to its trailing local name, then de-underscores it).

Source code in src/linkingtk/algorithms/ea/_simple_hhea_name.py

def entity_name_text(entity: Entity) -> str:
    """Cleaned label text for name-embedding, per the reference's own
    ``name.split("/")[-1].replace("_", " ")`` (strips a URI down to its
    trailing local name, then de-underscores it)."""
    label = label_texts(entity)[0] if entity.labels else entity.id
    return label.split("/")[-1].replace("_", " ").replace("\xa0", "")

whiten

whiten(
    vectors: NDArray[floating[Any]],
    kernel: NDArray[floating[Any]],
    bias: NDArray[floating[Any]],
) -> npt.NDArray[np.floating[Any]]

Applies a whitening transform and L2-normalizes each row.

Ports process_name_embedding.py's transform_and_normalize.

Source code in src/linkingtk/algorithms/ea/_simple_hhea_name.py

def whiten(
    vectors: npt.NDArray[np.floating[Any]],
    kernel: npt.NDArray[np.floating[Any]],
    bias: npt.NDArray[np.floating[Any]],
) -> npt.NDArray[np.floating[Any]]:
    """Applies a whitening transform and L2-normalizes each row.

    Ports `process_name_embedding.py`'s ``transform_and_normalize``.
    """
    transformed = (vectors + bias) @ kernel
    norms = np.linalg.norm(transformed, axis=1, keepdims=True)
    result: npt.NDArray[np.floating[Any]] = transformed / np.clip(norms, 1e-12, None)
    return result

_simple_hhea_structure

node2vec structural embeddings for SimpleHHEALinker.

Ports the biased-random-walk structural-embedding stage of https://github.com/DataArcTech/Simple-HHEA (feature_perprocessing/preproccess.py + longterm/{main.py,node2vec.py} + get_deep_emb.py) directly from its own hand-rolled alias-sampling walk code in longterm/node2vec.py -- not the third-party node2vec PyPI package this module first used.

That first attempt was a real, confirmed bug, not just a fidelity gap: longterm/node2vec.py's get_alias_edge uses a non-standard formula (weight * p for the "return to the previous node" term -- the file even has the textbook weight / p commented out directly above it, a deliberate inversion of the Grover & Leskovec convention). The node2vec PyPI package implements the standard (divide-by-p) formula, so copying the reference's literal p=1e-100 into it did the opposite of the intended "never immediately backtrack" -- confirmed directly: sampled walks degenerated into pure two-node oscillation (['d','c','d','c', ...]). Retuning p for the standard package's own convention (p=100) produced genuinely exploratory walks, but scored far worse on the real ICEWS-WIKI benchmark than the degenerate ones (Hits@1 0.037-0.041 vs. 0.262, across several configurations, including 10x more Word2Vec epochs and a 10x harder degree cap -- ruling out undertraining and hub-node dilution as the cause). The reference's actual formula isn't equivalent to either end of the standard package's p/q dial: at q=1, its elif branch (a neighbor that's also adjacent to the previous node) gets zeroed out too, alongside the return term -- a third walk character the standard package has no way to express. Hand-porting the exact formula below, rather than continuing to search for an equivalent parameterization, is the only way to reproduce it.

Ported near-verbatim: _alias_setup/_alias_draw (the classic alias-method sampling utilities), _precompute_transition_probs/ _get_alias_edge/_node2vec_walk/_simulate_walks (the reference's Graph.preprocess_transition_probs/get_alias_edge/node2vec_walk/ simulate_walks). Differences from the reference's own code, all mechanical, none behavioral:

• _alias_draw takes an explicit np.random.Generator instead of reading numpy's global random state (this module already seeds everything else that way).
• np.int (removed in modern numpy) becomes np.int64.
• Each node's sorted neighbor list is computed once and cached (_sorted_neighbors), not recomputed by a fresh sorted(G.neighbors(...)) call on every single walk step and every get_alias_edge call like the reference does -- pure performance (the ordering, and therefore every alias table's meaning, is identical either way; a real-graph hub node still at the max_degree cap would otherwise be re-sorted on every one of its many visits).

Two deliberate simplifications vs. the reference (flagged in #62's plan, not silent):

• Walks are entity-only, with no relation-id interleaving. The reference splices relation ids into each walk (entity-relation-entity-...) before handing it to Word2Vec, but the paper text itself never mentions this and it isn't ablated separately -- treated here as a reference-code-only embellishment on top of "a biased random walk method", not the core contribution.
• The "merge known-aligned entities into one graph node before walking" trick (the reference's node2same) is driven directly by the train_pairs SimpleHHEALinker.fit already receives as ground_truth, instead of reverse-engineering the reference's ref_ent_ids[:1500] file-position slice. Each pair's target id is remapped onto its source id directly (one fewer moving part than the reference's brand-new shared id per pair) -- functionally the same "one shared node" outcome.

One addition the reference doesn't have at all, found necessary against the real ICEWS-WIKI dataset (not a toy/synthetic graph): its combined graph has an extremely power-law-skewed degree distribution (median node degree 18, but the single highest-degree node -- "United States", ICEWS event data's most common actor -- has 273,317 edges; confirmed by direct measurement, not assumed). The alias-table transition-probability precompute is O(degree) per edge, so a hub node alone costs O(degree^2) -- tens of billions of operations for that one node, which never finished in practice with the original (uncapped) graph (confirmed: killed after several CPU-bound minutes stuck in exactly that precompute step, well before any walk was even simulated). cap_node_degree randomly subsamples any node's incident edges down to max_degree before walking -- a standard, well-precedented mitigation for random-walk embeddings on scale-free graphs (not something the reference needed to consider, or that a toy-scale unit-test fixture would ever surface).

build_structure_embeddings

build_structure_embeddings(
    entity_ids: list[str],
    triples: list[Triple],
    train_pairs: list[tuple[str, str]],
    *,
    dimensions: int = 64,
    walk_length: int = 80,
    num_walks: int = 10,
    p: float = 1e-100,
    q: float = 1.0,
    window: int = 10,
    epochs: int = 5,
    workers: int = 1,
    max_degree: int = 1000,
    random_state: int | None = None,
) -> dict[str, npt.NDArray[np.floating[Any]]]

One structural embedding per entity_ids, via a biased random walk + skip-gram.

p/q/walk_length/num_walks/window/epochs default to the reference's own values -- see this module's docstring for why p=1e-100 (a tiny, near-zero value) is correct here, ported from the reference's own non-standard alias-edge formula (weight * p, not weight / p), unlike the standard node2vec package's convention this module used at first. Entities with no edges at all in triples fall back to a random vector -- ports get_deep_emb.py's own except: np.random.random_sample(...) fallback for the same case.

Parameters:

Name	Type	Description	Default
workers	int	Worker threads for gensim's Word2Vec fit (real in-process threading, not separate OS processes -- safe to raise). Walk simulation itself is always single-threaded (the reference's own simulate_walks is too, absent an explicit --workers override it doesn't use for longterm/main.py's default single-process invocation).	1
max_degree	int	Caps every node's neighbor count before walking (see cap_node_degree) -- needed for real event-KG-scale graphs (e.g. ICEWS-WIKI), whose degree distribution is extreme enough that the transition-probability precompute is otherwise intractable. Irrelevant at toy/unit-test scale.	1000

Raises:

Type	Description
OptionalDependencyError	If networkx/gensim aren't installed.

Source code in src/linkingtk/algorithms/ea/_simple_hhea_structure.py

def build_structure_embeddings(
    entity_ids: list[str],
    triples: list[Triple],
    train_pairs: list[tuple[str, str]],
    *,
    dimensions: int = 64,
    walk_length: int = 80,
    num_walks: int = 10,
    p: float = 1e-100,
    q: float = 1.0,
    window: int = 10,
    epochs: int = 5,
    workers: int = 1,
    max_degree: int = 1000,
    random_state: int | None = None,
) -> dict[str, npt.NDArray[np.floating[Any]]]:
    """One structural embedding per ``entity_ids``, via a biased random walk + skip-gram.

    ``p``/``q``/``walk_length``/``num_walks``/``window``/``epochs`` default
    to the reference's own values -- see this module's docstring for why
    `p=1e-100` (a tiny, near-zero value) is correct *here*, ported from
    the reference's own non-standard alias-edge formula (`weight * p`,
    not `weight / p`), unlike the standard `node2vec` package's
    convention this module used at first. Entities with no edges at all
    in `triples` fall back to a random vector -- ports `get_deep_emb.py`'s
    own ``except: np.random.random_sample(...)`` fallback for the same
    case.

    Args:
        workers: Worker threads for gensim's `Word2Vec` fit (real
            in-process threading, not separate OS processes -- safe to
            raise). Walk simulation itself is always single-threaded (the
            reference's own `simulate_walks` is too, absent an explicit
            `--workers` override it doesn't use for `longterm/main.py`'s
            default single-process invocation).
        max_degree: Caps every node's neighbor count before walking (see
            [cap_node_degree][linkingtk.algorithms.ea._simple_hhea_structure.cap_node_degree])
            -- needed for real event-KG-scale graphs (e.g. ICEWS-WIKI),
            whose degree distribution is extreme enough that the
            transition-probability precompute is otherwise intractable.
            Irrelevant at toy/unit-test scale.

    Raises:
        OptionalDependencyError: If `networkx`/`gensim` aren't installed.
    """
    try:
        import networkx as nx
        from gensim.models import Word2Vec
    except ImportError as exc:
        raise OptionalDependencyError("SimpleHHEALinker", "kge") from exc

    merge_map = _merge_map(train_pairs)

    def representative(entity_id: str) -> str:
        return merge_map.get(entity_id, entity_id)

    edges: set[tuple[str, str]] = set()
    for subject, _, obj in triples:
        u, v = representative(subject), representative(obj)
        if u != v:
            edges.add((u, v) if u < v else (v, u))
    # Sorted, not `list(edges)` directly -- a plain set's iteration order
    # depends on PYTHONHASHSEED (randomized per process by default), which
    # would silently make `cap_node_degree`'s rng.permutation select
    # different *edges* run to run even for the same `random_state`
    # (confirmed: this was flaking a fully-seeded unit test before being
    # sorted here).
    rng = np.random.default_rng(random_state)
    capped_edges = cap_node_degree(sorted(edges), max_degree, rng)

    graph = nx.Graph()
    graph.add_edges_from(capped_edges, weight=1.0)

    embeddings: dict[str, npt.NDArray[np.floating[Any]]] = {}
    if graph.number_of_nodes() == 0:
        for entity_id in entity_ids:
            embeddings[entity_id] = rng.random(dimensions).astype(np.float32)
        return embeddings

    sorted_neighbors = {node: sorted(graph.neighbors(node)) for node in graph.nodes()}
    alias_nodes, alias_edges = _precompute_transition_probs(graph, sorted_neighbors, p, q)
    walks = _simulate_walks(
        graph, sorted_neighbors, alias_nodes, alias_edges, num_walks, walk_length, rng
    )

    model = Word2Vec(
        walks,
        vector_size=dimensions,
        window=window,
        min_count=0,
        sg=1,
        epochs=epochs,
        workers=workers,
        seed=random_state if random_state is not None else 1,
    )

    for entity_id in entity_ids:
        node = representative(entity_id)
        if node in model.wv:
            embeddings[entity_id] = np.asarray(model.wv[node], dtype=np.float32)
        else:
            embeddings[entity_id] = rng.random(dimensions).astype(np.float32)
    return embeddings

cap_node_degree

cap_node_degree(
    edges: list[tuple[str, str]],
    max_degree: int,
    rng: Generator,
) -> list[tuple[str, str]]

Randomly subsamples edges so no node ends up with more than max_degree.

Shuffles edges (so which ones survive isn't biased toward triple order/relation type), then greedily keeps each edge only while both endpoints are still under max_degree -- a node already at the cap loses only its "excess" edges, not a uniform fraction of all of them, so a rare hub node is trimmed while everything else is left alone.

Source code in src/linkingtk/algorithms/ea/_simple_hhea_structure.py

def cap_node_degree(
    edges: list[tuple[str, str]], max_degree: int, rng: np.random.Generator
) -> list[tuple[str, str]]:
    """Randomly subsamples `edges` so no node ends up with more than `max_degree`.

    Shuffles `edges` (so which ones survive isn't biased toward triple
    order/relation type), then greedily keeps each edge only while both
    endpoints are still under `max_degree` -- a node already at the cap
    loses only its "excess" edges, not a uniform fraction of all of them,
    so a rare hub node is trimmed while everything else is left alone.
    """
    order = rng.permutation(len(edges))
    counts: dict[str, int] = {}
    kept: list[tuple[str, str]] = []
    for index in order:
        u, v = edges[index]
        if counts.get(u, 0) >= max_degree or counts.get(v, 0) >= max_degree:
            continue
        kept.append((u, v))
        counts[u] = counts.get(u, 0) + 1
        counts[v] = counts.get(v, 0) + 1
    return kept

_simple_hhea_torch

Time2Vec + fusion model for SimpleHHEALinker.

Ports model.py from https://github.com/DataArcTech/Simple-HHEA: a small MLP that concatenates (name, time, structure) feature groups -- each already computed elsewhere (_simple_hhea_name.py, _simple_hhea_structure.py, build_time_histogram below) -- into one final per-entity embedding. Name/time/structure embeddings are held as fixed buffers (not trained); only the small linear/Time2Vec layers on top are trainable, matching the reference's plain (non-nn.Parameter) tensor attributes.

TIME_SPAN module-attribute

TIME_SPAN = 1 + 27 * 13

Number of month-bins: 1 catch-all pre-1995 bin, plus 27 years * 13 (the reference's own off-by-one month indexing, see _time_bin) monthly bins -- matches the reference's own 1+27*13, wide enough for ICEWS data (1995-2021ish).

CosineActivation

CosineActivation(in_features: int, out_features: int)

Bases: Module

Time2Vec's periodic + linear feature bank (Kazemi et al. 2019).

Source code in src/linkingtk/algorithms/ea/_simple_hhea_torch.py

def __init__(self, in_features: int, out_features: int) -> None:
    super().__init__()
    self.w0 = nn.Parameter(torch.randn(in_features, 1))
    self.b0 = nn.Parameter(torch.randn(1))
    self.w = nn.Parameter(torch.randn(in_features, out_features - 1))
    self.b = nn.Parameter(torch.randn(out_features - 1))

SimpleHHEAModel

SimpleHHEAModel(
    name_emb: NDArray[floating[Any]],
    time_emb: NDArray[floating[Any]] | None,
    structure_emb: NDArray[floating[Any]] | None,
    emb_size: int = 64,
    structure_size: int = 8,
    time_size: int = 8,
    device: device | str = "cpu",
)

Bases: Module

Fuses name/time/structure feature groups into one final embedding per entity.

Parameters:

Name	Type	Description	Default
name_emb	NDArray[floating[Any]]	(num_entities, name_dim) raw (whitened) name embeddings, fixed.	required
time_emb	NDArray[floating[Any]] | None	(num_entities, TIME_SPAN) month-activity histogram (see build_time_histogram), fixed. None disables the time branch.	required
structure_emb	NDArray[floating[Any]] | None	(num_entities, structure_dim) node2vec embeddings, fixed. None disables the structure branch.	required
emb_size	int	Final output embedding dimension.	64
structure_size	int	Post-projection size of the structure branch before concatenation.	8
time_size	int	Post-projection size of the time branch before concatenation.	8
device	device | str	Where to place the fixed feature tensors -- plain attributes (matching the reference), not registered buffers, so they don't move with a later .to(device) call on the model; set correctly up front instead.	'cpu'

Source code in src/linkingtk/algorithms/ea/_simple_hhea_torch.py

def __init__(
    self,
    name_emb: npt.NDArray[np.floating[Any]],
    time_emb: npt.NDArray[np.floating[Any]] | None,
    structure_emb: npt.NDArray[np.floating[Any]] | None,
    emb_size: int = 64,
    structure_size: int = 8,
    time_size: int = 8,
    device: torch.device | str = "cpu",
) -> None:
    super().__init__()
    self.use_time = time_emb is not None
    self.use_structure = structure_emb is not None

    linear_size = emb_size
    if self.use_time:
        linear_size += time_size
        assert time_emb is not None
        self.time_emb: torch.Tensor = torch.tensor(time_emb).float().to(device)
        self.fc_time_0 = nn.Linear(32, 32)
        self.fc_time = nn.Linear(32, time_size)
        self.time2vec = Time2Vec(hidden_dim=32)
        self.time_span_index: torch.Tensor = (
            torch.arange(TIME_SPAN).unsqueeze(1).float().to(device)
        )

    if self.use_structure:
        linear_size += structure_size
        assert structure_emb is not None
        self.structure_emb: torch.Tensor = torch.tensor(structure_emb).float().to(device)
        self.fc_structure_0 = nn.Linear(structure_emb.shape[-1], emb_size)
        self.fc_structure = nn.Linear(emb_size, structure_size)

    self.name_emb: torch.Tensor = torch.tensor(name_emb).float().to(device)
    self.fc_name_0 = nn.Linear(name_emb.shape[-1], emb_size)
    self.fc_name = nn.Linear(emb_size, emb_size)

    self.fc_final = nn.Linear(linear_size, emb_size)
    self.dropout = nn.Dropout(p=0.3)

build_time_histogram

build_time_histogram(
    entity_ids: list[str],
    temporal_triples: list[
        tuple[str, str, str, str | None, str | None]
    ],
) -> npt.NDArray[np.floating[Any]]

Per-entity month-activity histogram, (len(entity_ids), TIME_SPAN).

For each (subject, relation, object, start_label, end_label) temporal triple, increments the subject's histogram at every month bin in [start, end] (inclusive) if both resolve; a single bin if only one resolves; skipped if neither does. A resolved year before 1995 (ICEWS's real lower bound) falls into bin 0, the catch-all bin. One symmetric rule for every triple -- unlike the reference's own load_ent_time_matrix, whose two KG-side branches diverge (one increments two independent point-bins, the other fills a range) and whose range branch reads a stale time_y left over from an unrelated earlier loop (a genuine bug, not a deliberate per-side design choice -- confirmed by reading utils.py directly). Not ported here.

Source code in src/linkingtk/algorithms/ea/_simple_hhea_torch.py

def build_time_histogram(
    entity_ids: list[str], temporal_triples: list[tuple[str, str, str, str | None, str | None]]
) -> npt.NDArray[np.floating[Any]]:
    """Per-entity month-activity histogram, ``(len(entity_ids), TIME_SPAN)``.

    For each ``(subject, relation, object, start_label, end_label)``
    temporal triple, increments the *subject*'s histogram at every month
    bin in ``[start, end]`` (inclusive) if both resolve; a single bin if
    only one resolves; skipped if neither does. A resolved year before
    1995 (ICEWS's real lower bound) falls into bin ``0``, the catch-all
    bin. One symmetric rule for every triple -- unlike the reference's own
    `load_ent_time_matrix`, whose two KG-side branches diverge (one
    increments two independent point-bins, the other fills a range) and
    whose range branch reads a stale `time_y` left over from an unrelated
    earlier loop (a genuine bug, not a deliberate per-side design choice
    -- confirmed by reading `utils.py` directly). Not ported here.
    """
    index = {entity_id: row for row, entity_id in enumerate(entity_ids)}
    histogram = np.zeros((len(entity_ids), TIME_SPAN), dtype=np.float64)
    for subject, _relation, _obj, start_label, end_label in temporal_triples:
        row = index.get(subject)
        if row is None:
            continue
        start = _parse_time_label(start_label)
        end = _parse_time_label(end_label)
        if start is None and end is None:
            continue
        start = start or end
        end = end or start
        assert start is not None and end is not None
        start_bin = 0 if start[0] < 1995 else _time_bin(*start)
        end_bin = 0 if end[0] < 1995 else _time_bin(*end)
        if start_bin > end_bin:
            start_bin, end_bin = end_bin, start_bin
        histogram[row, start_bin : end_bin + 1] += 1
    return histogram

_simple_hhea_training

Margin-ranking training loop for SimpleHHEALinker.

Ports main_SimpleHHEA.py's train()/get_train_set()/l1() from https://github.com/DataArcTech/Simple-HHEA. Negatives are uniform-random entity ids, not hard-negative mining -- a simplification already present in the reference itself, not something to "improve" here.

build_training_quads

build_training_quads(
    train_pairs: NDArray[int64],
    batch_size: int,
    node_size: int,
    rng: Generator,
) -> npt.NDArray[np.int64]

Builds a fixed (batch_size, 4) array of (l, r, random_l, random_r) rows.

Ports get_train_set exactly, including its one real quirk: l/r are train_pairs repeated (not resampled) enough times to reach batch_size rows (ceil(batch_size / len(train_pairs)) copies, shuffled, then truncated to exactly batch_size) -- built once, not per epoch. The reference's own call site passes the node count as batch_size (not e.g. len(train_pairs)), so train_pairs gets repeated roughly node_size / len(train_pairs) times; the two random -negative columns are drawn once here too, then held fixed for the entire training run (train()'s epoch loop never redraws them).

Source code in src/linkingtk/algorithms/ea/_simple_hhea_training.py

def build_training_quads(
    train_pairs: npt.NDArray[np.int64], batch_size: int, node_size: int, rng: np.random.Generator
) -> npt.NDArray[np.int64]:
    """Builds a fixed ``(batch_size, 4)`` array of ``(l, r, random_l, random_r)`` rows.

    Ports `get_train_set` exactly, including its one real quirk: `l`/`r`
    are `train_pairs` *repeated* (not resampled) enough times to reach
    `batch_size` rows (`ceil(batch_size / len(train_pairs))` copies,
    shuffled, then truncated to exactly `batch_size`) -- built **once**,
    not per epoch. The reference's own call site passes the *node count*
    as `batch_size` (not e.g. `len(train_pairs)`), so `train_pairs` gets
    repeated roughly `node_size / len(train_pairs)` times; the two random
    -negative columns are drawn once here too, then held fixed for the
    entire training run (`train()`'s epoch loop never redraws them).
    """
    negative_ratio = batch_size // len(train_pairs) + 1
    repeated = np.tile(train_pairs, (negative_ratio, 1))
    rng.shuffle(repeated)
    repeated = repeated[:batch_size]
    negatives = rng.integers(0, node_size, size=repeated.shape)
    return np.concatenate([repeated, negatives], axis=-1)

train

train(
    model: SimpleHHEAModel,
    train_pairs: NDArray[int64],
    node_size: int,
    *,
    num_epochs: int = 1500,
    learning_rate: float = 0.01,
    weight_decay: float = 0.001,
    gamma: float = 1.0,
    random_state: int | None = None,
) -> None

Trains model in place via margin ranking loss.

One fixed set of training quads is built up front (see build_training_quads, batch_size=node_size matching the reference's own call). Every epoch: one full forward pass over the entire entity set (model() -- the reference's Simple_HHEA.forward() has no batching, every entity's fused embedding is recomputed every step), then sum(relu(gamma + d(l,r) - d(l,r')) + relu(gamma + d(l,r) - d(l',r))) / len(quads) where d is L1 distance -- ports main_SimpleHHEA.py's train() exactly (minus its retain_graph=True, unneeded here since each epoch builds a fresh graph from a fresh model() call).

Source code in src/linkingtk/algorithms/ea/_simple_hhea_training.py

def train(
    model: SimpleHHEAModel,
    train_pairs: npt.NDArray[np.int64],
    node_size: int,
    *,
    num_epochs: int = 1500,
    learning_rate: float = 0.01,
    weight_decay: float = 0.001,
    gamma: float = 1.0,
    random_state: int | None = None,
) -> None:
    """Trains ``model`` in place via margin ranking loss.

    One fixed set of training quads is built up front (see
    [build_training_quads][linkingtk.algorithms.ea._simple_hhea_training.build_training_quads],
    `batch_size=node_size` matching the reference's own call). Every
    epoch: one full forward pass over the entire entity set (`model()` --
    the reference's `Simple_HHEA.forward()` has no batching, every
    entity's fused embedding is recomputed every step), then
    ``sum(relu(gamma + d(l,r) - d(l,r')) + relu(gamma + d(l,r) - d(l',r)))
    / len(quads)`` where ``d`` is L1 distance -- ports `main_SimpleHHEA.py`'s
    `train()` exactly (minus its `retain_graph=True`, unneeded here since
    each epoch builds a fresh graph from a fresh `model()` call).
    """
    rng = np.random.default_rng(random_state)
    quads = build_training_quads(train_pairs, node_size, node_size, rng)
    device = model.name_emb.device
    indices = torch.from_numpy(quads).long().to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)

    for _epoch in range(num_epochs):
        model.train()
        optimizer.zero_grad()

        features = model()[indices]
        anchor_l, anchor_r = features[:, 0, :], features[:, 1, :]
        negative_l, negative_r = features[:, 2, :], features[:, 3, :]

        positive_distance = _l1_distance(anchor_l, anchor_r)
        loss = torch.sum(
            nn.functional.relu(gamma + positive_distance - _l1_distance(anchor_l, negative_r))
            + nn.functional.relu(gamma + positive_distance - _l1_distance(negative_l, anchor_r))
        ) / len(quads)

        loss.backward()  # type: ignore[no-untyped-call]
        optimizer.step()

ChatEA

chatea

ChatEA linker for Entity Alignment.

Ports ChatEA (Jiang et al. 2024, "Unlocking the Power of Large Language Models for Entity Alignment", https://aclanthology.org/2024.acl-long.408.pdf) from its reference implementation at https://github.com/DataArcTech/ChatEA -- an LLM re-ranker layered on top of an embedding-based EA method's own top-K candidates, rather than a standalone EA method. Filed as #22, a follow-up to #62 (SimpleHHEALinker), which is the reference's own base embedding method and this linker's intended base_linker.

See _chatea_context, _chatea_prompts, and _chatea_reasoning for the individual ported pieces and their documented fidelity decisions vs. the reference.

ChatEALinker

ChatEALinker(
    base_linker: EmbeddingLinker,
    client: LlmClient,
    top_k: int = 20,
    neigh_num: int = 5,
    desc_neigh_num: int = 25,
    desc_max_tokens: int = 80,
    threshold: float = 0.5,
    history_len: int = 3,
    relation_names: dict[str, str] | None = None,
    temporal_triples: list[TemporalTriple] | None = None,
    use_name: bool = True,
    use_desc: bool = True,
    use_struct: bool = True,
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

Re-ranks an already-fitted embedding linker's top-K candidates with an LLM.

Unlike most linkers in linkingtk.algorithms.ea, has no fit() -- base_linker must already be fitted before construction (its embeddings are what generates candidates), and the LLM itself needs no training. This mirrors LlmBaseLinker, which is also fit()-free for the same reason.

Parameters:

Name	Type	Description	Default
base_linker	EmbeddingLinker	An already-fitted linker exposing source_embedding/ target_embedding (see EmbeddingLinker) -- e.g. a fitted SimpleHHEALinker.	required
client	LlmClient	An LlmClient (e.g. from create_client). This class never imports openai/anthropic itself.	required
top_k	int	Candidates considered per source entity, ranked by base_linker's own cosine similarity. The reference's own default (20).	20
neigh_num	int	Neighbor triples shown per entity in the reasoning prompt (the reference's own --neigh 5 default for the main re-ranking pass).	5
desc_neigh_num	int	Neighbor triples shown per entity when generating its one-sentence description (the reference's own --neigh 25 default -- wider than neigh_num since description generation happens once per entity, not once per comparison).	25
desc_max_tokens	int	Forwarded to the description-generation LLM call.	80
threshold	float	If the base method's own top-2 candidate scores differ by more than this, accept the top-1 candidate immediately with no LLM call at all (the reference's own confidence shortcut).	0.5
history_len	int	Chat turns (user+assistant pairs) kept in the running conversation history within one source entity's re-ranking.	3
relation_names	dict[str, str] | None	relation_id -> display name, e.g. from IcewsWikiDataset.load_relation_labels. Falls back to the raw id for any relation not present, or when omitted entirely.	None
temporal_triples	list[TemporalTriple] | None	Optional (subject_id, relation_id, object_id, start_label, end_label) temporal facts (e.g. from IcewsWikiDataset.load_temporal_graphs). When given, used instead of graph for neighbor rendering (strictly more informative -- same ids, plus time) and enables the time-scoring dimension. None (default) disables the time dimension entirely, matching SimpleHHEALinker.fit's own convention for the same reference feature-toggle.	None
use_name	bool	Whether to score name similarity.	True
use_desc	bool	Whether to generate and score entity descriptions (requires use_name, see resolve_dimensions).	True
use_struct	bool	Whether to score structural (neighbor-tuple) similarity.	True
matching	Matcher	Strategy used to resolve scored candidates into final links, same as every other linker's matching argument.	DEFAULT_MATCHER

Note

blocking is accepted for interface compliance but not used -- candidates come from base_linker's own embedding similarity, not a BlockingStrategy, same pattern LlmBaseLinker/ StringSimilarityLinker already document for their own unused params. dataset2 must be a list[Entity], not an EntitySource -- candidate generation needs every target's embedding upfront.

Source code in src/linkingtk/algorithms/ea/chatea.py

def __init__(
    self,
    base_linker: EmbeddingLinker,
    client: LlmClient,
    top_k: int = 20,
    neigh_num: int = 5,
    desc_neigh_num: int = 25,
    desc_max_tokens: int = 80,
    threshold: float = 0.5,
    history_len: int = 3,
    relation_names: dict[str, str] | None = None,
    temporal_triples: list[TemporalTriple] | None = None,
    use_name: bool = True,
    use_desc: bool = True,
    use_struct: bool = True,
    matching: Matcher = DEFAULT_MATCHER,
) -> None:
    self.base_linker = base_linker
    self.client = client
    self.top_k = top_k
    self.neigh_num = neigh_num
    self.desc_neigh_num = desc_neigh_num
    self.desc_max_tokens = desc_max_tokens
    self.threshold = threshold
    self.history_len = history_len
    self.relation_names = relation_names or {}
    self.temporal_triples = temporal_triples
    self.use_name = use_name
    self.use_desc = use_desc
    self.use_struct = use_struct
    self.matching = matching

EmbeddingLinker

Bases: Protocol

What ChatEALinker needs from its base_linker -- the same shape linkingtk.eval.ranking's own private _ScoringLinker Protocol uses, duplicated locally since that one is module-private. Every linker in linkingtk.algorithms.ea (e.g. SimpleHHEALinker) already implements this.

source_embedding

source_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's source side.

Source code in src/linkingtk/algorithms/ea/chatea.py

def source_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score `entity_id` as a scored pair's source side."""

target_embedding

target_embedding(
    entity_id: str,
) -> npt.NDArray[np.floating[Any]]

Vector used to score entity_id as a scored pair's target side.

Source code in src/linkingtk/algorithms/ea/chatea.py

def target_embedding(self, entity_id: str) -> npt.NDArray[np.floating[Any]]:
    """Vector used to score `entity_id` as a scored pair's target side."""

_chatea_context

Per-entity context (neighbor triples, names, LLM-generated descriptions) for ChatEALinker's prompts.

Ports the neighbor-rendering half of the reference's tools_for_ChatEA.py's NeighborGenerator.get_neighbors (as NeighborIndex) and the description pre-compute of preobtain_description.py's generate_prompt/get_description (as generate_descriptions).

NeighborTriple module-attribute

NeighborTriple = tuple[
    str, str, str, "str | None", "str | None"
]

(subject_name, relation_name, object_name, start_label, end_label) -- time labels are None for non-temporal datasets/triples.

EntityContext dataclass

EntityContext(
    entity_id: str,
    name: str,
    description: str,
    neighbors: list[NeighborTriple] = list(),
)

One entity's rendered prompt content.

NeighborIndex

NeighborIndex(
    triples: list[Triple],
    temporal_triples: list[TemporalTriple] | None,
    entity_names: dict[str, str],
    relation_names: dict[str, str],
    neigh_num: int,
)

Looks up an entity's name and (capped) rendered neighbor triples.

Parameters:

Name	Type	Description	Default
triples	list[Triple]	Combined relational triples, used only when temporal_triples is None.	required
temporal_triples	list[TemporalTriple] | None	Combined temporal triples (preferred over triples when given, since they carry strictly more information -- the same relation/entity ids, plus time).	required
entity_names	dict[str, str]	entity_id -> display name (from Entity.labels).	required
relation_names	dict[str, str]	relation_id -> display name (e.g. from IcewsWikiDataset.load_relation_labels). Falls back to the raw id for any relation not present.	required
neigh_num	int	Max neighbor triples kept per entity (both as subject and object), mirroring the reference's own --neigh cap.	required

Source code in src/linkingtk/algorithms/ea/_chatea_context.py

def __init__(
    self,
    triples: list[Triple],
    temporal_triples: list[TemporalTriple] | None,
    entity_names: dict[str, str],
    relation_names: dict[str, str],
    neigh_num: int,
) -> None:
    self._entity_names = entity_names
    self._relation_names = relation_names
    self._neigh_num = neigh_num
    self._raw: dict[str, list[TemporalTriple]] = defaultdict(list)
    source: Iterable[TemporalTriple] = (
        temporal_triples
        if temporal_triples is not None
        else [(s, r, o, None, None) for s, r, o in triples]
    )
    for subject_id, relation_id, object_id, start, end in source:
        fact = (subject_id, relation_id, object_id, start, end)
        self._raw[subject_id].append(fact)
        self._raw[object_id].append(fact)

format_triple

format_triple(triple: NeighborTriple) -> str

Renders one neighbor triple as prompt text, e.g. "(A, rel, B, 2005-11, 2005-11)".

Source code in src/linkingtk/algorithms/ea/_chatea_context.py

def format_triple(triple: NeighborTriple) -> str:
    """Renders one neighbor triple as prompt text, e.g. ``"(A, rel, B, 2005-11, 2005-11)"``."""
    subject, relation, obj, start, end = triple
    if start is None and end is None:
        return f"({subject}, {relation}, {obj})"
    return f"({subject}, {relation}, {obj}, {start or '?'}, {end or '?'})"

generate_descriptions

generate_descriptions(
    client: LlmClient,
    entity_ids: Iterable[str],
    index: NeighborIndex,
    *,
    max_tokens: int = 80,
) -> dict[str, str]

One cached LLM-generated one-sentence description per id in entity_ids.

Ports preobtain_description.py's get_description -- called once per entity actually needed (a source entity or one of its top-k candidates), not once per comparison.

Source code in src/linkingtk/algorithms/ea/_chatea_context.py

def generate_descriptions(
    client: LlmClient,
    entity_ids: Iterable[str],
    index: NeighborIndex,
    *,
    max_tokens: int = 80,
) -> dict[str, str]:
    """One cached LLM-generated one-sentence description per id in `entity_ids`.

    Ports `preobtain_description.py`'s `get_description` -- called once
    per entity actually needed (a source entity or one of its top-k
    candidates), not once per comparison.
    """
    descriptions: dict[str, str] = {}
    for entity_id in entity_ids:
        if entity_id in descriptions:
            continue
        messages = _description_prompt(index.name(entity_id), index.neighbors(entity_id))
        try:
            response = client.complete_structured(
                messages, schema=_DESCRIPTION_SCHEMA, max_tokens=max_tokens
            )
            descriptions[entity_id] = str(response.get("description", "")).strip()
        except Exception:
            logger.warning(
                "ChatEALinker: description generation failed for entity %r, "
                "using empty description",
                entity_id,
                exc_info=True,
            )
            descriptions[entity_id] = ""
    return descriptions

_chatea_prompts

Prompt/scoring-schema construction for ChatEALinker.

Ports main_ChatEA.py's init_fields/init_prompt (the system prompt -- task framing plus a worked reasoning example) and generate_prompt/ ask_for_accuracy (the per-candidate reasoning prompt and the "rethinking" confirmation prompt) -- adapted to score via LlmClient.complete_structured against a fixed JSON schema instead of the reference's free-text + regex-parsed response (get_score), matching LlmBaseLinker's own existing convention in this repo.

The reference's use_code/class-definition text-assembly branching (init_prompt's combinatorial "how do I describe self.tuples in pseudo-code depending on which flags are off" logic) is simplified here to one plain description of whichever info types are enabled -- prompt- text cosmetics for ablations this port isn't running. The actual algorithm (candidate windowing, weighted scoring, thresholds, rethinking) lives in _chatea_reasoning and is ported faithfully.

ScoringDimensions dataclass

ScoringDimensions(
    keys: tuple[str, ...], weights: dict[str, float]
)

Which of the four score dimensions are active, and their (renormalized) weights.

build_candidate_prompt

build_candidate_prompt(
    main: EntityContext,
    candidate: EntityContext,
    candidate_names: list[str],
    ranked_so_far: list[tuple[str, float]],
    dims: ScoringDimensions,
) -> tuple[LlmMessage, str]

Ports generate_prompt. Returns (full_message, simple_prompt_text) -- simple_prompt_text is the compact form used for the running chat history (mirrors the reference's own prompt, simple_prompt pair).

Source code in src/linkingtk/algorithms/ea/_chatea_prompts.py

def build_candidate_prompt(
    main: EntityContext,
    candidate: EntityContext,
    candidate_names: list[str],
    ranked_so_far: list[tuple[str, float]],
    dims: ScoringDimensions,
) -> tuple[LlmMessage, str]:
    """Ports `generate_prompt`. Returns ``(full_message, simple_prompt_text)`` --
    `simple_prompt_text` is the compact form used for the running chat history
    (mirrors the reference's own ``prompt, simple_prompt`` pair)."""

    def render(entity: EntityContext) -> str:
        parts = [f"id={entity.entity_id!r}"]
        if "name" in dims.keys:
            parts.append(f"name={entity.name!r}")
        if "desc" in dims.keys:
            parts.append(f"description={entity.description!r}")
        if "struct" in dims.keys:
            triples_text = "; ".join(format_triple(t) for t in entity.neighbors)
            parts.append(f"tuples=[{triples_text}]")
        if "time" in dims.keys:
            time_text = "; ".join(f"{t[3] or '?'} to {t[4] or '?'}" for t in entity.neighbors)
            parts.append(f"time_info=[{time_text}]")
        return "Entity(" + ", ".join(parts) + ")"

    cand_list_text = ", ".join(candidate_names)
    ranked_text = ""
    if ranked_so_far:
        ranked_pairs = ", ".join(f"({name}, {score:.3f})" for name, score in ranked_so_far)
        ranked_text = f" Ranked so far (candidate, score): [{ranked_pairs}]."

    output_format = ", ".join(f"[{_FIELD_LABELS[k]}]" for k in dims.keys)
    full_content = (
        f"Candidate entities that may align with [Main Entity] {main.name}: "
        f"[{cand_list_text}].{ranked_text}\n"
        f"[Main Entity] = {render(main)}\n"
        f"[Candidate Entity] = {render(candidate)}\n"
        "Compared with other candidates, do [Main Entity] and [Candidate Entity] align "
        f"or match? Score {output_format} each as an integer 1-5."
    )
    simple_content = (
        f"Do [Main Entity] {main.name} and [Candidate Entity] {candidate.name} align or match?"
    )
    return LlmMessage(role="user", content=full_content), simple_content

build_rethink_prompt

build_rethink_prompt(
    main_name: str,
    candidate_names: list[str],
    ranked: list[tuple[str, float, int]],
) -> LlmMessage

Ports ask_for_accuracy's prompt: "are these alignments good enough?".

Source code in src/linkingtk/algorithms/ea/_chatea_prompts.py

def build_rethink_prompt(
    main_name: str, candidate_names: list[str], ranked: list[tuple[str, float, int]]
) -> LlmMessage:
    """Ports `ask_for_accuracy`'s prompt: "are these alignments good enough?"."""
    cand_list_text = ", ".join(candidate_names)
    aligned_text = ", ".join(
        f"(candidate={name}, score={score:.3f}, rank={rank})" for name, score, rank in ranked
    )
    content = (
        f"Candidate entities that may align with [Main Entity] {main_name}: [{cand_list_text}]. "
        f"Current ranked alignments: [{aligned_text}]. "
        "Compared with the candidate list, are these alignments good enough -- meaning the "
        "top-ranked candidate's score clears a reasonable threshold and is clearly higher than "
        "the rest? Answer with a boolean."
    )
    return LlmMessage(role="user", content=content)

build_system_prompt

build_system_prompt(dims: ScoringDimensions) -> str

Ports init_prompt: task framing plus one worked reasoning example.

Source code in src/linkingtk/algorithms/ea/_chatea_prompts.py

def build_system_prompt(dims: ScoringDimensions) -> str:
    """Ports `init_prompt`: task framing plus one worked reasoning example."""
    used_info = []
    if "name" in dims.keys:
        used_info.append("name")
    if "desc" in dims.keys:
        used_info.append("description")
    if "struct" in dims.keys:
        used_info.append("structure (neighboring knowledge tuples)")
    if "time" in dims.keys:
        used_info.append("temporal information")
    used_info.append("your own knowledge")

    example_steps = []
    if "name" in dims.keys:
        example_steps.append(
            "NAME SIMILARITY: 'Fudan University' and 'Fudan_University' are almost the "
            "same, so NAME SIMILARITY = 5 (very high)"
        )
    if "desc" in dims.keys:
        example_steps.append(
            "PROBABILITY OF DESCRIPTION POINTING SAME ENTITY: both descriptions describe "
            "a university in Shanghai founded in 1905, so this = 5 (very high)"
        )
    if "struct" in dims.keys:
        example_steps.append(
            "STRUCTURE SIMILARITY: 'China' is a common neighbor of both entities, so "
            "STRUCTURE SIMILARITY = 3 (medium)"
        )
    if "time" in dims.keys:
        example_steps.append(
            "TIME SIMILARITY: the candidate entity has no matching time information, so "
            "just assume TIME SIMILARITY = 2 (low)"
        )
    example = (
        "Worked example: Main Entity = Entity(id='8535', name='Fudan University'), "
        "Candidate Entity = Entity(id='24431', name='Fudan_University'). "
        "Reasoning step by step: " + "; ".join(example_steps) + "."
    )

    output_fields = ", ".join(f"[{_FIELD_LABELS[k]}]" for k in dims.keys)

    return (
        "You are a helpful assistant, helping align or match entities between two "
        f"knowledge graphs using {', '.join(used_info)}.\n\n"
        f"{example}\n\n"
        f"For every candidate entity, score {output_fields}, each an integer from "
        "1 (very low) to 5 (very high)."
    )

resolve_dimensions

resolve_dimensions(
    use_name: bool,
    use_desc: bool,
    use_struct: bool,
    use_time: bool,
) -> ScoringDimensions

Picks active dimensions and renormalizes _BASE_WEIGHTS over them.

Mirrors init_fields's scale_weights = [w/sum(weights) for w in weights]. The description dimension requires use_name too, same gating as the reference (if use_name and use_desc and use_code) -- "does this description point at the same entity as that name" presupposes a name was shown at all.

Source code in src/linkingtk/algorithms/ea/_chatea_prompts.py

def resolve_dimensions(
    use_name: bool, use_desc: bool, use_struct: bool, use_time: bool
) -> ScoringDimensions:
    """Picks active dimensions and renormalizes `_BASE_WEIGHTS` over them.

    Mirrors `init_fields`'s `scale_weights = [w/sum(weights) for w in weights]`.
    The description dimension requires `use_name` too, same gating as the
    reference (`if use_name and use_desc and use_code`) -- "does this
    description point at the same entity as that name" presupposes a name
    was shown at all.
    """
    enabled = {
        "name": use_name,
        "desc": use_desc and use_name,
        "struct": use_struct,
        "time": use_time,
    }
    keys = tuple(k for k in _DIMENSION_ORDER if enabled[k])
    total = sum(_BASE_WEIGHTS[k] for k in keys) or 1.0
    weights = {k: _BASE_WEIGHTS[k] / total for k in keys}
    return ScoringDimensions(keys=keys, weights=weights)

summarize_response

summarize_response(
    response: dict[str, Any], dims: ScoringDimensions
) -> str

A short text summary of a scoring response, used as the assistant's turn in the running chat history (mirrors the reference's own simple_response, keeping history compact across many calls).

Source code in src/linkingtk/algorithms/ea/_chatea_prompts.py

def summarize_response(response: dict[str, Any], dims: ScoringDimensions) -> str:
    """A short text summary of a scoring response, used as the assistant's
    turn in the running chat history (mirrors the reference's own
    ``simple_response``, keeping history compact across many calls)."""
    parts = [f"{_FIELD_LABELS[k]} = {response.get(f'{k}_score', 1)} out of 5" for k in dims.keys]
    return ", ".join(parts) + "."

weighted_score

weighted_score(
    response: dict[str, Any], dims: ScoringDimensions
) -> float

Combines a scoring response into one score in [0, 1].

Ports eval_alignment_for_evaluate's score += weights[j] * (sims[j] - 1) * 0.25.

Source code in src/linkingtk/algorithms/ea/_chatea_prompts.py

def weighted_score(response: dict[str, Any], dims: ScoringDimensions) -> float:
    """Combines a scoring response into one score in ``[0, 1]``.

    Ports `eval_alignment_for_evaluate`'s ``score += weights[j] * (sims[j] - 1) * 0.25``.
    """
    total = 0.0
    for key in dims.keys:
        raw = response.get(f"{key}_score", 1)
        clipped = min(5, max(1, int(raw)))
        total += dims.weights[key] * (clipped - 1) * 0.25
    return round(total, 4)

_chatea_reasoning

Per-source-entity iterative LLM re-ranking loop for ChatEALinker.

Ports main_ChatEA.py's eval_alignment_for_evaluate -- the confidence shortcut, widening candidate windows, weighted scoring, and the stop/continue/rethink control flow -- minus its ref_ent/base_rank ground-truth-peeking shortcut (the reference skips its LLM loop entirely when the true target isn't even in the top-20 candidates, since no re-ranking could ever reach it; that's pure benchmark-cost accounting using ground truth only evaluation code has access to). BaseLinker.link's general interface has no ground truth to peek at, so keeping that shortcut here would be a real interface violation, not a faithful port -- its cost-saving effect is instead reproduced at the benchmark-script layer, where ground truth is legitimately available (see examples/chatea_ea.py).

Structured JSON scoring (via _chatea_prompts) replaces the reference's regex-parsed free text, matching LlmBaseLinker's own convention.

rerank_candidates

rerank_candidates(
    client: LlmClient,
    system_prompt: str,
    main: EntityContext,
    candidates: list[tuple[EntityContext, float]],
    config: ReasoningConfig,
) -> dict[str, float]

Re-ranks candidates (base-method context + score, sorted descending by base score) for main, returning entity_id -> final weighted score for every candidate actually reasoned about by the LLM.

Ports eval_alignment_for_evaluate. An empty candidates returns {}.

Source code in src/linkingtk/algorithms/ea/_chatea_reasoning.py

def rerank_candidates(
    client: LlmClient,
    system_prompt: str,
    main: EntityContext,
    candidates: list[tuple[EntityContext, float]],
    config: ReasoningConfig,
) -> dict[str, float]:
    """Re-ranks `candidates` (base-method context + score, sorted descending
    by base score) for `main`, returning `entity_id -> final weighted score`
    for every candidate actually reasoned about by the LLM.

    Ports `eval_alignment_for_evaluate`. An empty `candidates` returns `{}`.
    """
    if not candidates:
        return {}

    if len(candidates) >= 2 and candidates[0][1] - candidates[1][1] > config.threshold:
        return {candidates[0][0].entity_id: 1.0}

    context_by_id = {ctx.entity_id: ctx for ctx, _ in candidates}
    candidate_names = [ctx.name for ctx, _ in candidates]
    system_messages = [LlmMessage(role="system", content=system_prompt)]
    chat_history: list[LlmMessage] = []
    scores: dict[str, float] = {}

    for window in config.search_windows:
        bounded_window = min(window, len(candidates))
        for candidate_ctx, _ in reversed(candidates[:bounded_window]):
            ranked_so_far = sorted(scores.items(), key=lambda item: -item[1])
            ranked_named = [(context_by_id[cid].name, score) for cid, score in ranked_so_far]
            full_message, simple_prompt = build_candidate_prompt(
                main, candidate_ctx, candidate_names, ranked_named, config.dims
            )
            messages = system_messages + chat_history + [full_message]
            try:
                response = client.complete_structured(
                    messages, schema=scoring_schema(config.dims), max_tokens=700
                )
            except Exception:
                logger.warning(
                    "ChatEALinker: scoring call failed for %r vs %r, skipping",
                    main.entity_id,
                    candidate_ctx.entity_id,
                    exc_info=True,
                )
                continue

            scores[candidate_ctx.entity_id] = weighted_score(response, config.dims)

            chat_history = chat_history + [
                LlmMessage(role="user", content=simple_prompt),
                LlmMessage(role="assistant", content=summarize_response(response, config.dims)),
            ]
            if len(chat_history) > config.history_len * 2:
                chat_history = chat_history[-config.history_len * 2 :]

        if not scores:
            continue

        ranked = sorted(scores.items(), key=lambda item: -item[1])
        top_id, top_score = ranked[0]
        if top_score < 0.4:
            good_enough = False
        elif top_score > 0.8:
            good_enough = True
        elif len(ranked) > 1 and top_score - ranked[1][1] > 0.5:
            good_enough = True
        else:
            good_enough = _ask_for_accuracy(
                client, system_messages, main, candidate_names, ranked, context_by_id
            )
        if good_enough:
            break

    return scores

Entity Linking

el

Entity Linking (EL) algorithms.

EL links named-entity mentions in text (with context, no description) to entries in a knowledge base (with description, no context). Related systems: spaCy's EntityLinker, ReFinED, BLINK.

BlinkEncoder

BlinkEncoder(
    mention_model_name: str = "distilbert-base-uncased",
    entity_model_name: str | None = None,
    embedding_dim: int = 256,
    max_length: int = 96,
    encode_batch_size: int = 64,
)

Bases: Module

Two independently-parameterized transformer bi-encoder for Entity Linking.

Satisfies TrainableModel: encode() dispatches each entity to a mention-side or entity-side sub-transformer (by whether entity.context is not None, mirroring ReFinEDEncoder's dispatch) -- unlike ReFinEDEncoder, the two sides share no weights at all, BLINK's real distinguishing idea.

Parameters:

Name	Type	Description	Default
mention_model_name	str	Hugging Face model id for the mention-side encoder, loaded via AutoModel/AutoTokenizer.from_pretrained. Defaults to distilbert-base-uncased, matching ReFinEDEncoder's precedent.	'distilbert-base-uncased'
entity_model_name	str | None	Hugging Face model id for the entity-side encoder. Defaults to the same as mention_model_name.	None
embedding_dim	int	Output embedding dimension for both sides, independent of either backbone's native hidden size.	256
max_length	int	Maximum token length per encoded text (truncated).	96
encode_batch_size	int	Maximum number of texts forwarded through a sub-transformer in one go; encode() chunks larger inputs internally and concatenates the results, so its own signature is unaffected. Needed at real-dataset scale -- unlike ReFinEDEncoder, which has never been called with more than a few thousand entities at once, a single Zeshel domain's entity dictionary can have 100K+ entries, and forwarding all of them through a transformer in one unbatched call exhausts GPU memory.	64

Source code in src/linkingtk/algorithms/el/blink.py

def __init__(
    self,
    mention_model_name: str = "distilbert-base-uncased",
    entity_model_name: str | None = None,
    embedding_dim: int = 256,
    max_length: int = 96,
    encode_batch_size: int = 64,
) -> None:
    super().__init__()
    self.mention_encoder = _SubEncoder(
        mention_model_name, embedding_dim, [_MENTION_START, _MENTION_END]
    )
    self.entity_encoder = _SubEncoder(
        entity_model_name or mention_model_name, embedding_dim, [_ENTITY_SEP]
    )
    self.max_length = max_length
    self.encode_batch_size = encode_batch_size

encode

encode(entities: list[Entity]) -> torch.Tensor

Return a (len(entities), embedding_dim) L2-normalized embedding tensor.

Mixes mention-side and entity-side entities in one call by encoding each group with its own sub-encoder and reassembling the result in the original order -- callers (e.g. Trainer) never need to split a batch themselves.

Source code in src/linkingtk/algorithms/el/blink.py

def encode(self, entities: list[Entity]) -> torch.Tensor:
    """Return a ``(len(entities), embedding_dim)`` L2-normalized embedding tensor.

    Mixes mention-side and entity-side entities in one call by
    encoding each group with its own sub-encoder and reassembling the
    result in the original order -- callers (e.g.
    [Trainer][linkingtk.train.trainer.Trainer]) never need to split a
    batch themselves.
    """
    mention_positions = [i for i, entity in enumerate(entities) if entity.context is not None]
    entity_positions = [i for i, entity in enumerate(entities) if entity.context is None]

    if not entity_positions:
        texts = [_mention_text(entities[i]) for i in mention_positions]
        return self.mention_encoder.encode_texts(texts, self.max_length, self.encode_batch_size)
    if not mention_positions:
        texts = [_entity_text(entities[i]) for i in entity_positions]
        return self.entity_encoder.encode_texts(texts, self.max_length, self.encode_batch_size)

    mention_emb = self.mention_encoder.encode_texts(
        [_mention_text(entities[i]) for i in mention_positions],
        self.max_length,
        self.encode_batch_size,
    )
    entity_emb = self.entity_encoder.encode_texts(
        [_entity_text(entities[i]) for i in entity_positions],
        self.max_length,
        self.encode_batch_size,
    )
    output = mention_emb.new_empty(len(entities), mention_emb.shape[-1])
    output[mention_positions] = mention_emb
    output[entity_positions] = entity_emb
    return output

BlinkLinker

BlinkLinker(
    mention_model_name: str = "distilbert-base-uncased",
    entity_model_name: str | None = None,
    embedding_dim: int = 256,
    max_length: int = 96,
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

Bi-encoder Entity Linking scored by a two-tower BlinkEncoder.

Like ReFinEDLinker, this has no fit() -- self.encoder is a plain TrainableModel; train it via Trainer directly, then call link():

linker = BlinkLinker()
trainer = Trainer(model=linker.encoder, args=args, train_data=pairs, eval_data=eval_pairs)
trainer.train()
results = linker.link(mentions, kb_entries)

link() runs (just poorly) on an untrained encoder too -- there's no _fitted flag, since nothing here can observe whether Trainer.train() was ever called on self.encoder.

Parameters:

Name	Type	Description	Default
mention_model_name	str	Forwarded to BlinkEncoder.	'distilbert-base-uncased'
entity_model_name	str | None	Forwarded to BlinkEncoder.	None
embedding_dim	int	Forwarded to BlinkEncoder.	256
max_length	int	Forwarded to BlinkEncoder.	96
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER

Source code in src/linkingtk/algorithms/el/blink.py

def __init__(
    self,
    mention_model_name: str = "distilbert-base-uncased",
    entity_model_name: str | None = None,
    embedding_dim: int = 256,
    max_length: int = 96,
    matching: Matcher = DEFAULT_MATCHER,
) -> None:
    self.encoder = BlinkEncoder(
        mention_model_name, entity_model_name, embedding_dim, max_length
    )
    self.matching = matching

score_candidates

score_candidates(
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]

Blocked candidates per source entity, scored but not yet matched.

Satisfies CandidateScorer -- what link() itself builds internally, exposed so LlmRerankerLinker (#23) can re-rank a narrowed top-k instead of every blocked pair.

Source code in src/linkingtk/algorithms/el/blink.py

def score_candidates(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]:
    """Blocked candidates per source entity, scored but not yet matched.

    Satisfies [CandidateScorer][linkingtk.algorithms.llm_reranker.CandidateScorer]
    -- what `link()` itself
    builds internally, exposed so
    [LlmRerankerLinker][linkingtk.algorithms.llm_reranker.LlmRerankerLinker]
    (#23) can re-rank a narrowed top-k instead of every blocked pair.
    """
    pairs = blocking.candidate_pairs(dataset1, dataset2)
    if not pairs:
        return {}

    mentions_by_id = {entity1.id: entity1 for entity1, _ in pairs}
    entities_by_id = {entity2.id: entity2 for _, entity2 in pairs}
    mention_ids = list(mentions_by_id)
    entity_ids = list(entities_by_id)

    with torch.no_grad():
        self.encoder.eval()
        mention_emb = self.encoder.encode([mentions_by_id[i] for i in mention_ids])
        entity_emb = self.encoder.encode([entities_by_id[i] for i in entity_ids])
        self.encoder.train()

    mention_row = {entity_id: row for row, entity_id in enumerate(mention_ids)}
    entity_row = {entity_id: row for row, entity_id in enumerate(entity_ids)}
    similarities = mention_emb @ entity_emb.T

    candidates_by_source: dict[str, list[tuple[str, float]]] = defaultdict(list)
    for entity1, entity2 in pairs:
        score = similarities[mention_row[entity1.id], entity_row[entity2.id]]
        candidates_by_source[entity1.id].append((entity2.id, float(score)))

    return candidates_by_source

ReFinEDEncoder

ReFinEDEncoder(
    model_name: str = "distilbert-base-uncased",
    embedding_dim: int = 256,
    max_length: int = 96,
)

Bases: Module

Shared/tied transformer bi-encoder for Entity Linking.

Satisfies TrainableModel: encode() runs both mentions and KB entries through the same transformer weights, distinguishing the two purely by how each entity's text is formatted (mention-span markers vs. label+description) -- ReFinED's real distinguishing idea versus BLINK's two independently encoded sides.

Parameters:

Name	Type	Description	Default
model_name	str	A Hugging Face model id, loaded via AutoModel/AutoTokenizer.from_pretrained. Defaults to distilbert-base-uncased, matching MultiKELinker's transformer-default precedent. Pass a tiny test model (e.g. hf-internal-testing/tiny-random-DistilBertModel, as test_multike.py already does) for fast, network-light tests.	'distilbert-base-uncased'
embedding_dim	int	Output embedding dimension, independent of the backbone's native hidden size.	256
max_length	int	Maximum token length per encoded text (truncated).	96

Source code in src/linkingtk/algorithms/el/refined.py

def __init__(
    self,
    model_name: str = "distilbert-base-uncased",
    embedding_dim: int = 256,
    max_length: int = 96,
) -> None:
    super().__init__()
    from transformers import AutoModel, AutoTokenizer

    self.tokenizer = AutoTokenizer.from_pretrained(model_name)
    self.tokenizer.add_special_tokens(
        {"additional_special_tokens": [_MENTION_START, _MENTION_END]}
    )
    self.transformer = AutoModel.from_pretrained(model_name)
    self.transformer.resize_token_embeddings(len(self.tokenizer))
    self.proj = nn.Linear(self.transformer.config.hidden_size, embedding_dim)
    self.max_length = max_length

encode

encode(entities: list[Entity]) -> torch.Tensor

Return a (len(entities), embedding_dim) L2-normalized embedding tensor.

Source code in src/linkingtk/algorithms/el/refined.py

def encode(self, entities: list[Entity]) -> torch.Tensor:
    """Return a ``(len(entities), embedding_dim)`` L2-normalized embedding tensor."""
    texts = [_format_entity(entity) for entity in entities]
    device = self.proj.weight.device
    encoded = self.tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=self.max_length,
        return_tensors="pt",
    ).to(device)
    hidden_states = self.transformer(**encoded).last_hidden_state
    pooled = _mean_pool(hidden_states, encoded["attention_mask"])
    return functional.normalize(self.proj(pooled), dim=-1)

ReFinEDLinker

ReFinEDLinker(
    model_name: str = "distilbert-base-uncased",
    embedding_dim: int = 256,
    max_length: int = 96,
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

Bi-encoder Entity Linking scored by a shared/tied ReFinEDEncoder.

Unlike every other linker in this repo, this has no fit() -- self.encoder is a plain TrainableModel; train it via Trainer directly (matching the documented Trainer(model=linker_model, ...) interface), then call link():

linker = ReFinEDLinker()
trainer = Trainer(model=linker.encoder, args=args, train_data=pairs, eval_data=eval_pairs)
trainer.train()
results = linker.link(mentions, kb_entries)

link() runs (just poorly) on an untrained encoder too -- there's no _fitted flag, since nothing here can observe whether Trainer.train() was ever called on self.encoder.

Parameters:

Name	Type	Description	Default
model_name	str	Forwarded to ReFinEDEncoder.	'distilbert-base-uncased'
embedding_dim	int	Forwarded to ReFinEDEncoder.	256
max_length	int	Forwarded to ReFinEDEncoder.	96
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER

Source code in src/linkingtk/algorithms/el/refined.py

def __init__(
    self,
    model_name: str = "distilbert-base-uncased",
    embedding_dim: int = 256,
    max_length: int = 96,
    matching: Matcher = DEFAULT_MATCHER,
) -> None:
    self.encoder = ReFinEDEncoder(model_name, embedding_dim, max_length)
    self.matching = matching

score_candidates

score_candidates(
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]

Blocked candidates per source entity, scored but not yet matched.

Satisfies CandidateScorer -- what link() itself builds internally, exposed so LlmRerankerLinker (#23) can re-rank a narrowed top-k instead of every blocked pair.

Source code in src/linkingtk/algorithms/el/refined.py

def score_candidates(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]:
    """Blocked candidates per source entity, scored but not yet matched.

    Satisfies [CandidateScorer][linkingtk.algorithms.llm_reranker.CandidateScorer]
    -- what `link()` itself
    builds internally, exposed so
    [LlmRerankerLinker][linkingtk.algorithms.llm_reranker.LlmRerankerLinker]
    (#23) can re-rank a narrowed top-k instead of every blocked pair.
    """
    pairs = blocking.candidate_pairs(dataset1, dataset2)
    if not pairs:
        return {}

    mentions_by_id = {entity1.id: entity1 for entity1, _ in pairs}
    entities_by_id = {entity2.id: entity2 for _, entity2 in pairs}
    mention_ids = list(mentions_by_id)
    entity_ids = list(entities_by_id)

    with torch.no_grad():
        self.encoder.eval()
        mention_emb = self.encoder.encode([mentions_by_id[i] for i in mention_ids])
        entity_emb = self.encoder.encode([entities_by_id[i] for i in entity_ids])
        self.encoder.train()

    mention_row = {entity_id: row for row, entity_id in enumerate(mention_ids)}
    entity_row = {entity_id: row for row, entity_id in enumerate(entity_ids)}
    similarities = mention_emb @ entity_emb.T

    candidates_by_source: dict[str, list[tuple[str, float]]] = defaultdict(list)
    for entity1, entity2 in pairs:
        score = similarities[mention_row[entity1.id], entity_row[entity2.id]]
        candidates_by_source[entity1.id].append((entity2.id, float(score)))

    return candidates_by_source

ReFinED

refined

ReFinED-style bi-encoder Entity Linking.

ReFinED (Ayoola et al., "ReFinED: An Efficient Zero-shot-capable Approach to End-to-End Entity Linking", 2022, https://github.com/amazon-science/ReFinED) is impractical to wrap directly: it isn't on PyPI (GitHub-zip install only), is pinned to Python 3.8 (this repo requires >=3.10), and needs GB-scale Wikipedia model downloads on first use. This module instead builds a custom bi-encoder that follows ReFinED's real distinguishing architectural idea -- a single shared/tied transformer encoder that scores both the mention-in-context and the candidate-entity representation through the same weights (unlike BLINK's two-stage separate-encoder-plus-cross-encoder-reranker design) -- trained via Trainer, not the original package.

See ReFinEDEncoder and ReFinEDLinker.

ReFinEDEncoder

ReFinEDEncoder(
    model_name: str = "distilbert-base-uncased",
    embedding_dim: int = 256,
    max_length: int = 96,
)

Bases: Module

Shared/tied transformer bi-encoder for Entity Linking.

Satisfies TrainableModel: encode() runs both mentions and KB entries through the same transformer weights, distinguishing the two purely by how each entity's text is formatted (mention-span markers vs. label+description) -- ReFinED's real distinguishing idea versus BLINK's two independently encoded sides.

Parameters:

Name	Type	Description	Default
model_name	str	A Hugging Face model id, loaded via AutoModel/AutoTokenizer.from_pretrained. Defaults to distilbert-base-uncased, matching MultiKELinker's transformer-default precedent. Pass a tiny test model (e.g. hf-internal-testing/tiny-random-DistilBertModel, as test_multike.py already does) for fast, network-light tests.	'distilbert-base-uncased'
embedding_dim	int	Output embedding dimension, independent of the backbone's native hidden size.	256
max_length	int	Maximum token length per encoded text (truncated).	96

Source code in src/linkingtk/algorithms/el/refined.py

def __init__(
    self,
    model_name: str = "distilbert-base-uncased",
    embedding_dim: int = 256,
    max_length: int = 96,
) -> None:
    super().__init__()
    from transformers import AutoModel, AutoTokenizer

    self.tokenizer = AutoTokenizer.from_pretrained(model_name)
    self.tokenizer.add_special_tokens(
        {"additional_special_tokens": [_MENTION_START, _MENTION_END]}
    )
    self.transformer = AutoModel.from_pretrained(model_name)
    self.transformer.resize_token_embeddings(len(self.tokenizer))
    self.proj = nn.Linear(self.transformer.config.hidden_size, embedding_dim)
    self.max_length = max_length

encode

encode(entities: list[Entity]) -> torch.Tensor

Return a (len(entities), embedding_dim) L2-normalized embedding tensor.

Source code in src/linkingtk/algorithms/el/refined.py

def encode(self, entities: list[Entity]) -> torch.Tensor:
    """Return a ``(len(entities), embedding_dim)`` L2-normalized embedding tensor."""
    texts = [_format_entity(entity) for entity in entities]
    device = self.proj.weight.device
    encoded = self.tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=self.max_length,
        return_tensors="pt",
    ).to(device)
    hidden_states = self.transformer(**encoded).last_hidden_state
    pooled = _mean_pool(hidden_states, encoded["attention_mask"])
    return functional.normalize(self.proj(pooled), dim=-1)

ReFinEDLinker

ReFinEDLinker(
    model_name: str = "distilbert-base-uncased",
    embedding_dim: int = 256,
    max_length: int = 96,
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

Bi-encoder Entity Linking scored by a shared/tied ReFinEDEncoder.

Unlike every other linker in this repo, this has no fit() -- self.encoder is a plain TrainableModel; train it via Trainer directly (matching the documented Trainer(model=linker_model, ...) interface), then call link():

linker = ReFinEDLinker()
trainer = Trainer(model=linker.encoder, args=args, train_data=pairs, eval_data=eval_pairs)
trainer.train()
results = linker.link(mentions, kb_entries)

link() runs (just poorly) on an untrained encoder too -- there's no _fitted flag, since nothing here can observe whether Trainer.train() was ever called on self.encoder.

Parameters:

Name	Type	Description	Default
model_name	str	Forwarded to ReFinEDEncoder.	'distilbert-base-uncased'
embedding_dim	int	Forwarded to ReFinEDEncoder.	256
max_length	int	Forwarded to ReFinEDEncoder.	96
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER

Source code in src/linkingtk/algorithms/el/refined.py

def __init__(
    self,
    model_name: str = "distilbert-base-uncased",
    embedding_dim: int = 256,
    max_length: int = 96,
    matching: Matcher = DEFAULT_MATCHER,
) -> None:
    self.encoder = ReFinEDEncoder(model_name, embedding_dim, max_length)
    self.matching = matching

score_candidates

score_candidates(
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]

Blocked candidates per source entity, scored but not yet matched.

Satisfies CandidateScorer -- what link() itself builds internally, exposed so LlmRerankerLinker (#23) can re-rank a narrowed top-k instead of every blocked pair.

Source code in src/linkingtk/algorithms/el/refined.py

def score_candidates(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]:
    """Blocked candidates per source entity, scored but not yet matched.

    Satisfies [CandidateScorer][linkingtk.algorithms.llm_reranker.CandidateScorer]
    -- what `link()` itself
    builds internally, exposed so
    [LlmRerankerLinker][linkingtk.algorithms.llm_reranker.LlmRerankerLinker]
    (#23) can re-rank a narrowed top-k instead of every blocked pair.
    """
    pairs = blocking.candidate_pairs(dataset1, dataset2)
    if not pairs:
        return {}

    mentions_by_id = {entity1.id: entity1 for entity1, _ in pairs}
    entities_by_id = {entity2.id: entity2 for _, entity2 in pairs}
    mention_ids = list(mentions_by_id)
    entity_ids = list(entities_by_id)

    with torch.no_grad():
        self.encoder.eval()
        mention_emb = self.encoder.encode([mentions_by_id[i] for i in mention_ids])
        entity_emb = self.encoder.encode([entities_by_id[i] for i in entity_ids])
        self.encoder.train()

    mention_row = {entity_id: row for row, entity_id in enumerate(mention_ids)}
    entity_row = {entity_id: row for row, entity_id in enumerate(entity_ids)}
    similarities = mention_emb @ entity_emb.T

    candidates_by_source: dict[str, list[tuple[str, float]]] = defaultdict(list)
    for entity1, entity2 in pairs:
        score = similarities[mention_row[entity1.id], entity_row[entity2.id]]
        candidates_by_source[entity1.id].append((entity2.id, float(score)))

    return candidates_by_source

BLINK

blink

BLINK-style bi-encoder Entity Linking.

BLINK (Wu et al., "Scalable Zero-shot Entity Linking with Dense Entity Retrieval", 2020, https://github.com/facebookresearch/BLINK) is impractical to wrap directly: it isn't on PyPI (GitHub-only), and its release is pinned to old transformers/Python versions with GB-scale Wikipedia model downloads. This module instead builds a custom bi-encoder that follows BLINK's real distinguishing architectural idea -- two independently-parameterized transformers, one for the mention-in- context, one for the candidate entity's title+description, each scored via [CLS]-pooled dot product (Section 4.1 of the paper) -- unlike ReFinEDLinker's single shared/tied encoder. Trained via Trainer, not the original package.

Candidate generation (this module's link()) only covers BLINK's first stage; the paper's second-stage cross-encoder reranker is out of scope here (see issue #23).

See BlinkEncoder and BlinkLinker.

BlinkEncoder

BlinkEncoder(
    mention_model_name: str = "distilbert-base-uncased",
    entity_model_name: str | None = None,
    embedding_dim: int = 256,
    max_length: int = 96,
    encode_batch_size: int = 64,
)

Bases: Module

Two independently-parameterized transformer bi-encoder for Entity Linking.

Satisfies TrainableModel: encode() dispatches each entity to a mention-side or entity-side sub-transformer (by whether entity.context is not None, mirroring ReFinEDEncoder's dispatch) -- unlike ReFinEDEncoder, the two sides share no weights at all, BLINK's real distinguishing idea.

Parameters:

Name	Type	Description	Default
mention_model_name	str	Hugging Face model id for the mention-side encoder, loaded via AutoModel/AutoTokenizer.from_pretrained. Defaults to distilbert-base-uncased, matching ReFinEDEncoder's precedent.	'distilbert-base-uncased'
entity_model_name	str | None	Hugging Face model id for the entity-side encoder. Defaults to the same as mention_model_name.	None
embedding_dim	int	Output embedding dimension for both sides, independent of either backbone's native hidden size.	256
max_length	int	Maximum token length per encoded text (truncated).	96
encode_batch_size	int	Maximum number of texts forwarded through a sub-transformer in one go; encode() chunks larger inputs internally and concatenates the results, so its own signature is unaffected. Needed at real-dataset scale -- unlike ReFinEDEncoder, which has never been called with more than a few thousand entities at once, a single Zeshel domain's entity dictionary can have 100K+ entries, and forwarding all of them through a transformer in one unbatched call exhausts GPU memory.	64

Source code in src/linkingtk/algorithms/el/blink.py

def __init__(
    self,
    mention_model_name: str = "distilbert-base-uncased",
    entity_model_name: str | None = None,
    embedding_dim: int = 256,
    max_length: int = 96,
    encode_batch_size: int = 64,
) -> None:
    super().__init__()
    self.mention_encoder = _SubEncoder(
        mention_model_name, embedding_dim, [_MENTION_START, _MENTION_END]
    )
    self.entity_encoder = _SubEncoder(
        entity_model_name or mention_model_name, embedding_dim, [_ENTITY_SEP]
    )
    self.max_length = max_length
    self.encode_batch_size = encode_batch_size

encode

encode(entities: list[Entity]) -> torch.Tensor

Return a (len(entities), embedding_dim) L2-normalized embedding tensor.

Mixes mention-side and entity-side entities in one call by encoding each group with its own sub-encoder and reassembling the result in the original order -- callers (e.g. Trainer) never need to split a batch themselves.

Source code in src/linkingtk/algorithms/el/blink.py

def encode(self, entities: list[Entity]) -> torch.Tensor:
    """Return a ``(len(entities), embedding_dim)`` L2-normalized embedding tensor.

    Mixes mention-side and entity-side entities in one call by
    encoding each group with its own sub-encoder and reassembling the
    result in the original order -- callers (e.g.
    [Trainer][linkingtk.train.trainer.Trainer]) never need to split a
    batch themselves.
    """
    mention_positions = [i for i, entity in enumerate(entities) if entity.context is not None]
    entity_positions = [i for i, entity in enumerate(entities) if entity.context is None]

    if not entity_positions:
        texts = [_mention_text(entities[i]) for i in mention_positions]
        return self.mention_encoder.encode_texts(texts, self.max_length, self.encode_batch_size)
    if not mention_positions:
        texts = [_entity_text(entities[i]) for i in entity_positions]
        return self.entity_encoder.encode_texts(texts, self.max_length, self.encode_batch_size)

    mention_emb = self.mention_encoder.encode_texts(
        [_mention_text(entities[i]) for i in mention_positions],
        self.max_length,
        self.encode_batch_size,
    )
    entity_emb = self.entity_encoder.encode_texts(
        [_entity_text(entities[i]) for i in entity_positions],
        self.max_length,
        self.encode_batch_size,
    )
    output = mention_emb.new_empty(len(entities), mention_emb.shape[-1])
    output[mention_positions] = mention_emb
    output[entity_positions] = entity_emb
    return output

BlinkLinker

BlinkLinker(
    mention_model_name: str = "distilbert-base-uncased",
    entity_model_name: str | None = None,
    embedding_dim: int = 256,
    max_length: int = 96,
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

Bi-encoder Entity Linking scored by a two-tower BlinkEncoder.

Like ReFinEDLinker, this has no fit() -- self.encoder is a plain TrainableModel; train it via Trainer directly, then call link():

linker = BlinkLinker()
trainer = Trainer(model=linker.encoder, args=args, train_data=pairs, eval_data=eval_pairs)
trainer.train()
results = linker.link(mentions, kb_entries)

link() runs (just poorly) on an untrained encoder too -- there's no _fitted flag, since nothing here can observe whether Trainer.train() was ever called on self.encoder.

Parameters:

Name	Type	Description	Default
mention_model_name	str	Forwarded to BlinkEncoder.	'distilbert-base-uncased'
entity_model_name	str | None	Forwarded to BlinkEncoder.	None
embedding_dim	int	Forwarded to BlinkEncoder.	256
max_length	int	Forwarded to BlinkEncoder.	96
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER

Source code in src/linkingtk/algorithms/el/blink.py

def __init__(
    self,
    mention_model_name: str = "distilbert-base-uncased",
    entity_model_name: str | None = None,
    embedding_dim: int = 256,
    max_length: int = 96,
    matching: Matcher = DEFAULT_MATCHER,
) -> None:
    self.encoder = BlinkEncoder(
        mention_model_name, entity_model_name, embedding_dim, max_length
    )
    self.matching = matching

score_candidates

score_candidates(
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]

Blocked candidates per source entity, scored but not yet matched.

Satisfies CandidateScorer -- what link() itself builds internally, exposed so LlmRerankerLinker (#23) can re-rank a narrowed top-k instead of every blocked pair.

Source code in src/linkingtk/algorithms/el/blink.py

def score_candidates(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]:
    """Blocked candidates per source entity, scored but not yet matched.

    Satisfies [CandidateScorer][linkingtk.algorithms.llm_reranker.CandidateScorer]
    -- what `link()` itself
    builds internally, exposed so
    [LlmRerankerLinker][linkingtk.algorithms.llm_reranker.LlmRerankerLinker]
    (#23) can re-rank a narrowed top-k instead of every blocked pair.
    """
    pairs = blocking.candidate_pairs(dataset1, dataset2)
    if not pairs:
        return {}

    mentions_by_id = {entity1.id: entity1 for entity1, _ in pairs}
    entities_by_id = {entity2.id: entity2 for _, entity2 in pairs}
    mention_ids = list(mentions_by_id)
    entity_ids = list(entities_by_id)

    with torch.no_grad():
        self.encoder.eval()
        mention_emb = self.encoder.encode([mentions_by_id[i] for i in mention_ids])
        entity_emb = self.encoder.encode([entities_by_id[i] for i in entity_ids])
        self.encoder.train()

    mention_row = {entity_id: row for row, entity_id in enumerate(mention_ids)}
    entity_row = {entity_id: row for row, entity_id in enumerate(entity_ids)}
    similarities = mention_emb @ entity_emb.T

    candidates_by_source: dict[str, list[tuple[str, float]]] = defaultdict(list)
    for entity1, entity2 in pairs:
        score = similarities[mention_row[entity1.id], entity_row[entity2.id]]
        candidates_by_source[entity1.id].append((entity2.id, float(score)))

    return candidates_by_source

Word Sense Disambiguation

wsd

Word Sense Disambiguation (WSD) algorithms.

WSD identifies the relevant dictionary sense for a mention in context. Similar to EL, but the target dataset is a dictionary rather than a KB. Related systems: Lesk, EWISER, GlossBERT, ESC.

EwiserEncoder

EwiserEncoder(
    model_name_or_path: str = "bert-large-cased",
    vocabulary: SenseVocabulary | None = None,
    decoder_hidden_dim: int = 512,
    adjacency: Tensor | None = None,
    structured_logits_trainable: bool = True,
    structured_logits_renormalize: bool = False,
    freeze_encoder: bool = True,
    dropout: float = 0.1,
    num_summed_layers: int = 4,
    max_length: int = 512,
    forward_batch_size: int = 16,
    output_embedding_init: Tensor | None = None,
)

Bases: Module

Frozen-BERT full-inventory sense classifier with a WordNet-graph output layer.

Parameters:

Name	Type	Description	Default
model_name_or_path	str	A Hugging Face encoder id, e.g. "bert-large-cased" (the published checkpoints' own encoder). Pass a tiny test model for fast, network-light tests.	'bert-large-cased'
vocabulary	SenseVocabulary | None	The SenseVocabulary the output layer classifies against.	None
decoder_hidden_dim	int	Width of the decoder's single hidden layer. Defaults to 512, matching the reference.	512
adjacency	Tensor | None	An optional sparse [len(vocabulary), len(vocabulary)] WordNet-relation adjacency (from build_relation_adjacency, or a checkpoint's own baked-in one via from_checkpoint). Wires up the graph-propagation step (StructuredLogits) if given; the model degrades to a plain (non-graph-aware) classifier if omitted.	None
structured_logits_trainable	bool	Whether the graph's edge weights are updated during training.	True
structured_logits_renormalize	bool	Forwarded to StructuredLogits.	False
freeze_encoder	bool	Whether the BERT encoder's parameters are frozen. Defaults to True, matching the paper -- exposed as a real constructor flag rather than hardcoded since nothing else in this package forbids fine-tuning it, but the default preserves fidelity to EWISER's own recipe (see the module docstring).	True
dropout	float	Decoder dropout rate. Checkpoint loading doesn't restore this (dropout has no learned parameters), so it only affects EwiserTrainer-driven training.	0.1
output_embedding_init	Tensor | None	An optional [len(vocabulary), decoder_hidden_dim] tensor to initialize decoder.logits.weight from, in place of its default random init -- see load_synset_centroid_vectors/ build_synset_centroid_vectors_from_lmms to build one. This is what makes EwiserTrainer.freeze_output_epochs meaningful -- freezing a pretrained output layer for a few epochs protects it from noisy early gradients (the paper's own motivation); freezing the layer's plain default random init has no such benefit.	None
num_summed_layers	int	The encoder input is the sum of the last num_summed_layers BERT hidden-state layers, not just the final layer -- confirmed against the reference's own TaggerModel.build_model, whose ``use_all_hidden=len(layers) 1readsargs.context_embeddings_type(a string, e.g."bert") where it evidently meant to readargs.context_embeddings_layers(the actual 4-entry layer list) -- sincelen("bert") == 4 > 1regardless of the real layer count, this makesuse_all_hidden**always** true for every BERT-backed checkpoint, silently overriding the separate, correctly-namedcontext_embeddings_use_all_hiddenconfig flag (which reads `False` on all three released checkpoints and would, if honored, sum only the single last layer). A reference quirk, not a reference bug we get to ignore: it *is* what these checkpoints were actually trained against, confirmed by tracing a checkpoint-reproduction gap down to this exact vector (norm ~82 summed vs. ~18 single-layer, cosine 0.69 apart -- unmistakably different vectors, not numerical noise). Defaults to4, matching every released checkpoint's owncontext_embeddings_layers`` length.	4
max_length	int	Maximum subword sequence length per sentence (truncated).	512
forward_batch_size	int	Number of distinct sentences score() batches through the encoder+decoder per forward pass -- unrelated to how many (mention, candidate_sense) pairs score() is called with (which may share far fewer distinct sentences, see score's own docstring).	16

Raises:

Type	Description
LinkingTKError	If vocabulary isn't given.

Source code in src/linkingtk/algorithms/wsd/ewiser.py

def __init__(
    self,
    model_name_or_path: str = "bert-large-cased",
    vocabulary: SenseVocabulary | None = None,
    decoder_hidden_dim: int = 512,
    adjacency: torch.Tensor | None = None,
    structured_logits_trainable: bool = True,
    structured_logits_renormalize: bool = False,
    freeze_encoder: bool = True,
    dropout: float = 0.1,
    num_summed_layers: int = 4,
    max_length: int = 512,
    forward_batch_size: int = 16,
    output_embedding_init: torch.Tensor | None = None,
) -> None:
    super().__init__()
    if vocabulary is None:
        raise LinkingTKError("EwiserEncoder requires a `vocabulary` (see SenseVocabulary).")
    from transformers import AutoModel, AutoTokenizer

    self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    self.encoder = AutoModel.from_pretrained(model_name_or_path)
    self.freeze_encoder = freeze_encoder
    for parameter in self.encoder.parameters():
        parameter.requires_grad = not freeze_encoder

    self.vocabulary = vocabulary
    self.num_summed_layers = num_summed_layers
    self.max_length = max_length
    self.forward_batch_size = forward_batch_size
    self.decoder = EwiserDecoder(
        input_dim=self.encoder.config.hidden_size,
        hidden_dim=decoder_hidden_dim,
        vocab_size=len(vocabulary),
        adjacency=adjacency,
        structured_logits_trainable=structured_logits_trainable,
        structured_logits_renormalize=structured_logits_renormalize,
        dropout=dropout,
    )
    if output_embedding_init is not None:
        expected_shape = (len(vocabulary), self.decoder.logits.weight.shape[1])
        if tuple(output_embedding_init.shape) != expected_shape:
            raise LinkingTKError(
                f"output_embedding_init has shape {tuple(output_embedding_init.shape)}, "
                f"expected {expected_shape} (len(vocabulary), decoder_hidden_dim)."
            )
        self.decoder.logits.weight.data = output_embedding_init.to(
            self.decoder.logits.weight.dtype
        )

from_checkpoint classmethod

from_checkpoint(
    checkpoint_path: str | Path,
    offsets_path: str | Path,
    lexicon: str = "omw-en:1.4",
    model_name_or_path: str = "bert-large-cased",
    forward_batch_size: int = 16,
) -> EwiserEncoder

Load one of the three EWISER checkpoints released alongside the paper.

Parameters:

Name	Type	Description	Default
checkpoint_path	str | Path	Local path to a checkpoint file (e.g. ewiser.semcor.pt) -- not bundled with this package.	required
offsets_path	str | Path	Local path to that checkpoint's own res/dictionaries/offsets.txt -- also not bundled (see SenseVocabulary.from_offsets_file).	required
lexicon	str	Forwarded to SenseVocabulary.from_offsets_file.	'omw-en:1.4'
model_name_or_path	str	The encoder to pair with the checkpoint's decoder weights. Defaults to "bert-large-cased", matching every released checkpoint's own recorded config.	'bert-large-cased'
forward_batch_size	int	Forwarded to EwiserEncoder.__init__.	16

Returns:

Type	Description
EwiserEncoder	An EwiserEncoder with the checkpoint's decoder weights (FFN
EwiserEncoder	graph adjacency) loaded, paired with a freshly downloaded
EwiserEncoder	frozen encoder (never part of the checkpoint -- see the module
EwiserEncoder	docstring).

Raises:

Type	Description
LinkingTKError	If offsets_path doesn't resolve to the same vocabulary size as the checkpoint's own output layer -- the surest sign it isn't the exact file this checkpoint was trained with.

Source code in src/linkingtk/algorithms/wsd/ewiser.py

@classmethod
def from_checkpoint(
    cls,
    checkpoint_path: str | Path,
    offsets_path: str | Path,
    lexicon: str = "omw-en:1.4",
    model_name_or_path: str = "bert-large-cased",
    forward_batch_size: int = 16,
) -> EwiserEncoder:
    """Load one of the three EWISER checkpoints released alongside the paper.

    Args:
        checkpoint_path: Local path to a checkpoint file (e.g.
            ``ewiser.semcor.pt``) -- not bundled with this package.
        offsets_path: Local path to that checkpoint's own
            ``res/dictionaries/offsets.txt`` -- also not bundled (see
            [SenseVocabulary.from_offsets_file][linkingtk.algorithms.wsd._ewiser_vocab.SenseVocabulary.from_offsets_file]).
        lexicon: Forwarded to `SenseVocabulary.from_offsets_file`.
        model_name_or_path: The encoder to pair with the checkpoint's
            decoder weights. Defaults to ``"bert-large-cased"``,
            matching every released checkpoint's own recorded config.
        forward_batch_size: Forwarded to `EwiserEncoder.__init__`.

    Returns:
        An `EwiserEncoder` with the checkpoint's decoder weights (FFN
        + graph adjacency) loaded, paired with a freshly downloaded
        frozen encoder (never part of the checkpoint -- see the module
        docstring).

    Raises:
        LinkingTKError: If `offsets_path` doesn't resolve to the same
            vocabulary size as the checkpoint's own output layer --
            the surest sign it isn't the exact file this checkpoint
            was trained with.
    """
    data = load_fairseq_checkpoint(checkpoint_path)
    args = data["args"]
    state_dict = data["model"]

    vocabulary = SenseVocabulary.from_offsets_file(offsets_path, lexicon=lexicon)
    vocab_size = int(state_dict["decoder.logits.weight"].shape[0])
    if len(vocabulary) != vocab_size:
        raise LinkingTKError(
            f"offsets_path resolved to a {len(vocabulary)}-entry vocabulary, but "
            f"the checkpoint's own output layer has {vocab_size} entries -- "
            "offsets_path must be the exact res/dictionaries/offsets.txt this "
            "checkpoint was trained with."
        )

    adjacency = torch.sparse_coo_tensor(
        state_dict["decoder.structured_logits.adjacency_pars.0"],
        state_dict["decoder.structured_logits.adjacency_pars.1"],
        size=tuple(state_dict["decoder.structured_logits.adjacency_pars.2"].tolist()),
    )

    encoder = cls(
        model_name_or_path=model_name_or_path,
        vocabulary=vocabulary,
        decoder_hidden_dim=int(state_dict["decoder.linears.0.weight"].shape[0]),
        adjacency=adjacency,
        structured_logits_trainable=bool(
            getattr(args, "decoder_structured_logits_trainable", True)
        ),
        forward_batch_size=forward_batch_size,
    )
    encoder.decoder.load_state_dict(
        {key[len("decoder.") :]: value for key, value in state_dict.items()}
    )
    return encoder

score

score(pairs: list[tuple[Entity, Entity]]) -> torch.Tensor

Return a (len(pairs),) tensor: each pair's raw sense-classification logit.

Unlike a cross-encoder's score, this does not run one forward pass per pair. It groups pairs by each mention's unique sentence (see mention_sentence_and_span), runs the encoder+decoder once per distinct sentence (batched forward_batch_size sentences at a time), and reads each pair's score off as logits[mention_word_index, candidate_sense_index] from that sentence's shared output. A candidate whose sense id isn't in vocabulary, or a mention whose span doesn't resolve to a word position, scores -inf (never spuriously wins a match).

Each chunk's [num_words, len(vocabulary)] logits are consumed (scores extracted) and released before the next chunk is encoded, rather than holding every distinct sentence's full logits matrix in memory for the whole call -- at len(vocabulary) ~117k, a large evaluation set's worth of these held simultaneously (thousands of sentences at once, e.g. UFSAC's "ALL" split) is enough to exhaust even a 24GB GPU; confirmed by hitting exactly that CUDA out of memory running examples/ewiser_reproduction.py before this was fixed.

Source code in src/linkingtk/algorithms/wsd/ewiser.py

def score(self, pairs: list[tuple[Entity, Entity]]) -> torch.Tensor:
    """Return a ``(len(pairs),)`` tensor: each pair's raw sense-classification logit.

    Unlike a cross-encoder's `score`, this does **not** run one
    forward pass per pair. It groups `pairs` by each mention's unique
    sentence (see `mention_sentence_and_span`), runs the
    encoder+decoder once per distinct sentence (batched
    `forward_batch_size` sentences at a time), and reads each pair's
    score off as ``logits[mention_word_index, candidate_sense_index]``
    from that sentence's shared output. A candidate whose sense id
    isn't in `vocabulary`, or a mention whose span doesn't resolve to
    a word position, scores ``-inf`` (never spuriously wins a match).

    Each chunk's ``[num_words, len(vocabulary)]`` logits are consumed
    (scores extracted) and released before the next chunk is encoded,
    rather than holding every distinct sentence's full logits matrix
    in memory for the whole call -- at `len(vocabulary)` ~117k, a
    large evaluation set's worth of these held simultaneously
    (thousands of sentences at once, e.g. UFSAC's "ALL" split) is
    enough to exhaust even a 24GB GPU; confirmed by hitting exactly
    that `CUDA out of memory` running `examples/ewiser_reproduction.py`
    before this was fixed.
    """
    if not pairs:
        return torch.empty(0)

    sentences: list[str] = []
    sentence_index: dict[str, int] = {}
    mention_word_index: dict[int, int | None] = {}
    pairs_by_sentence: dict[int, list[int]] = {}
    for pair_index, (mention, _sense) in enumerate(pairs):
        text, start, end = mention_sentence_and_span(mention)
        if text not in sentence_index:
            sentence_index[text] = len(sentences)
            sentences.append(text)
        if id(mention) not in mention_word_index:
            tokens = whitespace_tokenize_with_offsets(text)
            mention_word_index[id(mention)] = word_index_for_span(tokens, start, end)
        pairs_by_sentence.setdefault(sentence_index[text], []).append(pair_index)

    scores: list[torch.Tensor | None] = [None] * len(pairs)
    for start in range(0, len(sentences), self.forward_batch_size):
        chunk_sentences = sentences[start : start + self.forward_batch_size]
        chunk_logits = self._encode_chunk(chunk_sentences)
        for offset, logits in enumerate(chunk_logits):
            for pair_index in pairs_by_sentence.get(start + offset, []):
                mention, sense = pairs[pair_index]
                word_index = mention_word_index[id(mention)]
                sense_index = self.vocabulary.index_for(sense.id)
                if word_index is None or sense_index is None:
                    scores[pair_index] = logits.new_tensor(float("-inf"))
                else:
                    scores[pair_index] = logits[word_index, sense_index]

    assert all(score is not None for score in scores)  # every pair index was filled above
    return torch.stack(scores)  # type: ignore[arg-type]

EwiserLinker

EwiserLinker(
    model: EwiserEncoder,
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

WSD linking scored by an EwiserEncoder.

Unlike GlossBertLinker, this takes an already-constructed EwiserEncoder rather than a bare model_name_or_path string -- EwiserEncoder needs a SenseVocabulary (and usually an adjacency) at construction time, with no sane string-only default.

encoder = EwiserEncoder.from_checkpoint("ewiser.semcor.pt", "offsets.txt")
linker = EwiserLinker(encoder)
results = linker.link(mentions, senses)

Or, for from-scratch/continued training, construct EwiserEncoder directly and train it via EwiserTrainer.

Parameters:

Name	Type	Description	Default
model	EwiserEncoder	The EwiserEncoder to score candidates with.	required
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER

Source code in src/linkingtk/algorithms/wsd/ewiser.py

def __init__(self, model: EwiserEncoder, matching: Matcher = DEFAULT_MATCHER) -> None:
    self.model = model
    self.matching = matching

score_candidates

score_candidates(
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]

Blocked candidates per source entity, scored but not yet matched.

Satisfies CandidateScorer -- what link() itself builds internally, exposed so LlmRerankerLinker (#23) can re-rank a narrowed top-k instead of every blocked pair.

If dataset2 is a WnEntitySource whose lexicon differs from model.vocabulary's own, candidate ids are translated into the vocabulary's lexicon via synset_id_via_ili before scoring (see _translate_for_scoring's own docstring) -- without this, every candidate would fail SenseVocabulary.index_for and silently score -inf (#67). Candidate ids in the returned dict are still dataset2's own (untranslated) ids -- the translation is scoring-only, invisible to the caller.

Source code in src/linkingtk/algorithms/wsd/ewiser.py

def score_candidates(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]:
    """Blocked candidates per source entity, scored but not yet matched.

    Satisfies [CandidateScorer][linkingtk.algorithms.llm_reranker.CandidateScorer]
    -- what `link()` itself builds internally, exposed so
    [LlmRerankerLinker][linkingtk.algorithms.llm_reranker.LlmRerankerLinker]
    (#23) can re-rank a narrowed top-k instead of every blocked pair.

    If `dataset2` is a `WnEntitySource` whose `lexicon` differs from
    `model.vocabulary`'s own, candidate ids are translated into the
    vocabulary's lexicon via
    [synset_id_via_ili][linkingtk.sources.wn.synset_id_via_ili] before
    scoring (see `_translate_for_scoring`'s own docstring) -- without
    this, every candidate would fail
    [SenseVocabulary.index_for][linkingtk.algorithms.wsd._ewiser_vocab.SenseVocabulary.index_for]
    and silently score ``-inf`` (#67). Candidate ids in the returned
    dict are still `dataset2`'s own (untranslated) ids -- the
    translation is scoring-only, invisible to the caller.
    """
    pairs = list(blocking.candidate_pairs(dataset1, dataset2))
    if not pairs:
        return {}

    scoring_pairs = self._translate_for_scoring(pairs, dataset2)

    with torch.no_grad():
        self.model.eval()
        scores = self.model.score(scoring_pairs)
        self.model.train()

    candidates_by_source: dict[str, list[tuple[str, float]]] = defaultdict(list)
    for (entity1, entity2), score in zip(pairs, scores.tolist(), strict=True):
        candidates_by_source[entity1.id].append((entity2.id, score))

    return candidates_by_source

GlossBertEncoder

GlossBertEncoder(
    model_name_or_path: str = "bert-base-uncased",
    max_length: int = 512,
    forward_batch_size: int = 32,
)

Bases: Module

BertForSequenceClassification-based cross-encoder scorer for WSD.

Satisfies CrossEncoderModel. Unlike ReFinEDEncoder/ BlinkEncoder (generic AutoModel, any backbone), this is hardcoded to BERT -- GlossBERT's own reference implementation is BERT-specific, and this is exactly what lets model_name_or_path point at the paper's own published checkpoint directory as well as a fresh bert-base-uncased.

Parameters:

Name	Type	Description	Default
model_name_or_path	str	A Hugging Face model id, or a local directory (e.g. the paper's published Sent_CLS_WS checkpoint, extracted). Defaults to bert-base-uncased, matching the paper. Pass a tiny test model (e.g. hf-internal-testing/tiny-random-BertForSequenceClassification) for fast, network-light tests.	'bert-base-uncased'
max_length	int	Maximum token length for the tokenized (context, gloss) pair (truncated). The paper uses 512.	512
forward_batch_size	int	Chunk size score() tokenizes/forwards at once, regardless of how many pairs it's called with -- GlossBertLinker.link() and CrossEncoderTrainer's eval step both pass every one of a mention's (or dataset's) candidate pairs to score() in a single call, which at max_length=512 would exhaust GPU memory well before it reached even a few hundred real-corpus pairs if forwarded as one batch. Gradients still flow through every chunk normally (this only bounds how many pairs share one forward pass, not how many contribute to a loss) -- unrelated to TrainingArguments.batch_size, which bounds how many pairs one gradient step trains on.	32

Source code in src/linkingtk/algorithms/wsd/glossbert.py

def __init__(
    self,
    model_name_or_path: str = "bert-base-uncased",
    max_length: int = 512,
    forward_batch_size: int = 32,
) -> None:
    super().__init__()
    from transformers import BertForSequenceClassification, BertTokenizer

    self.tokenizer = BertTokenizer.from_pretrained(model_name_or_path)
    self.model = BertForSequenceClassification.from_pretrained(model_name_or_path, num_labels=2)
    self.max_length = max_length
    self.forward_batch_size = forward_batch_size

from_checkpoint classmethod

from_checkpoint(
    checkpoint_path: str | Path,
    max_length: int = 512,
    forward_batch_size: int = 32,
) -> GlossBertEncoder

Load the paper's own published checkpoint directory (e.g. Sent_CLS_WS, extracted).

Same entry point name as EwiserEncoder.from_checkpoint/ EscEncoder.from_checkpoint, for a consistent checkpoint-loading interface across WSD linkers (issue #64) -- unlike those two, no format translation is needed here (the paper's checkpoint is already a plain Hugging Face config.json/pytorch_model.bin/vocab.txt directory, exactly what model_name_or_path already accepts), so this is a thin wrapper around the constructor.

Parameters:

Name	Type	Description	Default
checkpoint_path	str | Path	Local path to the extracted checkpoint directory -- not bundled with this package.	required
max_length	int	Forwarded to GlossBertEncoder.__init__.	512
forward_batch_size	int	Forwarded to GlossBertEncoder.__init__.	32

Returns:

Type	Description
GlossBertEncoder	A GlossBertEncoder with the checkpoint's weights loaded.

Source code in src/linkingtk/algorithms/wsd/glossbert.py

@classmethod
def from_checkpoint(
    cls,
    checkpoint_path: str | Path,
    max_length: int = 512,
    forward_batch_size: int = 32,
) -> GlossBertEncoder:
    """Load the paper's own published checkpoint directory (e.g. `Sent_CLS_WS`, extracted).

    Same entry point name as
    [EwiserEncoder.from_checkpoint][linkingtk.algorithms.wsd.ewiser.EwiserEncoder.from_checkpoint]/
    [EscEncoder.from_checkpoint][linkingtk.algorithms.wsd.esc.EscEncoder.from_checkpoint],
    for a consistent checkpoint-loading interface across WSD linkers
    (issue #64) -- unlike those two, no format translation is needed
    here (the paper's checkpoint is already a plain Hugging Face
    `config.json`/`pytorch_model.bin`/`vocab.txt` directory, exactly
    what `model_name_or_path` already accepts), so this is a thin
    wrapper around the constructor.

    Args:
        checkpoint_path: Local path to the extracted checkpoint
            directory -- not bundled with this package.
        max_length: Forwarded to `GlossBertEncoder.__init__`.
        forward_batch_size: Forwarded to `GlossBertEncoder.__init__`.

    Returns:
        A `GlossBertEncoder` with the checkpoint's weights loaded.
    """
    return cls(
        model_name_or_path=str(checkpoint_path),
        max_length=max_length,
        forward_batch_size=forward_batch_size,
    )

score

score(pairs: list[tuple[Entity, Entity]]) -> torch.Tensor

Return a (len(pairs),) tensor: the class-1-minus-class-0 logit margin per pair.

Source code in src/linkingtk/algorithms/wsd/glossbert.py

def score(self, pairs: list[tuple[Entity, Entity]]) -> torch.Tensor:
    """Return a ``(len(pairs),)`` tensor: the class-1-minus-class-0 logit margin per pair."""
    margins = [
        self._score_chunk(pairs[start : start + self.forward_batch_size])
        for start in range(0, len(pairs), self.forward_batch_size)
    ]
    return torch.cat(margins) if margins else torch.empty(0)

GlossBertLinker

GlossBertLinker(
    model_name_or_path: str = "bert-base-uncased",
    max_length: int = 512,
    forward_batch_size: int = 32,
    matching: Matcher = DEFAULT_MATCHER,
    model: GlossBertEncoder | None = None,
)

Bases: BaseLinker

WSD linking scored by a GlossBertEncoder.

Unlike every other linker in this repo, this has no fit() -- self.model is a plain CrossEncoderModel; train it via CrossEncoderTrainer directly, then call link():

linker = GlossBertLinker()
trainer = CrossEncoderTrainer(model=linker.model, args=args, train_data=pairs)
trainer.train()
results = linker.link(mentions, senses)

Or skip training entirely and load the paper's own published checkpoint, either directly (GlossBertLinker(model_name_or_path="/path/to/Sent_CLS_WS")) or, matching EwiserLinker/ EscLinker's pattern, by passing an already-built encoder:

encoder = GlossBertEncoder.from_checkpoint("/path/to/Sent_CLS_WS")
linker = GlossBertLinker(model=encoder)

Parameters:

Name	Type	Description	Default
model_name_or_path	str	Forwarded to GlossBertEncoder. Ignored if model is given.	'bert-base-uncased'
max_length	int	Forwarded to GlossBertEncoder. Ignored if model is given.	512
forward_batch_size	int	Forwarded to GlossBertEncoder. Ignored if model is given.	32
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
model	GlossBertEncoder | None	An already-built GlossBertEncoder (e.g. from GlossBertEncoder.from_checkpoint), used as-is instead of constructing one from model_name_or_path.	None

Source code in src/linkingtk/algorithms/wsd/glossbert.py

def __init__(
    self,
    model_name_or_path: str = "bert-base-uncased",
    max_length: int = 512,
    forward_batch_size: int = 32,
    matching: Matcher = DEFAULT_MATCHER,
    model: GlossBertEncoder | None = None,
) -> None:
    self.model = (
        model
        if model is not None
        else GlossBertEncoder(model_name_or_path, max_length, forward_batch_size)
    )
    self.matching = matching

score_candidates

score_candidates(
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]

Blocked candidates per source entity, scored but not yet matched.

Satisfies CandidateScorer -- what link() itself builds internally, exposed so LlmRerankerLinker (#23) can re-rank a narrowed top-k instead of every blocked pair.

Source code in src/linkingtk/algorithms/wsd/glossbert.py

def score_candidates(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]:
    """Blocked candidates per source entity, scored but not yet matched.

    Satisfies [CandidateScorer][linkingtk.algorithms.llm_reranker.CandidateScorer]
    -- what `link()` itself builds internally, exposed so
    [LlmRerankerLinker][linkingtk.algorithms.llm_reranker.LlmRerankerLinker]
    (#23) can re-rank a narrowed top-k instead of every blocked pair.
    """
    pairs = list(blocking.candidate_pairs(dataset1, dataset2))
    if not pairs:
        return {}

    with torch.no_grad():
        self.model.eval()
        scores = self.model.score(pairs)
        self.model.train()

    candidates_by_source: dict[str, list[tuple[str, float]]] = defaultdict(list)
    for (entity1, entity2), score in zip(pairs, scores.tolist(), strict=True):
        candidates_by_source[entity1.id].append((entity2.id, score))

    return candidates_by_source

LeskLinker

LeskLinker()

Bases: StringSimilarityLinker

WSD linker that scores candidate senses by context/gloss word overlap.

A preconfigured StringSimilarityLinker (source_field="context", target_field="description", metric="word_overlap") — see that class for the ranking and alternatives behavior. Kept as its own named class since Lesk is a well-known algorithm in its own right.

Source code in src/linkingtk/algorithms/wsd/lesk.py

def __init__(self) -> None:
    super().__init__(source_field="context", target_field="description", metric="word_overlap")

lesk

Lesk-style Word Sense Disambiguation linker.

References

Lesk, M. (1986). Automatic sense disambiguation using machine readable dictionaries: how to tell a pine cone from an ice cream cone. In Proceedings of the 5th annual international conference on Systems documentation (SIGDOC '86).

LeskLinker

LeskLinker()

Bases: StringSimilarityLinker

WSD linker that scores candidate senses by context/gloss word overlap.

A preconfigured StringSimilarityLinker (source_field="context", target_field="description", metric="word_overlap") — see that class for the ranking and alternatives behavior. Kept as its own named class since Lesk is a well-known algorithm in its own right.

Source code in src/linkingtk/algorithms/wsd/lesk.py

def __init__(self) -> None:
    super().__init__(source_field="context", target_field="description", metric="word_overlap")

GlossBERT

glossbert

GlossBERT Word Sense Disambiguation.

GlossBERT (Huang et al., "GlossBERT: BERT for Word Sense Disambiguation with Gloss Knowledge", EMNLP 2019, https://github.com/HSLCY/GlossBERT) is a true cross-encoder: it concatenates a mention's context sentence and a candidate sense's gloss into one BERT input ([CLS] context [SEP] gloss [SEP]) and classifies the pair as matching or not, unlike ReFinEDLinker/ BlinkLinker's bi-encoders (independently encode() each side, then compare embeddings). Trained via CrossEncoderTrainer, built specifically for this shape.

This implements the paper's best-reported configuration, GlossBERT(Sent-CLS-WS), ported directly from the reference implementation's own data-preparation code (preparation/generate_sent_cls_ws.py), not just the paper text:

• Candidates are every WordNet sense of the mention's surface lemma, across every part of speech (not filtered to the mention's own gold POS) -- matches this repo's existing ExactMatch/EntitySource.search() design already.
• The gloss side of each pair is f"{mention_surface_text} : {gloss}" -- the mention's own surface form, not the sense's dictionary lemma, confirmed against the reference's wordnet/index.sense.gloss.
• "Weak supervision" (the -WS in the paper's best model) wraps the mention span in the context sentence with literal " quote characters -- not special tokens. Confirmed by the paper's own published checkpoint (state dict bert.* + classifier.weight (2, 768), vocab size 30522 -- stock bert-base-uncased, nothing added).
• The classification head is BertForSequenceClassification-shaped (num_labels=2, [CLS]-pooled, softmax cross-entropy in the original) -- score() returns the class-1-minus-class-0 logit margin, which is mathematically identical to that 2-class softmax cross-entropy (so CrossEncoderTrainer's generic binary-cross-entropy-on-the-margin loss trains it exactly the same way, without a GlossBERT-specific loss function).

Because the architecture matches exactly, GlossBertEncoder also loads the paper's own published checkpoint directly via GlossBertEncoder.from_checkpoint (same entry point name as EwiserEncoder.from_checkpoint/ EscEncoder.from_checkpoint, issue #64) -- see examples/glossbert_reproduction.py for a real evaluation against it.

See GlossBertEncoder and GlossBertLinker.

GlossBertEncoder

GlossBertEncoder(
    model_name_or_path: str = "bert-base-uncased",
    max_length: int = 512,
    forward_batch_size: int = 32,
)

Bases: Module

BertForSequenceClassification-based cross-encoder scorer for WSD.

Satisfies CrossEncoderModel. Unlike ReFinEDEncoder/ BlinkEncoder (generic AutoModel, any backbone), this is hardcoded to BERT -- GlossBERT's own reference implementation is BERT-specific, and this is exactly what lets model_name_or_path point at the paper's own published checkpoint directory as well as a fresh bert-base-uncased.

Parameters:

Name	Type	Description	Default
model_name_or_path	str	A Hugging Face model id, or a local directory (e.g. the paper's published Sent_CLS_WS checkpoint, extracted). Defaults to bert-base-uncased, matching the paper. Pass a tiny test model (e.g. hf-internal-testing/tiny-random-BertForSequenceClassification) for fast, network-light tests.	'bert-base-uncased'
max_length	int	Maximum token length for the tokenized (context, gloss) pair (truncated). The paper uses 512.	512
forward_batch_size	int	Chunk size score() tokenizes/forwards at once, regardless of how many pairs it's called with -- GlossBertLinker.link() and CrossEncoderTrainer's eval step both pass every one of a mention's (or dataset's) candidate pairs to score() in a single call, which at max_length=512 would exhaust GPU memory well before it reached even a few hundred real-corpus pairs if forwarded as one batch. Gradients still flow through every chunk normally (this only bounds how many pairs share one forward pass, not how many contribute to a loss) -- unrelated to TrainingArguments.batch_size, which bounds how many pairs one gradient step trains on.	32

Source code in src/linkingtk/algorithms/wsd/glossbert.py

def __init__(
    self,
    model_name_or_path: str = "bert-base-uncased",
    max_length: int = 512,
    forward_batch_size: int = 32,
) -> None:
    super().__init__()
    from transformers import BertForSequenceClassification, BertTokenizer

    self.tokenizer = BertTokenizer.from_pretrained(model_name_or_path)
    self.model = BertForSequenceClassification.from_pretrained(model_name_or_path, num_labels=2)
    self.max_length = max_length
    self.forward_batch_size = forward_batch_size

from_checkpoint classmethod

from_checkpoint(
    checkpoint_path: str | Path,
    max_length: int = 512,
    forward_batch_size: int = 32,
) -> GlossBertEncoder

Load the paper's own published checkpoint directory (e.g. Sent_CLS_WS, extracted).

Same entry point name as EwiserEncoder.from_checkpoint/ EscEncoder.from_checkpoint, for a consistent checkpoint-loading interface across WSD linkers (issue #64) -- unlike those two, no format translation is needed here (the paper's checkpoint is already a plain Hugging Face config.json/pytorch_model.bin/vocab.txt directory, exactly what model_name_or_path already accepts), so this is a thin wrapper around the constructor.

Parameters:

Name	Type	Description	Default
checkpoint_path	str | Path	Local path to the extracted checkpoint directory -- not bundled with this package.	required
max_length	int	Forwarded to GlossBertEncoder.__init__.	512
forward_batch_size	int	Forwarded to GlossBertEncoder.__init__.	32

Returns:

Type	Description
GlossBertEncoder	A GlossBertEncoder with the checkpoint's weights loaded.

Source code in src/linkingtk/algorithms/wsd/glossbert.py

@classmethod
def from_checkpoint(
    cls,
    checkpoint_path: str | Path,
    max_length: int = 512,
    forward_batch_size: int = 32,
) -> GlossBertEncoder:
    """Load the paper's own published checkpoint directory (e.g. `Sent_CLS_WS`, extracted).

    Same entry point name as
    [EwiserEncoder.from_checkpoint][linkingtk.algorithms.wsd.ewiser.EwiserEncoder.from_checkpoint]/
    [EscEncoder.from_checkpoint][linkingtk.algorithms.wsd.esc.EscEncoder.from_checkpoint],
    for a consistent checkpoint-loading interface across WSD linkers
    (issue #64) -- unlike those two, no format translation is needed
    here (the paper's checkpoint is already a plain Hugging Face
    `config.json`/`pytorch_model.bin`/`vocab.txt` directory, exactly
    what `model_name_or_path` already accepts), so this is a thin
    wrapper around the constructor.

    Args:
        checkpoint_path: Local path to the extracted checkpoint
            directory -- not bundled with this package.
        max_length: Forwarded to `GlossBertEncoder.__init__`.
        forward_batch_size: Forwarded to `GlossBertEncoder.__init__`.

    Returns:
        A `GlossBertEncoder` with the checkpoint's weights loaded.
    """
    return cls(
        model_name_or_path=str(checkpoint_path),
        max_length=max_length,
        forward_batch_size=forward_batch_size,
    )

score

score(pairs: list[tuple[Entity, Entity]]) -> torch.Tensor

Return a (len(pairs),) tensor: the class-1-minus-class-0 logit margin per pair.

Source code in src/linkingtk/algorithms/wsd/glossbert.py

def score(self, pairs: list[tuple[Entity, Entity]]) -> torch.Tensor:
    """Return a ``(len(pairs),)`` tensor: the class-1-minus-class-0 logit margin per pair."""
    margins = [
        self._score_chunk(pairs[start : start + self.forward_batch_size])
        for start in range(0, len(pairs), self.forward_batch_size)
    ]
    return torch.cat(margins) if margins else torch.empty(0)

GlossBertLinker

GlossBertLinker(
    model_name_or_path: str = "bert-base-uncased",
    max_length: int = 512,
    forward_batch_size: int = 32,
    matching: Matcher = DEFAULT_MATCHER,
    model: GlossBertEncoder | None = None,
)

Bases: BaseLinker

WSD linking scored by a GlossBertEncoder.

Unlike every other linker in this repo, this has no fit() -- self.model is a plain CrossEncoderModel; train it via CrossEncoderTrainer directly, then call link():

linker = GlossBertLinker()
trainer = CrossEncoderTrainer(model=linker.model, args=args, train_data=pairs)
trainer.train()
results = linker.link(mentions, senses)

Or skip training entirely and load the paper's own published checkpoint, either directly (GlossBertLinker(model_name_or_path="/path/to/Sent_CLS_WS")) or, matching EwiserLinker/ EscLinker's pattern, by passing an already-built encoder:

encoder = GlossBertEncoder.from_checkpoint("/path/to/Sent_CLS_WS")
linker = GlossBertLinker(model=encoder)

Parameters:

Name	Type	Description	Default
model_name_or_path	str	Forwarded to GlossBertEncoder. Ignored if model is given.	'bert-base-uncased'
max_length	int	Forwarded to GlossBertEncoder. Ignored if model is given.	512
forward_batch_size	int	Forwarded to GlossBertEncoder. Ignored if model is given.	32
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER
model	GlossBertEncoder | None	An already-built GlossBertEncoder (e.g. from GlossBertEncoder.from_checkpoint), used as-is instead of constructing one from model_name_or_path.	None

Source code in src/linkingtk/algorithms/wsd/glossbert.py

def __init__(
    self,
    model_name_or_path: str = "bert-base-uncased",
    max_length: int = 512,
    forward_batch_size: int = 32,
    matching: Matcher = DEFAULT_MATCHER,
    model: GlossBertEncoder | None = None,
) -> None:
    self.model = (
        model
        if model is not None
        else GlossBertEncoder(model_name_or_path, max_length, forward_batch_size)
    )
    self.matching = matching

score_candidates

score_candidates(
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]

Blocked candidates per source entity, scored but not yet matched.

Satisfies CandidateScorer -- what link() itself builds internally, exposed so LlmRerankerLinker (#23) can re-rank a narrowed top-k instead of every blocked pair.

Source code in src/linkingtk/algorithms/wsd/glossbert.py

def score_candidates(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]:
    """Blocked candidates per source entity, scored but not yet matched.

    Satisfies [CandidateScorer][linkingtk.algorithms.llm_reranker.CandidateScorer]
    -- what `link()` itself builds internally, exposed so
    [LlmRerankerLinker][linkingtk.algorithms.llm_reranker.LlmRerankerLinker]
    (#23) can re-rank a narrowed top-k instead of every blocked pair.
    """
    pairs = list(blocking.candidate_pairs(dataset1, dataset2))
    if not pairs:
        return {}

    with torch.no_grad():
        self.model.eval()
        scores = self.model.score(pairs)
        self.model.train()

    candidates_by_source: dict[str, list[tuple[str, float]]] = defaultdict(list)
    for (entity1, entity2), score in zip(pairs, scores.tolist(), strict=True):
        candidates_by_source[entity1.id].append((entity2.id, score))

    return candidates_by_source

EWISER

ewiser

EWISER Word Sense Disambiguation.

EWISER (Bevilacqua & Navigli, "Breaking Through the 80% Glass Ceiling: Raising the State of the Art in Word Sense Disambiguation by Incorporating Knowledge Graph Information", ACL 2020, https://github.com/SapienzaNLP/ewiser) is a full-inventory sense classifier, architecturally very different from GlossBertLinker's pairwise cross-encoder: it encodes a whole sentence once with a frozen BERT, projects every word's hidden state through a small FFN to raw logits over the entire WordNet sense inventory, then applies one linear WordNet-relation graph-propagation step to those logits (see StructuredLogits) before a mention's candidate senses are read off as indices into that one shared output vector -- there's no gloss text involved in the forward pass at all. EwiserEncoder.score(pairs) still returns the same (len(pairs),) -> torch.Tensor shape GlossBertEncoder.score does (so EwiserLinker's link() can reuse the same blocking -> score -> group -> match pattern every linker in this package uses), but computes it very differently internally -- see EwiserEncoder.score's own docstring.

Reimplemented from the paper and from the tensor shapes/keys observed directly in the three checkpoints released alongside it (not ported from EWISER's own source, which is CC-BY-NC-SA 4.0 licensed -- incompatible with redistributing literal code under this package's Apache-2.0 license). Confirmed by reading the reference's own model/criterion code directly (not guessed):

• The BERT encoder is never fine-tuned -- every released checkpoint's own state dict holds only 11 tensors, all in the small decoder (2-layer FFN + the graph module); the encoder's weights come straight from a frozen bert-large-cased, loaded fresh from Hugging Face, not from the checkpoint.
• The encoder input is the sum of the last 4 BERT hidden-state layers (context_embeddings_layers=[-4,-3,-2,-1]), not just last_hidden_state -- despite every released checkpoint's own recorded context_embeddings_use_all_hidden=False, which would suggest otherwise. Confirmed this flag is dead: the reference's own TaggerModel.build_model computes use_all_hidden from len(args.context_embeddings_type) > 1 (a variable-name mix-up -- context_embeddings_type is the string "bert", not the actual layer list), which for "bert" (length 4) is always True, unconditionally. See EwiserEncoder's own num_summed_layers docstring for how this was traced.
• Each whitespace-split word is subword-tokenized in isolation via a from-scratch WordPiece implementation (wordpiece_tokenize_words), not a single tokenizer(words, is_split_into_words=True) call -- confirmed, by tracing a real checkpoint-reproduction gap layer by layer down to one divergent subtoken, that is_split_into_words=True still runs BERT's punctuation-splitting pre-tokenizer on each word first (splitting e.g. UFSAC's own single-token "Oct." into separate "Oct"/"." pre-tokens before wordpiece), producing a standalone "." token where EWISER's own reference (wordpiece-tokenizing the literal string "Oct." directly, no pre-splitting) produces a "##." continuation piece instead -- a different vocabulary entry with an unrelated embedding. See wordpiece_tokenize_words's own docstring for the full diagnosis.
• EwiserEncoder.from_checkpoint reads a checkpoint's own baked-in sparse WordNet adjacency directly from its state dict -- no graph construction needed for checkpoint-based inference (see build_relation_adjacency for the from-scratch-training path that does need one). Unpickling a checkpoint file requires tolerating references to training-only classes from its original (pre-release) package name, qbert -- see load_fairseq_checkpoint.

See EwiserEncoder and EwiserLinker.

EwiserEncoder

EwiserEncoder(
    model_name_or_path: str = "bert-large-cased",
    vocabulary: SenseVocabulary | None = None,
    decoder_hidden_dim: int = 512,
    adjacency: Tensor | None = None,
    structured_logits_trainable: bool = True,
    structured_logits_renormalize: bool = False,
    freeze_encoder: bool = True,
    dropout: float = 0.1,
    num_summed_layers: int = 4,
    max_length: int = 512,
    forward_batch_size: int = 16,
    output_embedding_init: Tensor | None = None,
)

Bases: Module

Frozen-BERT full-inventory sense classifier with a WordNet-graph output layer.

Parameters:

Name	Type	Description	Default
model_name_or_path	str	A Hugging Face encoder id, e.g. "bert-large-cased" (the published checkpoints' own encoder). Pass a tiny test model for fast, network-light tests.	'bert-large-cased'
vocabulary	SenseVocabulary | None	The SenseVocabulary the output layer classifies against.	None
decoder_hidden_dim	int	Width of the decoder's single hidden layer. Defaults to 512, matching the reference.	512
adjacency	Tensor | None	An optional sparse [len(vocabulary), len(vocabulary)] WordNet-relation adjacency (from build_relation_adjacency, or a checkpoint's own baked-in one via from_checkpoint). Wires up the graph-propagation step (StructuredLogits) if given; the model degrades to a plain (non-graph-aware) classifier if omitted.	None
structured_logits_trainable	bool	Whether the graph's edge weights are updated during training.	True
structured_logits_renormalize	bool	Forwarded to StructuredLogits.	False
freeze_encoder	bool	Whether the BERT encoder's parameters are frozen. Defaults to True, matching the paper -- exposed as a real constructor flag rather than hardcoded since nothing else in this package forbids fine-tuning it, but the default preserves fidelity to EWISER's own recipe (see the module docstring).	True
dropout	float	Decoder dropout rate. Checkpoint loading doesn't restore this (dropout has no learned parameters), so it only affects EwiserTrainer-driven training.	0.1
output_embedding_init	Tensor | None	An optional [len(vocabulary), decoder_hidden_dim] tensor to initialize decoder.logits.weight from, in place of its default random init -- see load_synset_centroid_vectors/ build_synset_centroid_vectors_from_lmms to build one. This is what makes EwiserTrainer.freeze_output_epochs meaningful -- freezing a pretrained output layer for a few epochs protects it from noisy early gradients (the paper's own motivation); freezing the layer's plain default random init has no such benefit.	None
num_summed_layers	int	The encoder input is the sum of the last num_summed_layers BERT hidden-state layers, not just the final layer -- confirmed against the reference's own TaggerModel.build_model, whose ``use_all_hidden=len(layers) 1readsargs.context_embeddings_type(a string, e.g."bert") where it evidently meant to readargs.context_embeddings_layers(the actual 4-entry layer list) -- sincelen("bert") == 4 > 1regardless of the real layer count, this makesuse_all_hidden**always** true for every BERT-backed checkpoint, silently overriding the separate, correctly-namedcontext_embeddings_use_all_hiddenconfig flag (which reads `False` on all three released checkpoints and would, if honored, sum only the single last layer). A reference quirk, not a reference bug we get to ignore: it *is* what these checkpoints were actually trained against, confirmed by tracing a checkpoint-reproduction gap down to this exact vector (norm ~82 summed vs. ~18 single-layer, cosine 0.69 apart -- unmistakably different vectors, not numerical noise). Defaults to4, matching every released checkpoint's owncontext_embeddings_layers`` length.	4
max_length	int	Maximum subword sequence length per sentence (truncated).	512
forward_batch_size	int	Number of distinct sentences score() batches through the encoder+decoder per forward pass -- unrelated to how many (mention, candidate_sense) pairs score() is called with (which may share far fewer distinct sentences, see score's own docstring).	16

Raises:

Type	Description
LinkingTKError	If vocabulary isn't given.

Source code in src/linkingtk/algorithms/wsd/ewiser.py

def __init__(
    self,
    model_name_or_path: str = "bert-large-cased",
    vocabulary: SenseVocabulary | None = None,
    decoder_hidden_dim: int = 512,
    adjacency: torch.Tensor | None = None,
    structured_logits_trainable: bool = True,
    structured_logits_renormalize: bool = False,
    freeze_encoder: bool = True,
    dropout: float = 0.1,
    num_summed_layers: int = 4,
    max_length: int = 512,
    forward_batch_size: int = 16,
    output_embedding_init: torch.Tensor | None = None,
) -> None:
    super().__init__()
    if vocabulary is None:
        raise LinkingTKError("EwiserEncoder requires a `vocabulary` (see SenseVocabulary).")
    from transformers import AutoModel, AutoTokenizer

    self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    self.encoder = AutoModel.from_pretrained(model_name_or_path)
    self.freeze_encoder = freeze_encoder
    for parameter in self.encoder.parameters():
        parameter.requires_grad = not freeze_encoder

    self.vocabulary = vocabulary
    self.num_summed_layers = num_summed_layers
    self.max_length = max_length
    self.forward_batch_size = forward_batch_size
    self.decoder = EwiserDecoder(
        input_dim=self.encoder.config.hidden_size,
        hidden_dim=decoder_hidden_dim,
        vocab_size=len(vocabulary),
        adjacency=adjacency,
        structured_logits_trainable=structured_logits_trainable,
        structured_logits_renormalize=structured_logits_renormalize,
        dropout=dropout,
    )
    if output_embedding_init is not None:
        expected_shape = (len(vocabulary), self.decoder.logits.weight.shape[1])
        if tuple(output_embedding_init.shape) != expected_shape:
            raise LinkingTKError(
                f"output_embedding_init has shape {tuple(output_embedding_init.shape)}, "
                f"expected {expected_shape} (len(vocabulary), decoder_hidden_dim)."
            )
        self.decoder.logits.weight.data = output_embedding_init.to(
            self.decoder.logits.weight.dtype
        )

from_checkpoint classmethod

from_checkpoint(
    checkpoint_path: str | Path,
    offsets_path: str | Path,
    lexicon: str = "omw-en:1.4",
    model_name_or_path: str = "bert-large-cased",
    forward_batch_size: int = 16,
) -> EwiserEncoder

Load one of the three EWISER checkpoints released alongside the paper.

Parameters:

Name	Type	Description	Default
checkpoint_path	str | Path	Local path to a checkpoint file (e.g. ewiser.semcor.pt) -- not bundled with this package.	required
offsets_path	str | Path	Local path to that checkpoint's own res/dictionaries/offsets.txt -- also not bundled (see SenseVocabulary.from_offsets_file).	required
lexicon	str	Forwarded to SenseVocabulary.from_offsets_file.	'omw-en:1.4'
model_name_or_path	str	The encoder to pair with the checkpoint's decoder weights. Defaults to "bert-large-cased", matching every released checkpoint's own recorded config.	'bert-large-cased'
forward_batch_size	int	Forwarded to EwiserEncoder.__init__.	16

Returns:

Type	Description
EwiserEncoder	An EwiserEncoder with the checkpoint's decoder weights (FFN
EwiserEncoder	graph adjacency) loaded, paired with a freshly downloaded
EwiserEncoder	frozen encoder (never part of the checkpoint -- see the module
EwiserEncoder	docstring).

Raises:

Type	Description
LinkingTKError	If offsets_path doesn't resolve to the same vocabulary size as the checkpoint's own output layer -- the surest sign it isn't the exact file this checkpoint was trained with.

Source code in src/linkingtk/algorithms/wsd/ewiser.py

@classmethod
def from_checkpoint(
    cls,
    checkpoint_path: str | Path,
    offsets_path: str | Path,
    lexicon: str = "omw-en:1.4",
    model_name_or_path: str = "bert-large-cased",
    forward_batch_size: int = 16,
) -> EwiserEncoder:
    """Load one of the three EWISER checkpoints released alongside the paper.

    Args:
        checkpoint_path: Local path to a checkpoint file (e.g.
            ``ewiser.semcor.pt``) -- not bundled with this package.
        offsets_path: Local path to that checkpoint's own
            ``res/dictionaries/offsets.txt`` -- also not bundled (see
            [SenseVocabulary.from_offsets_file][linkingtk.algorithms.wsd._ewiser_vocab.SenseVocabulary.from_offsets_file]).
        lexicon: Forwarded to `SenseVocabulary.from_offsets_file`.
        model_name_or_path: The encoder to pair with the checkpoint's
            decoder weights. Defaults to ``"bert-large-cased"``,
            matching every released checkpoint's own recorded config.
        forward_batch_size: Forwarded to `EwiserEncoder.__init__`.

    Returns:
        An `EwiserEncoder` with the checkpoint's decoder weights (FFN
        + graph adjacency) loaded, paired with a freshly downloaded
        frozen encoder (never part of the checkpoint -- see the module
        docstring).

    Raises:
        LinkingTKError: If `offsets_path` doesn't resolve to the same
            vocabulary size as the checkpoint's own output layer --
            the surest sign it isn't the exact file this checkpoint
            was trained with.
    """
    data = load_fairseq_checkpoint(checkpoint_path)
    args = data["args"]
    state_dict = data["model"]

    vocabulary = SenseVocabulary.from_offsets_file(offsets_path, lexicon=lexicon)
    vocab_size = int(state_dict["decoder.logits.weight"].shape[0])
    if len(vocabulary) != vocab_size:
        raise LinkingTKError(
            f"offsets_path resolved to a {len(vocabulary)}-entry vocabulary, but "
            f"the checkpoint's own output layer has {vocab_size} entries -- "
            "offsets_path must be the exact res/dictionaries/offsets.txt this "
            "checkpoint was trained with."
        )

    adjacency = torch.sparse_coo_tensor(
        state_dict["decoder.structured_logits.adjacency_pars.0"],
        state_dict["decoder.structured_logits.adjacency_pars.1"],
        size=tuple(state_dict["decoder.structured_logits.adjacency_pars.2"].tolist()),
    )

    encoder = cls(
        model_name_or_path=model_name_or_path,
        vocabulary=vocabulary,
        decoder_hidden_dim=int(state_dict["decoder.linears.0.weight"].shape[0]),
        adjacency=adjacency,
        structured_logits_trainable=bool(
            getattr(args, "decoder_structured_logits_trainable", True)
        ),
        forward_batch_size=forward_batch_size,
    )
    encoder.decoder.load_state_dict(
        {key[len("decoder.") :]: value for key, value in state_dict.items()}
    )
    return encoder

score

score(pairs: list[tuple[Entity, Entity]]) -> torch.Tensor

Return a (len(pairs),) tensor: each pair's raw sense-classification logit.

Unlike a cross-encoder's score, this does not run one forward pass per pair. It groups pairs by each mention's unique sentence (see mention_sentence_and_span), runs the encoder+decoder once per distinct sentence (batched forward_batch_size sentences at a time), and reads each pair's score off as logits[mention_word_index, candidate_sense_index] from that sentence's shared output. A candidate whose sense id isn't in vocabulary, or a mention whose span doesn't resolve to a word position, scores -inf (never spuriously wins a match).

Each chunk's [num_words, len(vocabulary)] logits are consumed (scores extracted) and released before the next chunk is encoded, rather than holding every distinct sentence's full logits matrix in memory for the whole call -- at len(vocabulary) ~117k, a large evaluation set's worth of these held simultaneously (thousands of sentences at once, e.g. UFSAC's "ALL" split) is enough to exhaust even a 24GB GPU; confirmed by hitting exactly that CUDA out of memory running examples/ewiser_reproduction.py before this was fixed.

Source code in src/linkingtk/algorithms/wsd/ewiser.py

def score(self, pairs: list[tuple[Entity, Entity]]) -> torch.Tensor:
    """Return a ``(len(pairs),)`` tensor: each pair's raw sense-classification logit.

    Unlike a cross-encoder's `score`, this does **not** run one
    forward pass per pair. It groups `pairs` by each mention's unique
    sentence (see `mention_sentence_and_span`), runs the
    encoder+decoder once per distinct sentence (batched
    `forward_batch_size` sentences at a time), and reads each pair's
    score off as ``logits[mention_word_index, candidate_sense_index]``
    from that sentence's shared output. A candidate whose sense id
    isn't in `vocabulary`, or a mention whose span doesn't resolve to
    a word position, scores ``-inf`` (never spuriously wins a match).

    Each chunk's ``[num_words, len(vocabulary)]`` logits are consumed
    (scores extracted) and released before the next chunk is encoded,
    rather than holding every distinct sentence's full logits matrix
    in memory for the whole call -- at `len(vocabulary)` ~117k, a
    large evaluation set's worth of these held simultaneously
    (thousands of sentences at once, e.g. UFSAC's "ALL" split) is
    enough to exhaust even a 24GB GPU; confirmed by hitting exactly
    that `CUDA out of memory` running `examples/ewiser_reproduction.py`
    before this was fixed.
    """
    if not pairs:
        return torch.empty(0)

    sentences: list[str] = []
    sentence_index: dict[str, int] = {}
    mention_word_index: dict[int, int | None] = {}
    pairs_by_sentence: dict[int, list[int]] = {}
    for pair_index, (mention, _sense) in enumerate(pairs):
        text, start, end = mention_sentence_and_span(mention)
        if text not in sentence_index:
            sentence_index[text] = len(sentences)
            sentences.append(text)
        if id(mention) not in mention_word_index:
            tokens = whitespace_tokenize_with_offsets(text)
            mention_word_index[id(mention)] = word_index_for_span(tokens, start, end)
        pairs_by_sentence.setdefault(sentence_index[text], []).append(pair_index)

    scores: list[torch.Tensor | None] = [None] * len(pairs)
    for start in range(0, len(sentences), self.forward_batch_size):
        chunk_sentences = sentences[start : start + self.forward_batch_size]
        chunk_logits = self._encode_chunk(chunk_sentences)
        for offset, logits in enumerate(chunk_logits):
            for pair_index in pairs_by_sentence.get(start + offset, []):
                mention, sense = pairs[pair_index]
                word_index = mention_word_index[id(mention)]
                sense_index = self.vocabulary.index_for(sense.id)
                if word_index is None or sense_index is None:
                    scores[pair_index] = logits.new_tensor(float("-inf"))
                else:
                    scores[pair_index] = logits[word_index, sense_index]

    assert all(score is not None for score in scores)  # every pair index was filled above
    return torch.stack(scores)  # type: ignore[arg-type]

EwiserLinker

EwiserLinker(
    model: EwiserEncoder,
    matching: Matcher = DEFAULT_MATCHER,
)

Bases: BaseLinker

WSD linking scored by an EwiserEncoder.

Unlike GlossBertLinker, this takes an already-constructed EwiserEncoder rather than a bare model_name_or_path string -- EwiserEncoder needs a SenseVocabulary (and usually an adjacency) at construction time, with no sane string-only default.

encoder = EwiserEncoder.from_checkpoint("ewiser.semcor.pt", "offsets.txt")
linker = EwiserLinker(encoder)
results = linker.link(mentions, senses)

Or, for from-scratch/continued training, construct EwiserEncoder directly and train it via EwiserTrainer.

Parameters:

Name	Type	Description	Default
model	EwiserEncoder	The EwiserEncoder to score candidates with.	required
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER

Source code in src/linkingtk/algorithms/wsd/ewiser.py

def __init__(self, model: EwiserEncoder, matching: Matcher = DEFAULT_MATCHER) -> None:
    self.model = model
    self.matching = matching

score_candidates

score_candidates(
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]

Blocked candidates per source entity, scored but not yet matched.

Satisfies CandidateScorer -- what link() itself builds internally, exposed so LlmRerankerLinker (#23) can re-rank a narrowed top-k instead of every blocked pair.

If dataset2 is a WnEntitySource whose lexicon differs from model.vocabulary's own, candidate ids are translated into the vocabulary's lexicon via synset_id_via_ili before scoring (see _translate_for_scoring's own docstring) -- without this, every candidate would fail SenseVocabulary.index_for and silently score -inf (#67). Candidate ids in the returned dict are still dataset2's own (untranslated) ids -- the translation is scoring-only, invisible to the caller.

Source code in src/linkingtk/algorithms/wsd/ewiser.py

def score_candidates(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]:
    """Blocked candidates per source entity, scored but not yet matched.

    Satisfies [CandidateScorer][linkingtk.algorithms.llm_reranker.CandidateScorer]
    -- what `link()` itself builds internally, exposed so
    [LlmRerankerLinker][linkingtk.algorithms.llm_reranker.LlmRerankerLinker]
    (#23) can re-rank a narrowed top-k instead of every blocked pair.

    If `dataset2` is a `WnEntitySource` whose `lexicon` differs from
    `model.vocabulary`'s own, candidate ids are translated into the
    vocabulary's lexicon via
    [synset_id_via_ili][linkingtk.sources.wn.synset_id_via_ili] before
    scoring (see `_translate_for_scoring`'s own docstring) -- without
    this, every candidate would fail
    [SenseVocabulary.index_for][linkingtk.algorithms.wsd._ewiser_vocab.SenseVocabulary.index_for]
    and silently score ``-inf`` (#67). Candidate ids in the returned
    dict are still `dataset2`'s own (untranslated) ids -- the
    translation is scoring-only, invisible to the caller.
    """
    pairs = list(blocking.candidate_pairs(dataset1, dataset2))
    if not pairs:
        return {}

    scoring_pairs = self._translate_for_scoring(pairs, dataset2)

    with torch.no_grad():
        self.model.eval()
        scores = self.model.score(scoring_pairs)
        self.model.train()

    candidates_by_source: dict[str, list[tuple[str, float]]] = defaultdict(list)
    for (entity1, entity2), score in zip(pairs, scores.tolist(), strict=True):
        candidates_by_source[entity1.id].append((entity2.id, score))

    return candidates_by_source

_ewiser_vocab

Sense vocabulary: index <-> synset id, for EWISER's full-inventory output layer.

EwiserEncoder classifies every word against a fixed, closed vocabulary of WordNet senses rather than scoring one candidate at a time (unlike GlossBertEncoder's cross-encoder). SenseVocabulary is that fixed index space.

Two construction paths, since the index order matters differently depending on where it comes from:

• from_offsets_file loads EWISER's own published checkpoint dictionary (frequency-sorted over the original authors' training corpus) -- this order is not reproducible from first principles (it depends on corpus statistics we don't have), so checkpoint-compatible inference requires the literal file, supplied locally by the caller (mirrors how examples/glossbert_reproduction.py requires a local copy of GlossBERT's own published checkpoint directory -- neither file is redistributed by this package).
• from_wn builds a fresh vocabulary for from-scratch training, where any consistent order works -- nothing downstream depends on matching upstream's specific index assignment.

SenseVocabulary

SenseVocabulary(
    index_to_synset_id: list[str | None],
    lexicon: str | None = None,
)

index <-> synset id mapping for an EwiserEncoder's output layer.

Parameters:

Name	Type	Description	Default
index_to_synset_id	list[str | None]	The vocabulary itself, in index order. An entry of None marks a reserved/special slot (e.g. fairseq's <s>/<pad>/</s>/<unk>, or an offset that didn't resolve in the target lexicon) that never matches a real candidate sense.	required
lexicon	str | None	The wn lexicon specifier index_to_synset_id's entries are ids in (e.g. "omw-en:1.4"), or None if they aren't wn ids at all (e.g. from_wn's from-scratch training path). Used only for EwiserLinker's own mismatched-lexicon check against a WnEntitySource dataset2 -- see index_for's docstring for why that check exists.	None

Source code in src/linkingtk/algorithms/wsd/_ewiser_vocab.py

def __init__(self, index_to_synset_id: list[str | None], lexicon: str | None = None) -> None:
    self._index_to_synset_id = index_to_synset_id
    self.lexicon = lexicon
    self._synset_id_to_index = {
        synset_id: index
        for index, synset_id in enumerate(index_to_synset_id)
        if synset_id is not None
    }

from_offsets_file classmethod

from_offsets_file(
    path: str | Path,
    lexicon: str = "omw-en:1.4",
    nspecial: int = 4,
) -> SenseVocabulary

Load EWISER's own res/dictionaries/offsets.txt.

Each line is "<token> <frequency>"; token is either a WordNet 3.0 offset ("wn:02084071n") or a non-offset dummy symbol (e.g. fairseq's "madeupword0000" padding entry), which becomes a reserved (None) slot like the nspecial prefix. Offsets that don't resolve in lexicon (version skew) also become None rather than raising, since a handful of unresolvable entries shouldn't block loading the other ~117,000.

Parameters:

Name	Type	Description	Default
path	str | Path	Local path to offsets.txt (not bundled with this package -- see the module docstring).	required
lexicon	str	Passed to wn30_offset_to_synset_id.	'omw-en:1.4'
nspecial	int	Reserved slots prepended before the file's own entries, matching fairseq's convention (<s>, <pad>, </s>, <unk>) the checkpoints were trained with.	4

Returns:

Type	Description
SenseVocabulary	A SenseVocabulary of length nspecial + <lines in path>,
SenseVocabulary	matching a published checkpoint's decoder.logits.weight
SenseVocabulary	row count (117664 for the three EWISER checkpoints released
SenseVocabulary	alongside the paper).

Source code in src/linkingtk/algorithms/wsd/_ewiser_vocab.py

@classmethod
def from_offsets_file(
    cls, path: str | Path, lexicon: str = "omw-en:1.4", nspecial: int = 4
) -> SenseVocabulary:
    """Load EWISER's own ``res/dictionaries/offsets.txt``.

    Each line is ``"<token> <frequency>"``; `token` is either a
    WordNet 3.0 offset (``"wn:02084071n"``) or a non-offset dummy
    symbol (e.g. fairseq's ``"madeupword0000"`` padding entry), which
    becomes a reserved (``None``) slot like the `nspecial` prefix.
    Offsets that don't resolve in `lexicon` (version skew) also become
    ``None`` rather than raising, since a handful of unresolvable
    entries shouldn't block loading the other ~117,000.

    Args:
        path: Local path to `offsets.txt` (not bundled with this
            package -- see the module docstring).
        lexicon: Passed to
            [wn30_offset_to_synset_id][linkingtk.sources.wn.wn30_offset_to_synset_id].
        nspecial: Reserved slots prepended before the file's own
            entries, matching fairseq's convention (``<s>``, ``<pad>``,
            ``</s>``, ``<unk>``) the checkpoints were trained with.

    Returns:
        A `SenseVocabulary` of length ``nspecial + <lines in path>``,
        matching a published checkpoint's ``decoder.logits.weight``
        row count (117664 for the three EWISER checkpoints released
        alongside the paper).
    """
    index_to_synset_id: list[str | None] = [None] * nspecial
    for line in Path(path).read_text().splitlines():
        token = line.split(maxsplit=1)[0] if line.strip() else ""
        if token.startswith("wn:"):
            index_to_synset_id.append(wn30_offset_to_synset_id(token, lexicon=lexicon))
        else:
            index_to_synset_id.append(None)
    return cls(index_to_synset_id, lexicon=lexicon)

from_wn classmethod

from_wn(
    synset_ids: list[str], nspecial: int = 4
) -> SenseVocabulary

Build a fresh vocabulary for from-scratch training.

Parameters:

Name	Type	Description	Default
synset_ids	list[str]	The senses to index, e.g. every synset observed in a training corpus. Deduplicated and sorted for a deterministic (order-independent-correctness) index assignment.	required
nspecial	int	Reserved slots prepended before the real entries.	4

Returns:

Type	Description
SenseVocabulary	A SenseVocabulary of length nspecial + len(set(synset_ids)).

Source code in src/linkingtk/algorithms/wsd/_ewiser_vocab.py

@classmethod
def from_wn(cls, synset_ids: list[str], nspecial: int = 4) -> SenseVocabulary:
    """Build a fresh vocabulary for from-scratch training.

    Args:
        synset_ids: The senses to index, e.g. every synset observed in
            a training corpus. Deduplicated and sorted for a
            deterministic (order-independent-correctness) index
            assignment.
        nspecial: Reserved slots prepended before the real entries.

    Returns:
        A `SenseVocabulary` of length ``nspecial + len(set(synset_ids))``.
    """
    index_to_synset_id: list[str | None] = [None] * nspecial
    index_to_synset_id.extend(sorted(set(synset_ids)))
    return cls(index_to_synset_id)

index_for

index_for(synset_id: str) -> int | None

The output-layer index for synset_id, or None if it's not in this vocabulary.

A synset_id from the wrong lexicon silently returns None just like a genuinely absent one -- there's no way to tell "not in this WordNet's sense inventory" apart from "belongs to a different, non-overlapping id space" from a string lookup alone (confirmed real: from_offsets_file's default lexicon is "omw-en:1.4", whose ids look like "omw-en-00319912-v", but WnEntitySource's own default is "oewn:2021", whose ids look like "oewn-00319912-v" -- and per synset_id_to_wn30_offset's own docstring, oewn:*'s offsets have diverged from PWN 3.0, so even the numeric part rarely matches). Left to itself, every candidate from a mismatched WnEntitySource would fail this lookup, and EwiserEncoder.score would score every one of them -inf -- silently, with no exception (see #67). EwiserLinker avoids this by translating a WnEntitySource candidate's id into self.lexicon's id space via synset_id_via_ili before it ever reaches this method, whenever the two lexicons differ -- see EwiserLinker.score_candidates's own docstring.

Source code in src/linkingtk/algorithms/wsd/_ewiser_vocab.py

def index_for(self, synset_id: str) -> int | None:
    """The output-layer index for `synset_id`, or ``None`` if it's not in this vocabulary.

    A `synset_id` from the *wrong* lexicon silently returns ``None``
    just like a genuinely absent one -- there's no way to tell "not in
    this WordNet's sense inventory" apart from "belongs to a
    different, non-overlapping id space" from a string lookup alone
    (confirmed real: `from_offsets_file`'s default ``lexicon`` is
    ``"omw-en:1.4"``, whose ids look like ``"omw-en-00319912-v"``, but
    `WnEntitySource`'s own default is ``"oewn:2021"``, whose ids look
    like ``"oewn-00319912-v"`` -- and per
    [synset_id_to_wn30_offset][linkingtk.sources.wn.synset_id_to_wn30_offset]'s
    own docstring, `oewn:*`'s offsets have diverged from PWN 3.0, so
    even the numeric part rarely matches). Left to itself, every
    candidate from a mismatched `WnEntitySource` would fail this
    lookup, and `EwiserEncoder.score` would score every one of them
    ``-inf`` -- silently, with no exception (see #67). `EwiserLinker`
    avoids this by translating a `WnEntitySource` candidate's id into
    `self.lexicon`'s id space via
    [synset_id_via_ili][linkingtk.sources.wn.synset_id_via_ili] before
    it ever reaches this method, whenever the two lexicons differ --
    see `EwiserLinker.score_candidates`'s own docstring.
    """
    return self._synset_id_to_index.get(synset_id)

synset_id_for

synset_id_for(index: int) -> str | None

The synset id at index, or None for a reserved/unresolved slot.

Source code in src/linkingtk/algorithms/wsd/_ewiser_vocab.py

def synset_id_for(self, index: int) -> str | None:
    """The synset id at `index`, or ``None`` for a reserved/unresolved slot."""
    return self._index_to_synset_id[index]

_ewiser_graph

WordNet relation-graph construction for StructuredLogits.

Builds a sparse [V, V] adjacency over a SenseVocabulary from wn's own hypernym/similar/verb_group/derivation relations -- confirmed live against this repo's own installed lexicon that all four are real, populated relation types (derivation is sense-level; the other three are synset-level).

Reimplements the algorithm observed in EWISER's own graph builder (ResourceManager.make_adjacency_from_files), not its code (that source is CC-BY-NC-SA licensed, not redistributed here): for every node, gather every incoming edge, normalize their weights to sum to 1 over all predecessors first, then truncate to the top max_incoming by weight -- in that order. The surviving top-k weights generally do not re-sum to 1, which is a real, confirmed property of the reference's own behavior, not a bug in this reimplementation.

Edge direction matches the reference: a hypernym relation edge is added (hyponym_index, hypernym_index) -- the more general (and typically better-attested) hypernym's logit propagates down into its more specific hyponym, letting the model transfer confidence to senses it saw less often in training. similar/verb_group/derivation are only added in the direction each call reports -- confirmed directly that wn already stores all three reciprocally (both synsets/senses on a pair list each other), so the reverse direction is picked up for free when the loop below reaches the other node; adding it a second time here would double the edge weight.

Sourced natively via wn rather than EWISER's own BabelNet-id-keyed res/edges/*.tsv files, avoiding both a heavy extra resource dependency and the bnids_map.txt remapping step those files require. This means a from-scratch graph cannot be verified bit-exactly against EWISER's own (different relation coverage, different tie-breaking order) -- only structurally, on synthetic data (see tests/algorithms/wsd/test_ewiser_graph.py). Only relevant for EwiserTrainer's from-scratch/continued-training path -- checkpoint-based inference uses the checkpoint's own already-baked-in adjacency instead (see EwiserEncoder.from_checkpoint).

build_relation_adjacency

build_relation_adjacency(
    vocabulary: SenseVocabulary,
    lexicon: str = "omw-en:1.4",
    max_incoming: int = 5,
) -> torch.Tensor

Build a sparse [len(vocabulary), len(vocabulary)] WordNet relation adjacency.

Parameters:

Name	Type	Description	Default
vocabulary	SenseVocabulary	The sense inventory to build edges over -- only relation targets that are themselves in vocabulary become edges (a target outside vocabulary is silently dropped, the same "unresolved -> excluded, not an error" convention used elsewhere in this package, e.g. UfsacDataset).	required
lexicon	str	wn lexicon specifier to query relations against.	'omw-en:1.4'
max_incoming	int	Maximum incoming edges kept per node, by normalized weight (see module docstring for the normalize-then-truncate order).	5

Returns:

Type	Description
Tensor	A coalesced sparse COO tensor, A[child, parent] = weight.

Raises:

Type	Description
OptionalDependencyError	If wn isn't installed.

Source code in src/linkingtk/algorithms/wsd/_ewiser_graph.py

def build_relation_adjacency(
    vocabulary: SenseVocabulary,
    lexicon: str = "omw-en:1.4",
    max_incoming: int = 5,
) -> torch.Tensor:
    """Build a sparse ``[len(vocabulary), len(vocabulary)]`` WordNet relation adjacency.

    Args:
        vocabulary: The sense inventory to build edges over -- only
            relation targets that are themselves in `vocabulary` become
            edges (a target outside `vocabulary` is silently dropped, the
            same "unresolved -> excluded, not an error" convention used
            elsewhere in this package, e.g.
            [UfsacDataset][linkingtk.datasets.ufsac.UfsacDataset]).
        lexicon: `wn` lexicon specifier to query relations against.
        max_incoming: Maximum incoming edges kept per node, by
            normalized weight (see module docstring for the
            normalize-then-truncate order).

    Returns:
        A coalesced sparse COO tensor, ``A[child, parent] = weight``.

    Raises:
        OptionalDependencyError: If `wn` isn't installed.
    """
    try:
        import wn
    except ImportError as exc:
        raise OptionalDependencyError("build_relation_adjacency", "wn") from exc

    incoming: dict[int, dict[int, float]] = {}

    def add_edge(child_index: int, parent_index: int) -> None:
        if child_index == parent_index:
            return
        node = incoming.setdefault(child_index, {})
        node[parent_index] = node.get(parent_index, 0.0) + 1.0

    for index in range(len(vocabulary)):
        synset_id = vocabulary.synset_id_for(index)
        if synset_id is None:
            continue
        try:
            synset = wn.synset(synset_id, lexicon=lexicon)
        except wn.Error:
            logger.warning(
                "build_relation_adjacency: %s not found in %s, skipping", synset_id, lexicon
            )
            continue

        relations = synset.relations(*_SYNSET_RELATIONS)
        for relation_type in _SYNSET_RELATIONS:
            for target in relations.get(relation_type, []):
                target_index = vocabulary.index_for(target.id)
                if target_index is not None:
                    add_edge(index, target_index)

        for sense in synset.senses():
            for target_sense in sense.relations("derivation").get("derivation", []):
                target_index = vocabulary.index_for(target_sense.synset().id)
                if target_index is not None:
                    add_edge(index, target_index)

    return _to_sparse_tensor(incoming, size=len(vocabulary), max_incoming=max_incoming)

_ewiser_structured_logits

One linear WordNet-relation graph-propagation step over sense logits.

EWISER's central idea (Bevilacqua & Navigli, ACL 2020): instead of scoring each sense independently, boost every synset's raw classification logit by a weighted sum of its WordNet-neighbors' logits, so the model can transfer confidence from a frequently-labeled sense (e.g. a common hypernym) to a rarer one it's graph-adjacent to, even if that rarer sense was never seen labeled in training. This module reimplements that operation -- new_logits = (A @ old_logits) + old_logits for a sparse [V, V] adjacency A -- from the paper's description and the shape observed directly in the released checkpoints' state dicts, not ported from EWISER's own (CC-BY-NC-SA licensed) source.

Deliberately reimplemented with plain torch.sparse.mm/ torch.sparse_coo_tensor instead of the reference's own torch_scatter/torch_sparse dependency: those packages are unconditionally imported by the reference's equivalent module for what is, in substance, one sparse-dense matmul plus an optional row-sum normalization -- and they're notoriously fragile, tightly pinned to exact torch+CUDA build combinations. A deliberate deviation from upstream, not an oversight.

The trainable=True, renormalize=False path (the default, and what every released checkpoint and the reference's own bin/train-ewiser.sh actually use in both training stages) runs through _SparsePropagate below instead of a bare torch.sparse.mm call: torch.sparse.mm's default backward, when the sparse operand itself requires grad, materializes a dense [V, V] gradient rather than an O(nnz) one -- invisible at a toy vocabulary size, but confirmed directly (issue #58) to OOM at EWISER's own real ~117,664-entry WordNet vocabulary (a dense [117664, 117664] float32 gradient is ~51.6GB). _SparsePropagate computes the same values via an explicit, O(nnz) indexed gather-multiply-sum instead, verified against torch.autograd.gradcheck and a dense reference implementation, and measured at 3ms/0.4GB peak at the full ~117,664-entry / ~168K-edge scale (vs. an unfixed OOM). The renormalize=True combination is untouched (no released checkpoint or bin/train-ewiser.sh config uses it, so it isn't scale-verified the way the default path now is) and still goes through the original torch.sparse.mm path.

StructuredLogits

StructuredLogits(
    adjacency: Tensor,
    trainable: bool = False,
    renormalize: bool = False,
)

Bases: Module

Graph-propagation step: new_logits = (A @ old_logits) + old_logits.

Stores the sparse adjacency's indices/values/size as separate parameters under the same names (adjacency_pars.0/adjacency_pars.1/adjacency_pars.2) the released EWISER checkpoints use for their own equivalent module, so a checkpoint's state dict loads via plain load_state_dict(strict=True) with no key renaming.

Parameters:

Name	Type	Description	Default
adjacency	Tensor	A sparse (coalesced or not) [V, V] COO tensor, e.g. from build_relation_adjacency or loaded directly from a checkpoint's own baked-in adjacency.	required
trainable	bool	Whether the adjacency's edge weights are updated during training (the indices/shape are always fixed).	False
renormalize	bool	If True, additionally divide each node's propagated neighbor sum by that node's total incoming weight (computed via torch.sparse.sum, no torch_scatter needed). The three released checkpoints all use False (their own edge weights are pre-normalized at graph-build time, see build_relation_adjacency).	False

Source code in src/linkingtk/algorithms/wsd/_ewiser_structured_logits.py

def __init__(
    self,
    adjacency: torch.Tensor,
    trainable: bool = False,
    renormalize: bool = False,
) -> None:
    super().__init__()
    adjacency = adjacency.coalesce()
    self.adjacency_pars = nn.ParameterList(
        [
            nn.Parameter(adjacency.indices(), requires_grad=False),
            nn.Parameter(adjacency.values(), requires_grad=trainable),
            nn.Parameter(torch.tensor(adjacency.shape), requires_grad=False),
        ]
    )
    self.renormalize = renormalize

forward

forward(logits: Tensor) -> torch.Tensor

Propagate logits ([..., V]) one step through the WordNet graph.

Source code in src/linkingtk/algorithms/wsd/_ewiser_structured_logits.py

def forward(self, logits: torch.Tensor) -> torch.Tensor:
    """Propagate `logits` (``[..., V]``) one step through the WordNet graph."""
    vocab_size = logits.shape[-1]
    flat = logits.reshape(-1, vocab_size)
    indices, values, size = self.adjacency_pars
    if self.renormalize:
        adjacency = self._adjacency()
        neighbors = torch.sparse.mm(adjacency, flat.transpose(0, 1)).transpose(0, 1)
        row_sum = torch.sparse.sum(adjacency, dim=1).to_dense().clamp(min=1e-12)
        neighbors = neighbors / row_sum.unsqueeze(0)
    else:
        neighbors = _SparsePropagate.apply(  # type: ignore[no-untyped-call]
            indices, values, tuple(size.tolist()), flat
        )
    result: torch.Tensor = (neighbors + flat).reshape(logits.shape)
    return result

_ewiser_decoder

The trainable EWISER decoder head, and released-checkpoint loading.

Split out of ewiser.py to keep that module focused on the public EwiserEncoder/EwiserLinker API (see this package's ~300-line-per-file convention).

EwiserDecoder

EwiserDecoder(
    input_dim: int,
    hidden_dim: int,
    vocab_size: int,
    adjacency: Tensor | None = None,
    structured_logits_trainable: bool = True,
    structured_logits_renormalize: bool = False,
    dropout: float = 0.1,
)

Bases: Module

The small trainable head: BatchNorm -> Linear -> swish -> Linear + graph conv.

Submodule names (norm, linears, logits, structured_logits) match the released checkpoints' own state dict keys exactly (confirmed directly against a real checkpoint file), so a checkpoint's decoder weights load via plain load_state_dict(strict=True).

Source code in src/linkingtk/algorithms/wsd/_ewiser_decoder.py

def __init__(
    self,
    input_dim: int,
    hidden_dim: int,
    vocab_size: int,
    adjacency: torch.Tensor | None = None,
    structured_logits_trainable: bool = True,
    structured_logits_renormalize: bool = False,
    dropout: float = 0.1,
) -> None:
    super().__init__()
    self.norm = nn.BatchNorm1d(input_dim)
    self.linears = nn.ModuleList([nn.Linear(input_dim, hidden_dim)])
    self.dropout = nn.Dropout(dropout)
    self.activation = nn.SiLU()  # swish(x) = x * sigmoid(x), confirmed against the reference
    self.logits = nn.Linear(hidden_dim, vocab_size, bias=False)
    self.structured_logits = (
        StructuredLogits(
            adjacency,
            trainable=structured_logits_trainable,
            renormalize=structured_logits_renormalize,
        )
        if adjacency is not None
        else None
    )

load_fairseq_checkpoint

load_fairseq_checkpoint(path: str | Path) -> dict[str, Any]

Load a released EWISER .pt checkpoint file.

Its pickle stream references training-only bookkeeping classes from its original (pre-release) package name, qbert, and from fairseq (e.g. qbert.fairseq_ext.meters.SumMeter, fairseq.meters.AverageMeter) -- none of which this package depends on or needs (only data["model"]'s plain tensors and data["args"]'s plain argparse.Namespace are used), but standard unpickling fails the entire load if even one referenced class can't be resolved. Works around this with a permissive Unpickler.find_class that substitutes an inert placeholder class for anything unresolvable instead of failing -- confirmed safe by inspecting exactly which classes get requested against a real checkpoint file (only optimizer/meter training-history objects, never anything under data["model"]).

Source code in src/linkingtk/algorithms/wsd/_ewiser_decoder.py

def load_fairseq_checkpoint(path: str | Path) -> dict[str, Any]:
    """Load a released EWISER ``.pt`` checkpoint file.

    Its pickle stream references training-only bookkeeping classes from
    its original (pre-release) package name, ``qbert``, and from
    ``fairseq`` (e.g. ``qbert.fairseq_ext.meters.SumMeter``,
    ``fairseq.meters.AverageMeter``) -- none of which this package depends
    on or needs (only `data["model"]`'s plain tensors and `data["args"]`'s
    plain `argparse.Namespace` are used), but standard unpickling fails the
    *entire* load if even one referenced class can't be resolved. Works
    around this with a permissive `Unpickler.find_class` that substitutes
    an inert placeholder class for anything unresolvable instead of
    failing -- confirmed safe by inspecting exactly which classes get
    requested against a real checkpoint file (only optimizer/meter
    training-history objects, never anything under `data["model"]`).
    """

    class _PermissiveUnpickler(pickle.Unpickler):
        def find_class(self, module: str, name: str) -> Any:
            try:
                return super().find_class(module, name)
            except Exception:
                return type(name, (), {})

    pickle_module = types.ModuleType("_ewiser_permissive_pickle")
    pickle_module.Unpickler = _PermissiveUnpickler  # type: ignore[attr-defined]
    pickle_module.load = pickle.load  # type: ignore[attr-defined]

    data: dict[str, Any] = torch.load(
        str(path), map_location="cpu", weights_only=False, pickle_module=pickle_module
    )
    return data

_ewiser_sense_embeddings

Pretrained sense-embedding loaders for EwiserEncoder's output layer.

EWISER's own strongest results (Section 4.2, "Output Embeddings", of the paper) come from initializing decoder.logits.weight -- the sense output embedding layer -- from off-the-shelf sense vectors rather than a random init, then optionally freezing it for a few epochs (EwiserTrainer.freeze_output_epochs). Two off-the-shelf vector sets are named in the paper: LMMS 2048d (Loureiro and Jorge, 2019) and SensEmBERT+LMMS (Scarlini et al., 2020, enhancing LMMS with BabelNet+Wikipedia for nouns, backed off to LMMS for verbs/adjectives/adverbs -- SensEmBERT's own vectors are "in the same space as LMMS", per the paper). The paper's own recorded pipeline for both (confirmed directly from its Section 4.2 text, not from the reference's CC-BY-NC-SA-licensed scripts): normalize each raw sense vector to unit length, take the centroid of every sense in a synset, reduce to 512d with truncated SVD, normalize again.

Two loaders, for two different starting points:

• load_synset_centroid_vectors loads an already-built, already-reduced synset-vector file -- e.g. the literal res/embeddings/sensembert+lmms.svd512.synset-centroid.vec the reference's own training script uses (confirmed against bin/train-ewiser.sh's $EMBEDDINGS variable), a GloVe-style file keyed by WordNet 3.0 offset tokens ("wn:02084071n", same convention SenseVocabulary.from_offsets_file already parses). This is the fastest path to the paper's own best configuration when such a file is available -- no SVD or centroiding needed at load time.
• build_synset_centroid_vectors_from_lmms builds the same shape of tensor from LMMS's own raw, sensekey-keyed release files (also GloVe-style, no header -- confirmed directly against a real downloaded LMMS release) -- for users who only have the raw per-sensekey vectors, not a pre-built synset file.

Neither loader ingests SensEmBERT's own raw (pre-combination) vector files -- SensEmBERT was never open-sourced and its own site (sensembert.org) is no longer reachable, so there is no real file or confirmed format to build and verify a loader against. Use load_synset_centroid_vectors with a pre-built SensEmBERT+LMMS file (as the reference's own release provides) to get the paper's full best configuration; build_synset_centroid_vectors_from_lmms alone only covers the LMMS-only configuration (still a real, evaluated row in the paper's own Table 2, just not its best one).

build_synset_centroid_vectors_from_lmms

build_synset_centroid_vectors_from_lmms(
    path: str | Path,
    vocabulary: SenseVocabulary,
    base: Tensor,
    target_dim: int | None = None,
    lexicon: str = "omw-en:1.4",
    progress: bool = True,
) -> tuple[torch.Tensor, int]

Build a synset-centroid init tensor from LMMS's own raw sensekey vectors.

Reimplements the paper's own recorded pipeline (Section 4.2, "Output Embeddings" -- not ported from the reference's CC-BY-NC-SA-licensed bin/get_centroids.py/bin/reduce_dims.py, though independently confirmed to agree with them): each raw sense vector is L2-normalized, every sense sharing a synset (via sensekey_to_synset_id) is averaged into that synset's centroid, then (if target_dim differs from the raw dimensionality) reduced with sklearn.decomposition.TruncatedSVD and L2-normalized again -- the paper states the final vectors are unit-length, and SVD projection doesn't preserve norms, so the second normalization is necessary to actually reach that state, not optional polish.

This resolves every sensekey in path via sensekey_to_synset_id (one wn lookup each) -- a real LMMS release has ~100K-200K sensekeys, so this is a slow, one-time vocabulary-init pass, not something to call per training step.

Parameters:

Name	Type	Description	Default
path	str | Path	A GloVe-style plaintext file, no header, one line per sensekey: "<sensekey> <float> <float> ...\n" (LMMS's own release format -- confirmed directly against a real downloaded release).	required
vocabulary	SenseVocabulary	The SenseVocabulary to align rows against.	required
base	Tensor	A [len(vocabulary), dim] tensor to fill in place, same contract as load_synset_centroid_vectors's base.	required
target_dim	int | None	Reduce to this many dimensions via truncated SVD if given and different from the file's own raw dimensionality (the paper reduces LMMS's raw 2048d to 512d). None (the default) keeps the raw dimensionality as-is.	None
lexicon	str	Passed to sensekey_to_synset_id.	'omw-en:1.4'
progress	bool	Show a tqdm progress bar.	True

Returns:

Type	Description
Tensor	(base, matched), same shape of result as
int	load_synset_centroid_vectors.

Source code in src/linkingtk/algorithms/wsd/_ewiser_sense_embeddings.py

def build_synset_centroid_vectors_from_lmms(
    path: str | Path,
    vocabulary: SenseVocabulary,
    base: torch.Tensor,
    target_dim: int | None = None,
    lexicon: str = "omw-en:1.4",
    progress: bool = True,
) -> tuple[torch.Tensor, int]:
    """Build a synset-centroid init tensor from LMMS's own raw sensekey vectors.

    Reimplements the paper's own recorded pipeline (Section 4.2, "Output
    Embeddings" -- not ported from the reference's CC-BY-NC-SA-licensed
    ``bin/get_centroids.py``/``bin/reduce_dims.py``, though independently
    confirmed to agree with them): each raw sense vector is L2-normalized,
    every sense sharing a synset (via `sensekey_to_synset_id`) is averaged
    into that synset's centroid, then (if `target_dim` differs from the
    raw dimensionality) reduced with `sklearn.decomposition.TruncatedSVD`
    and L2-normalized again -- the paper states the *final* vectors are
    unit-length, and SVD projection doesn't preserve norms, so the second
    normalization is necessary to actually reach that state, not optional
    polish.

    This resolves every sensekey in `path` via `sensekey_to_synset_id`
    (one `wn` lookup each) -- a real LMMS release has ~100K-200K
    sensekeys, so this is a slow, one-time vocabulary-init pass, not
    something to call per training step.

    Args:
        path: A GloVe-style plaintext file, no header, one line per
            sensekey: ``"<sensekey> <float> <float> ...\\n"`` (LMMS's own
            release format -- confirmed directly against a real
            downloaded release).
        vocabulary: The `SenseVocabulary` to align rows against.
        base: A ``[len(vocabulary), dim]`` tensor to fill in place, same
            contract as `load_synset_centroid_vectors`'s `base`.
        target_dim: Reduce to this many dimensions via truncated SVD if
            given and different from the file's own raw dimensionality
            (the paper reduces LMMS's raw 2048d to 512d). `None` (the
            default) keeps the raw dimensionality as-is.
        lexicon: Passed to `sensekey_to_synset_id`.
        progress: Show a `tqdm` progress bar.

    Returns:
        ``(base, matched)``, same shape of result as
        `load_synset_centroid_vectors`.
    """
    from tqdm import tqdm

    vectors_by_synset: dict[str, list[torch.Tensor]] = {}
    with Path(path).open() as handle:
        lines = tqdm(handle, desc=str(path), unit=" lines", disable=not progress)
        for line in lines:
            line = line.rstrip("\n")
            if not line:
                continue
            sense_key, *floats = line.split(" ")
            synset_id = sensekey_to_synset_id(sense_key, lexicon=lexicon)
            if synset_id is None:
                continue
            vector = torch.tensor([float(x) for x in floats])
            vectors_by_synset.setdefault(synset_id, []).append(_l2_normalize(vector))

    centroids = {
        synset_id: _l2_normalize(torch.stack(vectors).mean(0))
        for synset_id, vectors in vectors_by_synset.items()
    }

    if centroids and target_dim is not None:
        raw_dim = next(iter(centroids.values())).shape[0]
        if target_dim != raw_dim:
            centroids = _reduce_dims(centroids, target_dim)

    matched = _fill_from_synset_vectors(base, vocabulary, centroids)
    logger.info(
        "build_synset_centroid_vectors_from_lmms: %d/%d vocabulary entries matched (%s)",
        matched,
        len(vocabulary),
        path,
    )
    return base, matched

load_synset_centroid_vectors

load_synset_centroid_vectors(
    path: str | Path,
    vocabulary: SenseVocabulary,
    base: Tensor,
    lexicon: str = "omw-en:1.4",
    progress: bool = True,
) -> tuple[torch.Tensor, int]

Load a pre-built, offset-keyed synset-vector file into base, in place.

Parameters:

Name	Type	Description	Default
path	str | Path	A GloVe-style plaintext file, no header, one line per synset: "<token> <float> <float> ...\n", where token is a WordNet 3.0 offset ("wn:02084071n") -- the same format and token convention as offsets.txt (SenseVocabulary.from_offsets_file), and the format of the reference's own res/embeddings/*.svd512.synset-centroid.vec resource files. A line whose token isn't in that shape (e.g. a dummy padding entry) is skipped, not an error.	required
vocabulary	SenseVocabulary	The SenseVocabulary to align rows against.	required
base	Tensor	A [len(vocabulary), dim] tensor to fill in place -- typically a freshly-constructed decoder's own decoder.logits.weight.data (or a clone of it), so that vocabulary entries this file has no vector for keep whatever init base already had rather than being zeroed.	required
lexicon	str	Passed to wn30_offset_to_synset_id.	'omw-en:1.4'
progress	bool	Show a tqdm progress bar (this file is typically ~100K+ lines).	True

Returns:

Type	Description
Tensor	(base, matched) -- base (mutated in place, also returned
int	for chaining) and the number of vocabulary entries a vector was
tuple[Tensor, int]	found for.

Source code in src/linkingtk/algorithms/wsd/_ewiser_sense_embeddings.py

def load_synset_centroid_vectors(
    path: str | Path,
    vocabulary: SenseVocabulary,
    base: torch.Tensor,
    lexicon: str = "omw-en:1.4",
    progress: bool = True,
) -> tuple[torch.Tensor, int]:
    """Load a pre-built, offset-keyed synset-vector file into `base`, in place.

    Args:
        path: A GloVe-style plaintext file, no header, one line per
            synset: ``"<token> <float> <float> ...\\n"``, where `token`
            is a WordNet 3.0 offset (``"wn:02084071n"``) -- the same
            format and token convention as `offsets.txt`
            (`SenseVocabulary.from_offsets_file`), and the format of the
            reference's own ``res/embeddings/*.svd512.synset-centroid.vec``
            resource files. A line whose token isn't in that shape (e.g.
            a dummy padding entry) is skipped, not an error.
        vocabulary: The `SenseVocabulary` to align rows against.
        base: A ``[len(vocabulary), dim]`` tensor to fill in place --
            typically a freshly-constructed decoder's own
            ``decoder.logits.weight.data`` (or a clone of it), so that
            vocabulary entries this file has no vector for keep whatever
            init `base` already had rather than being zeroed.
        lexicon: Passed to `wn30_offset_to_synset_id`.
        progress: Show a `tqdm` progress bar (this file is typically
            ~100K+ lines).

    Returns:
        ``(base, matched)`` -- `base` (mutated in place, also returned
        for chaining) and the number of `vocabulary` entries a vector was
        found for.
    """
    from tqdm import tqdm

    vectors_by_synset: dict[str, torch.Tensor] = {}
    with Path(path).open() as handle:
        lines = tqdm(handle, desc=str(path), unit=" lines", disable=not progress)
        for line in lines:
            line = line.rstrip("\n")
            if not line:
                continue
            token, *floats = line.split(" ")
            if not token.startswith("wn:"):
                continue
            synset_id = wn30_offset_to_synset_id(token, lexicon=lexicon)
            if synset_id is None:
                continue
            vectors_by_synset[synset_id] = torch.tensor([float(x) for x in floats])

    matched = _fill_from_synset_vectors(base, vocabulary, vectors_by_synset)
    logger.info(
        "load_synset_centroid_vectors: %d/%d vocabulary entries matched (%s)",
        matched,
        len(vocabulary),
        path,
    )
    return base, matched

_ewiser_text

Sentence/word alignment helpers for EwiserEncoder.

EWISER classifies one vector per whitespace word (mean-pooled from BERT subwords), unlike GlossBertEncoder, which never needs word-level positions. These helpers resolve a mention's character span to a word index, subword-tokenize each whitespace word in isolation (see wordpiece_tokenize_words's docstring for why this can't just be a high-level tokenizer call), and pool subword hidden states back up to word level.

mean_pool_subwords

mean_pool_subwords(
    hidden: Tensor,
    word_ids_batch: list[list[int | None]],
    num_words: list[int],
) -> list[torch.Tensor]

Mean-pool hidden ([batch, seq, H]) subword vectors back to one vector per word.

word_ids_batch[b][t] is the word index subtoken t of sequence b belongs to (None for special/padding tokens), from the tokenizer's own fast word_ids() alignment.

Source code in src/linkingtk/algorithms/wsd/_ewiser_text.py

def mean_pool_subwords(
    hidden: torch.Tensor, word_ids_batch: list[list[int | None]], num_words: list[int]
) -> list[torch.Tensor]:
    """Mean-pool `hidden` (``[batch, seq, H]``) subword vectors back to one vector per word.

    `word_ids_batch[b][t]` is the word index subtoken `t` of sequence `b`
    belongs to (``None`` for special/padding tokens), from the tokenizer's
    own fast `word_ids()` alignment.
    """
    pooled: list[torch.Tensor] = []
    for batch_index, word_ids in enumerate(word_ids_batch):
        n = num_words[batch_index]
        real = [(t, wid) for t, wid in enumerate(word_ids) if wid is not None and wid < n]
        positions = torch.tensor([t for t, _wid in real], dtype=torch.long, device=hidden.device)
        word_index = torch.tensor([wid for _t, wid in real], dtype=torch.long, device=hidden.device)
        vectors = hidden[batch_index, positions]

        sums = hidden.new_zeros(n, hidden.shape[-1]).index_add_(0, word_index, vectors)
        counts = hidden.new_zeros(n).index_add_(
            0, word_index, torch.ones_like(word_index, dtype=hidden.dtype)
        )
        pooled.append(sums / counts.clamp(min=1.0).unsqueeze(-1))
    return pooled

mention_sentence_and_span

mention_sentence_and_span(
    entity: Entity,
) -> tuple[str, int, int]

The sentence text and mention character span to run EwiserEncoder against.

Uses context's (text, start, end) offsets when it's a ContextWithSpan (see linkingtk.core.entity) (the shape SemCorDataset/ UfsacDataset mentions carry -- every mention from the same sentence shares an identical text, which is what lets EwiserEncoder.score batch by unique sentence instead of per pair). Falls back to treating the mention's own label as a standalone one-word "sentence" when context carries no span.

Source code in src/linkingtk/algorithms/wsd/_ewiser_text.py

def mention_sentence_and_span(entity: Entity) -> tuple[str, int, int]:
    """The sentence text and mention character span to run `EwiserEncoder` against.

    Uses `context`'s ``(text, start, end)`` offsets when it's a
    `ContextWithSpan` (see [linkingtk.core.entity][]) (the shape
    [SemCorDataset][linkingtk.datasets.semcor.SemCorDataset]/
    [UfsacDataset][linkingtk.datasets.ufsac.UfsacDataset] mentions carry --
    every mention from the same sentence shares an identical `text`, which
    is what lets `EwiserEncoder.score` batch by unique sentence instead of
    per pair). Falls back to treating the mention's own label as a
    standalone one-word "sentence" when `context` carries no span.
    """
    context = entity.context
    if isinstance(context, tuple):
        return context
    text = context if isinstance(context, str) else " ".join(label_texts(entity))
    return text, 0, len(text)

whitespace_tokenize_with_offsets

whitespace_tokenize_with_offsets(
    text: str,
) -> list[tuple[str, int, int]]

Split text on whitespace, keeping each word's character span.

Source code in src/linkingtk/algorithms/wsd/_ewiser_text.py

def whitespace_tokenize_with_offsets(text: str) -> list[tuple[str, int, int]]:
    """Split `text` on whitespace, keeping each word's character span."""
    tokens: list[tuple[str, int, int]] = []
    start: int | None = None
    for index, char in enumerate(text):
        if char.isspace():
            if start is not None:
                tokens.append((text[start:index], start, index))
                start = None
        elif start is None:
            start = index
    if start is not None:
        tokens.append((text[start:], start, len(text)))
    return tokens

word_index_for_span

word_index_for_span(
    tokens: list[tuple[str, int, int]], start: int, end: int
) -> int | None

The index of tokens containing character span [start, end), or None.

Prefers an exact containment match (tok_start <= start < tok_end); falls back to the token with the greatest character overlap for a span that doesn't land cleanly on a whitespace-token boundary.

Source code in src/linkingtk/algorithms/wsd/_ewiser_text.py

def word_index_for_span(tokens: list[tuple[str, int, int]], start: int, end: int) -> int | None:
    """The index of `tokens` containing character span ``[start, end)``, or ``None``.

    Prefers an exact containment match (`tok_start <= start < tok_end`);
    falls back to the token with the greatest character overlap for a
    span that doesn't land cleanly on a whitespace-token boundary.
    """
    for index, (_word, tok_start, tok_end) in enumerate(tokens):
        if tok_start <= start < tok_end:
            return index

    best_index: int | None = None
    best_overlap = 0
    for index, (_word, tok_start, tok_end) in enumerate(tokens):
        overlap = min(end, tok_end) - max(start, tok_start)
        if overlap > best_overlap:
            best_index, best_overlap = index, overlap
    return best_index

wordpiece_tokenize_words

wordpiece_tokenize_words(
    tokenizer: PreTrainedTokenizerBase, words: list[str]
) -> tuple[list[str], list[int | None]]

Subword-tokenize each whitespace words entry in isolation, returning ([CLS] + subwords + [SEP], word_ids) (word_ids[t] is the index into words subtoken t came from, None for [CLS]/[SEP]).

Deliberately not tokenizer(words, is_split_into_words=True): that high-level call still runs BERT's own punctuation-splitting pre-tokenizer on each provided "word" before wordpiece, so a word containing internal punctuation (e.g. UFSAC's own single-token "Oct.") gets split into two independent pre-tokens ("Oct", ".") before wordpiece ever runs, producing a standalone "." token -- not the "##." continuation-piece EWISER's own reference produces by wordpiece-tokenizing the literal string "Oct." directly, with no pre-splitting step at all. Confirmed directly: this single-token divergence ("." vs "##." are different vocabulary entries with unrelated embeddings) was traced, layer by layer, as the exact source of a checkpoint-reproduction gap against EWISER's own published numbers -- affects any UFSAC word with internal punctuation (abbreviations like "Oct."/"Sept.", likely also "U.S.", possessives, etc.), a small fraction of tokens but enough to measurably shift results once fed through the checkpoint's own BatchNorm (calibrated tightly enough to the reference's exact activations that even one wrong subtoken's ripple through attention measurably shifts other positions' representations).

Implements BERT's WordPiece algorithm directly (greedy longest-match against the tokenizer's own vocab) rather than depending on any tokenizer-internal wordpiece_tokenizer attribute -- not exposed by every transformers version/backend.

Source code in src/linkingtk/algorithms/wsd/_ewiser_text.py

def wordpiece_tokenize_words(
    tokenizer: PreTrainedTokenizerBase, words: list[str]
) -> tuple[list[str], list[int | None]]:
    """Subword-tokenize each whitespace `words` entry in isolation, returning
    ``([CLS] + subwords + [SEP], word_ids)`` (`word_ids[t]` is the index into
    `words` subtoken `t` came from, `None` for `[CLS]`/`[SEP]`).

    Deliberately **not** ``tokenizer(words, is_split_into_words=True)``:
    that high-level call still runs BERT's own punctuation-splitting
    pre-tokenizer on each provided "word" before wordpiece, so a word
    containing internal punctuation (e.g. UFSAC's own single-token
    ``"Oct."``) gets split into two independent pre-tokens (``"Oct"``,
    ``"."``) before wordpiece ever runs, producing a *standalone* ``"."``
    token -- not the ``"##."`` continuation-piece EWISER's own
    reference produces by wordpiece-tokenizing the literal string
    ``"Oct."`` directly, with no pre-splitting step at all. Confirmed
    directly: this single-token divergence (`"."` vs `"##."` are
    different vocabulary entries with unrelated embeddings) was traced,
    layer by layer, as the exact source of a checkpoint-reproduction gap
    against EWISER's own published numbers -- affects any UFSAC word with
    internal punctuation (abbreviations like "Oct."/"Sept.", likely also
    "U.S.", possessives, etc.), a small fraction of tokens but enough to
    measurably shift results once fed through the checkpoint's own
    BatchNorm (calibrated tightly enough to the reference's exact
    activations that even one wrong subtoken's ripple through attention
    measurably shifts other positions' representations).

    Implements BERT's WordPiece algorithm directly (greedy longest-match
    against the tokenizer's own vocab) rather than depending on any
    tokenizer-internal ``wordpiece_tokenizer`` attribute -- not exposed by
    every `transformers` version/backend.
    """
    vocab = tokenizer.get_vocab()
    unk_token = tokenizer.unk_token
    subwords: list[str] = []
    word_ids: list[int | None] = []
    for word_index, word in enumerate(words):
        for piece in _wordpiece_tokenize_one_word(word, vocab, unk_token):
            subwords.append(piece)
            word_ids.append(word_index)
    tokens = [tokenizer.cls_token, *subwords, tokenizer.sep_token]
    return tokens, [None, *word_ids, None]

ESC

esc

ESC Word Sense Disambiguation.

ESC (Barba, Pasini & Navigli, "ESC: Redesigning WSD with Extractive Sense Comprehension", NAACL 2021, https://github.com/SapienzaNLP/esc) reframes WSD as extractive span comprehension: a mention's context sentence (with the target word bracketed by literal <classify>/</classify> markers, see insert_classify_markers) is paired with every candidate sense's gloss, concatenated into one sequence, and a standard extractive-QA head (transformers.AutoModelForQuestionAnswering) predicts which candidate's gloss span answers the "question" -- unlike GlossBertLinker's one-forward-pass-per-(mention, candidate) cross-encoder, or EwiserLinker's one-forward -pass-per-sentence full-inventory classifier: ESC is one forward pass per mention, jointly observing every one of that mention's candidates at once.

Reimplemented from the paper and from understanding the reference's own source (esc/esc_pl_module.py, esc/utils/definitions_tokenizer.py), not ported from it -- the reference repo is CC-BY-NC-SA 4.0, incompatible with redistributing literal code under this package's Apache-2.0 license (same constraint already documented in EwiserEncoder's module docstring for the EWISER port). Confirmed directly against the paper's own published SemCor checkpoint (its Lightning hyper_parameters/state_dict, not just the paper text):

• transformer_model: facebook/bart-large, squad_head: False, use_special_tokens: False, use_pmask: False -- the checkpoint uses the plain AutoModelForQuestionAnswering path, not any of the reference's custom XLNet/SQuAD-head/special-token variants. Its state_dict holds only qa_model.model.* (a stock BART encoder+decoder) and qa_model.qa_outputs.{weight,bias} (a Linear(hidden_size, 2) start/end head).
• transformers' BartForQuestionAnswering.forward runs the full encoder-decoder stack (not encoder-only) -- decoder_input_ids defaults to a right-shift of input_ids when omitted -- and applies the QA head to the decoder's last hidden state. A real architectural detail of this checkpoint, not an incidental implementation choice: simply calling AutoModelForQuestionAnswering.from_pretrained(...) with input_ids/attention_mask (and, for training, start_positions/end_positions) reproduces it exactly, with .loss computed internally as HF's own (start_loss + end_loss) / 2 cross-entropy -- no hand-rolled loss needed (see EscTrainer).
• Candidate scoring uses the reference's own "probabilistic" prediction type (the README's own example command, and what its headline numbers are reported under): start_logprob[start_token] + end_logprob[end_token] for each candidate's resolved gloss span, both taken from a log_softmax over the full (padded) sequence.

See EscEncoder and EscLinker.

EscEncoder

EscEncoder(
    model_name_or_path: str = "facebook/bart-large",
    max_length: int = 1024,
    forward_batch_size: int = 8,
)

Bases: Module

Extractive-QA cross-encoder scorer for WSD.

Unlike GlossBertEncoder (one forward pass per (mention, candidate) pair), score() runs one forward pass per mention, jointly observing every one of that mention's candidates present in pairs -- see score's own docstring.

Parameters:

Name	Type	Description	Default
model_name_or_path	str	A Hugging Face question-answering model id (e.g. "facebook/bart-large", the paper's own backbone), or a local directory. Pass a tiny test model for fast, network-light tests.	'facebook/bart-large'
max_length	int	Maximum joint-sequence token length (truncated).	1024
forward_batch_size	int	Number of distinct mentions score() batches through the encoder per forward pass -- unrelated to how many (mention, candidate_sense) pairs score() is called with (which may share far fewer distinct mentions).	8

Source code in src/linkingtk/algorithms/wsd/esc.py

def __init__(
    self,
    model_name_or_path: str = "facebook/bart-large",
    max_length: int = 1024,
    forward_batch_size: int = 8,
) -> None:
    super().__init__()
    from transformers import AutoModelForQuestionAnswering, AutoTokenizer

    self.tokenizer = AutoTokenizer.from_pretrained(
        model_name_or_path, use_fast=True, add_prefix_space=True
    )
    self.qa_model = AutoModelForQuestionAnswering.from_pretrained(model_name_or_path)
    self.max_length = max_length
    self.forward_batch_size = forward_batch_size

forward

forward(
    input_ids: Tensor,
    attention_mask: Tensor,
    start_positions: Tensor | None = None,
    end_positions: Tensor | None = None,
) -> Any

Thin pass-through to the wrapped QA model -- used by EscTrainer for its .loss (HF's own (start_loss + end_loss) / 2 cross-entropy) and by score() for its .start_logits/.end_logits.

Source code in src/linkingtk/algorithms/wsd/esc.py

def forward(
    self,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    start_positions: torch.Tensor | None = None,
    end_positions: torch.Tensor | None = None,
) -> Any:
    """Thin pass-through to the wrapped QA model -- used by
    [EscTrainer][linkingtk.train.esc_trainer.EscTrainer] for its
    `.loss` (HF's own `(start_loss + end_loss) / 2` cross-entropy) and
    by `score()` for its `.start_logits`/`.end_logits`.
    """
    return self.qa_model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        start_positions=start_positions,
        end_positions=end_positions,
    )

from_checkpoint classmethod

from_checkpoint(
    checkpoint_path: str | Path,
    model_name_or_path: str = "facebook/bart-large",
    max_length: int = 1024,
    forward_batch_size: int = 8,
) -> EscEncoder

Load the paper's own published Lightning checkpoint (e.g. the SemCor checkpoint linked from the reference README).

Parameters:

Name	Type	Description	Default
checkpoint_path	str | Path	Local path to a .ckpt file.	required
model_name_or_path	str	The backbone architecture to construct before loading the checkpoint's weights -- must match what the checkpoint was trained with (its own hyper_parameters["transformer_model"]; every published ESC checkpoint uses "facebook/bart-large").	'facebook/bart-large'
max_length	int	Forwarded to EscEncoder.__init__.	1024
forward_batch_size	int	Forwarded to EscEncoder.__init__.	8

Returns:

Type	Description
EscEncoder	An EscEncoder with the checkpoint's QA-head weights loaded.

Source code in src/linkingtk/algorithms/wsd/esc.py

@classmethod
def from_checkpoint(
    cls,
    checkpoint_path: str | Path,
    model_name_or_path: str = "facebook/bart-large",
    max_length: int = 1024,
    forward_batch_size: int = 8,
) -> EscEncoder:
    """Load the paper's own published Lightning checkpoint (e.g. the
    SemCor checkpoint linked from the reference README).

    Args:
        checkpoint_path: Local path to a `.ckpt` file.
        model_name_or_path: The backbone architecture to construct
            before loading the checkpoint's weights -- must match what
            the checkpoint was trained with (its own
            `hyper_parameters["transformer_model"]`; every published
            ESC checkpoint uses ``"facebook/bart-large"``).
        max_length: Forwarded to `EscEncoder.__init__`.
        forward_batch_size: Forwarded to `EscEncoder.__init__`.

    Returns:
        An `EscEncoder` with the checkpoint's QA-head weights loaded.
    """
    _stub_pytorch_lightning()
    checkpoint = torch.load(str(checkpoint_path), map_location="cpu", weights_only=False)
    state_dict = {
        key[len("qa_model.") :]: value
        for key, value in checkpoint["state_dict"].items()
        if key.startswith("qa_model.")
    }
    encoder = cls(
        model_name_or_path=model_name_or_path,
        max_length=max_length,
        forward_batch_size=forward_batch_size,
    )
    encoder.qa_model.load_state_dict(state_dict)
    return encoder

score

score(pairs: list[tuple[Entity, Entity]]) -> torch.Tensor

Return a (len(pairs),) tensor: each pair's start_logprob + end_logprob for its candidate's resolved gloss span.

Unlike a cross-encoder's score, this does not run one forward pass per pair: it groups pairs by mention, builds one joint sequence per mention (context sentence + every one of that mention's candidates' glosses, concatenated -- see build_joint_sequence), and runs the encoder once per mention (batched forward_batch_size mentions at a time). A candidate whose gloss span can't be resolved within the (possibly truncated) joint sequence scores -inf (never spuriously wins a match) -- same convention EwiserEncoder.score uses for an unresolved candidate.

Source code in src/linkingtk/algorithms/wsd/esc.py

def score(self, pairs: list[tuple[Entity, Entity]]) -> torch.Tensor:
    """Return a ``(len(pairs),)`` tensor: each pair's ``start_logprob +
    end_logprob`` for its candidate's resolved gloss span.

    Unlike a cross-encoder's `score`, this does **not** run one forward
    pass per pair: it groups `pairs` by mention, builds one joint
    sequence per mention (context sentence + every one of that
    mention's candidates' glosses, concatenated -- see
    `build_joint_sequence`), and runs the encoder once per mention
    (batched `forward_batch_size` mentions at a time). A candidate
    whose gloss span can't be resolved within the (possibly truncated)
    joint sequence scores ``-inf`` (never spuriously wins a match) --
    same convention `EwiserEncoder.score` uses for an unresolved
    candidate.
    """
    if not pairs:
        return torch.empty(0)

    mention_order: dict[int, int] = {}
    mentions: list[Entity] = []
    candidates_by_mention: dict[int, dict[str, Entity]] = {}
    pairs_by_mention: dict[int, list[int]] = {}
    for pair_index, (mention, sense) in enumerate(pairs):
        key = id(mention)
        if key not in mention_order:
            mention_order[key] = len(mentions)
            mentions.append(mention)
            candidates_by_mention[key] = {}
        candidates_by_mention[key].setdefault(sense.id, sense)
        pairs_by_mention.setdefault(mention_order[key], []).append(pair_index)

    scores: list[torch.Tensor | None] = [None] * len(pairs)
    for chunk_start in range(0, len(mentions), self.forward_batch_size):
        chunk_mentions = mentions[chunk_start : chunk_start + self.forward_batch_size]
        chunk_candidates = [
            list(candidates_by_mention[id(mention)].values()) for mention in chunk_mentions
        ]
        chunk_start_logprobs, chunk_end_logprobs, chunk_spans = self._encode_chunk(
            chunk_mentions, chunk_candidates
        )
        for offset in range(len(chunk_mentions)):
            mention_index = chunk_start + offset
            candidate_index_by_id = {
                sense.id: index for index, sense in enumerate(chunk_candidates[offset])
            }
            for pair_index in pairs_by_mention.get(mention_index, []):
                _mention, sense = pairs[pair_index]
                candidate_index = candidate_index_by_id[sense.id]
                span = chunk_spans[offset][candidate_index]
                if span is None:
                    scores[pair_index] = chunk_start_logprobs[offset].new_tensor(float("-inf"))
                else:
                    start_token, end_token = span
                    scores[pair_index] = (
                        chunk_start_logprobs[offset][start_token]
                        + chunk_end_logprobs[offset][end_token]
                    )

    assert all(score is not None for score in scores)  # every pair index was filled above
    return torch.stack(scores)  # type: ignore[arg-type]

EscLinker

EscLinker(
    model: EscEncoder, matching: Matcher = DEFAULT_MATCHER
)

Bases: BaseLinker

WSD linking scored by an EscEncoder.

encoder = EscEncoder.from_checkpoint("escher_semcor_best.ckpt")
linker = EscLinker(encoder)
results = linker.link(mentions, senses)

Or, for from-scratch/continued training, construct EscEncoder directly and train it via EscTrainer.

Parameters:

Name	Type	Description	Default
model	EscEncoder	The EscEncoder to score candidates with.	required
matching	Matcher	Strategy used to resolve scored candidates into final links. Defaults to GreedyMatcher.	DEFAULT_MATCHER

Source code in src/linkingtk/algorithms/wsd/esc.py

def __init__(self, model: EscEncoder, matching: Matcher = DEFAULT_MATCHER) -> None:
    self.model = model
    self.matching = matching

score_candidates

score_candidates(
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]

Blocked candidates per source entity, scored but not yet matched.

Satisfies CandidateScorer -- what link() itself builds internally, exposed so LlmRerankerLinker (#23) can re-rank a narrowed top-k instead of every blocked pair.

Source code in src/linkingtk/algorithms/wsd/esc.py

def score_candidates(
    self,
    dataset1: list[Entity],
    dataset2: list[Entity] | EntitySource,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> dict[str, list[tuple[str, float]]]:
    """Blocked candidates per source entity, scored but not yet matched.

    Satisfies [CandidateScorer][linkingtk.algorithms.llm_reranker.CandidateScorer]
    -- what `link()` itself builds internally, exposed so
    [LlmRerankerLinker][linkingtk.algorithms.llm_reranker.LlmRerankerLinker]
    (#23) can re-rank a narrowed top-k instead of every blocked pair.
    """
    pairs = list(blocking.candidate_pairs(dataset1, dataset2))
    if not pairs:
        return {}

    with torch.no_grad():
        self.model.eval()
        scores = self.model.score(pairs)
        self.model.train()

    candidates_by_source: dict[str, list[tuple[str, float]]] = defaultdict(list)
    for (entity1, entity2), score in zip(pairs, scores.tolist(), strict=True):
        candidates_by_source[entity1.id].append((entity2.id, score))

    return candidates_by_source

_esc_text

Joint-sequence construction helpers for EscEncoder.

ESC scores a mention's candidate senses with a single extractive-QA forward pass per mention: the context sentence (with the target word bracketed by literal <classify>/</classify> markers) followed by every candidate's gloss, concatenated, tokenized as one (sequence1, sequence2) pair. The correct candidate is identified by a token span within sequence2 -- these helpers build that joint sequence and resolve each candidate's span in it.

Reimplemented from understanding the reference's own esc/utils/definitions_tokenizer.py (prepare_sample_without_st), not copied -- the reference repo is CC-BY-NC-SA 4.0, incompatible with redistributing literal code under this package's Apache-2.0 license (see EwiserEncoder's module docstring for the same constraint on the EWISER port). Uses the fast tokenizer's own Encoding.sequence_ids property to tell sequence1's tokens from sequence2's, rather than the reference's own heuristic of scanning for two consecutive (0, 0)-offset tokens in a row -- a more direct, less fragile way to find the same boundary.

build_joint_sequence

build_joint_sequence(
    tokenizer: PreTrainedTokenizerBase,
    context_sentence: str,
    candidate_glosses: list[str],
    max_length: int,
) -> tuple[
    torch.Tensor, torch.Tensor, list[tuple[int, int] | None]
]

Tokenize (context_sentence, " ".join(candidate_glosses)) as one pair sequence and resolve each candidate's (start_token, end_token) span within it.

Returns (input_ids, attention_mask, spans) for one sequence (batching/padding happens at the caller, matching EwiserEncoder._encode_chunk's existing pattern) -- input_ids/ attention_mask are 1-D. spans[i] is None if candidate i's gloss span couldn't be resolved to token boundaries, or falls outside a sequence truncated to max_length (the caller treats this the same way EwiserEncoder.score treats an unresolved candidate: unscoreable, not a crash).

Source code in src/linkingtk/algorithms/wsd/_esc_text.py

def build_joint_sequence(
    tokenizer: PreTrainedTokenizerBase,
    context_sentence: str,
    candidate_glosses: list[str],
    max_length: int,
) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int] | None]]:
    """Tokenize ``(context_sentence, " ".join(candidate_glosses))`` as one
    pair sequence and resolve each candidate's ``(start_token, end_token)``
    span within it.

    Returns ``(input_ids, attention_mask, spans)`` for **one** sequence
    (batching/padding happens at the caller, matching
    `EwiserEncoder._encode_chunk`'s existing pattern) -- `input_ids`/
    `attention_mask` are 1-D. `spans[i]` is `None` if candidate `i`'s gloss
    span couldn't be resolved to token boundaries, or falls outside a
    sequence truncated to `max_length` (the caller treats this the same
    way `EwiserEncoder.score` treats an unresolved candidate: unscoreable,
    not a crash).
    """
    definitions_seq = " ".join(candidate_glosses)
    definitions_offsets: list[tuple[int, int]] = []
    cursor = 0
    for gloss in candidate_glosses:
        start = cursor
        end = start + len(gloss)
        definitions_offsets.append((start, end))
        cursor = end + 1  # +1 for the joining space

    encoded = tokenizer(
        context_sentence,
        definitions_seq,
        return_tensors="pt",
        truncation=True,
        max_length=max_length,
    )
    input_ids = encoded["input_ids"][0]
    attention_mask = encoded["attention_mask"][0]
    encoding = encoded.encodings[0]
    offsets = encoding.offsets
    sequence_ids = encoding.sequence_ids

    char_start_to_token: dict[int, int] = {}
    char_end_to_token: dict[int, int] = {}
    for token_index, (sequence_id, (char_start, char_end)) in enumerate(
        zip(sequence_ids, offsets, strict=True)
    ):
        if sequence_id != 1:
            continue
        char_start_to_token.setdefault(char_start, token_index)
        char_end_to_token.setdefault(char_end, token_index)

    spans: list[tuple[int, int] | None] = []
    for gloss_start, gloss_end in definitions_offsets:
        start_token = char_start_to_token.get(gloss_start)
        end_token = char_end_to_token.get(gloss_end)
        if start_token is None or end_token is None:
            spans.append(None)
        else:
            spans.append((start_token, end_token))

    return input_ids, attention_mask, spans

candidate_gloss

candidate_gloss(sense: Entity) -> str

Definition text, capitalized, with a trailing period -- matches the reference's own gloss formatting (synset.definition().capitalize() + ".").

Source code in src/linkingtk/algorithms/wsd/_esc_text.py

def candidate_gloss(sense: Entity) -> str:
    """``Definition text, capitalized, with a trailing period`` -- matches
    the reference's own gloss formatting (`synset.definition().capitalize()
    + "."`)."""
    text = description_text(sense).strip()
    if not text:
        return "."
    text = text[0].upper() + text[1:]
    return text if text.endswith(".") else f"{text}."

insert_classify_markers

insert_classify_markers(
    text: str, start: int, end: int
) -> str

Bracket text's [start, end) mention span with literal <classify>/</classify> markers (plain text, not added special tokens -- matches the published checkpoint's own use_special_tokens=False training configuration).

Source code in src/linkingtk/algorithms/wsd/_esc_text.py

def insert_classify_markers(text: str, start: int, end: int) -> str:
    """Bracket `text`'s ``[start, end)`` mention span with literal
    ``<classify>``/``</classify>`` markers (plain text, not added special
    tokens -- matches the published checkpoint's own
    ``use_special_tokens=False`` training configuration).
    """
    return f"{text[:start]}{CLASSIFY_START} {text[start:end]} {CLASSIFY_END}{text[end:]}"

pad_encoded_sequences

pad_encoded_sequences(
    sequences: list[tuple[Tensor, Tensor]],
    pad_token_id: int,
) -> tuple[torch.Tensor, torch.Tensor]

Right-pad a list of (input_ids, attention_mask) 1-D tensor pairs (e.g. from repeated build_joint_sequence calls) into one batch.

Shared by EscEncoder's scoring path and EscTrainer's training path, so both batch mentions of differing joint-sequence length the same way.

Source code in src/linkingtk/algorithms/wsd/_esc_text.py

def pad_encoded_sequences(
    sequences: list[tuple[torch.Tensor, torch.Tensor]], pad_token_id: int
) -> tuple[torch.Tensor, torch.Tensor]:
    """Right-pad a list of ``(input_ids, attention_mask)`` 1-D tensor pairs
    (e.g. from repeated `build_joint_sequence` calls) into one batch.

    Shared by [EscEncoder][linkingtk.algorithms.wsd.esc.EscEncoder]'s
    scoring path and [EscTrainer][linkingtk.train.esc_trainer.EscTrainer]'s
    training path, so both batch mentions of differing joint-sequence
    length the same way.
    """
    max_len = max(ids.size(0) for ids, _mask in sequences)
    batch_input_ids = torch.full((len(sequences), max_len), pad_token_id, dtype=torch.long)
    batch_attention_mask = torch.zeros((len(sequences), max_len), dtype=torch.long)
    for row, (ids, mask) in enumerate(sequences):
        batch_input_ids[row, : ids.size(0)] = ids
        batch_attention_mask[row, : mask.size(0)] = mask
    return batch_input_ids, batch_attention_mask

Word Sense Alignment

wsa

Word Sense Alignment (WSA) algorithms.

WSA finds equivalent senses between two dictionary entries, either monolingually or across languages. Related systems: NAISC.