NiFi scripts: anon doc script improvement (allow modelpack path).

CogStack · Nov 27, 2023 · 84a18f3 · 84a18f3
1 parent dd3af8c
commit 84a18f3
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/nifi/user-scripts/anonymise_doc.py b/nifi/user-scripts/anonymise_doc.py
@@ -1,9 +1,14 @@
 from medcat.utils.ner.deid import DeIdModel
 import sys
-
+import os
 input_text = sys.stdin.read()
 
-model_pack_path = "/opt/models/de_id_base.zip"
+model_pack_path = os.environ.get("MODEL_PACK_PATH", "/opt/models/de_id_base.zip")
+
+for arg in sys.argv:
+    _arg = arg.split("=", 1)
+    if _arg[0] == "model_pack_path":
+        model_pack_path = _arg[1]
 
 cat = DeIdModel.load_model_pack(model_pack_path)