diff --git a/rerankers/utils.py b/rerankers/utils.py index cf6aaaf..d7d4e99 100644 --- a/rerankers/utils.py +++ b/rerankers/utils.py @@ -93,6 +93,11 @@ def prep_docs( "'None' doc_ids detected, reverting to auto-generated integer ids..." ) doc_ids = list(range(len(docs))) + # Write the generated ids back onto the documents (the string-input + # path below already does this); otherwise every returned Document + # keeps doc_id=None and get_score_by_docid/get_result_by_docid break. + for doc, doc_id in zip(docs, doc_ids): + doc.doc_id = doc_id if metadata is not None: if docs[0].metadata is not None: diff --git a/tests/test_results.py b/tests/test_results.py index e0d7deb..be02f55 100644 --- a/tests/test_results.py +++ b/tests/test_results.py @@ -31,3 +31,13 @@ def test_result_validation_error(): with pytest.raises(ValueError) as excinfo: Result(document=Document(doc_id=2, text="Doc 2")) assert "Either score or rank must be provided." in str(excinfo.value) + + +def test_prep_docs_document_input_autogenerates_ids(): + from rerankers.utils import prep_docs + + # Document objects with no explicit doc_id must receive auto-generated + # integer ids so get_score_by_docid / get_result_by_docid can resolve them. + docs = [Document(text="a"), Document(text="b"), Document(text="c")] + out = prep_docs(docs) + assert [d.doc_id for d in out] == [0, 1, 2]