🧑‍💻 add --test flag

Signed-off-by: Peter <[email protected]>

app.py

@@ -179,16 +179,29 @@ def get_parser():
         default=False,
         help="turn on verbose logging",
     )
+    parser.add_argument(
+        "--test",
+        action="store_true",
+        default=False,
+        help="load the smallest model for simple testing",
+    )
+
     return parser
 
 
 if __name__ == "__main__":
     args = get_parser().parse_args()
     default_model = str(args.model)
+    test = args.test
+    if test:
+        logging.info("loading the smallest model for testing")
+        default_model = "ethzanalytics/distilgpt2-tiny-conversational"
+
     model_loc = Path(default_model)  # if the model is a path, use it
     basic_sc = args.basic_sc  # whether to use the baseline spellchecker
     gram_model = str(args.gram_model)
     device = 0 if torch.cuda.is_available() else -1
+
     print(f"CUDA avail is {torch.cuda.is_available()}")
 
     my_chatbot = (