Add ICLMultiChoiceTaskDataset.max_sequence_length property

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- Added `ICLMultiChoiceTaskDataset.max_sequence_length` property.
+
 ## [v0.1.0](https://github.com/allenai/OLMo-in-loop-evals/releases/tag/v0.1.0) - 2024-10-28
 
 ### Added

src/olmo_eval/tasks.py

@@ -57,6 +57,7 @@ def __init__(
 
         # prep examples
         self.prep_examples()
+        self._max_sequence_length: Optional[int] = None
 
     def __getitem__(self, index):
         return self.samples[index]
@@ -147,6 +148,16 @@ def pad_tokens_until_max(self, tokens, max_len=2048):
 
             return tokens
 
+    @property
+    def max_sequence_length(self) -> int:
+        if self._max_sequence_length is None:
+            max_seq_len = 0
+            for sample in self.samples:
+                if len(sample["query"]) > max_seq_len:
+                    max_seq_len = len(sample["query"])
+            self._max_sequence_length = max_seq_len
+        return self._max_sequence_length
+
     def collate_fn(self, data):
         # pad to max length
         # 'ctx', 'continuation', 'query' can all have variable length
@@ -1391,6 +1402,8 @@ def __init__(
         # prep examples
         self.prep_examples()
 
+        self._max_sequence_length: Optional[int] = None
+
     def prep_examples(self):
         current_doc_id_offset = 0
         max_doc_id = 0

src/test/tasks_test.py

@@ -13,6 +13,7 @@ def test_build_task(label: str):
     )
     task = build_task(label, tokenizer)
     assert len(task) >= 2
+    assert task.max_sequence_length > 0
     instance1, instance2 = task[0], task[1]
     batch = task.collate_fn([instance1, instance2])
     assert isinstance(batch, dict)