Model Path Parameter

silverlightning926 · Jun 12, 2024 · f7232da · f7232da
1 parent 1e29c6c
commit f7232da
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/src/build_model.py b/src/build_model.py
@@ -4,12 +4,14 @@
 
 import os
 
+from main import MODEL_PATH
+
 
 def _build_model():
 
-    if os.path.exists('model.keras'):
+    if os.path.exists(MODEL_PATH):
         print('Model already exists. Loading model.')
-        model = load_model('model.keras')
+        model = load_model(MODEL_PATH)
 
     else:
 

diff --git a/src/main.py b/src/main.py
@@ -1,6 +1,8 @@
 from load_data import getData
 from build_model import build_model
 
+MODEL_PATH = 'model.keras'
+
 
 def main():
     train_data, test_data = getData()
@@ -11,7 +13,7 @@ def main():
 
     model.evaluate(test_data)
 
-    model.save('model.h5')
+    model.save(MODEL_PATH)
 
 
 if __name__ == "__main__":