Update app.py

app.py

@@ -33,6 +33,8 @@ repetition_penalty=1.7
 
 #llm = LLM(model_name, max_model_len=4096)
 
+system_prompt = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nTu es Albert, l'agent conversationnel des services publics qui peut décrire des documents de référence ou aider à des tâches de rédaction<|eot_id|>"
+
 #Vector search over the database
 def vector_search(sentence_query):
 
@@ -64,24 +66,19 @@ class StopOnTokens(StoppingCriteria):
                 return True
         return False
 
+
 def predict(message, history):
     text = vector_search(message)
     message = message + "\n\n### Source ###\n" + text
     history_transformer_format = history + [[message, ""]]
-
-    messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])
-                for item in history_transformer_format])
-
-    return messages
-
-def predict_alt(message, history):
-    history_transformer_format = history + [[message, ""]]
     stop = StopOnTokens()
 
-    messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])
+    messages = "".join(["".join(["<|start_header_id|>user<|end_header_id|>\n\n"+item[0], "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"+item[1]])
                 for item in history_transformer_format])
 
-    model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
+    messages = system_prompt + messages
+
+    """"model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
     streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         model_inputs,
@@ -101,7 +98,8 @@ def predict_alt(message, history):
     for new_token in streamer:
         if new_token != '<':
             partial_message += new_token
-            yield partial_message
+            yield partial_message"""
+    return messages
 
 # Define the Gradio interface
 title = "Tchap"