Fix prepare_huggingface_generation_config

app.py

@@ -149,7 +149,6 @@ def prepare_huggingface_generation_config(generation_config):
 
     # I'm not sure what is the right behavior here, but it is better to be explicit
     for name, params in GENERATION_CONFIG_PARAMS.items():
-        # Checking for START to examine the a slider parameters only
         if (
             "START" in params
             and params["SAMPLING"]
@@ -159,21 +158,21 @@ def prepare_huggingface_generation_config(generation_config):
             if generation_config[name] == params["DEFAULT"]:
                 generation_config[name] = None
             else:
-                assert generation_config["do_sample"]
+                assert generation_config.get("do_sample", False)
 
-    # TODO: refactor this part
-    if generation_config["is_chat"]:
+    # FIX: use .get() to avoid KeyError
+    if generation_config.get("is_chat", False):
         generation_config["max_tokens"] = generation_config.pop("max_new_tokens")
-
         generation_config["stop"] = generation_config.pop("stop_sequences")
         del generation_config["do_sample"]
         del generation_config["top_k"]
 
-    is_chat = generation_config.pop("is_chat")
+    is_chat = generation_config.pop("is_chat", False)
 
     return generation_config, is_chat
 
 
+
 def escape_markdown(text):
     escape_dict = {
         "*": r"\*",