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:
            # Разбиваем текст на части, например, по 500 символов
            text_parts = [text[i:i + 500] for i in range(0, len(text), 500)]
            gpt_results = []
            for part in text_parts:
                gpt_result = self.model_gpt(part, max_length=150)
                gpt_results.append(gpt_result[0]['generated_text'])
            return "\n".join(gpt_results)  # Объединяем ответы для всех частей
        elif step == 2:
            # ... (код обработки с помощью T5)
        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()