import gradio as gr
from transformers import pipeline

# Загрузка модели для генерации текста
text_generator = pipeline("text-generation", model="WhiteRabbitNeo/WhiteRabbitNeo-13B-v1")

# Функция для генерации текста через модель Transformers
def generate_text(prompt):
    # Генерация текста на основе введенного пользователем промпта
    generated_text = ""
    while not generated_text.strip():  # Проверяем, что генерируемый текст не пустой
        generated_text = text_generator(prompt, max_length=50, do_sample=True)[0]['generated_text']
    return generated_text

# Создание интерфейса Gradio
iface = gr.Interface(
    fn=generate_text,
    inputs="text",
    outputs="text",
    title="Chat with Transformers",
    description="Start a conversation with the model. Enter your message and see the response.",
    allow_flagging=False  # Отключение возможности пользователей помечать результаты
)

# Запуск интерфейса
iface.launch(share=True)  # Разрешение на публикацию интерфейса в Интернете