renamed the interface to UltraForKnowledgeGraphReasoning

README.md

@@ -44,7 +44,7 @@ from modeling import UltraLinkPrediction
 from ultra.datasets import CoDExSmall
 from ultra.eval import test
 
-model = UltraLinkPrediction.from_pretrained("mgalkin/ultra_3g")
+model = UltraForKnowledgeGraphReasoning.from_pretrained("mgalkin/ultra_3g")
 dataset = CoDExSmall(root="./datasets/")
 test(model, mode="test", dataset=dataset, gpus=None)
 # Expected results for ULTRA 3g

modeling.py

@@ -1,7 +1,7 @@
 import os
 import sys
 from transformers import PretrainedConfig, PreTrainedModel
-sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+#sys.path.append(os.path.dirname(os.path.dirname(__file__)))
 from ultra.models import Ultra
 from ultra.datasets import WN18RR, CoDExSmall, FB15k237, FB15k237Inductive
 from ultra.eval import test
@@ -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_3g")
+    model = UltraForKnowledgeGraphReasoning.from_pretrained("mgalkin/ultra_3g")
     dataset = CoDExSmall(root="./datasets/")
     test(model, mode="test", dataset=dataset, gpus=None)
     # mrr:      0.472035