renaming and small tweaks

Author: alkidbaci

retrieval_aug_predictors/extend_kg.py retrieval_aug_predictors/enrich_kg.pyretrieval_aug_predictors/extend_kg.py renamed to retrieval_aug_predictors/enrich_kg.py

@@ -25,8 +25,9 @@ def run(args):
         triples = f.readlines()
 
     for prd in predictions:
-        triple = prd[0] + "\t" + prd[1] + " \t" + prd[2][0][0]
-        triples.append(triple + "\n")
+        triple = prd[0] + "\t" + prd[1] + " \t" + prd[2][0][0]+ "\n"
+        if triple not in triples:
+            triples.append(triple)
 
     with open(args.dataset_dir + "/" + args.kg_out, "w") as out:
         out.writelines(triples)
@@ -40,7 +41,7 @@ def run(args):
     parser.add_argument("--pred_out", type=str, default=None,
                         help="Name of the output file where the predictions will be saved.")
 
-    parser.add_argument("--kg_out", type=str, default="extended_train.txt",
+    parser.add_argument("--kg_out", type=str, default="enriched_train.txt",
                         help="Name of the output file where the extended train set will be saved.")
 
     run(parser.parse_args())

retrieval_aug_predictors/models/RALP.py retrieval_aug_predictors/models/CATC.pyretrieval_aug_predictors/models/RALP.py renamed to retrieval_aug_predictors/models/CATC.py

@@ -4,7 +4,8 @@
 from retrieval_aug_predictors.models import KG, AbstractBaseLinkPredictorClass
 
 
-class RALP(AbstractBaseLinkPredictorClass):
+class CATC(AbstractBaseLinkPredictorClass):
+    """Context-Aware Triple Completion"""
     def __init__(self, knowledge_graph: KG = None,
                  name="ralp-1.0",
                  base_url="http://tentris-ml.cs.upb.de:8501/v1",

retrieval_aug_predictors/models/demir_ensemble.py

@@ -88,8 +88,9 @@ def forward(self, subject, predicate, few_shot_examples) -> List[Tuple[str, floa
             return []
 
 
-class DemirEnsemble(AbstractBaseLinkPredictorClass):
-    """Ensemble approach combining multiple prediction strategies"""
+class RALP(AbstractBaseLinkPredictorClass):
+    """Retrieval-Augmented Link Prediction (ex DemirEnsemble).
+    Ensemble approach combining multiple prediction strategies"""
 
     def __init__(self, knowledge_graph, base_url, api_key, temperature, seed, llm_model, use_val: bool = False):
         super().__init__(knowledge_graph, name="DemirEnsemble")
@@ -288,7 +289,7 @@ def process(scores, storage, **kwargs):
         results:dict = evaluate_lp_k_vs_all(model=model, triple_idx=kg.test_set[:args.eval_size],
                              er_vocab=kg.er_vocab, info='Eval KvsAll Starts', batch_size=args.batch_size)
     else:
-        x = kg.test_set[:, [0, 1]]
+        x = kg.train_set[:, [0, 1]]
         results = model.get_predicted_triples(x, args.k)
 
     print("Results: {}".format(results))