import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Загрузка модели и токенизатора
model_name = "deepseek/r1"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def inference(prompt, temperature, top_p, max_length):
    try:
        # Токенизация входного текста
        inputs = tokenizer(prompt, return_tensors="np")

        # Генерация ответа
        outputs = model.generate(
            inputs.input_ids,
            max_length=max_length,
            temperature=temperature,
            top_p=top_p,
            do_sample=True
        )

        # Декодирование ответа
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response

    except Exception as e:
        return f"Ошибка: {str(e)}"

# Создание интерфейса Gradio
with gr.Blocks() as demo:
    gr.Markdown("# DeepSeek-r1 Model")

    with gr.Row():
        input_text = gr.Textbox(label="Входной текст", placeholder="Введите ваш текст здесь...")
        output_text = gr.Textbox(label="Выходной текст", placeholder="Ответ модели появится здесь...")

    with gr.Row():
        max_length = gr.Slider(label="Max Length", minimum=1, maximum=1000, value=500)
        temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, value=1.0)
        top_p = gr.Slider(label="Top P", minimum=0.1, maximum=1.0, value=1.0)

    run_button = gr.Button("Run")

    # Подключение функции к кнопке
    run_button.click(
        fn=inference,
        inputs=[input_text, temperature, top_p, max_length],
        outputs=output_text
    )

# Запуск приложения
if __name__ == "__main__":
    demo.launch()