import gradio as gr
from ctransformers import AutoModelForCausalLM

model_name = "lmstudio-community/Devstral-Small-2505-GGUF"
model_file = "devstral-small-2505.Q4_K_M.gguf"  # Выберем версию с квантизацией Q4_K_M для экономии памяти

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    model_file=model_file,
    model_type="mistral",
    gpu_layers=50,  # Используем GPU насколько возможно
    context_length=4096  # Максимальный контекст
)

def generate_text(prompt, max_tokens=512, temperature=0.7, top_p=0.9):
    # Форматируем запрос в стиле Mistral
    formatted_prompt = f"<s>[INST] {prompt} [/INST]"
    
    # Генерируем ответ
    response = model(
        formatted_prompt,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=1.1,
        stream=False
    )
    
    # Удаляем исходный запрос из ответа
    return response.replace(formatted_prompt, "").strip()

custom_css = """
:root {
    --primary-color: #4F46E5;
    --secondary-color: #6366F1;
    --background-color: #F9FAFB;
    --surface-color: #FFFFFF;
    --text-color: #1F2937;
    --border-radius: 10px;
}

body {
    background-color: var(--background-color);
}

.container {
    max-width: 900px;
    margin: auto;
    padding-top: 1.5rem;
}

.title {
    text-align: center;
    color: var(--primary-color);
    font-size: 2.2rem;
    font-weight: 700;
    margin-bottom: 0.5rem;
}

.subtitle {
    text-align: center;
    color: var(--text-color);
    opacity: 0.8;
    margin-bottom: 2rem;
}

footer {display: none !important;}

.gradio-container {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
}

.gr-button {
    border-radius: var(--border-radius) !important;
}

.gr-button-primary {
    background-color: var(--primary-color) !important;
}

.gr-input, .gr-textarea {
    border-radius: var(--border-radius) !important;
    border: 1px solid #E5E7EB !important;
}

.gr-box {
    border-radius: var(--border-radius) !important;
    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06) !important;
    background-color: var(--surface-color) !important;
    padding: 1.5rem !important;
}

.advanced-options {
    margin-top: 1rem;
    padding: 1rem;
    border-radius: var(--border-radius);
    background: #F3F4F6;
}

.footer-text {
    text-align: center;
    margin-top: 1rem;
    color: var(--text-color);
    opacity: 0.7;
    font-size: 0.9rem;
}
"""

with gr.Blocks(css=custom_css) as demo:
    with gr.Column(elem_classes="container"):
        gr.Markdown("# Devstral Code Assistant", elem_classes="title")
        gr.Markdown("Powered by Devstral-Small-2505 - Specialized for code generation", elem_classes="subtitle")
        
        with gr.Box():
            prompt = gr.Textbox(
                placeholder="Write a function in Python to implement a binary search tree", 
                label="Your Request",
                lines=5
            )
            
            with gr.Row():
                submit_btn = gr.Button("Generate Code", variant="primary", scale=2)
                clear_btn = gr.Button("Clear", scale=1)
            
            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    with gr.Column():
                        max_tokens = gr.Slider(
                            minimum=64, maximum=2048, value=512, step=64,
                            label="Maximum Output Length"
                        )
                    with gr.Column():
                        temperature = gr.Slider(
                            minimum=0.1, maximum=1.0, value=0.7, step=0.1,
                            label="Temperature (Creativity)"
                        )
                with gr.Row():
                    with gr.Column():
                        top_p = gr.Slider(
                            minimum=0.1, maximum=1.0, value=0.9, step=0.05,
                            label="Top-p (Nucleus Sampling)"
                        )
            
            output = gr.Textbox(
                label="Generated Code", 
                lines=12,
                show_copy_button=True
            )
        
        gr.Markdown(
            "⚡ Optimized for code generation and technical tasks", 
            elem_classes="footer-text"
        )
    
    submit_btn.click(
        generate_text, 
        inputs=[prompt, max_tokens, temperature, top_p], 
        outputs=output
    )
    clear_btn.click(lambda: "", None, prompt)

demo.launch()