Update app.py

app.py

@@ -31,16 +31,27 @@ def llama_o1_template(data):
     text = template.format(content=data)
     return text
 
+def format_response(response):
+    response = response.replace('<start_of_father_id>','')
+    response = response.replace('<end_of_father_id><start_of_local_id>','👉')
+    response = response.replace('<end_of_local_id><start_of_thought>',', ')
+    response = response.replace('<end_of_thought><start_of_rating>','')
+    response = response.replace('<end_of_rating>','')
+    response = response.replace('<positive_rating>','👍')
+    response = response.replace('<negative_rating>','👎')
+
 def generate_text(message, history, max_tokens=512, temperature=0.9, top_p=0.95):
     temp = ""
     input_texts = [llama_o1_template(message)]
     input_texts = [input_text.replace('<|end_of_text|>','') for input_text in input_texts]
     #print(f"input_texts[0]: {input_texts[0]}")
     inputs = model.tokenize(input_texts[0].encode('utf-8'))
-    for token in model.generate(inputs, top_p=top_p, temp=temperature):
+    tokens = model.generate(inputs, top_p=top_p, temp=temperature):
         #print(f"token: {token}")
-        text = model.detokenize([token])
+    texts = model.detokenize(tokens,special=True)
         #print(f"text detok: {text}")
+    texts = format_response(texts)
+    for text in texts:
         temp += text.decode('utf-8')
         yield temp