Update app.py

app.py

@@ -1,296 +1,226 @@
-# ---------------------------------------------------------------
-#  Fathom-R1-14B ZeroGPU chat-demo  (Gradio Blocks)
-# ---------------------------------------------------------------
-
+# app.py  –  Gradio chatbot for FractalAIResearch/Fathom-R1-14B
+# ---------------------------------------------------------------------
 import gradio as gr
 import spaces
-import torch, re, uuid, tiktoken
-from transformers import (AutoModelForCausalLM,
-                          AutoTokenizer,
-                          TextIteratorStreamer)
+import torch
 from threading import Thread
-
-# ────────────────────────────────────────────────────────────────
-# 1.  Load the model on the single GPU supplied by ZeroGPU
-#    (4-bit to stay well below the 24 GB VRAM of an A10G)
-# ────────────────────────────────────────────────────────────────
-model_name = "FractalAIResearch/Fathom-R1-14B"
-
-try:
-    # 1-line 4-bit loading (needs bitsandbytes, already in HF Space image)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        device_map="auto",
-        load_in_4bit=True,
-        trust_remote_code=True
-    )
-except RuntimeError:
-    # fallback to fp16 if 4-bit isn’t available
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        trust_remote_code=True
-    )
-
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-device = next(model.parameters()).device      # usually  cuda:0
-
-
-# ────────────────────────────────────────────────────────────────
-# 2.  Helpers
-# ────────────────────────────────────────────────────────────────
-def format_math(text: str) -> str:
-    "Replace [...]/\\(...\\) with $$...$$ for nicer math rendering"
-    text = re.sub(r"\[(.*?)\]", r"$$\1$$", text, flags=re.DOTALL)
-    return text.replace(r"\(", "$").replace(r"\)", "$")
-
-
-def generate_conversation_id() -> str:
-    return str(uuid.uuid4())[:8]
-
-
-# tiktoken – we just keep it to count tokens during streaming
-enc = tiktoken.encoding_for_model("gpt-3.5-turbo")
-
-
-# Build a prompt that Fathom-R1 understands
-BOS, SEP, EOS = "<|im_start|>", "<|im_sep|>", "<|im_end|>"
-
-system_message = (
-    "Your role as an assistant involves thoroughly exploring questions "
-    "through a systematic thinking process before providing the final "
-    "precise and accurate solutions. …"   # same text you used before
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
 )
 
+# ---------------------------------------------------------------------
+# 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,
+)
 
-def build_prompt(history, user_msg: str) -> str:
-    prompt = f"{BOS}system{SEP}{system_message}{EOS}"
-    for m in history:
-        role = m["role"]
-        prompt += f"{BOS}{role}{SEP}{m['content']}{EOS}"
-    prompt += f"{BOS}user{SEP}{user_msg}{EOS}{BOS}assistant{SEP}"
-    return prompt
-
+print("✅ Model loaded")
 
-# ────────────────────────────────────────────────────────────────
-# 3.  Generation (runs on the GPU for 60 s max per call)
-# ────────────────────────────────────────────────────────────────
-@spaces.GPU(duration=60)
-def generate_response(user_message,
-                      max_tokens,
-                      temperature,
-                      top_p,
-                      history_state):
+# ---------------------------------------------------------------------
+# 2.  Helper: build a prompt with the tokenizer’s chat_template
+# ---------------------------------------------------------------------
+def build_chat_prompt(history, user_message, system_message):
     """
-    Takes exactly the same signature the rest of the UI expects:
-    returns   (visible_chatbot,  history_state)
+    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,
+    )
+
+# ---------------------------------------------------------------------
+# 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
 
-    prompt = build_prompt(history_state, user_message)
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    # 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."
+    )
+
+    prompt = build_chat_prompt(history_state, user_message, system_message)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    streamer = TextIteratorStreamer(tokenizer,
-                                    skip_prompt=True,
-                                    skip_special_tokens=True)
+    # Stream tokens as they come
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
 
-    gen_kwargs = dict(
+    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),
-        do_sample=True,
-        eos_token_id=tokenizer.eos_token_id,
-        pad_token_id=tokenizer.eos_token_id,
-        streamer=streamer
+        repetition_penalty=float(repetition_penalty),
+        streamer=streamer,
     )
 
-    # run generate in a background thread – lets us stream tokens
-    Thread(target=model.generate, kwargs=gen_kwargs).start()
+    # Run generate in a background thread so the UI stays responsive
+    Thread(target=model.generate, kwargs=generation_kwargs).start()
 
     assistant_response = ""
     new_history = history_state + [
         {"role": "user", "content": user_message},
-        {"role": "assistant", "content": ""}
+        {"role": "assistant", "content": ""},
     ]
 
-    # live-stream tokens to the UI
-    tokens_seen = 0
-    token_budget = int(max_tokens)
-
-    for new_tok in streamer:
-        assistant_response += new_tok
-        tokens_seen += len(enc.encode(new_tok))
-        new_history[-1]["content"] = format_math(assistant_response.strip())
+    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
-        if tokens_seen >= token_budget:
-            break
 
-    # final return
     yield new_history, new_history
 
-
-# ────────────────────────────────────────────────────────────────
-# 4.  Demo UI  – identical to your current one
-# ────────────────────────────────────────────────────────────────
+# ---------------------------------------------------------------------
+# 4.  Example questions (unchanged)
+# ---------------------------------------------------------------------
 example_messages = {
-    "IIT-JEE 2024 Mathematics": (
-        "A student appears for a quiz consisting of only true-false type "
-        "questions and answers all the questions. …"
-    ),
-    "IIT-JEE 2025         Physics": (
-        "A person sitting inside an elevator performs a weighing experiment …"
-    ),
-    "Goldman Sachs Interview Puzzle": (
-        "Four friends need to cross a dangerous bridge at night …"
-    ),
-    "IIT-JEE 2025 Mathematics": (
-        "Let S be the set of all seven-digit numbers that can be formed …"
-    )
+    "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.",
 }
 
+# ---------------------------------------------------------------------
+# 5.  Gradio UI (identical to the original, just lower default max_tokens)
+# ---------------------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    # session-scoped states
-    conversations_state = gr.State({})
-    current_convo_id   = gr.State(generate_conversation_id())
-    history_state      = gr.State([])
-
-    # Header
-    gr.HTML(
+    gr.Markdown(
         """
-        <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>
+        # 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
         """
     )
 
-    # Sidebar
-    with gr.Sidebar():
-        gr.Markdown("## Conversations")
-        conversation_selector = gr.Radio(choices=[], label="Select Conversation", interactive=True)
-        new_convo_button = gr.Button("New Conversation ➕")
+    history_state = gr.State([])
 
     with gr.Row():
+        # Settings panel
         with gr.Column(scale=1):
-            # intro text
-            gr.Markdown(
-                """
-                Welcome to the Fathom R1 14B Chatbot, developed by **Fractal AI Research**!  
-                This model excels at reasoning tasks in mathematics and science …
-
-                Once you close this demo window, all currently saved conversations will be lost.
-                """
-            )
-
-            # Settings
             gr.Markdown("### Settings")
-            max_tokens_slider = gr.Slider(6144, 32768, step=1024, value=16384, label="Max Tokens")
-            with gr.Accordion("Advanced Settings", open=True):
-                temperature_slider = gr.Slider(0.1, 2.0, value=0.6, label="Temperature")
-                top_p_slider       = gr.Slider(0.1, 1.0, value=0.95, label="Top-p")
-
-            gr.Markdown(
-                """
-                We sincerely acknowledge [VIDraft](https://huggingface.co/VIDraft) …
-                """
+            max_tokens_slider = gr.Slider(
+                minimum=64, maximum=4096, step=256, value=1024, label="Max Tokens"
             )
-
+            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
         with gr.Column(scale=4):
-            chatbot = gr.Chatbot(label="Chat", type="messages", height=520)
+            chatbot = gr.Chatbot(label="Chat", type="messages")
             with gr.Row():
-                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)
-
-            # examples
+                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("IIT-JEE 2025 Mathematics")
-                example2_button = gr.Button("IIT-JEE 2025         Physics")
-                example3_button = gr.Button("Goldman Sachs Interview Puzzle")
-                example4_button = gr.Button("IIT-JEE 2024 Mathematics")
-
-    # ───────── conversation-management helpers ──────────────────
-    def update_conversation_list(conversations):
-        return [conversations[cid]["title"] for cid in conversations]
-
-    def start_new_conversation(conversations):
-        new_id = generate_conversation_id()
-        conversations[new_id] = {"title": f"New Conversation {new_id}", "messages": []}
-        return new_id, [], gr.update(choices=update_conversation_list(conversations),
-                                     value=conversations[new_id]["title"]), conversations
-
-    def load_conversation(selected_title, conversations):
-        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
+                example1_button = gr.Button("Math reasoning")
+                example2_button = gr.Button("Logic puzzle")
+                example3_button = gr.Button("Physics problem")
 
-    # main “send” wrapper: keeps conversations dict in sync
-    def send_message(user_message, max_tokens, temperature, top_p,
-                     convo_id, history, conversations):
-        if convo_id not in conversations:
-            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"] = " ".join(user_message.strip().split()[:5])
-
-        # call the streamer generator and forward its yields
-        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(conversations),
-                             value=conversations[convo_id]["title"]),
-                   conversations)
-
-    # ───────── UI → functions wiring ────────────────────────────
+    # Button wiring
     submit_button.click(
-        fn=send_message,
-        inputs=[user_input, max_tokens_slider, temperature_slider, top_p_slider,
-                current_convo_id, history_state, conversations_state],
-        outputs=[chatbot, history_state, conversation_selector, conversations_state],
-        concurrency_limit=16
+        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
+        outputs=user_input,
     )
 
-    clear_button.click(fn=lambda: ([], []), inputs=None,
-                       outputs=[chatbot, history_state])
-
-    new_convo_button.click(fn=start_new_conversation,
-                           inputs=[conversations_state],
-                           outputs=[current_convo_id, history_state,
-                                    conversation_selector, conversations_state])
-
-    conversation_selector.change(fn=load_conversation,
-                                 inputs=[conversation_selector, conversations_state],
-                                 outputs=[current_convo_id, history_state, chatbot])
+    clear_button.click(
+        fn=lambda: ([], []),
+        inputs=None,
+        outputs=[chatbot, history_state],
+    )
 
-    # example buttons
-    example1_button.click(lambda: gr.update(value=example_messages["IIT-JEE 2025 Mathematics"]),
-                          None, user_input)
-    example2_button.click(lambda: gr.update(value=example_messages["IIT-JEE 2025         Physics"]),
-                          None, user_input)
-    example3_button.click(lambda: gr.update(value=example_messages["Goldman Sachs Interview Puzzle"]),
-                          None, user_input)
-    example4_button.click(lambda: gr.update(value=example_messages["IIT-JEE 2024 Mathematics"]),
-                          None, user_input)
+    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,
+    )
 
-# ────────────────────────────────────────────────────────────────
-# 5.  Launch
-# ────────────────────────────────────────────────────────────────
+# ---------------------------------------------------------------------
+# 6.  Launch
+# ---------------------------------------------------------------------
 if __name__ == "__main__":
-    demo.queue().launch(share=True, ssr_mode=False)
+    demo.launch(ssr_mode=False)