app.py

import gradio as gr
from huggingface_hub import InferenceClient

# Define available models (update with your actual model IDs)
model_list = {
    "Safe LM": "HuggingFaceH4/zephyr-7b-beta",  # Replace with your Safe LM model ID
    "Zephyr Beta": "HuggingFaceH4/zephyr-7b-beta",
    "Another Model": "HuggingFaceH4/zephyr-7b-beta"
}

def respond(message, history, system_message, max_tokens, temperature, top_p, selected_model):
    # Create an InferenceClient for the selected model
    client = InferenceClient(model_list.get(selected_model, "HuggingFaceH4/zephyr-7b-beta"))
    
    # Build conversation messages for the client
    messages = [{"role": "system", "content": system_message}]
    for user_msg, assistant_msg in history:
        if user_msg:  # Only add non-empty messages
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:  # Only add non-empty messages
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})
    
    response = ""
    
    # Stream the response from the client
    for token_message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = token_message.choices[0].delta.content
        if token is not None:  # Handle potential None values
            response += token
            yield response

# Custom CSS for branding with blue, teal and gold accents
css = """
body { 
    background-color: #f5f7fa; 
    font-family: 'Inter', 'Segoe UI', sans-serif;
}
.gradio-container { 
    background-color: #FFFFFF; 
    border-radius: 16px; 
    box-shadow: 0 4px 12px rgba(0,0,0,0.1); 
    max-width: 95%;
    margin: 20px auto;
}
.app-header {
    background: linear-gradient(135deg, #0a76d8, #36b3c2);
    padding: 20px;
    border-radius: 16px 16px 0 0;
    position: relative;
    color: white;
    overflow: hidden;
}
.app-header::before {
    content: "🔒";
    position: absolute;
    font-size: 120px;
    opacity: 0.1;
    right: -20px;
    top: -30px;
    transform: rotate(15deg);
}
.app-title {
    font-size: 28px;
    font-weight: 800;
    margin: 0;
    display: flex;
    align-items: center;
}
.app-title .safe {
    font-weight: 900;
    margin-right: 5px;
}
.app-title .lm {
    font-weight: 300;
    letter-spacing: 1px;
}
.app-title .shield {
    margin-right: 10px;
    font-size: 24px;
}
.app-subtitle {
    font-size: 14px;
    opacity: 0.9;
    margin-top: 5px;
}
.golden-accent {
    border-top: 3px solid #e6c200;
    margin: 0;
}
button, .gradio-button {
    background: linear-gradient(to right, #0a76d8, #36b3c2) !important;
    color: white !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 10px 20px !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
}
button:hover, .gradio-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 8px rgba(10, 118, 216, 0.3) !important;
}
.gradio-dropdown, .gradio-slider, .gradio-textbox, .gradio-checkbox {
    border-radius: 12px !important;
    border: 1px solid #e1e5eb !important;
}
.model-select {
    border-left: 4px solid #e6c200 !important;
    padding-left: 10px !important;
}
.footer {
    text-align: center;
    color: #666;
    font-size: 12px;
    margin-top: 20px;
    padding: 10px;
}
"""

with gr.Blocks(css=css) as demo:
    # Custom header with branding
    with gr.HTML():
        gr.HTML("""
        <div class="app-header">
            <h1 class="app-title">
                <span class="shield">🔒</span>
                <span class="safe">Safe</span>
                <span class="lm">Playground</span>
            </h1>
            <p class="app-subtitle">Responsible AI for everyone</p>
        </div>
        <hr class="golden-accent">
        """)
        
    with gr.Row():
        # Left sidebar: Model selector
        with gr.Column(scale=1):
            gr.Markdown("## Models")
            model_dropdown = gr.Dropdown(
                choices=list(model_list.keys()),
                label="Select Model",
                value="Safe LM",
                elem_classes=["model-select"]
            )
            
            # Settings
            gr.Markdown("### Settings")
            system_message = gr.Textbox(
                value="You are a friendly and safe assistant.",
                label="System Message",
                lines=2
            )
            max_tokens_slider = gr.Slider(
                minimum=1, maximum=2048, value=512, step=1, 
                label="Max New Tokens"
            )
            temperature_slider = gr.Slider(
                minimum=0.1, maximum=4.0, value=0.7, step=0.1, 
                label="Temperature"
            )
            top_p_slider = gr.Slider(
                minimum=0.1, maximum=1.0, value=0.95, step=0.05, 
                label="Top-p (nucleus sampling)"
            )
            
        # Main area: Chat interface
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                label="Conversation", 
                show_label=True,
                avatar_images=("👤", "🔒"),
                height=500
            )
            with gr.Row():
                user_input = gr.Textbox(
                    placeholder="Type your message here...", 
                    label="Your Message",
                    show_label=False,
                    scale=9
                )
                send_button = gr.Button(
                    "Send", 
                    scale=1,
                    variant="primary"
                )
            
            with gr.Row():
                clear_button = gr.Button("Clear Chat")
    
    # Footer with branding
    with gr.HTML():
        gr.HTML("""
        <div class="footer">
            <p>Safe Playground™ | Responsible AI Technology | © 2025</p>
        </div>
        """)
    
    # Fix 1: Correct event handling for the chatbot interface
    def user(user_message, history):
        # Return the user's message and add it to history
        return "", history + [[user_message, None]]
    
    def bot(history, system_message, max_tokens, temperature, top_p, selected_model):
        # Get the last user message from history
        user_message = history[-1][0]
        # Call respond function with the message
        response_generator = respond(
            user_message, 
            history[:-1],  # Pass history without the current message
            system_message, 
            max_tokens, 
            temperature, 
            top_p, 
            selected_model
        )
        # Update history as responses come in
        for response in response_generator:
            history[-1][1] = response
            yield history
    
    # Wire up the event chain
    user_input.submit(
        user,
        [user_input, chatbot],
        [user_input, chatbot],
        queue=False
    ).then(
        bot,
        [chatbot, system_message, max_tokens_slider, temperature_slider, top_p_slider, model_dropdown],
        [chatbot]
    )
    
    send_button.click(
        user,
        [user_input, chatbot],
        [user_input, chatbot],
        queue=False
    ).then(
        bot,
        [chatbot, system_message, max_tokens_slider, temperature_slider, top_p_slider, model_dropdown],
        [chatbot]
    )
    
    # Clear the chat history
    clear_button.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch()