Update app.py

app.py

@@ -34,13 +34,12 @@ def format_llama_messages(messages):
     prompt = []
 
     if messages[0]["role"] == "system":
-        content = "".join(["<<SYS>>\n", messages[0]["content"], "\n<</SYS>>\n\n", messages[1]["content"]])
-        messages = [{"role": messages[1]["role"], "content": content}] + messages[2:]
+        content = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"+messages[0]["content"]+ "<|eot_id|>\n\n"
+        prompt.append(content)
+    for user, answer in zip(messages[1::2], messages[2::2]):
+        prompt.append("<|start_header_id|>user<|end_header_id|>"+ (user["content"]).strip()+ "<|eot_id|><|start_header_id|>assistant<|end_header_id|>"+ (answer["content"]).strip()+ "<|eot_id|>")
 
-    for user, answer in zip(messages[::2], messages[1::2]):
-        prompt.extend(["<s>", "[INST] ", (user["content"]).strip(), " [/INST] ", (answer["content"]).strip(), "</s>"])
-
-    prompt.extend(["<s>", "[INST] ", (messages[-1]["content"]).strip(), " [/INST] "])
+    prompt.append("<|start_header_id|>user<|end_header_id|>"+ (messages[-1]["content"]).strip()+ "<|eot_id|>")
 
     return "".join(prompt)
 
@@ -64,7 +63,14 @@ def handle_query(query):
     response = requests.request("POST", API_URL, headers=headers, data=data)
     try:
         data=json.loads(response.content.decode("utf-8"))
-        return data[0]['generated_text'].split('[/INST] ')[1]
+        result=data[0]['generated_text']
+        last_comma_index = result.rfind('assistant<|end_header_id|>')
+
+        # If a comma is found, slice the string from the last comma onwards
+        if last_comma_index != -1:
+            result = result[last_comma_index + 26:].strip()
+            
+        return result
     except:
         return response