switched to 3B o1

app.py

@@ -5,18 +5,20 @@ import spaces
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
-client = InferenceClient("meditsolutions/Llama-3.2-SUN-2.5B-chat")
+client = InferenceClient("mkurman/llama-3.2-MEDIT-3B-o1")
 
 @spaces.GPU
 def respond(
     message,
     history: list[tuple[str, str]],
-    system_message,
+    # system_message,
     max_tokens,
     temperature,
     top_p,
 ):
-    messages = [{"role": "system", "content": system_message}]
+    # messages = [{"role": "system", "content": system_message}]
+
+    messages = []
 
     for val in history:
         if val[0]:
@@ -27,17 +29,33 @@ def respond(
     messages.append({"role": "user", "content": message})
 
     response = ""
+    counter = 0
 
     for message in client.chat_completion(
         messages,
         max_tokens=max_tokens,
         stream=True,
-        temperature=temperature,
+        temperature=temperature if temperature > 0.0 else None,
         top_p=top_p,
+        logprobs=True,
+        stop=['</Output>']
     ):
-        token = message.choices[0].delta.content
+        text = message.choices[0].delta.content
+
+        token = message.choices[0].logprobs.content[0].token
+
+        if token in ['<Thought>', '<|python_tag|>']:
+            text = '## Thinking:\n\n' + text
+
+        if token == '<Output>' and counter > 0:
+            text = '## Output:\n\n' + text
+
+        if token == '</Output>':
+            yield None
+            break
 
-        response += token
+        response += text
+        counter += 1
         yield response
 
 
@@ -46,12 +64,12 @@ For information on how to customize the ChatInterface, peruse the gradio docs: h
 """
 demo = gr.ChatInterface(
     respond,
-    title="LLama 3.2 MedIT SUN 2.5B by MedIT Solutions",
+    title="LLama 3.2 MedIT 3B o1",
     description="Built with Llama. Please note that this model is not a source of knowledge and is not intended to provide 100% accurate answers. If the model provides answers that are generally considered inappropriate, please contact me.",
     additional_inputs=[
-        gr.Textbox(value="You are a helpful, smart, kind, and efficient AI assistant. You always fulfill the user's requests to the best of your ability. You always think and reflect before providing final answers in a step-by-step manner.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        # gr.Textbox(value="You are a helpful, smart, kind, and efficient AI assistant. You always fulfill the user's requests to the best of your ability. You always think and reflect before providing final answers in a step-by-step manner.", label="System message"),
+        gr.Slider(minimum=1, maximum=4096, value=2048, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.0, maximum=4.0, value=0.0, step=0.1, label="Temperature"),
         gr.Slider(
             minimum=0.1,
             maximum=1.0,