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

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,
)

def predict(message, history):
    # Формирование чата из истории и нового сообщения
    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=256,  # Увеличено для более длинных ответов
        do_sample=True,
        temperature=0.7,
        top_p=0.95,
    )
    
    # Декодирование результата
    response = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
    
    return response.strip()

# Настройка интерфейса 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",
    examples=[
        "Привет! Как дела?",
        "Расскажи мне о квантовой физике",
        "Какие интересные книги ты можешь порекомендовать?",
    ],
    cache_examples=True,
    retry_btn="Повторить",
    undo_btn="Отменить последнее",
    clear_btn="Очистить",
)

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