Update app.py

app.py

@@ -7,6 +7,8 @@ import random
 import time
 import re
 from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, TextStreamer
+import transformers
+
 
 
 # Set an environment variable
@@ -45,10 +47,10 @@ def inference(query):
 
     tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda")
     outputs = model.generate(tokenized_chat, **generation_params)
-    decoded_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    decoded_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=False)
     assistant_response = decoded_outputs[0].split("<|im_start|>assistant\n")[-1].strip()
-    response_ = assistant_response.replace('<|im_end|>', "")
-    return response_
+    # response_ = assistant_response.replace('<|im_end|>', "")
+    return assistant_response
     # outputs = model.generate(tokenized_chat, **generation_params, streamer=streamer)
     # return outputs