update data gen/def

data/build_drop_compact.py

@@ -3,8 +3,6 @@ import gc
 from datasets import Dataset, load_dataset
 from tqdm import tqdm
 
-from ctx_to_lora.data.processing import closed_qa_prompting
-
 if __name__ == "__main__":
     ds_name = "ucinlp/drop"
 
@@ -16,7 +14,7 @@ if __name__ == "__main__":
             ctx = sample["passage"]
             if ctx not in ctx_qa_dict:
                 ctx_qa_dict[ctx] = {"prompts": [], "responses": []}
-            question = closed_qa_prompting(sample["question"])
+            question = sample["question"]
             answer = sample["answers_spans"]["spans"][0]
             ctx_qa_dict[ctx]["prompts"].append(question)
             ctx_qa_dict[ctx]["responses"].append(answer)

data/build_ropes_compact.py

@@ -3,8 +3,6 @@ import gc
 from datasets import Dataset, load_dataset
 from tqdm import tqdm
 
-from ctx_to_lora.data.preprocessing_fn import closed_qa_prompting
-
 if __name__ == "__main__":
     ds_name = "allenai/ropes"
 
@@ -18,7 +16,7 @@ if __name__ == "__main__":
             bg_txt = sample["background"]
             situation_txt = sample["situation"]
             ctx = ctx_template.format(background=bg_txt, situation=situation_txt)
-            q = closed_qa_prompting(sample["question"])
+            q = sample["question"]
             if ctx not in ctx_qa_dict:
                 ctx_qa_dict[ctx] = {"prompts": [], "responses": []}
             ctx_qa_dict[ctx]["prompts"].append(q)

data/build_squad_compact.py

@@ -3,8 +3,6 @@ import gc
 from datasets import Dataset, load_dataset
 from tqdm import tqdm
 
-from ctx_to_lora.data.preprocessing_fn import closed_qa_prompting
-
 if __name__ == "__main__":
     ds_name = "data/raw_datasets/squad"
 
@@ -16,7 +14,7 @@ if __name__ == "__main__":
             ctx = sample["context"]
             if ctx not in ctx_qa_dict:
                 ctx_qa_dict[ctx] = {"prompts": [], "responses": []}
-            question = closed_qa_prompting(sample["question"])
+            question = sample["question"]
             answer = sample["answers"]["text"][0]
             ctx_qa_dict[ctx]["prompts"].append(question)
             ctx_qa_dict[ctx]["responses"].append(answer)

data/generate_negative_natural_questions.py

@@ -1,44 +1,66 @@
 import os
 
-from datasets import Dataset, load_dataset
-from tqdm import tqdm
+from datasets import load_dataset
+
+
+def process_sample(sample):
+    """Process a single sample and return the processed data or None if invalid."""
+    ctx = " ".join(
+        [
+            token
+            for is_html, token in zip(
+                sample["document"]["tokens"]["is_html"],
+                sample["document"]["tokens"]["token"],
+            )
+            if not is_html
+        ]
+    ).strip()
+
+    if not ctx:
+        return {"prompt": None, "context": None, "answer": None}
+
+    answers = []
+    for answer in sample["annotations"]["short_answers"]:
+        answers += answer["text"]
+
+    if not answers:
+        return {"prompt": None, "context": None, "answer": None}
+
+    answer = answers[0].strip()
+    if not answer:
+        return {"prompt": None, "context": None, "answer": None}
+
+    prompt = sample["question"]["text"].capitalize().strip() + " ?"
+    return {"prompt": prompt, "context": ctx, "answer": answer.capitalize()}
+
+
+def shift_contexts(examples):
+    """Shift contexts by one position to create negative examples."""
+    contexts = examples["context"]
+    shifted_contexts = contexts[1:] + contexts[:1]
+    examples["context"] = shifted_contexts
+    return examples
+
 
 if __name__ == "__main__":
     ds = load_dataset("google-research-datasets/natural_questions")
     ds = ds.shuffle(seed=42)
-    for split in ["train", "validation"]:
-        out = []
-        for i, sample in enumerate(tqdm(ds[split])):
-            ctx = " ".join(
-                [
-                    token
-                    for is_html, token in zip(
-                        sample["document"]["tokens"]["is_html"],
-                        sample["document"]["tokens"]["token"],
-                    )
-                    if not is_html
-                ]
-            ).strip()
-            if not ctx:
-                continue
-            answers = []
-            for answer in sample["annotations"]["short_answers"]:
-                answers += answer["text"]
-            if not answers:
-                continue
-            answer = answers[0].strip()
-            if not answer:
-                continue
-            prompt = sample["question"]["text"].capitalize().strip() + " ?"
-            out.append(dict(prompt=prompt, context=ctx, answer=answer.capitalize()))
-            # if i >= 10:
-            #     break
-        shifted_ctxs = [d["context"] for d in out]
-        shifted_ctxs = shifted_ctxs[1:] + shifted_ctxs[:1]
-        for d, shifted_ctx in zip(out, shifted_ctxs):
-            d["context"] = shifted_ctx
-        new_ds = Dataset.from_list(out)
-        print(new_ds)
+
+    for split in ["validation", "train"]:
+        # Process samples in parallel using map
+        processed_ds = ds[split].map(
+            process_sample, remove_columns=ds[split].column_names, num_proc=16
+        )
+
+        # Filter out None values
+        processed_ds = processed_ds.filter(lambda x: x["context"] is not None)
+
+        # Shift contexts to create negative examples
+        processed_ds = processed_ds.map(
+            shift_contexts, batched=True, batch_size=len(processed_ds)
+        )
+
+        print(processed_ds)
         save_dir = "data/raw_datasets/negative_natural_questions"
         os.makedirs(save_dir, exist_ok=True)
-        new_ds.to_json(f"{save_dir}/{split}.jsonl")
+        processed_ds.to_json(f"{save_dir}/{split}.jsonl")

src/ctx_to_lora/data/definitions.py

@@ -537,6 +537,13 @@ DS_KWARGS = {
             split="test",
         ),
     ),
+    "bookqa": dict(
+        train=dict(
+            path="parquet",
+            data_files="data/raw_datasets/fw_qa_v3/min_0_to_20000/booksum_level_0.parquet",
+            split="train",
+        )
+    ),
     "gov_report": dict(
         train=dict(
             path="ccdv/govreport-summarization",
@@ -551,6 +558,13 @@ DS_KWARGS = {
             split="test",
         ),
     ),
+    "gov_report_qa": dict(
+        train=dict(
+            path="parquet",
+            data_files="data/raw_datasets/fw_qa_v3/min_0_to_20000/gov_report_level_0.parquet",
+            split="train",
+        )
+    ),
     # "wikitext-2": dict(
     #     train=dict(
     #         path="EleutherAI/wikitext_document_level",
@@ -684,7 +698,7 @@ DS_KWARGS = {
         )
     ),
     "codeparrot": dict(
-        dict(train=dict(path="transformersbook/codeparrot", split="train"))
+        dict(train=dict(path="transformersbook/codeparrot", split="train[:2000000]"))
     ),
     "programming_books": dict(
         dict(train=dict(path="open-phi/programming_books_llama", split="train"))