import gradio as gr
import transformers_gradio

# Load the models
demo = gr.load(name="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B", src=transformers_gradio.registry)
demo = gr.load(name="deepseek-ai/DeepSeek-R1", src=transformers_gradio.registry)
demo = gr.load(name="deepseek-ai/DeepSeek-R1-Zero", src=transformers_gradio.registry)

# Disable API names for all functions
for fn in demo.fns.values():
    fn.api_name = False

# Define the optional parameters section
def create_optional_parameters():
    with gr.Accordion("Optional Parameters (Click to Expand)", open=False):
        system_message = gr.Textbox(
            label="System Message", 
            value="You are a friendly Chatbot created by balione.com",
            lines=2
        )
        max_new_tokens = gr.Slider(minimum=1, maximum=4000, value=200, label="Max New Tokens")
        temperature = gr.Slider(minimum=0.10, maximum=4.00, value=0.70, label="Temperature")
        top_p = gr.Slider(minimum=0.10, maximum=1.00, value=0.90, label="Top-p (nucleus sampling)")
    return system_message, max_new_tokens, temperature, top_p

# Define the main interface
def chat_interface(user_input, system_message, max_new_tokens, temperature, top_p):
    # Placeholder response - integrate with actual model here
    response = f"""**System Message**: {system_message}
**Your Input**: {user_input}
**Parameters Used**:
- Max New Tokens: {max_new_tokens}
- Temperature: {temperature}
- Top-p: {top_p}

*Note: Actual model integration required for real responses*"""
    return response

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# DeepSeek Chatbot\nCreated by [ruslanmv.com](https://ruslanmv.com/)")
    
    with gr.Row():
        with gr.Column():
            user_input = gr.Textbox(label="Your Message", placeholder="Type your message here...", lines=3)
            submit_button = gr.Button("Submit", variant="primary")
        with gr.Column():
            output = gr.Markdown(label="Chatbot Response")
    
    # Add the optional parameters section
    system_message, max_new_tokens, temperature, top_p = create_optional_parameters()
    
    # Link the submit button to the chat interface
    submit_button.click(
        chat_interface,
        inputs=[user_input, system_message, max_new_tokens, temperature, top_p],
        outputs=output
    )

# Launch the demo
if __name__ == "__main__":
    demo.launch()