import gradio as gr
from ctransformers import AutoModelForCausalLM
import os

# Глобальная переменная для модели
model = None

def load_model():
    global model
    try:
        print("Начинаем загрузку модели из Hub...")
        
        # Загружаем файл модели из Hub
        model_path = hf_hub_download(
            repo_id="AugustLight/LLight-3.2-3B-Instruct",
            filename="Llight.Q8_0.gguf",
            repo_type="model"
        )
        
        print(f"Модель загружена в: {model_path}")
        
        # Инициализируем модель
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            model_type="llama",
            context_length=2048,
            gpu_layers=0,
            threads=4
        )
        print("Модель успешно инициализирована!")
        return model
        
    except Exception as e:
        print(f"Подробная ошибка при загрузке модели: {str(e)}")
        raise e

def respond(message, history, system_message, max_new_tokens, temperature, top_p):
    try:
        global model
        if model is None:
            model = load_model()
            
        # Формируем контекст из истории
        context = f"{system_message}\n\n"
        for user_msg, assistant_msg in history:
            context += f"User: {user_msg}\nAssistant: {assistant_msg}\n"
        context += f"User: {message}\nAssistant: "

        print(f"Генерируем ответ для контекста длиной {len(context)} символов")
        
        # Генерируем ответ
        response = model(
            context,
            max_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            stop=["User:", "\n\n", "<|endoftext|>"]
        )
        
        print(f"Ответ сгенерирован успешно, длина: {len(response)}")
        return response.strip()
        
    except Exception as e:
        error_msg = f"Произошла ошибка: {str(e)}"
        print(error_msg)
        return error_msg

# Создаем интерфейс
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(
            value="Ты дружелюбный и полезный ассистент. Ты всегда отвечаешь кратко и по делу.", 
            label="System message"
        ),
        gr.Slider(
            minimum=1, 
            maximum=2048, 
            value=512, 
            step=1, 
            label="Max new tokens"
        ),
        gr.Slider(
            minimum=0.1, 
            maximum=2.0, 
            value=0.7, 
            step=0.1, 
            label="Temperature"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)"
        ),
    ],
    title="GGUF Chat Model",
    description="Чат с GGUF моделью (Llight.Q8_0.gguf)",
    examples=[
        ["Привет! Как дела?"],
        ["Расскажи мне о себе"],
        ["Что ты умеешь делать?"]
    ],
    cache_examples=False
)

# Запускаем приложение
if __name__ == "__main__":
    # Пробуем загрузить модель при запуске
    try:
        print("Инициализация приложения...")
        model = load_model()
        print("Модель загружена успешно при старте")
    except Exception as e:
        print(f"Ошибка при инициализации: {str(e)}")
    
    demo.launch()