import gradio as gr
from huggingface_hub import InferenceClient

# Создаем клиент для модели
client = InferenceClient(model="mistralai/Mamba-Codestral-7B-v0.1")

def generate_code(prompt):
    response = client.text_generation(prompt, max_new_tokens=512)
    return response

# Создаем интерфейс Gradio
with gr.Blocks() as demo:
    gr.Markdown("## Mamba-Codestral-7B Code Generator")
    with gr.Row():
        prompt_input = gr.Textbox(label="Введите запрос", placeholder="Например: Напишите функцию сортировки на Python")
        submit_button = gr.Button("Сгенерировать код")
    output = gr.Code(label="Результат")
    
    submit_button.click(generate_code, inputs=prompt_input, outputs=output)

# Запуск интерфейса
demo.launch()