from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration, pipeline
from datetime import datetime

# Загрузка модели и токенизатора RAG
model_name = "facebook/rag-token-base"
tokenizer = RagTokenizer.from_pretrained(model_name)
retriever = RagRetriever.from_pretrained(model_name, index_name="exact", passages_path="datasets", trust_remote_code=True)
model = RagSequenceForGeneration.from_pretrained(model_name, retriever=retriever)

print("Модели загрузились в:", datetime.now())

# Создание пайплайна для задачи вопрос-ответ
qa_pipeline = pipeline("rag-sequence", model=model, tokenizer=tokenizer)

# Пример использования пайплайна для заданий на тему ридеров и электронных книжек
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 = "What are the advantages of e-readers?"

# Задание вопроса
input_text = f"{question} {context}"

# Получение ответа
result = qa_pipeline(input_text, max_length=50, num_beams=5, early_stopping=True)

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