"DB100K",
"OpenEA",
"Countries",
"WD50KT",
"Wikidata5M",
"PharmKG8k",
"PharmKG",
# Utilities
"dataset_resolver",
"get_dataset",
"has_dataset",
]
logger = logging.getLogger(__name__)
dataset_resolver = ClassResolver.from_entrypoint(group="pykeen.datasets",
base=Dataset)
if not dataset_resolver.lookup_dict:
raise RuntimeError(
dedent("""\
Datasets have been loaded with entrypoints since PyKEEN v1.0.5, which is now a
very old version of PyKEEN.
If you simply use `python3 -m pip install --upgrade pykeen`, the entrypoints will
not be reloaded. Instead, please reinstall PyKEEN using the following commands:
$ python3 -m pip uninstall pykeen
$ python3 -m pip install pykeen
If you are on Kaggle or Google Colab, please follow these instructions:
https://pykeen.readthedocs.io/en/stable/installation.html#google-colab-and-kaggle-users
"DistMultInteraction",
"DistMAInteraction",
"ERMLPInteraction",
"ERMLPEInteraction",
"HolEInteraction",
"KG2EInteraction",
"MuREInteraction",
"NTNInteraction",
"PairREInteraction",
"ProjEInteraction",
"RESCALInteraction",
"RotatEInteraction",
"SimplEInteraction",
"SEInteraction",
"TorusEInteraction",
"TransDInteraction",
"TransEInteraction",
"TransFInteraction",
"TransHInteraction",
"TransRInteraction",
"TripleREInteraction",
"TuckerInteraction",
"UMInteraction",
]
representation_resolver: ClassResolver[
Representation] = ClassResolver.from_subclasses(
base=Representation,
default=Embedding,
)
x_e: Optional[torch.FloatTensor] = None,
) -> torch.FloatTensor: # noqa: D102
if message is None or x_e is None:
raise ValueError(
f"{self.__class__.__name__} requires message and x_e.")
# view for heads
message_ = message.view(message.shape[0], self.num_heads, -1)
# compute attention coefficients, shape: (num_edges, num_heads)
alpha = self.activation(
torch.einsum(
"ihd,hd->ih",
torch.cat(
[
message_,
x_e[target].view(target.shape[0], self.num_heads, -1),
],
dim=-1,
),
self.weight,
))
# TODO we can use scatter_softmax from torch_scatter directly, kept this if we can rewrite it w/o scatter
alpha = softmax(alpha, index=target, num_nodes=x_e.shape[0], dim=0)
alpha = self.dropout(alpha)
return (message_ * alpha.view(-1, self.num_heads, 1)).view(
-1, self.num_heads * self.attention_dim)
edge_weight_resolver = ClassResolver.from_subclasses(
base=EdgeWeighting, default=SymmetricEdgeWeighting)
if self.stopper.should_evaluate(epoch):
# TODO how to pass inductive mode
if self.stopper.should_stop(epoch):
self.training_loop._should_stop = True
# Since the model is also used within the stopper, its graph and cache have to be cleared
self.model._free_graph_and_cache()
# When the stopper obtained a new best epoch, this model has to be saved for reconstruction
if self.stopper.best_epoch != self.last_best_epoch and self.best_epoch_model_file_path is not None:
self.training_loop._save_state(
path=self.best_epoch_model_file_path,
triples_factory=self.triples_factory)
self.last_best_epoch = epoch
callback_resolver: ClassResolver[
TrainingCallback] = ClassResolver.from_subclasses(base=TrainingCallback, )
#: A hint for constructing a :class:`MultiTrainingCallback`
TrainingCallbackHint = OneOrSequence[HintOrType[TrainingCallback]]
TrainingCallbackKwargsHint = OneOrSequence[OptionalKwargs]
class MultiTrainingCallback(TrainingCallback):
"""A wrapper for calling multiple training callbacks together."""
#: A collection of callbacks
callbacks: List[TrainingCallback]
def __init__(
self,
callbacks: TrainingCallbackHint = None,
There are two other major considerations when randomly sampling negative triples: the random sampling
strategy and the filtering of positive triples. A full guide on negative sampling with the SLCWA can be
found in :mod:`pykeen.sampling`. The following chart from [ali2020a]_ demonstrates the different potential
triples considered in LCWA vs. sLCWA based on the given true triples (in red):
.. image:: ../img/training_approaches.png
:alt: Troubleshooting Image 2
""" # noqa:E501
from class_resolver import ClassResolver
from .callbacks import TrainingCallback # noqa: F401
from .lcwa import LCWATrainingLoop # noqa: F401
from .slcwa import SLCWATrainingLoop # noqa: F401
from .training_loop import NonFiniteLossError, TrainingLoop # noqa: F401
__all__ = [
"TrainingLoop",
"SLCWATrainingLoop",
"LCWATrainingLoop",
"NonFiniteLossError",
"training_loop_resolver",
"TrainingCallback",
]
training_loop_resolver = ClassResolver.from_subclasses(
base=TrainingLoop, # type: ignore
default=SLCWATrainingLoop,
)
x=x,
source=target,
target=source,
edge_type=edge_type,
edge_weights=edge_weights,
accumulator=y,
)
if self.bias is not None:
y = y + self.bias
# activation
if self.activation is not None:
y = self.activation(y)
return y
decomposition_resolver = ClassResolver.from_subclasses(
base=Decomposition, default=BasesDecomposition)
class RGCNRepresentation(Representation):
r"""Entity representations enriched by R-GCN.
The GCN employed by the entity encoder is adapted to include typed edges.
The forward pass of the GCN is defined by:
.. math::
\textbf{e}_{i}^{l+1} = \sigma \left( \sum_{r \in \mathcal{R}}\sum_{j\in \mathcal{N}_{i}^{r}}
\frac{1}{c_{i,r}} \textbf{W}_{r}^{l} \textbf{e}_{j}^{l} + \textbf{W}_{0}^{l} \textbf{e}_{i}^{l}\right)
where $\mathcal{N}_{i}^{r}$ is the set of neighbors of node $i$ that are connected to
$i$ by relation $r$, $c_{i,r}$ is a fixed normalization constant (but it can also be introduced as an additional
def __init__(
self,
regularizers: Iterable[Regularizer],
total_weight: float = 1.0,
apply_only_once: bool = False,
):
super().__init__(weight=total_weight, apply_only_once=apply_only_once)
self.regularizers = nn.ModuleList(regularizers)
for r in self.regularizers:
if isinstance(r, NoRegularizer):
raise TypeError("Can not combine a no-op regularizer")
self.register_buffer(
name="normalization_factor",
tensor=torch.as_tensor(
sum(r.weight for r in self.regularizers),
).reciprocal(),
)
@property
def normalize(self): # noqa: D102
return any(r.normalize for r in self.regularizers)
def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102
return self.normalization_factor * sum(r.weight * r.forward(x) for r in self.regularizers)
regularizer_resolver: ClassResolver[Regularizer] = ClassResolver.from_subclasses(
base=Regularizer,
default=NoRegularizer,
)
increasing = False
def __call__(self, ranks: np.ndarray, num_candidates: Optional[np.ndarray] = None) -> float: # noqa: D102
return super().__call__(ranks=ranks) / self.expected_value(num_candidates=num_candidates)
@parse_docdata
class AdjustedArithmeticMeanRankIndex(ArithmeticMeanRank):
"""The adjusted arithmetic mean rank index (AMRI).
---
link: https://arxiv.org/abs/2002.06914
description: The re-indexed adjusted mean rank (AAMR)
"""
name = "Adjusted Arithmetic Mean Rank Index (AAMRI)"
value_range = ValueRange(lower=-1, lower_inclusive=True, upper=1, upper_inclusive=True)
synonyms = ("adjusted_mean_rank_index", "amri", "aamri")
increasing = True
supported_rank_types = (RANK_REALISTIC,)
needs_candidates = True
def __call__(self, ranks: np.ndarray, num_candidates: Optional[np.ndarray] = None) -> float: # noqa: D102
return 1.0 - (super().__call__(ranks=ranks) - 1.0) / (self.expected_value(num_candidates=num_candidates) - 1.0)
rank_based_metric_resolver: ClassResolver[RankBasedMetric] = ClassResolver.from_subclasses(
base=RankBasedMetric,
default=InverseHarmonicMeanRank, # mrr
)
def __call__(
self,
edge_index: numpy.ndarray,
known_anchors: Optional[numpy.ndarray] = None,
) -> numpy.ndarray: # noqa: D102
anchors = known_anchors or None
for selection in self.selections:
anchors = selection(edge_index=edge_index, known_anchors=anchors)
return anchors
anchor_selection_resolver: ClassResolver[
AnchorSelection] = ClassResolver.from_subclasses(
base=AnchorSelection,
default=DegreeAnchorSelection,
skip={SingleSelection},
)
class AnchorSearcher:
"""A method for finding the closest anchors."""
@abstractmethod
def __call__(self, edge_index: numpy.ndarray, anchors: numpy.ndarray,
k: int) -> numpy.ndarray:
"""
Find the $k$ closest anchor nodes for each entity.
:param edge_index: shape: (2, m)
the edge index
:param anchors: shape: (a,)
"NeptuneResultTracker",
"WANDBResultTracker",
"JSONResultTracker",
"CSVResultTracker",
"PythonResultTracker",
"TensorBoardResultTracker",
"ConsoleResultTracker",
# Utilities
"tracker_resolver",
"TrackerHint",
"resolve_result_trackers",
]
tracker_resolver: ClassResolver[ResultTracker] = ClassResolver.from_subclasses(
base=ResultTracker,
default=ResultTracker,
skip={FileResultTracker, MultiResultTracker},
)
def resolve_result_trackers(
result_tracker: OneOrManyHintOrType[ResultTracker] = None,
result_tracker_kwargs: OneOrManyOptionalKwargs = None,
) -> MultiResultTracker:
"""Resolve and compose result trackers.
:param result_tracker: Either none (will result in a Python result tracker),
a single tracker (as either a class, instance, or string for class name), or a list
of trackers (as either a class, instance, or string for class name
:param result_tracker_kwargs: Either none (will use all defaults), a single dictionary
(will be used for all trackers), or a list of dictionaries with the same length
# -*- coding: utf-8 -*-
"""Evaluation."""
from class_resolver import ClassResolver
from .classification_evaluator import ClassificationEvaluator, ClassificationMetricResults
from .evaluator import Evaluator, MetricResults, evaluate
from .rank_based_evaluator import RankBasedEvaluator, RankBasedMetricResults
__all__ = [
"evaluate",
"Evaluator",
"MetricResults",
"RankBasedEvaluator",
"RankBasedMetricResults",
"ClassificationEvaluator",
"ClassificationMetricResults",
"evaluator_resolver",
"metric_resolver",
]
evaluator_resolver: ClassResolver[Evaluator] = ClassResolver.from_subclasses(
base=Evaluator,
default=RankBasedEvaluator,
)
metric_resolver: ClassResolver[MetricResults] = ClassResolver.from_subclasses(
MetricResults)
def forward(self, a: torch.FloatTensor, b: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102
return self.__class__.func(a, b)
class SubtractionCompositionModule(FunctionalCompositionModule):
"""Composition by element-wise subtraction."""
func = torch.sub
class MultiplicationCompositionModule(FunctionalCompositionModule):
"""Composition by element-wise multiplication."""
func = torch.mul
class CircularCorrelationCompositionModule(FunctionalCompositionModule):
"""Composition by circular correlation via :func:`pykeen.nn.functional.circular_correlation`."""
func = circular_correlation
composition_resolver = ClassResolver.from_subclasses(
CompositionModule,
default=MultiplicationCompositionModule,
skip={
FunctionalCompositionModule,
},
)
from pykeen.pipeline import pipeline
results = pipeline(
dataset='YAGO3-10',
model='PairRE',
training_loop='sLCWA',
negative_sampler='bernoulli',
)
""" # noqa
from class_resolver import ClassResolver
from .basic_negative_sampler import BasicNegativeSampler
from .bernoulli_negative_sampler import BernoulliNegativeSampler
from .negative_sampler import NegativeSampler
from .pseudo_type import PseudoTypedNegativeSampler
__all__ = [
"NegativeSampler",
"BasicNegativeSampler",
"BernoulliNegativeSampler",
"PseudoTypedNegativeSampler",
# Utils
"negative_sampler_resolver",
]
negative_sampler_resolver = ClassResolver.from_subclasses(
NegativeSampler,
default=BasicNegativeSampler,
)
"TuckER",
"UM",
# Inductive Models
"InductiveNodePiece",
"InductiveNodePieceGNN",
# Evaluation-only models
"SoftInverseTripleBaseline",
"MarginalDistributionBaseline",
# Utils
"model_resolver",
"make_model",
"make_model_cls",
]
model_resolver = ClassResolver.from_subclasses(
base=Model,
skip={
# Abstract Models
_NewAbstractModel,
# We might be able to relax this later
ERModel,
LiteralModel,
# baseline models behave differently
EvaluationOnlyModel,
*get_subclasses(EvaluationOnlyModel),
# Old style models should never be looked up
_OldAbstractModel,
EntityRelationEmbeddingModel,
},
)
... loss='marginranking',
... loss_kwargs=dict(margin=1),
... training_loop='slcwa',
... training_kwargs=dict(num_epochs=100, batch_size=128),
... negative_sampler='basic',
... negative_sampler_kwargs=dict(num_negs_per_pos=1),
... evaluator_kwargs=dict(filtered=True),
... evaluation_kwargs=dict(batch_size=128),
... stopper='early',
... stopper_kwargs=dict(frequency=5, patience=2, relative_delta=0.002),
... )
"""
from class_resolver import ClassResolver
from .early_stopping import EarlyStopper, StopperCallback # noqa: F401
from .stopper import NopStopper, Stopper
__all__ = [
"Stopper",
"NopStopper",
"EarlyStopper",
# Utils
"stopper_resolver",
]
stopper_resolver = ClassResolver.from_subclasses(
Stopper,
default=NopStopper,
)
lower=-1.0,
upper=1.0,
description="A balanced measure applicable even with class imbalance",
link="https://en.wikipedia.org/wiki/Phi_coefficient",
)
# TODO there's something wrong with this, so add it later
# classifier_annotator.higher(
# rmc.pr_auc_score,
# name="AUC-PR",
# description="Area Under the Precision-Recall Curve",
# link="https://rexmex.readthedocs.io/en/latest/modules/root.html#rexmex.metrics.classification.pr_auc_score",
# )
classification_metric_resolver: ClassResolver[ClassificationMetric] = ClassResolver(
list(classifier_annotator.metrics.values()),
base=ClassificationMetric,
)
def _check():
"""Check that all functions in the classification module are annotated."""
for func in rmc.__dict__.values():
if not inspect.isfunction(func):
continue
parameters = inspect.signature(func).parameters
if "y_true" not in parameters or "y_score" not in parameters:
continue
if func in EXCLUDE:
if func in classifier_annotator.metrics:
raise ValueError(f"should not include {func.__name__}")
continue