WIP: adding new models in this file

dice-group · Mar 21, 2024 · de48fca · de48fca
1 parent c65abc2
commit de48fca
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/dicee/static_funcs.py b/dicee/static_funcs.py
@@ -3,7 +3,7 @@
 import datetime
 from typing import Tuple, List
 from .models import CMult, Pyke, DistMult, KeciBase, Keci, TransE, DeCaL,\
-    ComplEx, AConEx, AConvO, AConvQ, ConvQ, ConvO, ConEx, QMult, OMult, Shallom, LFMult
+    ComplEx, AConEx, AConvO, AConvQ, ConvQ, ConvO, ConEx, QMult, OMult, Shallom, LFMult, FMult, PolyMult
 from .models.pykeen_models import PykeenKGE
 from .models.transformers import BytE
 import time
@@ -421,6 +421,12 @@ def intialize_model(args: dict,verbose=0) -> Tuple[object, str]:
     elif model_name == 'DeCaL':
         model =DeCaL(args=args)
         form_of_labelling = 'EntityPrediction'
+    elif model_name == 'FMult':
+        model =FMult(args=args)
+        form_of_labelling = 'EntityPrediction'
+    elif model_name == 'PolyMult':
+        model =PolyMult(args=args)
+        form_of_labelling = 'EntityPrediction'
     else:
         raise ValueError(f"--model_name: {model_name} is not found.")
     return model, form_of_labelling