Update app_chat.py

app_chat.py

@@ -18,8 +18,11 @@ DESCRIPTION = """\
 """
 
 model_id = "nvidia/Hymba-1.5B-Instruct"
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
+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
 
 # class StoppingCriteriaSub(StoppingCriteria):
@@ -49,11 +52,11 @@ def generate(
 ) -> Iterator[str]:
     conversation = []
     if system_prompt:
-        conversation.append({"role": "System", "content": system_prompt})
+        conversation.append({"role": "system", "content": system_prompt})
     conversation += chat_history
-    conversation.append({"role": "User", "content": message})
+    conversation.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(conversation,  tokenize=True, add_generation_prompt=True, return_tensors="pt")
+    input_ids = tokenizer.apply_chat_template(conversation,  tokenize=True, add_generation_prompt=True, return_tensors="pt").to('cuda')
     
     # stopping_criteria = StoppingCriteriaList([StopStringCriteria(tokenizer=tokenizer, stop_strings="</s>")])