# app.py
# =============
# This is a complete app.py file for deploying the MTSAIR/Cotype-Nano model using Gradio and Hugging Face Transformers with chat and token streaming functionality, advanced settings, and English interface.

import gradio as gr
from transformers import pipeline

# Load the model and pipeline
model_name = "MTSAIR/Cotype-Nano"
pipe = pipeline("text-generation", model=model_name, device="cpu")

# Define the system prompt
system_prompt = {"role": "system", "content": "You are an AI assistant. Your task is to generate a detailed and comprehensive response."}

# Define the Gradio interface
def generate_response(history, user_input, temperature, max_new_tokens):
    messages = [system_prompt] + history + [{"role": "user", "content": user_input}]
    response = pipe(messages, max_new_tokens=max_new_tokens, temperature=temperature, return_full_text=False)
    generated_text = response[0]['generated_text']
    history.append({"role": "user", "content": user_input})
    history.append({"role": "assistant", "content": generated_text})
    return history, ""

# Function to clear chat history
def clear_chat():
    return [], ""

# Custom CSS for styling
custom_css = """
    #chatbot {
        height: 400px;
        overflow-y: auto;
        border: 1px solid #ccc;
        padding: 10px;
        border-radius: 5px;
    }
    .gr-button {
        margin-top: 10px;
        padding: 10px 20px;
        font-size: 16px;
        border-radius: 5px;
    }
    .gr-textbox {
        margin-top: 10px;
        padding: 10px;
        font-size: 16px;
        border-radius: 5px;
    }
    .gr-slider {
        margin-top: 20px;
    }
"""

# Create the Gradio interface
with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("## Cotype-Nano Text Generation Chat")

    chatbot = gr.Chatbot([], elem_id="chatbot", type='messages')

    with gr.Row():
        txt = gr.Textbox(
            show_label=False,
            placeholder="Type your message here...",
        )

        send_btn = gr.Button("Send")

    with gr.Row():
        clear_btn = gr.Button("Clear Chat")

    with gr.Row():
        temperature_slider = gr.Slider(0, 1, 0.7, step=0.1, label="Temperature")
        max_new_tokens_slider = gr.Slider(1, 1000, 100, step=1, label="Max New Tokens")

    send_btn.click(generate_response, [chatbot, txt, temperature_slider, max_new_tokens_slider], [chatbot, txt])
    txt.submit(generate_response, [chatbot, txt, temperature_slider, max_new_tokens_slider], [chatbot, txt])
    clear_btn.click(clear_chat, outputs=[chatbot, txt])

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