Update app.py

app.py

@@ -1,8 +1,5 @@
-import spaces
-
 #######################
 '''
-Name: Phine Inference
 License: MIT
 '''
 #######################
@@ -21,6 +18,7 @@ import transformers
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import random
+import spaces
 import re
 
 def cut_text_after_last_token(text, token):
@@ -60,18 +58,17 @@ class _SentinelTokenStoppingCriteria(transformers.StoppingCriteria):
 
 
 
-model_path = 'freecs/phine-2-v0'
+model_path = 'freecs/ArtificialThinker-Phi2'
 
-device = "cuda"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
 model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, load_in_4bit=False, torch_dtype=torch.float16).to(device) #remove .to() if load_in_4/8bit = True
 
-sys_message = "You are an AI assistant named Phine developed by FreeCS.org. You are polite and smart." #System Message
-@spaces.GPU(enable_queue=True)
-def phine(message, history, temperature, top_p, top_k, repetition_penalty):
 
+@spaces.GPU(enable_queue=True)
+def phine(message, history, temperature, top_p, top_k, repetition_penalty, sys_message):
 
 
     n = 0
@@ -83,7 +80,9 @@ def phine(message, history, temperature, top_p, top_k, repetition_penalty):
             if n%2 == 0:
               context+=f"""\n<|prompt|>{h}\n"""
             else:
-              context+=f"""<|response|>{h}"""
+              pattern = re.compile(r'<details>.*?</details>')
+              result = re.sub(pattern, '', h)
+              context+=f"""<|response|>{result}"""
             n+=1
     else:
 
@@ -92,7 +91,7 @@ def phine(message, history, temperature, top_p, top_k, repetition_penalty):
 
 
 
-    prompt = f"""\n<|system|>{sys_message}"""+context+"\n<|prompt|>"+message+"<|endoftext|>\n<|response|>"
+    prompt = f"""\n<|system|>{sys_message}"""+context+"\n<|prompt|>"+message+"<|endoftext|>\n<|reasoning|>"
     tokenized = tokenizer(prompt, return_tensors="pt").to(device)
 
 
@@ -106,28 +105,44 @@ def phine(message, history, temperature, top_p, top_k, repetition_penalty):
             starting_idx=tokenized.input_ids.shape[-1])
     ])
 
-        
+
     token = model.generate(**tokenized,
                         stopping_criteria=stopping_criteria_list,
                         do_sample=True,
                         max_length=2048, temperature=temperature, top_p=top_p, top_k = top_k, repetition_penalty = repetition_penalty
                            )
 
-    completion = tokenizer.decode(token[0], skip_special_tokens=False)
+    completion = tokenizer.decode(token[0], skip_special_tokens=True)
+
+    token = "<|reasoning|>"
+    reasoning = cut_text_after_last_token(completion, token)
+    prompt = f"""\n<|system|>{sys_message}"""+context+"\n<|prompt|>"+message+"\n<|reasoning|>"+reasoning+"\n<|response|>"
+
+    tokenized = tokenizer(prompt, return_tensors="pt").to(device)
+    token = model.generate(**tokenized,
+                        stopping_criteria=stopping_criteria_list,
+                        do_sample=True,
+                        max_length=2048, temperature=temperature, top_p=top_p, top_k = top_k, repetition_penalty = repetition_penalty
+                           )
+    completion = tokenizer.decode(token[0], skip_special_tokens=True)
     token = "<|response|>"
-    res = cut_text_after_last_token(completion, token)
-    return res.replace('<|endoftext|>', '')
+    response = cut_text_after_last_token(completion, token)
+
+
+    res = f"""<details><summary>Reasoning</summary>{reasoning}</details>\n\n{response}"""
+
+    return res
 
 
 demo = gr.ChatInterface(phine,
-    title="Phine Demo",
-    description="Demo of [Phine 2](https://huggingface.co/freecs/phine-2-v0). We are not affiliated with and do not endorse Phine 2.",
-    additional_inputs=[
-        gr.Slider(0.1, 2.0, label="temperature", value=0.5),
-        gr.Slider(0.1, 2.0, label="Top P", value=0.9),
-        gr.Slider(1, 500, label="Top K", value=50),
-        gr.Slider(0.1, 2.0, label="Repetition Penalty", value=1.15)
-    ]
-)
+                          additional_inputs=[
+                              gr.Slider(0.1, 2.0, label="temperature", value=0.5),
+                              gr.Slider(0.1, 2.0, label="Top P", value=0.9),
+                              gr.Slider(1, 500, label="Top K", value=50),
+                              gr.Slider(0.1, 2.0, label="Repetition Penalty", value=1.1),
+                              gr.Textbox(label="System Prompt",max_lines=1,interactive=True, value="You are an AI assistant named Phine developed by FreeCS.org. You are polite and smart.")
+                          ]
+                          )
+
 if __name__ == "__main__":
-    demo.queue().launch(share=True, debug=True) #If debug=True causes problems you can set it to False
+    demo.queue().launch(share=True, debug=True) #If debug=True causes problems you can set it to False