Update app_chat.py

app_chat.py

@@ -19,16 +19,14 @@ DEFAULT_MAX_NEW_TOKENS = 512
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 DESCRIPTION = """\
-# Hymba-1.5B chat
-
+# Hymba-1.5B-Instruct chat
+Feel free to chat with our model! More details: [Paper](https://arxiv.org/abs/2411.13676), [Model card](https://huggingface.co/nvidia/Hymba-1.5B-Instruct), [GitHub](https://github.com/NVlabs/hymba).
 """
 
 model_id = "nvidia/Hymba-1.5B-Instruct"
 model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", trust_remote_code=True)
 model = model.cuda().to(torch.bfloat16)
 model.compile()
-#model.to('cuda')
-#model.eval()
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 tokenizer.chat_template = "{{'<extra_id_0>System'}}{% for message in messages %}{% if message['role'] == 'system' %}{{'\n' + message['content'].strip()}}{% if tools or contexts %}{{'\n'}}{% endif %}{% endif %}{% endfor %}{% if tools %}{% for tool in tools %}{{ '\n<tool> ' + tool|tojson + ' </tool>' }}{% endfor %}{% endif %}{% if contexts %}{% if tools %}{{'\n'}}{% endif %}{% for context in contexts %}{{ '\n<context> ' + context.strip() + ' </context>' }}{% endfor %}{% endif %}{{'\n\n'}}{% for message in messages %}{% if message['role'] == 'user' %}{{ '<extra_id_1>User\n' + message['content'].strip() + '\n' }}{% elif message['role'] == 'assistant' %}{{ '<extra_id_1>Assistant\n' + message['content'].strip() + '\n' }}{% elif message['role'] == 'tool' %}{{ '<extra_id_1>Tool\n' + message['content'].strip() + '\n' }}{% endif %}{% endfor %}{%- if add_generation_prompt %}{{'<extra_id_1>Assistant\n'}}{%- endif %}"
 #tokenizer.use_default_system_prompt = False
@@ -73,7 +71,7 @@ def generate(
         gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
     input_ids = input_ids.to(model.device)
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=False)
+    streamer = TextIteratorStreamer(tokenizer, timeout=1.0, skip_prompt=True, skip_special_tokens=False)
     generate_kwargs = dict(
         {"input_ids": input_ids},
         streamer=streamer,