Update app.py

app.py

@@ -34,12 +34,12 @@ def predict(message, history):
     stop = StopOnTokens()
 
     # Formatting the input for the model.
-    messages = "</s>".join(["</s>".join(["\n<|user|>:" + item[0], "\n<|assistant|>:" + item[1]])
-                        for item in history_transformer_format])
-    model_inputs = tokenizer([messages], return_tensors="pt")
+    system_prompt = "<|im_start|>system\nYou are Phixtral, a helpful AI assistant.<|im_end|>"
+    messages = system_prompt + "".join(["".join(["\n<|im_start|>user\n" + item[0], "<|im_end|>\n<|im_start|>assistant\n" + item[1]]) for item in history_transformer_format])
+    input_ids = tokenizer([messages], return_tensors="pt").to('cuda')
     streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
-        model_inputs,
+        input_ids,
         streamer=streamer,
         max_new_tokens=1024,
         do_sample=True,