Update app.py

app.py

@@ -1,57 +1,108 @@
+# app.py
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
 import gradio as gr
+import spaces
 
-model_id = "FractalAIResearch/Fathom-R1-14B"  # or your HF repo path
+MODEL_NAME = "FractalAIResearch/Fathom-R1-14B"
 
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        device_map="auto",
-        torch_dtype=torch.bfloat16,
-        trust_remote_code=True,
-        low_cpu_mem_usage=True
-    )
-    return model, tokenizer
-
-model, tokenizer = load_model()
-
-def generate_response(message, history):
-    prompt = ""
-    for user, bot in history:
-        prompt += f"<|user|>\n{user.strip()}\n<|assistant|>\n{bot.strip()}\n"
-    prompt += f"<|user|>\n{message.strip()}\n<|assistant|>\n"
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        **inputs,
-        streamer=streamer,
-        max_new_tokens=16384,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9,
+@spaces.GPU
+class Chatbot:
+    def __init__(self):
+        print("⏳ Loading model...")
+        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True,
+        )
+        self.model.eval()
+        print("✅ Model loaded!")
+
+    def chat(self, messages, temperature, max_new_tokens, top_p, repetition_penalty):
+        # Format messages into prompt
+        prompt = self._format_messages(messages)
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.model.device)
+
+        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = dict(
+            input_ids=input_ids,
+            streamer=streamer,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            top_p=top_p,
+            temperature=temperature,
+            repetition_penalty=repetition_penalty,
+        )
+
+        thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        response = ""
+        for token in streamer:
+            response += token
+            yield response
+
+    def _format_messages(self, messages):
+        prompt = ""
+        for msg in messages:
+            if msg["role"] == "user":
+                prompt += f"<|user|>\n{msg['content'].strip()}\n"
+            elif msg["role"] == "assistant":
+                prompt += f"<|assistant|>\n{msg['content'].strip()}\n"
+        prompt += "<|assistant|>\n"
+        return prompt
+
+chatbot = Chatbot()
+
+# Chat state management
+def user_submit(user_message, history):
+    history = history + [{"role": "user", "content": user_message}, {"role": "assistant", "content": ""}]
+    return "", history, gr.update(visible=True)
+
+def generate(history, temperature, max_new_tokens, top_p, repetition_penalty):
+    response_gen = chatbot.chat(
+        history,
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
     )
-    thread = Thread(target=model.generate, kwargs=generate_kwargs)
-    thread.start()
-    for new_text in streamer:
-        yield new_text
-
-with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 Chat with Fathom-R1 14B")
-    chatbot = gr.Chatbot(show_copy_button=True)
-    msg = gr.Textbox(placeholder="Ask me anything...", container=False)
+    partial = ""
+    for chunk in response_gen:
+        partial = chunk
+        history[-1]["content"] = partial
+        yield history, history
+
+def reset():
+    return [], []
+
+with gr.Blocks(css="footer {display: none !important;}") as demo:
+    gr.Markdown("<h1 align='center'>🧠 Fathom R1 14B Chatbot</h1>")
+    chatbot_ui = gr.Chatbot([], elem_id="chatbot", height=500, bubble_full_width=False)
     state = gr.State([])
 
-    def user_submit(message, history):
-        history = history + [[message, ""]]
-        return "", history
+    with gr.Row():
+        with gr.Column(scale=6):
+            txt = gr.Textbox(placeholder="Ask a math question...", label="Your Message")
+        with gr.Column(scale=1):
+            submit = gr.Button("Submit", variant="primary")
+            clear = gr.Button("Clear")
 
-    msg.submit(user_submit, [msg, state], [msg, state]).then(
-        generate_response, [msg, state], chatbot
-    )
+    with gr.Accordion("Advanced settings", open=False):
+        temperature = gr.Slider(0.1, 1.5, value=0.7, label="Temperature")
+        max_new_tokens = gr.Slider(64, 2048, step=64, value=512, label="Max New Tokens")
+        top_p = gr.Slider(0.1, 1.0, value=0.95, label="Top-p")
+        repetition_penalty = gr.Slider(1.0, 2.0, value=1.1, label="Repetition Penalty")
+
+    submit.click(user_submit, [txt, state], [txt, state, chatbot_ui], queue=False)\
+          .then(generate, [state, temperature, max_new_tokens, top_p, repetition_penalty], [chatbot_ui, state])
+
+    txt.submit(user_submit, [txt, state], [txt, state, chatbot_ui], queue=False)\
+       .then(generate, [state, temperature, max_new_tokens, top_p, repetition_penalty], [chatbot_ui, state])
+
+    clear.click(reset, outputs=[chatbot_ui, state])
 
-if __name__ == "__main__":
-    demo.queue().launch()
+demo.queue().launch()