add booksum and gov_report

configs/pretrain_all_xl_and_sum.yaml

@@ -0,0 +1,63 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
+label_names: ["labels"]
+# eval_on_start: True
+# eval_strategy: "steps"
+# eval_steps: 500
+# save_strategy: "no"
+# # save_steps: 500
+# logging_strategy: "steps"
+# logging_steps: 100
+# use_liger_kernel: true
+# remove_unused_columns: false
+
+# needed to avoid OOM by compute the metrics batch by batch
+# w/o this the trainer stores logits of all sample in memory...
+# batch_eval_metrics: true
+
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 8
+max_val_samples_per_ds: 1000
+# optim: schedule_free_adamw
+
+learning_rate: 0.00002
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 5
+weight_decay: 0.01
+# 
+warmup_steps: 100
+
+dataloader_prefetch_factor: 8
+dataloader_num_workers: 8
+# LoRA
+lora_r: 8
+lora_dropout: 0.05
+target_modules:
+  - down_proj
+  - up_proj
+
+# data
+train_ds_names:
+- fw_qa_xl
+- ctx_qa
+- pwc
+- hotpot_qa
+- squad
+- drop
+- narrativeqa
+- quoref
+- ropes
+- synthetic_convqa
+- booksum
+- gov_report
+
+val_ds_names:
+- fw_qa_xl
+- ctx_qa
+- pwc
+- hotpot_qa
+- squad
+
+load_best_model_at_end: true
+metric_for_best_model: eval_pwc_loss

src/ctx_to_lora/data_utils.py

@@ -132,6 +132,26 @@ DS_KWARGS = {
             split="train",
         ),
     ),
+    "booksum": dict(
+        train=dict(
+            path="kmfoda/booksum",
+            split="train+validation+test",
+        )
+    ),
+    "gov_report": dict(
+        train=dict(
+            path="ccdv/govreport-summarization",
+            split="train",
+        ),
+        validation=dict(
+            path="ccdv/govreport-summarization",
+            split="validation",
+        ),
+        test=dict(
+            path="ccdv/govreport-summarization",
+            split="test",
+        ),
+    ),
     "gsm8k": dict(
         train=dict(
             path="openai/gsm8k",
@@ -574,6 +594,24 @@ def get_preprocessing_fn(ds_name: str) -> Callable[[dict[str, Any]], dict[str, A
                 "response": response,
             }
 
+    elif ds_name == "booksum":
+
+        def f(sample):
+            return {
+                "context": sample["chapter"],
+                "prompt": "Summarize the provided text.",
+                "response": sample["summary_text"],
+            }
+
+    elif ds_name == "gov_report":
+
+        def f(sample):
+            return {
+                "context": sample["report"],
+                "prompt": "Summarize the provided text.",
+                "response": sample["summary"],
+            }
+
     elif ds_name == "openmathintx-2":
 
         def f(sample):