Update app.py

app.py

@@ -1,226 +1,221 @@
-# app.py  –  Gradio chatbot for FractalAIResearch/Fathom-R1-14B
-# ---------------------------------------------------------------------
 import gradio as gr
 import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 from threading import Thread
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TextIteratorStreamer,
-)
+import re
+import uuid
 
-# ---------------------------------------------------------------------
-# 1.  Model & tokenizer
-# ---------------------------------------------------------------------
-MODEL_NAME = "FractalAIResearch/Fathom-R1-14B"
-
-print("⏳ Loading model … (this may take a couple of minutes)")
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    device_map="auto",          # dispatch across any available device(s)
-    trust_remote_code=True,     # Fathom uses custom modelling code
-    low_cpu_mem_usage=True,
-)
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_NAME,
-    trust_remote_code=True,
-)
+# Load model and tokenizer
+our_model_path = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-print("✅ Model loaded")
+our_model = AutoModelForCausalLM.from_pretrained(our_model_path, device_map="auto", torch_dtype="auto")
+our_tokenizer = AutoTokenizer.from_pretrained(our_model_path)
 
-# ---------------------------------------------------------------------
-# 2.  Helper: build a prompt with the tokenizer’s chat_template
-# ---------------------------------------------------------------------
-def build_chat_prompt(history, user_message, system_message):
-    """
-    history        : list[dict(role, content)]
-    user_message   : str
-    system_message : str
-    returns a single prompt string (not tokenised)
-    """
-    msgs = []
-    if system_message:
-        msgs.append({"role": "system", "content": system_message})
-    msgs.extend(history)
-    msgs.append({"role": "user", "content": user_message})
-
-    return tokenizer.apply_chat_template(
-        msgs,
-        tokenize=False,          # return pure text
-        add_generation_prompt=True,
-    )
+def format_math(text):
+    text = re.sub(r"\[(.*?)\]", r"$$\1$$", text, flags=re.DOTALL)
+    text = text.replace(r"\(", "$").replace(r"\)", "$")
+    return text
 
-# ---------------------------------------------------------------------
-# 3.  Generation endpoint
-# ---------------------------------------------------------------------
-@spaces.GPU(duration=60)        # short GPU reservation if available
-def generate_response(
-    user_message,
-    max_tokens,
-    temperature,
-    top_k,
-    top_p,
-    repetition_penalty,
-    history_state,
-):
-    # Empty input → nothing to do
-    if not user_message.strip():
-        return history_state, history_state
+# Global dictionary to store all conversations: {id: {"title": str, "messages": list}}
+conversations = {}
 
-    # System prompt (kept from your Phi-4 version)
-    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: "
-        "<think> … </think> and Solution."
-    )
+def generate_conversation_id():
+    return str(uuid.uuid4())[:8]
 
-    prompt = build_chat_prompt(history_state, user_message, system_message)
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    # Stream tokens as they come
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=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_k=int(top_k),
-        top_p=float(top_p),
-        repetition_penalty=float(repetition_penalty),
-        streamer=streamer,
-    )
+@spaces.GPU(duration=60)
+def generate_response(user_message, max_tokens, temperature, top_p, history_state):
+    if not user_message.strip():
+        return history_state, history_state
 
-    # Run generate in a background thread so the UI stays responsive
-    Thread(target=model.generate, kwargs=generation_kwargs).start()
+    model = our_model
+    tokenizer = our_tokenizer
+    start_tag = "<|im_start|>"
+    sep_tag = "<|im_sep|>"
+    end_tag = "<|im_end|>"
+
+    system_message = "Your role as an assistant..."
+    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}"
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    generation_kwargs = {
+        "input_ids": inputs["input_ids"],
+        "attention_mask": inputs["attention_mask"],
+        "max_new_tokens": int(max_tokens),
+        "do_sample": True,
+        "temperature": temperature,
+        "top_k": 50,
+        "top_p": top_p,
+        "repetition_penalty": 1.0,
+        "pad_token_id": tokenizer.eos_token_id,
+        "streamer": streamer,
+    }
+
+    try:
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+    except Exception:
+        yield history_state + [{"role": "user", "content": user_message}, {"role": "assistant", "content": "⚠️ Generation failed."}], history_state
+        return
 
     assistant_response = ""
     new_history = history_state + [
         {"role": "user", "content": user_message},
-        {"role": "assistant", "content": ""},
+        {"role": "assistant", "content": ""}
     ]
 
-    for token in streamer:
-        # strip any stray special tokens the model may output
-        cleaned = (
-            token.replace("<|im_start|>", "")
-                 .replace("<|im_end|>", "")
-                 .replace("<|im_sep|>", "")
-        )
-        assistant_response += cleaned
-        new_history[-1]["content"] = assistant_response.strip()
-        yield new_history, new_history
+    try:
+        for new_token in streamer:
+            if "<|end" in new_token:
+                continue
+            cleaned_token = new_token.replace("<|im_start|>", "").replace("<|im_sep|>", "").replace("<|im_end|>", "")
+            assistant_response += cleaned_token
+            new_history[-1]["content"] = assistant_response.strip()
+            yield new_history, new_history
+    except Exception:
+        pass
 
     yield new_history, new_history
 
-# ---------------------------------------------------------------------
-# 4.  Example questions (unchanged)
-# ---------------------------------------------------------------------
+
 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.",
+    "JEE Main 2025 Combinatorics": "From all the English alphabets, five letters are chosen and are arranged in alphabetical order. The total number of ways, in which the middle letter is 'M', is?",
+    "JEE Main 2025 Coordinate Geometry": "A circle \\(C\\) of radius 2 lies in the second quadrant and touches both the coordinate axes. Let \\(r\\) be the radius of a circle that has centre at the point \\((2, 5)\\) and intersects the circle \\(C\\) at exactly two points. If the set of all possible values of \\(r\\) is the interval \\((\\alpha, \\beta)\\), then \\(3\\beta - 2\\alpha\\) is?",
+    "JEE Main 2025 Probability & Statistics": "A coin is tossed three times. Let \(X\) denote the number of times a tail follows a head. If \\(\\mu\\) and \\(\\sigma^2\\) denote the mean and variance of \\(X\\), then the value of \\(64(\\mu + \\sigma^2)\\) is?",
+    "JEE Main 2025 Laws of Motion": "A massless spring gets elongated by amount x_1 under a tension of 5 N . Its elongation is x_2 under the tension of 7 N . For the elongation of 5x_1 - 2x_2 , the tension in the spring will be?"
 }
 
-# ---------------------------------------------------------------------
-# 5.  Gradio UI (identical to the original, just lower default max_tokens)
-# ---------------------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # Fathom-R1-14B Chatbot  
-        The model excels at multi-step reasoning in mathematics, logic, and science.
-
-        It returns two sections:\n
-        1. **<think>** – detailed chain-of-thought (reasoning)\n
-        2. **Solution** – concise, final answer
-        """
-    )
+    # Global heading stays at top
+    #gr.Markdown("# Ramanujan Ganit R1 14B V1 Chatbot")
+    gr.HTML(
+    """
+    <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" alt="Fractal AI Logo" style="height: 48px;">
+        </div>
+        <h1 style="margin: 0;">Ramanujan Ganit R1 14B V1 Chatbot</h1>
+    </div>
+    """
+)
+
+    with gr.Sidebar():
+        gr.Markdown("## Conversations")
+        conversation_selector = gr.Radio(choices=[], label="Select Conversation", interactive=True)
+        new_convo_button = gr.Button("New Conversation ➕")
 
+    current_convo_id = gr.State(generate_conversation_id())
     history_state = gr.State([])
 
     with gr.Row():
-        # Settings panel
         with gr.Column(scale=1):
-            gr.Markdown("### Settings")
-            max_tokens_slider = gr.Slider(
-                minimum=64, maximum=4096, step=256, value=1024, label="Max Tokens"
+            # INTRO TEXT MOVED HERE
+            gr.Markdown(
+                """
+                Welcome to the Ramanujan Ganit R1 14B V1 Chatbot, developed by Fractal AI Research! 
+                
+                Our model excels at reasoning tasks in mathematics and science.  
+                
+                Try the example problems below from JEE Main 2025 or type in your own problems to see how our model breaks down complex reasoning problems.
+
+                Please note that once you close this demo window, all currently saved conversations will be lost.
+                """
             )
-            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"
-                )
-
-        # Chat area
+
+            gr.Markdown("### Settings")
+            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")
+
+            # New acknowledgment line at bottom
+            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.Column(scale=4):
-            chatbot = gr.Chatbot(label="Chat", type="messages")
+            #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)
+                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=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")
+                example1_button = gr.Button("JEE Main 2025\nCombinatorics")
+                example2_button = gr.Button("JEE Main 2025\nCoordinate Geometry")
+                example3_button = gr.Button("JEE Main 2025\nProbability & Statistics")
+                example4_button = gr.Button("JEE Main 2025\nLaws of Motion")
+
+    def update_conversation_list():
+        return [conversations[cid]["title"] for cid in conversations]
+
+    def start_new_conversation():
+        new_id = generate_conversation_id()
+        conversations[new_id] = {"title": f"New Conversation {new_id}", "messages": []}
+        return new_id, [], gr.update(choices=update_conversation_list(), value=conversations[new_id]["title"])
+
+    def load_conversation(selected_title):
+        for cid, convo in conversations.items():
+            if convo["title"] == selected_title:
+                return cid, convo["messages"], convo["messages"]
+        return current_convo_id.value, history_state.value, history_state.value
+
+    def send_message(user_message, max_tokens, temperature, top_p, convo_id, history):
+        if convo_id not in conversations:
+            #title = user_message.strip().split("\n")[0][:40]
+            title = " ".join(user_message.strip().split()[:5])
+            conversations[convo_id] = {"title": title, "messages": history}
+        if conversations[convo_id]["title"].startswith("New Conversation"):
+            #conversations[convo_id]["title"] = user_message.strip().split("\n")[0][:40]
+            conversations[convo_id]["title"] = " ".join(user_message.strip().split()[:5])
+        for updated_history, new_history in generate_response(user_message, max_tokens, temperature, top_p, history):
+            conversations[convo_id]["messages"] = new_history
+            yield updated_history, new_history, gr.update(choices=update_conversation_list(), value=conversations[convo_id]["title"])
 
-    # Button wiring
     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],
+        fn=send_message,
+        inputs=[user_input, max_tokens_slider, temperature_slider, top_p_slider, current_convo_id, history_state],
+        outputs=[chatbot, history_state, conversation_selector]
     ).then(
         fn=lambda: gr.update(value=""),
         inputs=None,
-        outputs=user_input,
+        outputs=user_input
     )
 
     clear_button.click(
         fn=lambda: ([], []),
         inputs=None,
-        outputs=[chatbot, history_state],
+        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"]),
+    new_convo_button.click(
+        fn=start_new_conversation,
         inputs=None,
-        outputs=user_input,
+        outputs=[current_convo_id, history_state, conversation_selector]
     )
-    example3_button.click(
-        fn=lambda: gr.update(value=example_messages["Physics problem"]),
-        inputs=None,
-        outputs=user_input,
+
+    conversation_selector.change(
+        fn=load_conversation,
+        inputs=conversation_selector,
+        outputs=[current_convo_id, history_state, chatbot]
     )
 
-# ---------------------------------------------------------------------
-# 6.  Launch
-# ---------------------------------------------------------------------
-if __name__ == "__main__":
-    demo.launch(ssr_mode=False)
+    example1_button.click(fn=lambda: gr.update(value=example_messages["JEE Main 2025 Combinatorics"]), inputs=None, outputs=user_input)
+    example2_button.click(fn=lambda: gr.update(value=example_messages["JEE Main 2025 Coordinate Geometry"]), inputs=None, outputs=user_input)
+    example3_button.click(fn=lambda: gr.update(value=example_messages["JEE Main 2025 Probability & Statistics"]), inputs=None, outputs=user_input)
+    example4_button.click(fn=lambda: gr.update(value=example_messages["JEE Main 2025 Laws of Motion"]), inputs=None, outputs=user_input)
+
+demo.launch(share=True, ssr_mode=False)