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

model_id = "aeonium/Aeonium-v1.1-Chat-4B"

# Загрузка токенизатора и модели
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

@spaces.GPU
def predict(message, history, max_tokens, temperature, top_p):
    # Формирование чата из истории и нового сообщения
    chat = [{"role": "user" if i % 2 == 0 else "assistant", "content": m} 
            for i, (m, _) in enumerate(history)] + [{"role": "user", "content": message}]
    
    # Применение шаблона чата
    prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
    
    # Кодирование входных данных
    inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
    
    # Генерация ответа
    outputs = model.generate(
        input_ids=inputs, 
        max_new_tokens=max_tokens,
        do_sample=True,
        temperature=temperature,
        top_p=top_p,
    )
    
    # Декодирование результата
    response = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
    
    return response.strip().replace("assistant", "", 1)

# Настройка интерфейса Gradio
iface = gr.ChatInterface(
    predict,
    chatbot=gr.Chatbot(height=600),
    textbox=gr.Textbox(placeholder="Введите ваше сообщение здесь...", container=False, scale=7),
    title="Чат с Aeonium v1.1",
    description="Это чат-интерфейс для модели Aeonium v1.1 Chat 4B. Задавайте вопросы и получайте ответы!",
    theme="soft",
    retry_btn="Повторить",
    undo_btn="Отменить последнее",
    clear_btn="Очистить",
    additional_inputs=[
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Максимальное количество новых токенов"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Температура"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
    ],
)

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