From 415d1a5ba3b6f25572e0c0b2d1a1a33c0365e358 Mon Sep 17 00:00:00 2001 From: Ben Browning Date: Thu, 25 Jul 2024 15:31:45 -0400 Subject: [PATCH] Incorporate knowledge generation context selection improvements This incorporates fixes from https://github.com/aakankshaduggal/sdg/pull/21 to the way we select the expanded contexts during the knowledge data generation, cleaning up the logic around how we select which other documents to include in the expanded context during edge cases where we have low numbers of unique documents Co-authored-by: abhi1092 Signed-off-by: Ben Browning --- src/instructlab/sdg/datamixing.py | 59 ++++++++++++++++------------- tests/test_datamixing.py | 62 +++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+), 26 deletions(-) create mode 100644 tests/test_datamixing.py diff --git a/src/instructlab/sdg/datamixing.py b/src/instructlab/sdg/datamixing.py index 1e9ab439..44e12741 100644 --- a/src/instructlab/sdg/datamixing.py +++ b/src/instructlab/sdg/datamixing.py @@ -262,41 +262,48 @@ def _add_extra_contexts_to_samples(ds: Dataset, p, num_doc_in_context=4): `keep_context_separate` equal to True. When this finishes, the `context` column is removed from the dataset and all context moved to the user messages. + + This is inspired by the concepts of Retrieval Augmented FineTuning (RAFT) + from https://arxiv.org/abs/2403.10131 """ - all_context = ds["context"] - all_context = [ - " ".join(e.split(" ")[: random.randint(100, 500)]) for e in all_context - ] - ds = ds.add_column("row_idx", range(ds.num_rows)) + all_context = list(set(ds["context"])) def __pick_documents(rec, p): - # Loop until we find enough other documents to add to the context - # for this document. Exit the loop early if we have fewer total - # documents than the number of documents we want in our context - # so that we don't end up looping forever. This handles edge - # cases where the number of generated instructions is very low, - # like in CI or user's testing small sizes. - while True: - selected_docs = random.choices(range(ds.num_rows), k=num_doc_in_context) - if ds.num_rows <= num_doc_in_context: - break - if rec["row_idx"] not in selected_docs: - break - if random.uniform(0, 1) < p: - docs = [ - all_context[idx] for idx in selected_docs[: num_doc_in_context - 1] - ] + [rec["context"]] - # rec['indicator'] ='golden' + answer_document = [rec["context"]] + selected_docs = [e for e in all_context if e != answer_document] + if len(selected_docs) > 0: + if len(selected_docs) < num_doc_in_context: + logger.debug( + f"Number of unique documents is {len(selected_docs)} which is less than {num_doc_in_context}. Using all the documents in the expanded context." + ) + if random.uniform(0, 1) < p: + # golden/answer + distractor documents + docs = ( + random.sample(selected_docs, k=num_doc_in_context) + if len(selected_docs) >= num_doc_in_context + else selected_docs + [answer_document] + ) + else: + # distractor documents + docs = ( + random.sample(selected_docs, k=num_doc_in_context) + if len(selected_docs) >= num_doc_in_context + else selected_docs + ) else: - docs = [all_context[idx] for idx in selected_docs] - # rec['indicator'] = 'distractor' + logger.warning( + "Only 1 unique document found. Disabling expanded context injection, which may lead to poorer knowledge retention results." + ) + docs = [answer_document] random.shuffle(docs) docs = "\n".join(([f"Document:\n{e}\n\n" for idx, e in enumerate(docs)])) - user_idx, user_msg = [ + user_idx_msgs = [ (idx, rec_msg) for idx, rec_msg in enumerate(rec["messages"]) if rec_msg["role"] == "user" - ][0] + ] + assert len(user_idx_msgs) > 0, "No user role found in dataset" + user_idx, user_msg = user_idx_msgs[0] user_inst = user_msg["content"] rec["messages"][user_idx]["content"] = f"{docs}\n\n{user_inst}" rec["messages"] = rec["messages"] diff --git a/tests/test_datamixing.py b/tests/test_datamixing.py new file mode 100644 index 00000000..48101390 --- /dev/null +++ b/tests/test_datamixing.py @@ -0,0 +1,62 @@ +""" +Unit tests for the top-level datamixing module. +""" + +# Third Party +from datasets import Dataset + +# First Party +from instructlab.sdg.datamixing import _add_extra_contexts_to_samples + + +def _fake_context(msg_id): + return { + "context": f"context {msg_id}", + "id": msg_id, + "messages": [{"role": "user", "content": f"user content {msg_id}"}], + "metadata": '{"dataset": []}', + } + + +def test_add_extra_contexts_to_samples_with_one_sample(): + """ + Test _add_extra_contexts_to_samples doesn't error out when + given only one sample + """ + samples = Dataset.from_list([_fake_context("abc123")]) + dataset = _add_extra_contexts_to_samples(samples, p=0.4) + assert len(dataset) == 1 + + +def test_add_extra_contexts_to_samples_with_two_samples(): + """ + Test _add_extra_contexts_to_samples doesn't error out when + given only two samples + """ + samples = Dataset.from_list( + [ + _fake_context("abc123"), + _fake_context("bcd234"), + ] + ) + dataset = _add_extra_contexts_to_samples(samples, p=0.4) + assert len(dataset) == 2 + + +def test_add_extra_contexts_to_samples_with_six_samples(): + """ + Test _add_extra_contexts_to_samples doesn't error out when + given more samples + """ + samples = Dataset.from_list( + [ + _fake_context("s1"), + _fake_context("s2"), + _fake_context("s3"), + _fake_context("s4"), + _fake_context("s5"), + _fake_context("s6"), + ] + ) + dataset = _add_extra_contexts_to_samples(samples, p=0.4) + assert len(dataset) == 6