from transformers import T5Tokenizer, AutoModelForSeq2SeqLM

# Загрузка токенизатора и модели
tokenizer = T5Tokenizer.from_pretrained("D:/MyAI/IGM/googleflan-t5-large/")
model = AutoModelForSeq2SeqLM.from_pretrained("D:/MyAI/IGM/googleflan-t5-large/")

def generate_answer(input_text):
    # Формируем запрос к модели с явным указанием задачи
    prompt = f"{input_text}"
    
    # Подготовка текста для модели
    inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True)
    
    # Генерация ответа
    outputs = model.generate(
        inputs, 
        max_length=400,  # Увеличим максимальную длину вывода
        num_beams=30,     # Используем beam search с 5 путями
        early_stopping=True, 
        repetition_penalty=0.9,  # Добавим штраф за повторения
        no_repeat_ngram_size=15   # Предотвращаем повторение трехграмм
    )
    
    # Декодирование и вывод текста
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return answer

print("Запуск чата с моделью T5. Введите 'exit' для выхода.")

while True:
    input_text = input()
    if input_text.lower() == "exit":
        break
    
    # Вывод запроса
    print(f"Вы: {input_text}")
    
    # Генерация ответа
    answer = generate_answer(input_text)
    
    # Вывод ответа
    print(f"AI: {answer}")