renamed the interface to UltraForKnowledgeGraphReasoning

README.md

@@ -39,10 +39,10 @@ ULTRA performs **link prediction** (KG completion): given a query `(head, relati
 * Run **zero-shot inference** on any graph:
 
 ```python
-from modeling import UltraLinkPrediction
+from modeling import UltraForKnowledgeGraphReasoning
 from ultra.datasets import CoDExSmall
 from ultra.eval import test
-model = UltraLinkPrediction.from_pretrained("mgalkin/ultra_3g")
+model = UltraForKnowledgeGraphReasoning.from_pretrained("mgalkin/ultra_4g")
 dataset = CoDExSmall(root="./datasets/")
 test(model, mode="test", dataset=dataset, gpus=None)
 # Expected results for ULTRA 4g

modeling.py

@@ -39,7 +39,7 @@ class UltraConfig(PretrainedConfig):
 
         super().__init__(**kwargs)
 
-class UltraLinkPrediction(PreTrainedModel):
+class UltraForKnowledgeGraphReasoning(PreTrainedModel):
 
     config_class = UltraConfig
 
@@ -59,7 +59,7 @@ class UltraLinkPrediction(PreTrainedModel):
 
 if __name__ == "__main__":
 
-    model = UltraLinkPrediction.from_pretrained("mgalkin/ultra_4g")
+    model = UltraForKnowledgeGraphReasoning.from_pretrained("mgalkin/ultra_4g")
     dataset = CoDExSmall(root="./datasets/")
     test(model, mode="test", dataset=dataset, gpus=None)
     # mrr:      0.463971