|
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration, pipeline |
|
from datetime import datetime |
|
|
|
|
|
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']) |
|
|
|
|