updated the conditions
@@ -57,7 +57,13 @@ class ReciprocalRelationsModel(KgeModel):
def penalty(self, **kwargs):
penalty_result = self._base_model.penalty(**kwargs)
- if "batch" in kwargs and "triples" in kwargs["batch"]:
+
+ # Penalize the reciprocal relations when performing weighted regularization
+ is_weighted = self.get_p_embedder().get_option("regularize_args.weighted")
+ regularize = self.get_p_embedder().regularize
+ regularize_weight = self.get_p_embedder().get_option("regularize_weight")
+
+ if is_weighted and (regularize != "") and (regularize_weight != 0.0):
triples = kwargs["batch"]["triples"].to(self.config.get("job.device"))
reciprocal_indexes = triples[:, 1] + self.dataset.num_relations()
penalty_result += self.get_p_embedder().penalty(
