Update app.py

app.py

@@ -3,177 +3,143 @@ import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 from threading import Thread
+import re
+import uuid
 
-phi4_model_path = "microsoft/Phi-4-reasoning-plus"
+# ------------- MODEL LOADING -------------
+model_name = "FractalAIResearch/Fathom-R1-14B"  # CHANGE this if your org name or model is different
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-phi4_model = AutoModelForCausalLM.from_pretrained(phi4_model_path, device_map="auto", torch_dtype="auto")
-phi4_tokenizer = AutoTokenizer.from_pretrained(phi4_model_path)
-
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# -------- MATH FORMATTING UTILITIES --------
+def format_math(text):
+    text = re.sub(r"\[(.*?)\]", r"$$\1$$", text, flags=re.DOTALL)
+    text = text.replace(r"\(", "$").replace(r"\)", "$")
+    return text
+
+def generate_conversation_id():
+    return str(uuid.uuid4())[:8]
+
+# --------- PROMPT CONSTRUCTION ---------
+def construct_prompt(history_state, user_message, system_message):
+    # Adjust custom tags as per your Fathom R1 model prompt format. This is a typical chat prompt pattern.
+    prompt = f"[SYSTEM]\n{system_message.strip()}\n"
+    for m in history_state:
+        if m["role"] == "user":
+            prompt += f"[USER]\n{m['content']}\n"
+        elif m["role"] == "assistant":
+            prompt += f"[ASSISTANT]\n{m['content']}\n"
+    prompt += f"[USER]\n{user_message.strip()}\n[ASSISTANT]\n"
+    return prompt
+
+# --------- SYSTEM PROMPT SETUP ---------
+system_message = (
+    "Your role as an assistant involves thoroughly exploring questions through a systematic thinking process before providing the final precise and accurate solutions. "
+    "Please structure your response into two main sections: Thought and Solution using the specified format: <think> ... </think> ... {Solution section}. "
+    "In the Thought section, detail your reasoning in steps. Then, systematically present the final solution that is logical and concise."
+)
+
+# ----------- GENERATION FUNCTION -----------
 @spaces.GPU(duration=60)
-def generate_response(user_message, max_tokens, temperature, top_k, top_p, repetition_penalty, history_state):
+def generate_response(
+    user_message, max_tokens, temperature, top_p, history_state
+):
     if not user_message.strip():
         return history_state, history_state
-        
-    # Phi-4 model settings
-    model = phi4_model
-    tokenizer = phi4_tokenizer
-    start_tag = "<|im_start|>"
-    sep_tag = "<|im_sep|>"
-    end_tag = "<|im_end|>"
-
-    # Recommended prompt settings by Microsoft
-    system_message = "Your role as an assistant involves thoroughly exploring questions through a systematic thinking process before providing the final precise and accurate solutions. This requires engaging in a comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, backtracing, and iteration to develop well-considered thinking process. Please structure your response into two main sections: Thought and Solution using the specified format: <think> {Thought section} </think> {Solution section}. In the Thought section, detail your reasoning process in steps. Each step should include detailed considerations such as analysing questions, summarizing relevant findings, brainstorming new ideas, verifying the accuracy of the current steps, refining any errors, and revisiting previous steps. In the Solution section, based on various attempts, explorations, and reflections from the Thought section, systematically present the final solution that you deem correct. The Solution section should be logical, accurate, and concise and detail necessary steps needed to reach the conclusion. Now, try to solve the following question through the above guidelines:"
-    prompt = f"{start_tag}system{sep_tag}{system_message}{end_tag}"
-    for message in history_state:
-        if message["role"] == "user":
-            prompt += f"{start_tag}user{sep_tag}{message['content']}{end_tag}"
-        elif message["role"] == "assistant" and message["content"]:
-            prompt += f"{start_tag}assistant{sep_tag}{message['content']}{end_tag}"
-    prompt += f"{start_tag}user{sep_tag}{user_message}{end_tag}{start_tag}assistant{sep_tag}"
 
+    prompt = construct_prompt(history_state, user_message, system_message)
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
-
-    do_sample = not (temperature == 1.0 and top_k >= 100 and top_p == 1.0)
-
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
-
-    # sampling techniques
-    generation_kwargs = {
-        "input_ids": inputs["input_ids"],
-        "attention_mask": inputs["attention_mask"],
-        "max_new_tokens": int(max_tokens),
-        "do_sample": True,
-        "temperature": 0.8,
-        "top_k": int(top_k),
-        "top_p": 0.95,
-        "repetition_penalty": repetition_penalty,
-        "streamer": streamer,
-    }
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    generation_kwargs = dict(
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        max_new_tokens=int(max_tokens),
+        do_sample=True,
+        temperature=float(temperature),
+        top_p=float(top_p),
+        streamer=streamer,
+    )
 
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
-    # Stream the response
     assistant_response = ""
     new_history = history_state + [
         {"role": "user", "content": user_message},
         {"role": "assistant", "content": ""}
     ]
     for new_token in streamer:
-        cleaned_token = new_token.replace("<|im_start|>", "").replace("<|im_sep|>", "").replace("<|im_end|>", "")
-        assistant_response += cleaned_token
+        assistant_response += new_token
+        # Optionally, you can call format_math() on assistant_response
         new_history[-1]["content"] = assistant_response.strip()
         yield new_history, new_history
 
     yield new_history, new_history
 
+# ------------ EXAMPLE MESSAGES --------------
 example_messages = {
-    "Math reasoning": "If a rectangular prism has a length of 6 cm, a width of 4 cm, and a height of 5 cm, what is the length of the longest line segment that can be drawn from one vertex to another?",
-    "Logic puzzle": "Four people (Alex, Blake, Casey, and Dana) each have a different favorite color (red, blue, green, yellow) and a different favorite fruit (apple, banana, cherry, date). Given the following clues: 1) The person who likes red doesn't like dates. 2) Alex likes yellow. 3) The person who likes blue likes cherries. 4) Blake doesn't like apples or bananas. 5) Casey doesn't like yellow or green. Who likes what color and what fruit?",
-    "Physics problem": "A ball is thrown upward with an initial velocity of 15 m/s from a height of 2 meters above the ground. Assuming the acceleration due to gravity is 9.8 m/s², determine: 1) The maximum height the ball reaches. 2) The total time the ball is in the air before hitting the ground. 3) The velocity with which the ball hits the ground."
+    "IIT-JEE 2024 Mathematics": "A student appears for a quiz consisting of only true-false type questions and answers all the questions. The student knows the answers of some questions and guesses the answers for the remaining questions. Whenever the student knows the answer of a question, he gives the correct answer. Assume that probability of the student giving the correct answer for a question, given that he has guessed it, is $\\frac{1}{2}$. Also assume that the probability of the answer for a question being guessed, given that the student's answer is correct, is $\\frac{1}{6}$. Then the probability that the student knows the answer of a randomly chosen question is?",
+    "IIT-JEE 2025         Physics": "A person sitting inside an elevator performs a weighing experiment with an object of mass 50 kg. Suppose that the variation of the height 𝑦 (in m) of the elevator, from the ground, with time 𝑡 (in s) is given by 𝑦 = 8 [1 + sin ( 2𝜋𝑡/𝑇 )], where 𝑇 = 40𝜋 s. Taking acceleration due to gravity, 𝑔 = 10 m/s^2 , the maximum variation of the object’s weight (in N) as observed in the experiment is ?",
+    "Goldman Sachs Interview Puzzle": "Four friends need to cross a dangerous bridge at night. Unfortunately, they have only one torch and the bridge is too dangerous to cross without one. The bridge is only strong enough to support two people at a time. Not all people take the same time to cross the bridge. Times for each person: 1 min, 2 mins, 7 mins and 10 mins. What is the shortest time needed for all four of them to cross the bridge?",
+    "IIT-JEE 2025 Mathematics": "Let 𝑆 be the set of all seven-digit numbers that can be formed using the digits 0, 1 and 2. For example, 2210222 is in 𝑆, but 0210222 is NOT in 𝑆.Then the number of elements 𝑥 in 𝑆 such that at least one of the digits 0 and 1 appears exactly twice in 𝑥, is ?"
 }
 
+# --------------- UI (GRADIO) ----------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
+    conversations_state = gr.State({})
+    current_convo_id = gr.State(generate_conversation_id())
+    history_state = gr.State([])
+
+    gr.HTML(
         """
-        # Phi-4-reasoning-plus Chatbot 
-        Welcome to the Phi-4-reasoning-plus Chatbot! This model excels at multi-step reasoning tasks in mathematics, logic, and science.
-        
-        The model will provide responses with two sections:
-        1. **<think>**: A detailed step-by-step reasoning process showing its work
-        2. **Solution**: A concise, accurate final answer based on the reasoning
-        
-        Try the example problems below to see how the model breaks down complex reasoning problems.
+        <div style="display: flex; align-items: center; gap: 16px; margin-bottom: 1em;">
+            <div style="background-color: black; padding: 6px; border-radius: 8px;">
+                <img src="https://framerusercontent.com/images/j0KjQQyrUfkFw4NwSaxQOLAoBU.png" style="height: 48px;">
+            </div>
+            <h1 style="margin: 0;">Fathom R1 14B Chatbot</h1>
+        </div>
         """
     )
-    
-    history_state = gr.State([])
+
+    with gr.Sidebar():
+        gr.Markdown("## Conversations")
+        conversation_selector = gr.Radio(choices=[], label="Select Conversation", interactive=True)
+        new_convo_button = gr.Button("New Conversation ➕")
 
     with gr.Row():
         with gr.Column(scale=1):
+            gr.Markdown(
+                """
+                Welcome to the Fathom R1 14B Chatbot, developed by Fractal AI Research! 
+                
+                Our model excels at reasoning tasks in mathematics and science. Given that our model has been optimised for tasks requiring critical thinking, it might overthink for simple chat queries.
+
+                To check out our GitHub repository, click [here](https://github.com/FractalAIResearchLabs/Fathom-R1) 
+                
+                For training recipe details on how this model was built, please check [here](https://huggingface.co/FractalAIResearch/Fathom-R1-14B)  
+                
+                Try the example problems below from various popular entrance examinations and interviews or type in your own problems to see how our model breaks down and solves complex reasoning problems.
+
+                NOTE: Once you close this demo window, all currently saved conversations will be lost.
+                """
+            )
             gr.Markdown("### Settings")
-            max_tokens_slider = gr.Slider(
-                minimum=64,
-                maximum=32768,
-                step=1024,
-                value=4096,
-                label="Max Tokens"
+            max_tokens_slider = gr.Slider(minimum=6144, maximum=32768, step=1024, value=16384, label="Max Tokens")
+            with gr.Accordion("Advanced Settings", open=True):
+                temperature_slider = gr.Slider(minimum=0.1, maximum=2.0, value=0.6, label="Temperature")
+                top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-p")
+            gr.Markdown(
+                """
+                We sincerely acknowledge [VIDraft](https://huggingface.co/VIDraft) for their Phi 4 Reasoning Plus [space](https://huggingface.co/spaces/VIDraft/phi-4-reasoning-plus), which served as the starting point for this demo.
+                """
             )
-            with gr.Accordion("Advanced Settings", open=False):
-                temperature_slider = gr.Slider(
-                    minimum=0.1,
-                    maximum=2.0,
-                    value=0.8,
-                    label="Temperature"
-                )
-                top_k_slider = gr.Slider(
-                    minimum=1,
-                    maximum=100,
-                    step=1,
-                    value=50,
-                    label="Top-k"
-                )
-                top_p_slider = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.0,
-                    value=0.95,
-                    label="Top-p"
-                )
-                repetition_penalty_slider = gr.Slider(
-                    minimum=1.0,
-                    maximum=2.0,
-                    value=1.0,
-                    label="Repetition Penalty"
-                )
-        
         with gr.Column(scale=4):
-            chatbot = gr.Chatbot(label="Chat", type="messages")
+            chatbot = gr.Chatbot(label="Chat", type="messages", height=520)
             with gr.Row():
-                user_input = gr.Textbox(
-                    label="Your message",
-                    placeholder="Type your message here...",
-                    scale=3
-                )
-                submit_button = gr.Button("Send", variant="primary", scale=1)
-                clear_button = gr.Button("Clear", scale=1)
-            gr.Markdown("**Try these examples:**")
-            with gr.Row():
-                example1_button = gr.Button("Math reasoning")
-                example2_button = gr.Button("Logic puzzle")
-                example3_button = gr.Button("Physics problem")
-
-    submit_button.click(
-        fn=generate_response,
-        inputs=[user_input, max_tokens_slider, temperature_slider, top_k_slider, top_p_slider, repetition_penalty_slider, history_state],
-        outputs=[chatbot, history_state]
-    ).then(
-        fn=lambda: gr.update(value=""),
-        inputs=None,
-        outputs=user_input
-    )
-
-    clear_button.click(
-        fn=lambda: ([], []),
-        inputs=None,
-        outputs=[chatbot, history_state]
-    )
-
-    example1_button.click(
-        fn=lambda: gr.update(value=example_messages["Math reasoning"]),
-        inputs=None,
-        outputs=user_input
-    )
-    example2_button.click(
-        fn=lambda: gr.update(value=example_messages["Logic puzzle"]),
-        inputs=None,
-        outputs=user_input
-    )
-    example3_button.click(
-        fn=lambda: gr.update(value=example_messages["Physics problem"]),
-        inputs=None,
-        outputs=user_input
-    )
-
-demo.launch(ssr_mode=False)
+                user_input = gr.Textbox(label="User Input", placeholder="Type your question here...", lines=3, scale=8)
+                with gr.Column():
+                    submit_button = gr.Button("Send", variant="primary", scale