def test_legacy_hf_index_retriever_save_and_from_pretrained(self):
retriever = self.get_dummy_legacy_index_retriever()
with tempfile.TemporaryDirectory() as tmp_dirname:
retriever.save_pretrained(tmp_dirname)
retriever = RagRetriever.from_pretrained(tmp_dirname)
self.assertIsInstance(retriever, RagRetriever)
hidden_states = np.array([
np.ones(self.retrieval_vector_size),
-np.ones(self.retrieval_vector_size)
],
dtype=np.float32)
out = retriever.retrieve(hidden_states, n_docs=1)
self.assertTrue(out is not None)
def test_canonical_hf_index_retriever_save_and_from_pretrained(self):
retriever = self.get_dummy_canonical_hf_index_retriever()
with tempfile.TemporaryDirectory() as tmp_dirname:
with patch("transformers.models.rag.retrieval_rag.load_dataset"
) as mock_load_dataset:
mock_load_dataset.return_value = self.get_dummy_dataset()
retriever.save_pretrained(tmp_dirname)
retriever = RagRetriever.from_pretrained(tmp_dirname)
self.assertIsInstance(retriever, RagRetriever)
hidden_states = np.array([
np.ones(self.retrieval_vector_size),
-np.ones(self.retrieval_vector_size)
],
dtype=np.float32)
out = retriever.retrieve(hidden_states, n_docs=1)
self.assertTrue(out is not None)