...

epdb/management/commands/create_ml_models.py

@@ -2,7 +2,9 @@ from django.conf import settings as s
 from django.core.management.base import BaseCommand
 from django.db import transaction
 
-from epdb.models import MLRelativeReasoning, EnviFormer, Package
+from epdb.models import EnviFormer, MLRelativeReasoning
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 class Command(BaseCommand):
@@ -75,11 +77,13 @@ class Command(BaseCommand):
             return packages
 
         # Iteratively create models in options["model_names"]
-        print(f"Creating models: {options['model_names']}\n"
-              f"Data packages: {options['data_packages']}\n"
-              f"Rule Packages (only for MLRR): {options['rule_packages']}\n"
-              f"Eval Packages: {options['eval_packages']}\n"
-              f"Threshold: {options['threshold']:.2f}")
+        print(
+            f"Creating models: {options['model_names']}\n"
+            f"Data packages: {options['data_packages']}\n"
+            f"Rule Packages (only for MLRR): {options['rule_packages']}\n"
+            f"Eval Packages: {options['eval_packages']}\n"
+            f"Threshold: {options['threshold']:.2f}"
+        )
         data_packages = decode_packages(options["data_packages"])
         eval_packages = decode_packages(options["eval_packages"])
         rule_packages = decode_packages(options["rule_packages"])
@@ -90,10 +94,10 @@ class Command(BaseCommand):
                     pack,
                     data_packages=data_packages,
                     eval_packages=eval_packages,
-                    threshold=options['threshold'],
+                    threshold=options["threshold"],
                     name=f"EnviFormer - {', '.join(options['data_packages'])} - T{options['threshold']:.2f}",
                     description=f"EnviFormer transformer trained on {options['data_packages']} "
-                                f"evaluated on {options['eval_packages']}.",
+                    f"evaluated on {options['eval_packages']}.",
                 )
             elif model_name == "mlrr":
                 model = MLRelativeReasoning.create(
@@ -101,10 +105,10 @@ class Command(BaseCommand):
                     rule_packages=rule_packages,
                     data_packages=data_packages,
                     eval_packages=eval_packages,
-                    threshold=options['threshold'],
+                    threshold=options["threshold"],
                     name=f"ECC - {', '.join(options['data_packages'])} - T{options['threshold']:.2f}",
                     description=f"ML Relative Reasoning trained on {options['data_packages']} with rules from "
-                                f"{options['rule_packages']} and evaluated on {options['eval_packages']}.",
+                    f"{options['rule_packages']} and evaluated on {options['eval_packages']}.",
                 )
             else:
                 raise ValueError(f"Cannot create model of type {model_name}, unknown model type")

epdb/management/commands/load_enviformer.py

@@ -8,7 +8,9 @@ from django.conf import settings as s
 from django.core.management.base import BaseCommand
 from django.db import transaction
 
-from epdb.models import EnviFormer, Package
+from epdb.models import EnviFormer
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 class Command(BaseCommand):

epdb/management/commands/localize_urls.py

@@ -1,8 +1,8 @@
 from django.apps import apps
+from django.conf import settings as s
 from django.core.management.base import BaseCommand
-
-from django.db.models import F, Value, TextField, JSONField
-from django.db.models.functions import Replace, Cast
+from django.db.models import F, JSONField, TextField, Value
+from django.db.models.functions import Cast, Replace
 
 from epdb.models import EnviPathModel
 
@@ -23,10 +23,13 @@ class Command(BaseCommand):
         )
 
     def handle(self, *args, **options):
+        Package = s.GET_PACKAGE_MODEL()
+        print("Localizing urls for Package")
+        Package.objects.update(url=Replace(F("url"), Value(options["old"]), Value(options["new"])))
+
         MODELS = [
             "User",
             "Group",
-            "Package",
             "Compound",
             "CompoundStructure",
             "Pathway",