import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, BitsAndBytesConfig
import torch
from threading import Thread

# === Настройка квантизации 4-bit ===
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

model_id = "dphn/dolphin-2.9.1-llama-3-8b"

print("Загрузка токенизатора...")
tokenizer = AutoTokenizer.from_pretrained(model_id)

print("Загрузка модели в 4-bit...")
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=quantization_config,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)

# Настраиваем паддинг (важно для чата)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"

# Стример для потоковой генерации
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

def chat(message, history):
    # Формируем историю в формате messages
    messages = []
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    # Применяем шаблон чата
    input_text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    # Токенизируем
    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

    # Запускаем генерацию в отдельном потоке
    generation_kwargs = dict(
        inputs,
        streamer=streamer,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    # Потоковая отдача ответа
    generated_text = ""
    for new_text in streamer:
        generated_text += new_text
        yield generated_text

    thread.join()

# === Gradio интерфейс ===
with gr.Blocks(title="Dolphin 2.9.1 Llama 3 8B (4-bit)") as demo:
    gr.Markdown(
        """
        # Dolphin 2.9.1 Llama 3 8B (4-bit Q4_K_M)
        Лёгкая версия с квантизацией — работает на 16 ГБ RAM!  
        Нажмите **Enter** или кнопку **Отправить**.
        """
    )
    
    chatbot = gr.Chatbot(height=600)
    
    with gr.Row():
        msg = gr.Textbox(
            label="Ваше сообщение",
            placeholder="Введите сообщение и нажмите Enter или кнопку 'Отправить'...",
            lines=3,
            scale=4
        )
        submit_btn = gr.Button("Отправить", variant="primary", scale=1)

    clear = gr.Button("Очистить чат")

    # Обработка отправки: и по кнопке, и по Enter
    submit_event = submit_btn.click(
        fn=chat,
        inputs=[msg, chatbot],
        outputs=chatbot
    ).then(
        fn=lambda: "",  # Очистка поля ввода
        inputs=None,
        outputs=msg
    )

    msg.submit(
        fn=chat,
        inputs=[msg, chatbot],
        outputs=chatbot
    ).then(
        fn=lambda: "",
        inputs=None,
        outputs=msg
    )

    clear.click(
        fn=lambda: ([], ""),  # Очистка чата и поля ввода
        inputs=None,
        outputs=[chatbot, msg]
    )

# Запуск
if __name__ == "__main__":
    demo.launch(share=False, inbrowser=True)