Fix model path

gatewayd-io · Oct 17, 2024 · f0f4e7d · f0f4e7d
1 parent 1311692
commit f0f4e7d
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/training/test_train.py b/training/test_train.py
@@ -35,14 +35,18 @@
 def model(request):
     # Load dataset
     data = None
+    prefix = ""
     try:
         data = pd.read_csv(request.param["dataset"])
     except FileNotFoundError:
         # Check if the dataset is in the parent directory
-        data = pd.read_csv("../" + request.param["dataset"])
+        prefix = "../"
+        data = pd.read_csv(prefix + request.param["dataset"])
 
     # Load TF model from SavedModel
-    sqli_model = TFSMLayer(request.param["model_path"], call_endpoint="serving_default")
+    sqli_model = TFSMLayer(
+        prefix + request.param["model_path"], call_endpoint="serving_default"
+    )
 
     # Tokenize the sample
     tokenizer = Tokenizer(num_words=MAX_WORDS, filters="")