SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Add dataset_revision argument to RagConfig (#29610)

Browse Source 
* add arg

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar
2024-03-14 16:48:11 +01:00
committed by GitHub
parent 956f44f11a
commit 2cc3cc835f
4 changed files with 34 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tests/models/rag/test_modeling_tf_rag.py
View File

@@ -590,6 +590,7 @@ class TFRagModelIntegrationTests(unittest.TestCase):
use_dummy_dataset=True,
retrieval_vector_size=768,
retrieval_batch_size=8,
dataset_revision="b24a417",
)
@slow
@@ -794,7 +795,9 @@ class TFRagModelIntegrationTests(unittest.TestCase):
@slow
def test_rag_token_greedy_search(self):
tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True)
retriever = RagRetriever.from_pretrained(
"facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True, dataset_revision="b24a417"
)
rag_token = TFRagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever)
# check first two questions
@@ -828,7 +831,9 @@ class TFRagModelIntegrationTests(unittest.TestCase):
def test_rag_token_generate_batch(self):
# NOTE: gold labels comes from num_beam=4, so this is effectively beam-search test
tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True)
retriever = RagRetriever.from_pretrained(
"facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True, dataset_revision="b24a417"
)
rag_token = TFRagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever)
input_dict = tokenizer(
@@ -871,7 +876,10 @@ class TFRagModelIntegrationTests(unittest.TestCase):
def test_rag_sequence_generate_batch(self):
tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
retriever = RagRetriever.from_pretrained(
"facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True
"facebook/rag-sequence-nq",
index_name="exact",
use_dummy_dataset=True,
dataset_revision="b24a417",
)
rag_sequence = TFRagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", retriever=retriever)
@@ -908,7 +916,7 @@ class TFRagModelIntegrationTests(unittest.TestCase):
def test_rag_sequence_generate_batch_from_context_input_ids(self):
tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
retriever = RagRetriever.from_pretrained(
"facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True
"facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True, dataset_revision="b24a417"
)
rag_sequence = TFRagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", retriever=retriever)
input_dict = tokenizer(
@@ -976,6 +984,7 @@ class TFRagModelSaveLoadTests(unittest.TestCase):
use_dummy_dataset=True,
retrieval_vector_size=768,
retrieval_batch_size=8,
dataset_revision="b24a417",
)
@slow