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

[RAG] Fix retrieval offset in RAG's HfIndex and better integration tests (#7372)

Browse Source 
* Fix retrieval offset in RAG's HfIndex

* update slow tests

* style

* fix new test

* style

* add better tests

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
This commit is contained in:
Quentin Lhoest
2020-09-25 16:12:46 +02:00
committed by GitHub
parent 571c7a11c1
commit cf1c88e092
4 changed files with 138 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_retrieval_rag.py
View File

@@ -166,7 +166,7 @@ class RagRetrieverTest(TestCase):
self.assertEqual(len(doc_dicts[0]["id"]), n_docs)
self.assertEqual(doc_dicts[0]["id"][0], "1") # max inner product is reached with second doc
self.assertEqual(doc_dicts[1]["id"][0], "0") # max inner product is reached with first doc
self.assertListEqual(list(doc_ids), [1, 0])
self.assertListEqual(doc_ids.tolist(), [[1], [0]])
def test_legacy_index_retriever_retrieve(self):
n_docs = 1
@@ -181,7 +181,7 @@ class RagRetrieverTest(TestCase):
self.assertEqual(len(doc_dicts[0]["text"]), n_docs)
self.assertEqual(doc_dicts[0]["text"][0], "bar") # max inner product is reached with second doc
self.assertEqual(doc_dicts[1]["text"][0], "foo") # max inner product is reached with first doc
self.assertListEqual(list(doc_ids), [1, 0])
self.assertListEqual(doc_ids.tolist(), [[1], [0]])
@require_torch
def test_hf_index_retriever_call(self):