import gradio as gr
from transformers import pipeline

# Завантаження оптимізованої моделі
def load_model():
    try:
        # Використовуємо компактну модель для економії ресурсів
        generator = pipeline("text-generation", model="EleutherAI/gpt-neo-125M", device_map="auto")
        return generator
    except Exception as e:
        return f"Помилка завантаження моделі: {e}"

# Ініціалізація моделі
generator = load_model()
conversation_history = []

def chat_with_npc(user_input):
    global conversation_history
    if isinstance(generator, str):
        return generator
    try:
        # Додаємо новий запит до історії
        conversation_history.append({"role": "user", "content": user_input})
        
        # Формуємо контекст
        context = "\n".join([f"{msg['role']}: {msg['content']}" for msg in conversation_history])
        
        # Генеруємо відповідь
        response = generator(context, max_length=300, num_return_sequences=1, do_sample=True, temperature=0.7, top_p=0.9)
        bot_reply = response[0]["generated_text"]
        
        # Додаємо відповідь до історії
        conversation_history.append({"role": "assistant", "content": bot_reply})
        return bot_reply
    except Exception as e:
        return f"Помилка генерації відповіді: {e}"

    
# Функція для генерації відповіді
def chat_with_npc(user_input):
    if isinstance(generator, str):  # Якщо модель не завантажилась
        return generator
    try:
        # Налаштування для оптимізації пам'яті
        response = generator(
            user_input,
            max_length=300,  # Довгий запит
            num_return_sequences=1,  # Генеруємо лише один результат
            do_sample=True,
            temperature=0.7,  # Контроль творчості
            top_p=0.9  # Відсікання малоймовірних результатів
        )
        return response[0]["generated_text"]
    except Exception as e:
        return f"Помилка генерації відповіді: {e}"

# Інтерфейс Gradio
interface = gr.Interface(
    fn=chat_with_npc,
    inputs="text",
    outputs="text",
    title="Alt Помічник",
    description="Оптимізований віртуальний помічник для запитів з великою кількістю слів.",
    examples=[
        ["Як можна розвивати духовний потенціал у сучасному світі?"],
        ["Що таке NPC і яка його роль у симуляції життя?"]
    ]
)

# Запуск програми
if __name__ == "__main__":
    interface.launch()