Update app.py

app.py

@@ -1,64 +1,175 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-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("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", 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.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-
-if __name__ == "__main__":
-    demo.launch()
+import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import torch
+from threading import Thread
+import re
+
+phi4_model_path = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
+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)
+
+def format_math(text):
+    text = re.sub(r"\[(.*?)\]", r"$$\1$$", text, flags=re.DOTALL)
+    text = text.replace(r"\(", "$").replace(r"\)", "$")
+    return text
+
+@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
+
+    model = phi4_model
+    tokenizer = phi4_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,  # Fixed default
+        "top_p": top_p,
+        "repetition_penalty": 1.0,  # Fixed default
+        "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": ""}
+    ]
+
+    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
+
+example_messages = {
+    "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?",
+    "Co-ordinate 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?",
+    "Prob-Stats": "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?"
+}
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # Ramanujan Ganit R1 14B V1 Chatbot 
+    
+        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.
+        """
+    )
+
+    history_state = gr.State([])
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            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=False):
+                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"
+                )
+
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(label="Chat", type="messages")
+            with gr.Row():
+                user_input = gr.Textbox(
+                    label="User Input",
+                    placeholder="Type your question here...",
+                    scale=3
+                )
+                submit_button = gr.Button("Send", variant="primary", scale=1)
+                #stop_button = gr.Button("Stop", variant="stop", scale=1, interactive=True)
+                clear_button = gr.Button("Clear", scale=1)
+            gr.Markdown("**Try these examples:**")
+            with gr.Row():
+                example1_button = gr.Button("Combinatorics")
+                example2_button = gr.Button("Co-ordinate Geometry")
+                example3_button = gr.Button("Prob-Stats")
+
+    submit_button.click(
+        fn=generate_response,
+        inputs=[user_input, max_tokens_slider, temperature_slider, top_p_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["Combinatorics"]),
+        inputs=None,
+        outputs=user_input
+    )
+    example2_button.click(
+        fn=lambda: gr.update(value=example_messages["Co-ordinate Geometry"]),
+        inputs=None,
+        outputs=user_input
+    )
+    example3_button.click(
+        fn=lambda: gr.update(value=example_messages["Prob-Stats"]),
+        inputs=None,
+        outputs=user_input
+    )
+
+demo.launch(ssr_mode=False)