Add model input validation for run.py

Signed-off-by: Anna Jung (VMware) <[email protected]>
vmware-archive · Dec 1, 2021 · 2189fe1 · 2189fe1
1 parent 7c401c7
commit 2189fe1
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/ml-conversational-analytic-tool/run.py b/ml-conversational-analytic-tool/run.py
@@ -51,6 +51,10 @@ def run(annotated_filename, dataset_filename, outcome, encoding_type, model_type
     parser.add_argument('-pad', action='store_true', default=False, help='Pad total length of each pull')
 
     args = parser.parse_args()
+
+    if args.model != 'CNN' and args.model != 'LSTM':
+        raise Exception("Model must be either CNN or LSTM")
+
     encodingType = 'role'
     if not args.roleRelevant:
         encodingType = 'role-agnostic'