Ensemble of Chain of Thought Models leading to the best link prediction result on Countries

Author: Demirrr

dicee/knowledge_graph.py

@@ -3,6 +3,8 @@
 import sys
 import pandas as pd
 import polars as pl
+import numpy as np
+
 class KG:
     """ Knowledge Graph """
 
@@ -137,8 +139,22 @@ def exists(self,h:str,r:str,t:str):
         return ((self.raw_train_set == pd.Series(row_to_check)).all(axis=1)).any()
 
     def __iter__(self):
-        for h, r, t in self.raw_train_set.to_numpy().tolist():
+        if self.raw_train_set is not None:
+            graph=self.raw_train_set.to_numpy()
+        elif self.train_set is not None:
+            assert isinstance(self.train_set,np.ndarray)
+            graph=self.train_set
+        else:
+            raise RuntimeError(f"Dataset {self.dataset_dir} and {self.raw_train_set} & {self.train_set} are None")
+        assert graph.shape[0]>=0 and graph.shape[1]==3, "Invalid graph shape!"
+
+        if hasattr(self,"idx_to_entity") is False:
+            self.idx_to_entity = self.entity_to_idx.set_index(self.entity_to_idx.index)['entity'].to_dict()
+            self.idx_to_relations = self.relation_to_idx.set_index(self.relation_to_idx.index)['relation'].to_dict()
+
+        for h, r, t in graph.tolist():
             yield self.idx_to_entity[h], self.idx_to_relations[r], self.idx_to_entity[t]
+
     def __len__(self):
         return len(self.raw_train_set)
 

dicee/read_preprocess_save_load_kg/preprocess.py

@@ -311,7 +311,6 @@ def preprocess_with_pandas(self) -> None:
         self.kg.raw_test_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprocal,
                                                          eval_model=self.kg.eval_model,
                                                          df=self.kg.raw_test_set, info="Test")
-
         # (2) Construct integer indexing for entities and relations.
         self.sequential_vocabulary_construction()
         self.kg.num_entities, self.kg.num_relations = len(self.kg.entity_to_idx), len(self.kg.relation_to_idx)

retrieval_aug_predictors/models/Demir.py

@@ -14,7 +14,7 @@
     "H@10": 0.9583333333333334,
     "MRR": 0.7921296296296297
 }
-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.7083333333333334,
     "H@3": 0.9583333333333334,

retrieval_aug_predictors/models/demir_ensemble.py

@@ -6,21 +6,20 @@
     "H@10": 1.0,
     "MRR": 1.0
 }
-
 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@1": 1.0,
+    "H@3": 1.0,
     "H@10": 1.0,
-    "MRR": 0.9666666666666667
+    "MRR": 1.0
 }
 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,
+    "H@1": 0.9166666666666666,
+    "H@3": 1.0,
     "H@10": 1.0,
-    "MRR": 0.9249999999999999
-}
+    "MRR": 0.951388888888889
+}(
 """
 
 import dspy
@@ -49,7 +48,7 @@ class MultiLabelLinkPredictionWithScores(dspy.Signature):
 class MultiLabelLinkPredictor(dspy.Module):
     def __init__(self):
         super().__init__()
-        self.predictor = dspy.Predict(MultiLabelLinkPredictionWithScores)
+        self.predictor = dspy.ChainOfThought(MultiLabelLinkPredictionWithScores)
 
     def forward(self, subject, predicate, few_shot_examples) -> List[Tuple[str, float]]:
         # Format examples more structured with clearer JSON expectations
@@ -96,16 +95,13 @@ def __init__(self, knowledge_graph, base_url, api_key, temperature, seed, llm_mo
         super().__init__(knowledge_graph, name="DemirEnsemble")
         self.temperature = temperature
         self.seed = seed
-
-        # Create multiple LLM models with different parameters
-        self.lm_high_temp = dspy.LM(model=f"openai/{llm_model}", api_key=api_key,
-                                    api_base=base_url, seed=seed, temperature=0.7,
-                                    cache=True, cache_in_memory=True)
-
-        self.lm_low_temp = dspy.LM(model=f"openai/{llm_model}", api_key=api_key,
-                                   api_base=base_url, seed=seed, temperature=0.1,
-                                   cache=True, cache_in_memory=True)
-
+        # () Initialize ensemble.
+        self.ensemble=[]
+        for i in range(0, 9):
+            temperature_coefficient=i*0.1
+            self.ensemble.append(dspy.LM(model=f"openai/{llm_model}", api_key=api_key,
+                    api_base=base_url, seed=seed, temperature=temperature_coefficient,
+                    cache=True, cache_in_memory=True))
         # Initialize data same as original
         self.train_set = [(self.idx_to_entity[idx_h],
                            self.idx_to_relation[idx_r],
@@ -136,14 +132,14 @@ def __init__(self, knowledge_graph, base_url, api_key, temperature, seed, llm_mo
     def _create_ensemble_predictors(self):
         """Create multiple predictors with different configurations"""
         predictors = []
-
-        # Standard predictor
-        dspy.configure(lm=self.lm_low_temp)
-        predictors.append(MultiLabelLinkPredictor())
+        for i in self.ensemble:
+            # Standard predictor
+            dspy.configure(lm=i)
+            predictors.append(MultiLabelLinkPredictor())
 
         # Diverse predictor (high temp)
-        dspy.configure(lm=self.lm_high_temp)
-        predictors.append(MultiLabelLinkPredictor())
+        #dspy.configure(lm=self.lm_high_temp)
+        #predictors.append(MultiLabelLinkPredictor())
 
         return predictors