Merge pull request #272 from kylebgorman/predict

Expands list of separate features models
CUNY-CL · Dec 1, 2024 · d32a240 · d32a240
2 parents 44af54f + 8c6a029
commit d32a240
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/yoyodyne/predict.py b/yoyodyne/predict.py
@@ -36,9 +36,13 @@ def get_datamodule_from_argparse_args(
         data.DataModule.
     """
     separate_features = args.features_col != 0 and args.arch in [
-        "pointer_generator_rnn",
+        "hard_attention_gru",
+        "hard_attention_lstm",
+        "pointer_generator_gru",
+        "pointer_generator_lstm",
         "pointer_generator_transformer",
-        "transducer",
+        "transducer_grm",
+        "transducer_lstm",
     ]
     index = data.Index.read(args.model_dir)
     return data.DataModule(