34 lines
1.5 KiB
Python
Executable File
34 lines
1.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# Option 1 - all in one forward pass
|
|
|
|
tokenizer = RagTokenForGenerationizer.from_pretrained("facebook/rag")
|
|
retriever = RagRetriever.from_pretrained("facebook/rag")
|
|
rag = RagForSequenceGeneration.from_pretrained("facebook/rage", retriever=retriever)
|
|
|
|
# IMPORTANT these input_ids are encoded by the question encoder tokenizer!!!
|
|
question_encoder_input_ids = tokenizer("How many people live in Paris?", return_pt="pt").input_ids
|
|
|
|
# 1 step
|
|
answer_ids = rag.generate(input_ids=question_encoder_input_ids)
|
|
|
|
|
|
# Option 2 - step by step manually
|
|
tokenizer = RagTokenForGenerationizer.from_pretrained("facebook/rag")
|
|
retriever = RagRetriever.from_pretrained("facebook/rag")
|
|
rag = RagForSequenceGeneration.from_pretrained("facebook/rag") # Note now there is no retriever used as an input
|
|
|
|
# IMPORTANT these input_ids are encoded by the question encoder tokenizer!!!
|
|
question_encoder_input_ids = tokenizer("How many people live in Paris?", return_pt="pt").input_ids
|
|
|
|
# 3 steps
|
|
question_encoder_hidden_states = rag.question_encoder(question_encoder_input_ids)
|
|
|
|
# retriever takes care of:
|
|
# - decoding question encoder input ids to encoder input str
|
|
# - retrieving the document encodings
|
|
# - decoding document encodings
|
|
# - post processing: concat doc string and encoder input str
|
|
# - encoding concat to generator context input ids
|
|
context_input_ids, document_scores = rag.retriever(question_encoder_input_ids, question_encoder_hidden_states)
|
|
answer_ids = rag.generate(context_input_ids=context_input_ids, document_scores=document_scores)
|