Helmet dataset: max_seq_lenght = None by default

Author: mseeger

keys_values/data/base.py

@@ -38,15 +38,14 @@
 class LongContextDataset(Dataset):
     """
     Base class for some datasets we define here.
-
     """
 
     def __init__(
         self,
         data: List[Dict[str, str]],
         tokenizer: Tokenizer,
         prompt_style: Union[str, PromptStyle],
-        max_seq_length: int = -1,
+        max_seq_length: Optional[int] = None,
         transform: Optional[Callable[[Dict[str, str]], Dict[str, str]]] = None,
     ) -> None:
         self.data = data

keys_values/data/helmet.py

@@ -90,7 +90,7 @@ def __init__(
                 validation. The rest is used for training.
             ignore_index: Value used to mask prompt positions in the labels.
             max_seq_length: Sequences longer than this (in tokens) are filtered
-                out. Defaults to no filtering (``100000``).
+                out. Defaults to no filtering.
             seed: Random seed for the train/val split.
             metadata_dir: If given, sequence lengths for every case are stored
                 in a JSON metadata file in this directory so that subsequent
@@ -204,7 +204,7 @@ def _transform(
                 new_seq_lengths.append(seq_length)
             else:
                 seq_length = seq_lengths[idx]
-            if seq_length > self.max_seq_length:
+            if self.max_seq_length is not None and seq_length > self.max_seq_length:
                 continue
             output = instance["output"]
             results.append(
@@ -215,10 +215,7 @@ def _transform(
                 }
             )
         final_seq_lengths = new_seq_lengths if seq_lengths is None else seq_lengths
-        print(
-            f"Kept {len(results)} of {len(dataset)} {split} records "
-            f"(<= {self.max_seq_length} tokens)"
-        )
+        print(f"Kept {len(results)} of {len(dataset)} {split} records")
         return results, final_seq_lengths, needs_store
 
     def _get_seq_lengths(

keys_values/data/longbench_v2.py

@@ -83,7 +83,7 @@ def __init__(
         mask_prompt: bool = True,
         val_split_fraction: float = 0.1,
         ignore_index: int = -100,
-        max_seq_length: Optional[int] = None,
+        max_seq_length: Optional[int] = 100000,
         seed: int = 42,
         repo_id: str = "THUDM/LongBench-v2",
         access_token: Optional[str] = None,
@@ -388,7 +388,7 @@ def smart_lastrec_info(self, tokenizer: HFTokenizer) -> SmartInitialInformation:
 
 def filter_and_transform(
     dataset: Any,
-    max_seq_length: int,
+    max_seq_length: Optional[int],
     tokenizer: Tokenizer,
     seq_lengths: Optional[List[int]],
     head_model: str,
@@ -414,7 +414,10 @@ def filter_and_transform(
     test_results: RawDatasetType = []
     num_used = 0
     num_total = 0
-    print(f"\nProcessing dataset, filtering out records with > {max_seq_length} tokens")
+    if max_seq_length is not None:
+        print(f"\nProcessing dataset, filtering out records with > {max_seq_length} tokens")
+    else:
+        print(f"\nProcessing dataset")
     if seq_lengths is None:
         # Show progress bar: This takes a while
         data_iter = tqdm(dataset)
@@ -445,7 +448,7 @@ def filter_and_transform(
             new_seq_lengths.append(seq_length)
         else:
             seq_length = seq_lengths[idx]
-        if seq_length <= max_seq_length:
+        if max_seq_length is None or seq_length <= max_seq_length:
             num_used += 1
             train_results.append(
                 {
@@ -465,8 +468,8 @@ def filter_and_transform(
                     "num_tokens_instruction": seq_length,
                 }
             )
-    print(f"\nKept {num_used} of {num_total} records: {max_seq_length} tokens or less")
-    if test_set_tag == "rest":
+    print(f"\nKept {num_used} of {num_total} records")
+    if test_set_tag == "rest" and test_results:
         # Sort by increasing length
         test_results = sorted(
             test_results,

keys_values/data/module.py

@@ -53,7 +53,6 @@ class SequenceLengthFilteredDataModule(DataModule):
     use a :class:`EvaluationDataLoader`, which returns batches coupled with
     tasks. Here, we try to form micro batches with sequences of similar length,
     but there is no concept of macro batches.
-
     """
 
     def __init__(
@@ -75,7 +74,7 @@ def __init__(
             ignore_index: The index to use for elements to be ignored in the
                 label.
             max_seq_length: Sequences longer than this number of tokens are
-                filtered out. Defaults to 100000.
+                filtered out.
             seed: The random seed for creating the train/val splits and shuffling
                 the dataset.
             trainloader_longest_first: If set, :meth:`train_dataloader` returns
@@ -94,7 +93,7 @@ def __init__(
         self.mask_prompt = mask_prompt
         self.val_split_fraction = val_split_fraction
         self.ignore_index = ignore_index
-        self.max_seq_length = 100000 if max_seq_length is None else max_seq_length
+        self.max_seq_length = max_seq_length
         self.seed = seed
         self.head_model = None
         self._is_sequence_classification = None
@@ -268,7 +267,7 @@ def setup(self, stage: str = "") -> None:
                 ),
                 tokenizer=self.tokenizer,
                 prompt_style=Default(),
-                max_seq_length=-1,
+                max_seq_length=None,
             )
         else:
             test_kwargs = None

keys_values/data/sequence_classification.py

@@ -56,7 +56,7 @@ def __init__(
         tokenizer: Tokenizer,
         prompt_style: Union[str, PromptStyle],
         class_labels: Iterable[str],
-        max_seq_length: int = -1,
+        max_seq_length: Optional[int] = None,
         transform: Optional[Callable[[Dict[str, str]], Dict[str, str]]] = None,
     ) -> None:
         super().__init__(
@@ -108,8 +108,9 @@ def __getitem__(self, idx: int) -> Dict[str, Any]:
         if self.transform is not None:
             example = self.transform(example)
         prompt = self.prompt_style.apply(prompt=example["instruction"], **example)
+        max_length = -1 if self.max_seq_length is None else self.max_seq_length
         encoded_prompt = self.tokenizer.encode(
-            prompt, bos=False, eos=True, max_length=self.max_seq_length
+            prompt, bos=False, eos=True, max_length=max_length,
         )
         token_counts = {"raw_plus_prompt_template": len(encoded_prompt)}
         raw_count = example.get("num_tokens_instruction")

keys_values/data/sft_dataset.py

@@ -55,7 +55,7 @@ def __init__(
         data: List[Dict[str, str]],
         tokenizer: Tokenizer,
         prompt_style: Union[str, PromptStyle],
-        max_seq_length: int = -1,
+        max_seq_length: Optional[int] = None,
         mask_prompt: bool = True,
         ignore_index: int = -100,
         transform: Optional[Callable[[Dict[str, str]], Dict[str, str]]] = None,
@@ -84,9 +84,10 @@ def __getitem__(self, idx: int) -> Dict[str, Any]:
         if self.transform is not None:
             example = self.transform(example)
         prompt = self.prompt_style.apply(prompt=example["instruction"], **example)
+        max_length = -1 if self.max_seq_length is None else self.max_seq_length
         encoded_prompt = self.tokenizer.encode(
             prompt,
-            max_length=self.max_seq_length,
+            max_length=max_length,
         )
         targets = example["output"]
         if isinstance(targets, list):
@@ -99,15 +100,14 @@ def __getitem__(self, idx: int) -> Dict[str, Any]:
             _targets,
             bos=False,
             eos=True,
-            max_length=self.max_seq_length,
+            max_length=max_length,
         )
         encoded_prompt_and_response = torch.cat(
             (encoded_prompt, encoded_response)
         ).type(torch.int64)
-        msl = self.max_seq_length
-        if 0 < msl < len(encoded_prompt_and_response):
-            encoded_prompt_and_response = encoded_prompt_and_response[:msl]
-            encoded_prompt_and_response[msl - 1] = self.tokenizer.eos_id
+        if 0 < max_length < len(encoded_prompt_and_response):
+            encoded_prompt_and_response = encoded_prompt_and_response[:max_length]
+            encoded_prompt_and_response[max_length - 1] = self.tokenizer.eos_id
 
         # The labels are the full prompt with response, but with the prompt masked out
         labels = encoded_prompt_and_response.clone()