from transformers import pipeline
import gradio as gr

class TextProcessor:
    """
    Класс для обработки текста с использованием моделей GPT-Neo и T5.
    """
    def __init__(self, api_key=None):
        """
        Инициализирует объект TextProcessor.
        """
        self.api_key = api_key  # API-ключ, если он нужен для авторизации
        self.model_gpt = pipeline("text-generation", model="EleutherAI/gpt-neo-125M")
        self.model_t5 = pipeline("text2text-generation", model="t5-base")

    def process_text(self, step, text):
        """
        Обрабатывает текст с использованием выбранной модели.
        """
        if step == 1:
            # Пример генерации текста с помощью GPT-Neo
            gpt_result = self.model_gpt(text, max_length=100)
            return gpt_result
        elif step == 2:
            # Пример генерации текста с помощью T5 (например, для резюмирования)
            t5_result = self.model_t5(text)
            return t5_result
        else:
            return "Unknown step"

# Создание интерфейса Gradio
processor = TextProcessor()

def process_step(step, text):
    """
    Функция для обработки шага в Gradio.
    """
    result = processor.process_text(step, text)
    return result

iface = gr.Interface(
    fn=process_step,
    inputs=[
        gr.Radio(choices=[1, 2], label="Шаг"),  # Изменили  количество    вариантов
        gr.Textbox(lines=3, label="Текст"),
    ],
    outputs="text",
    title="Обработка Текста",
    description="Выберите шаг и введите текст.",
)
iface.launch()