import gradio as gr
from gradio_client import Client

# Создаем клиент для взаимодействия с API
client = Client("llamameta/Pixtral-Large-Instruct-2411")

# Функция для обработки сообщений чат-бота
def chatbot_response(message, chat_history, system_message, max_tokens, temperature, top_p):
    # Добавляем сообщение пользователя в историю
    chat_history.append((message, ""))

    # Генерируем ответ от чат-бота
    result = client.predict(
        message=message,
        system_message=system_message,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        api_name="/chat"
    )

    # Добавляем ответ чат-бота в историю
    chat_history[-1] = (message, result)

    return chat_history, chat_history

# Определяем интерфейс Gradio
iface = gr.Interface(
    fn=chatbot_response,
    inputs=[
        gr.inputs.Textbox(lines=2, placeholder="Введите ваше сообщение здесь...", label="Сообщение"),
        gr.inputs.State(initial_value=[]),  # Состояние для хранения истории сообщений
        gr.inputs.Textbox(lines=2, default="You are a friendly Chatbot created by balianone.com", label="Системное сообщение"),
        gr.inputs.Slider(minimum=100, maximum=18000, default=18000, label="Максимальное количество новых токенов"),
        gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.7, label="Температура"),
        gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.95, label="Top-p (нуклеарная сэмплирование)")
    ],
    outputs=[
        gr.outputs.Chatbot(label="Чат"),
        gr.outputs.State()  # Состояние для хранения истории сообщений
    ],
    title="Chatbot с использованием Pixtral-Large-Instruct",
    description="Простой чат-бот, использующий модель Pixtral-Large-Instruct для генерации ответов."
)

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