import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Загрузка модели и токенизатора
model_name = "Qwen/Qwen2.5-Coder-32B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Функция для генерации ответа
def generate_response(system_message, chat_history, user_input):
    if not chat_history:
        chat_history.append({"role": "system", "content": system_message})

    # Конкатенируем все сообщения для контекста
    conversation = " ".join([f"{entry['role']}: {entry['content']}" for entry in chat_history])
    conversation += f" user: {user_input}"

    # Токенизация ввода
    inputs = tokenizer(conversation, return_tensors="pt").to(model.device)

    # Генерация ответа
    with torch.no_grad():
        outputs = model.generate(inputs["input_ids"], max_length=512, num_return_sequences=1)

    # Декодирование и получение текста ответа
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Извлечение только нового ответа
    response = generated_text.split("assistant:")[-1].strip()

    # Добавляем пользовательский ввод и ответ модели в историю чата
    chat_history.append({"role": "user", "content": user_input})
    chat_history.append({"role": "assistant", "content": response})

    return chat_history

# Создаем интерфейс
with gr.Blocks() as demo:
    gr.Markdown("## Чат-бот для помощи в кодировании")

    system_message = gr.Textbox(placeholder="Введите системное сообщение...", label="Системное сообщение")
    chatbot = gr.Chatbot(label="Чат с ботом", type='messages')
    user_input = gr.Textbox(placeholder="Введите ваш вопрос или код здесь...", label="Ваш ввод")
    submit_btn = gr.Button("Отправить")

    submit_btn.click(generate_response, inputs=[system_message, chatbot, user_input], outputs=chatbot)

demo.launch()