add len info eval + triviaqa_retrieved

src/ctx_to_lora/data_utils.py

@@ -2,15 +2,13 @@ import logging
 import hashlib
 import json
 import random
-from functools import partial
 from os import path
-from glob import glob
 from typing import Any, Callable, Iterator, Optional
 
 
 import datasets
 import numpy as np
-from datasets import load_dataset, IterableDataset
+from datasets import load_dataset
 from transformers import PreTrainedTokenizerBase
 
 from ctx_to_lora.training_utils import TRAINING_TASK
@@ -329,7 +327,7 @@ def add_negative_prompt_fn(samples):
         prompts.append(prompt)
         responses.append(response)
 
-    logger.debug(f"Adding negative prompt...")
+    logger.debug("Adding negative prompt...")
     logger.debug(f"# unique contexts: {len(unique_contexts)}")
 
     # remove one last sample if the number of samples is odd
@@ -580,7 +578,8 @@ def construct_and_tokenize_ctx_qa(
     )
     if split == "train":
         # drop?
-        # currently the above function would raise if the chat is too long
+        # base model cant handle these samples naturally
+        # should skip
         ...
     else:
         tokenized_ds = tokenized_ds.map(
@@ -594,6 +593,13 @@ def construct_and_tokenize_ctx_qa(
             batch_size=100_000,
             num_proc=num_proc,
         )
+        tokenized_ds = tokenized_ds.map(
+            add_length_info,
+            fn_kwargs={"columns": ["input_ids"]},
+            batched=True,
+            batch_size=100_000,
+            num_proc=num_proc,
+        )
 
     # for use_kl_loss, we need "chat_ids" and "chat_attn_mask"
     if use_kl_loss:
@@ -645,6 +651,13 @@ def construct_and_tokenize_ctx_qa(
                 batch_size=100_000,
                 num_proc=num_proc,
             )
+            tokenized_ds = tokenized_ds.map(
+                add_length_info,
+                fn_kwargs={"columns": ["ctx_ids"]},
+                batched=True,
+                batch_size=100_000,
+                num_proc=num_proc,
+            )
     tokenized_ds = tokenized_ds.remove_columns(
         ["messages", "chat", "context", "prompt", "response"]
     )
@@ -885,6 +898,12 @@ def tokenize_chat_messages(
     return conversation_ids
 
 
+def add_length_info(samples: dict[str, any], columns: list[str]) -> dict[str, any]:
+    for col in columns:
+        samples[f"{col}_len"] = [len(t) for t in samples[col]]
+    return samples
+
+
 def truncate_middle_if_too_long(
     samples: dict[str, any],
     max_length: int,

src/ctx_to_lora/eval.py

@@ -259,7 +259,13 @@ def decode_test_result(test_dataset, test_result, tokenizer, ctx_tokenizer):
             d["context"] = ctx_tokenizer.decode(
                 sample["ctx_ids"], skip_special_tokens=True
             )
+        for k in sample:
+            if k.endswith("_len"):
+                d[k] = sample[k].item()
         out.append(d)
+    # sort samples by length if possible
+    len_key = "ctx_ids_len" if "ctx_ids_len" in sample else "input_ids_len"
+    sorted(out, key=lambda x: x[len_key])
 
     return out
 
@@ -297,6 +303,40 @@ def eval_generation(
             for k, v in qa_f1_metric.items():
                 eval_result.metrics[f"{split_name}_{k}"] = v
 
+        # Group by input length and compute metrics for each group
+        length_bins = [(0, 511), (512, 1023), (1024, 2047), (2048, 4095), (4096, 8192)]
+        # Ensure all keys for length metrics are present, even if a bin is empty
+        for low, high in length_bins:
+            eval_result.metrics[f"{split_name}_rouge1.f1_len_{low}-{high}"] = "None"
+            eval_result.metrics[f"{split_name}_rougeL.f1_len_{low}-{high}"] = "None"
+            if ds_name in CLOSED_QA_DATASETS:
+                eval_result.metrics[f"{split_name}_qa_f1_len_{low}-{high}"] = "None"
+
+        grouped_texts = defaultdict(lambda: {"generated": [], "label": [], "count": 0})
+        for txt in decoded_txts:
+            len_key = "ctx_ids_len" if "ctx_ids_len" in txt else "input_ids_len"
+            input_len = txt[len_key]
+            for low, high in length_bins:
+                if low <= input_len <= high:
+                    group_key = f"{low}-{high}"
+                    grouped_texts[group_key]["generated"].append(txt["generated"])
+                    grouped_texts[group_key]["label"].append(txt["label"])
+                    grouped_texts[group_key]["count"] += 1
+                    break
+
+        for group_key, data in grouped_texts.items():
+            if data["count"] > 0:
+                group_rouge_metrics = compute_rouge(data["generated"], data["label"])
+                for k, v in group_rouge_metrics.items():
+                    eval_result.metrics[f"{split_name}_{k}_len_{group_key}"] = v
+
+                if ds_name in CLOSED_QA_DATASETS:
+                    group_qa_f1_metric = compute_qa_f1_score(
+                        data["generated"], data["label"]
+                    )
+                    for k, v in group_qa_f1_metric.items():
+                        eval_result.metrics[f"{split_name}_{k}_len_{group_key}"] = v
+
         save_generated_text(
             decoded_txts,
             split=split_name,

src/ctx_to_lora/self_generate_data.py

@@ -0,0 +1,5 @@
+# when generating dont forget to remove dup contexts
+# merge QAs from dup contexts to construct
+# a doc with QAs at the end
+# maybe mix with some that doesn;t have QAs
+# TODO: implement