app.py

import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
import torch
from threading import Thread
import re
import uuid
from openai import OpenAI

client = OpenAI(
    base_url="https://a7g1ajqixo23revq.us-east-1.aws.endpoints.huggingface.cloud/v1/",
    api_key="hf_XXXXX"
)

def format_math(text):
    text = re.sub(r"\[(.*?)\]", r"$$\1$$", text, flags=re.DOTALL)
    text = text.replace(r"\(", "$").replace(r"\)", "$")
    return text

# Global dictionary to store all conversations: {id: {"title": str, "messages": list}}
conversations = {}

def generate_conversation_id():
    return str(uuid.uuid4())[:8]


import tiktoken                               
enc = tiktoken.encoding_for_model("gpt-3.5-turbo")  # any OpenAI encoding works

def generate_response(user_message,
                      max_tokens,
                      temperature,
                      top_p,
                      history_state):
    if not user_message.strip():
        return history_state, history_state

    system_message = "Your role as an assistant..."
    messages = [{"role": "system", "content": system_message}]
    for m in history_state:
        messages.append({"role": m["role"], "content": m["content"]})
    messages.append({"role": "user", "content": user_message})

    try:
        response = client.chat.completions.create(
            model="tgi",
            messages=messages,
            max_tokens=int(max_tokens),        # server-side limit
            temperature=temperature,
            top_p=top_p,
            stream=True
        )
    except Exception as e:
        print(f"[ERROR] OpenAI API call failed: {e}")
        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": ""}
    ]

    token_budget = int(max_tokens)
    tokens_seen = 0

    try:
        for chunk in response:
            if (not chunk.choices
                or not chunk.choices[0].delta
                or not chunk.choices[0].delta.content):
                continue

            token_text = chunk.choices[0].delta.content
            assistant_response += token_text
            # count how many tokens that piece is worth
            tokens_seen += len(enc.encode(token_text))

            new_history[-1]["content"] = assistant_response.strip()
            yield new_history, new_history

            if tokens_seen >= token_budget:
                break                            # stop the local loop
    except Exception:
        pass

    yield new_history, new_history


example_messages = {
    "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?"
}

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    # 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 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():
        with gr.Column(scale=1):
            # INTRO TEXT MOVED HERE
            gr.Markdown(
                """
                Welcome to the Ramanujan Ganit R1 14B 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.

                NOTE: that once you close this demo window, all currently saved conversations will be lost.
                """
            )

            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", height=520)
            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)
            gr.Markdown("**Try these examples:**")
            with gr.Row():
                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"])

    submit_button.click(
        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
    )

    clear_button.click(
        fn=lambda: ([], []),
        inputs=None,
        outputs=[chatbot, history_state]
    )

    new_convo_button.click(
        fn=start_new_conversation,
        inputs=None,
        outputs=[current_convo_id, history_state, conversation_selector]
    )

    conversation_selector.change(
        fn=load_conversation,
        inputs=conversation_selector,
        outputs=[current_convo_id, history_state, chatbot]
    )

    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)

if __name__ == "__main__":
    demo.queue(                                       # turn the queue on
        concurrency_count = 8,                       # how many jobs run simultaneously
        max_size         = 40,                       # 40-6 = 34 jobs can wait in line
        api_open         = False                    # no public /queue/status endpoint
    ).launch(
        share     = True,
        ssr_mode  = False
    )