Ensemble updated

Author: Demirrr

retrieval_aug_predictors/__init__.py

@@ -4,4 +4,4 @@
 from .models.RALP import RALP
 from .models.RCL import RCL
 from .models.Demir import Demir
-
+from .models.demir_ensemble import DemirEnsemble

retrieval_aug_predictors/models/demir_ensemble.py

@@ -1,20 +1,20 @@
 """
-python -m retrieval_aug_predictors.models.Demir --dataset_dir KGs/Countries-S1 --out "countries_s1_results.json" && cat countries_s1_results.json
+python -m retrieval_aug_predictors.models.demir_ensemble --dataset_dir KGs/Countries-S1 --out "countries_s1_results.json" && cat countries_s1_results.json
 {
     "H@1": 1.0,
     "H@3": 1.0,
     "H@10": 1.0,
     "MRR": 1.0
 }
 
-python -m retrieval_aug_predictors.models.Demir --dataset_dir KGs/Countries-S2 --out "countries_s2_results.json" && cat countries_s2_results.json
+python -m retrieval_aug_predictors.models.demir_ensemble --dataset_dir KGs/Countries-S2 --out "countries_s2_results.json" && cat countries_s2_results.json
 {
     "H@1": 0.9583333333333334,
     "H@3": 0.9583333333333334,
     "H@10": 1.0,
     "MRR": 0.9666666666666667
 }
-python -m retrieval_aug_predictors.models.Demir --dataset_dir KGs/Countries-S3 --out "countries_s3_results.json" && cat countries_s3_results.json
+python -m retrieval_aug_predictors.models.demir_ensemble --dataset_dir KGs/Countries-S3 --out "countries_s3_results.json" && cat countries_s3_results.json
 {
     "H@1": 0.875,
     "H@3": 0.9583333333333334,
@@ -231,14 +231,17 @@ def forward_k_vs_all(self, x: torch.LongTensor) -> torch.FloatTensor:
         return torch.FloatTensor(batch_predictions)
     def forward_triples(self, x: torch.LongTensor) -> torch.FloatTensor:
         raise NotImplementedError("RCL needs to implement it")
+
+
 # test the dspy model -> remove later
 if __name__ == "__main__":
     args=parser.parse_args()
     # Important: add_reciprocal=False in KvsAll implies that inverse relation has been introduced.
     # Therefore, The link prediction results are based on the missing tail rankings only!
     kg = KG(dataset_dir=args.dataset_dir, separator="\s+", eval_model=args.eval_model, add_reciprocal=False)
     sanity_checking(args,kg)
-    model = DemirEnsemble(knowledge_graph=kg, base_url=args.base_url, api_key=args.api_key, llm_model=args.llm_model_name, temperature=args.temperature, seed=args.seed)
+    model = DemirEnsemble(knowledge_graph=kg, base_url=args.base_url, api_key=args.api_key,
+                              llm_model=args.llm_model_name, temperature=args.temperature, seed=args.seed)
     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)
     if args.out and results: