from transformers import MT5ForConditionalGeneration, MT5Tokenizer

# Укажите путь к директории, где находятся model.safetensors, spiece.model и config.json
model_directory = './'

# Загрузка токенизатора из локальной директории
tokenizer = MT5Tokenizer.from_pretrained(pretrained_model_name_or_path=model_directory, sp_model_kwargs={"model_file": f"{model_directory}/spiece.model"})

# Загрузка модели из локальной директории
model = MT5ForConditionalGeneration.from_pretrained(pretrained_model_name_or_path=model_directory, local_files_only=True)

# Пример использования модели

# You are a PocketBook Reader bot. You only talk about PocketBook Reader and 
# you are not allowed to compare PocketBook with other brands like Kindle, Tolino, Kobo, etc.


# Вы бот PocketBook Reader. Вы говорите только о PocketBook Reader, и вам не разрешается сравнивать PocketBook с другими брендами, такими как Kindle, Tolino, Kobo и т. д.
#context = "Контекст, на основе которого нужно ответить на вопрос."
context = "Вы бот PocketBook Reader. Вы говорите только о PocketBook Reader, и вам не разрешается сравнивать PocketBook с другими брендами, такими как Kindle, Tolino, Kobo и т. д."
question = "Как включить PocketBook?"
# question = "Какой вопрос нужно задать?"

# Форматирование входных данных для модели
input_text = f"question: {question} context: {context}"
input_ids = tokenizer(input_text, return_tensors="pt").input_ids

# Генерация ответа
outputs = model.generate(input_ids)
answer = tokenizer.decode(outputs[0], skip_special_tokens=True)

print(answer)