Update app.py

app.py

@@ -20,15 +20,19 @@ model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_c
 
 @spaces.GPU(duration=120)
 def predict(input_text, history):
-    conv = ""
+    chat = []
     for item in history:
-        conv += f"<|user|>\n{item[0]}</s>\n<|assistant|>\n{item[1]}</s>\n"
-    conv += f"<|user|>\n{input_text}</s>\n"
+        chat.append({"role": "user", "content": item[0]})
+        if item[1] is not None:
+            chat.append({"role": "assistant", "content": item[1]})
+    chat.append({"role": "user", "content": message})
+    conv = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
     
     inputs = tokenizer(conv, return_tensors="pt").to("cuda")
     outputs = model.generate(**inputs, max_new_tokens=512)
     
     generated_text = tokenizer.batch_decode(outputs)[0]
-    return generated_text.split("<|assistant|>")[-1].strip()
+    generated_text = generated_text.split("<|assistant|>")[-1]
+    return generated_text
 
 gr.ChatInterface(predict, theme="soft").launch()