import gradio as gr
from transformers import pipeline, AutoTokenizer
import torch

# Функция для загрузки текста из файла
def load_text_from_file(file):
    with open(file.name, "r", encoding="utf-8") as f:
        text = f.read()
    return text

# Функция для генерации ответа на основе промта и текста
def generate_response(prompt, text):
    # Объединяем текст и промт
    full_prompt = f"{text}\n\n{prompt}"
    
    # Загружаем модель и токенизатор
    model_name = "ai-forever/rugpt3small_based_on_gpt2"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    # Инициализируем пайплайн для генерации текста
    generator = pipeline(
        'text-generation',
        model=model_name,
        tokenizer=tokenizer,
        device=0 if torch.cuda.is_available() else -1
    )
    
    # Генерация текста
    output = generator(
        full_prompt,
        max_new_tokens=100,  # Увеличим количество токенов для более длинного ответа
        do_sample=True,
        top_p=0.9,
        top_k=50,
        temperature=0.7,
        num_return_sequences=1,
        pad_token_id=tokenizer.eos_token_id,
        repetition_penalty=1.2,
        no_repeat_ngram_size=2,
        truncation=True
    )
    
    # Извлекаем сгенерированный текст
    generated_text = output[0]['generated_text']
    
    # Получаем только новую часть текста (после промпта)
    response = generated_text[len(full_prompt):].strip()
    
    return response

# Функция для обработки ввода и вывода
def process_input(file, prompt):
    # Загружаем текст из файла
    text = load_text_from_file(file)
    
    # Генерация ответа
    response = generate_response(prompt, text)
    
    return response

# Создаем интерфейс Gradio
def create_interface():
    # Описание интерфейса
    description = """
    Загрузите текстовый файл (.txt) и введите промт. Модель GPT на основе загруженного текста сгенерирует ответ.
    """
    
    # Создаем интерфейс
    interface = gr.Interface(
        fn=process_input,  # Функция для обработки ввода
        inputs=[
            gr.File(label="Загрузите текстовый файл (.txt)"),  # Поле для загрузки файла
            gr.Textbox(label="Введите промт", lines=3)  # Поле для ввода промта
        ],
        outputs=gr.Textbox(label="Ответ модели", lines=10),  # Поле для вывода ответа
        title="Анализ текста с помощью GPT",
        description=description
    )
    
    return interface

# Запуск приложения
if __name__ == "__main__":
    # Создаем интерфейс
    interface = create_interface()
    
    # Запускаем Gradio
    interface.launch()