argparse included

Author: Demirrr

retrieval_aug_predictors/arguments.py

@@ -0,0 +1,27 @@
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument("--dataset_dir", type=str, default="/home/cdemir/Desktop/Softwares/dice-embeddings/KGs/Countries-S1", help="Path to dataset.")
+parser.add_argument("--model", type=str, default="Demir", help="Model name to use for link prediction.",
+                    choices=["Demir", "GCL", "RCL", "RALP"])
+parser.add_argument("--base_url", type=str, default="http://harebell.cs.upb.de:8501/v1",
+                    choices=["http://harebell.cs.upb.de:8501/v1", "http://tentris-ml.cs.upb.de:8502/v1"],
+                    help="Base URL for the OpenAI client.")
+parser.add_argument("--llm_model_name", type=str, default="tentris", help="Model name of the LLM to use.")
+parser.add_argument("--temperature", type=float, default=0.0, help="Temperature hyperparameter for LLM calls.")
+parser.add_argument("--api_key", type=str, default=None, help="API key for the OpenAI client. If left to None, "
+                                                              "it will look at the environment variable named "
+                                                              "TENTRIS_TOKEN from a local .env file.")
+parser.add_argument("--eval_size", type=int, default=None,
+                    help="Amount of triples from the test set to evaluate. "
+                         "Leave it None to include all triples on the test set.")
+parser.add_argument("--eval_model", type=str, default="train_value_test",
+                    help="Type of evaluation model.")
+parser.add_argument("--batch_size", type=int, default=1)
+parser.add_argument("--chunk_size", type=int, default=1)
+parser.add_argument("--seed", type=int, default=42)
+parser.add_argument("--num_of_hops", type=int, default=1,
+                    help="Number of hops to use to extract a subgraph around an entity.")
+parser.add_argument("--max_relation_examples", type=int, default=2000,
+                    help="Maximum number of relation examples to include in RCL context.")
+parser.add_argument("--exclude_source", action="store_true",
+                    help="Exclude triples with the same source entity in RCL context.")

retrieval_aug_predictors/models/Demir.py

@@ -4,31 +4,18 @@
 from typing import List, Tuple
 from retrieval_aug_predictors.models import KG, AbstractBaseLinkPredictorClass
 from openai import OpenAI
-
-# 1. Define the Signature
-class KGLikelihood(dspy.Signature):
-    """Assess the likelihood that a triple (subject, predicate, candidate_object) is true,
-    given some context triples. Output a score between 0.0 and 1.0."""
-
-    context = dspy.InputField(desc="Known knowledge graph triples.")
-    subject = dspy.InputField(desc="The subject entity.")
-    predicate = dspy.InputField(desc="The relationship type.")
-    candidate_object = dspy.InputField(desc="The candidate object entity to score.")
-
-    score = dspy.OutputField(desc="A likelihood score between 0.0 and 1.0.")
-
+from collections import OrderedDict
+from retrieval_aug_predictors.arguments import parser
+from retrieval_aug_predictors.utils import sanity_checking
+from dicee.evaluator import evaluate_lp, evaluate_lp_k_vs_all
+from dotenv import load_dotenv
+load_dotenv()
 
 class MultiLabelLinkPredictionWithScores(dspy.Signature):
-    """Given a subject entity and a predicate (relation), predict a list of
-    object entities that satisfy the relation, along with a likelihood score for each.
-    Use the provided examples as a guide.
-    Output a JSON formatted list of objects, where each object has an 'entity' (string)
-    and a 'score' (float between 0.0 and 1.0) key."""
-
     examples = dspy.InputField(
         desc="Few-shot examples of (subject, predicate) -> [{'entity': entity1, 'score': score1}, ...].")
-    subject = dspy.InputField(desc="The subject entity.")
-    predicate = dspy.InputField(desc="The relationship type.")
+    subject:str = dspy.InputField(desc="The subject entity.")
+    predicate:str = dspy.InputField(desc="The relationship type.")
 
     # Updated OutputField requesting JSON
     objects_with_scores = dspy.OutputField(
@@ -48,6 +35,32 @@ def forward(self, subject, predicate, few_shot_examples)->List[Tuple[str, float]
         return [ (i["entity"],i["score"])for i in json.loads(dspy_pred.objects_with_scores)]
 
 class Demir(AbstractBaseLinkPredictorClass):
+    def __init__(self,knowledge_graph, base_url,api_key,temperature, seed,llm_model,use_val:bool=False):
+        super().__init__(knowledge_graph,name="Demir")
+        self.temperature = temperature
+        self.seed = seed
+        self.lm = dspy.LM(model=f"openai/{llm_model}", api_key=api_key,
+                          api_base=base_url,
+                          seed=seed,
+                          temperature=temperature,
+                          cache=True,cache_in_memory=True)
+        dspy.configure(lm=self.lm)
+        self.train_set: List[Tuple[str]] = [(self.idx_to_entity[idx_h],
+                                             self.idx_to_relation[idx_r],
+                                             self.idx_to_entity[idx_t]) for idx_h, idx_r, idx_t in
+                                            self.kg.train_set.tolist()]
+        # Validation dataset
+        self.val_set: List[Tuple[str]] = [(self.idx_to_entity[idx_h],
+                                           self.idx_to_relation[idx_r],
+                                           self.idx_to_entity[idx_t]) for idx_h, idx_r, idx_t in
+                                          self.kg.valid_set.tolist()]
+        self.triples = self.train_set + self.val_set if use_val else self.train_set
+
+        self.entity_relation_to_entities=dict()
+        for s,p,o in self.triples:
+            self.entity_relation_to_entities.setdefault((s,p),[]).append(o)
+        self.scoring_func = MultiLabelLinkPredictor()
+
     def forward_triples(self, x: torch.LongTensor) -> torch.FloatTensor:
         raise NotImplementedError("RCL needs to implement it")
     def forward_k_vs_all(self, x: torch.LongTensor) -> torch.FloatTensor:
@@ -70,35 +83,19 @@ def forward_k_vs_all(self, x: torch.LongTensor) -> torch.FloatTensor:
             batch_predictions.append(scores)
         return torch.FloatTensor(batch_predictions)
 
-    def __init__(self,knowledge_graph, base_url,api_key,temperature, seed,llm_model,use_val:bool=False):
-        super().__init__(knowledge_graph,name="Demir")
-        self.client = OpenAI(base_url=base_url, api_key=api_key)
-        self.temperature = temperature
-        self.seed = seed
 
-        self.lm = dspy.LM(model=f"openai/{llm_model}", api_key=api_key,
-                          api_base=base_url,
-                          seed=seed,
-                          temperature=temperature,
-                          cache=True,cache_in_memory=True,
-                          kwargs={"extra_body":{"truncate_prompt_tokens": 32_000}})
-        dspy.configure(lm=self.lm)
-        self.train_set: List[Tuple[str]] = [(self.idx_to_entity[idx_h],
-                                             self.idx_to_relation[idx_r],
-                                             self.idx_to_entity[idx_t]) for idx_h, idx_r, idx_t in
-                                            self.kg.train_set.tolist()]
-        # Validation dataset
-        self.val_set: List[Tuple[str]] = [(self.idx_to_entity[idx_h],
-                                           self.idx_to_relation[idx_r],
-                                           self.idx_to_entity[idx_t]) for idx_h, idx_r, idx_t in
-                                          self.kg.valid_set.tolist()]
-        self.triples = self.train_set + self.val_set if use_val else self.train_set
+# 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!
+    print(args)
+    kg = KG(dataset_dir=args.dataset_dir, separator="\s+", eval_model=args.eval_model, add_reciprocal=False)
 
-        self.entity_relation_to_entities=dict()
-        from collections import OrderedDict
-        for s,p,o in self.triples:
-            self.entity_relation_to_entities.setdefault((s,p),[]).append(o)
+    sanity_checking(args,kg)
 
-        # 4. Instantiate your predictor
-        self.scoring_func = MultiLabelLinkPredictor()
-        self.entities:List[str]=list(sorted(self.entity_to_idx.keys()))
+    model = Demir(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)
+    print(results)

retrieval_aug_predictors/utils.py

@@ -0,0 +1,9 @@
+import os
+def sanity_checking(args,kg):
+    if args.eval_size is not None:
+        assert len(kg.test_set) >= args.eval_size, (f"Evaluation size cant be greater than the "
+                                                    f"total amount of triples in the test set: {len(kg.test_set)}")
+    else:
+        args.eval_size = len(kg.test_set)
+    if args.api_key is None:
+        args.api_key = os.environ.get("TENTRIS_TOKEN")