Update app_chat.py

app_chat.py

@@ -19,8 +19,10 @@ DESCRIPTION = """\
 
 model_id = "nvidia/Hymba-1.5B-Instruct"
 model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", model_dtype="bfloat16", trust_remote_code=True)
-model.to('cuda')
-model.eval()
+model = model.cuda().to(torch.bfloat16)
+
+#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
@@ -56,7 +58,7 @@ def generate(
     conversation += chat_history
     conversation.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(conversation,  tokenize=True, add_generation_prompt=True, return_tensors="pt").to('cuda')
+    input_ids = tokenizer.apply_chat_template(conversation,  tokenize=True, add_generation_prompt=True, return_tensors="pt").to('cuda').to(torch.bfloat16)
     
     # stopping_criteria = StoppingCriteriaList([StopStringCriteria(tokenizer=tokenizer, stop_strings="</s>")])