from transformers import MT5Tokenizer, MT5ForConditionalGeneration, pipeline
from datetime import datetime

# Путь для сохранения модели
save_directory = "./"

# Загрузка сохраненной модели
loaded_tokenizer = MT5Tokenizer.from_pretrained(save_directory, legacy=False)
loaded_model = MT5ForConditionalGeneration.from_pretrained(save_directory)
print("Модели загрузились в :", datetime.now())
# exit(0)
# Создание пайплайна для задачи вопрос-ответ
qa_pipeline = pipeline("text2text-generation", model=loaded_model, tokenizer=loaded_tokenizer)

# Пример использования пайплайна для заданий на тему ридеров и электронных книжек
context = "Электронные ридеры это устройства для чтения электронных книг. Они обычно имеют экран на базе электронных чернил и позволяют хранить большое количество книг в одном устройстве."
context = "Electronic readers are devices for reading electronic books. They usually have an electronic ink screen and allow you to store a large number of books in one device."
question = "Какие преимущества электронных ридеров?"
question = "What are the advantages of e-readers?"
# Задание вопроса
# result = qa_pipeline(f"question: {question} context: {context}")


input_text = f"question: {question} context: {context}"
result = qa_pipeline(input_text, max_length=50, num_beams=5, early_stopping=True)


print("Ответ:", result[0]['generated_text'])