wip for multipack pretraining

* update datasets to cut down the warnings

* set versions for tokenizers and gradio

* upgrade transformers to latest version

requirements.txt

@@ -2,14 +2,15 @@
 auto-gptq==0.5.1
 packaging
 peft==0.6.0
-transformers==4.35.1
+transformers==4.35.2
+tokenizers==0.15.0
 bitsandbytes>=0.41.1
 accelerate==0.24.1
 deepspeed
 addict
 fire
 PyYAML>=6.0
-datasets>=2.14.0
+datasets>=2.15.0
 flash-attn==2.3.3
 sentencepiece
 wandb
@@ -29,7 +30,7 @@ scikit-learn==1.2.2
 pynvml
 art
 fschat==0.2.29
-gradio
+gradio==3.50.2
 tensorboard
 
 # remote filesystems

src/axolotl/prompt_strategies/completion.py

@@ -1,7 +1,6 @@
 """
 Basic completion text
 """
-import json
 from collections import defaultdict
 from typing import Any, Dict, Generator, Optional, Tuple
 
@@ -65,19 +64,6 @@ class CompletionPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
         return next(iter(self.prompter.build_prompt(instruction, input, response)))
 
 
-class CompletionJSONPromptTokenizationStrategy(CompletionPromptTokenizingStrategy):
-    """
-    Strategy to return the stringified JSON of the entire row as the training data
-    """
-
-    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
-        return (
-            json.dumps(prompt),
-            "",
-            "",
-        )
-
-
 class CompletionPrompter:
     """
     Prompter for completion
@@ -96,7 +82,7 @@ def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
     strat = CompletionPromptTokenizingStrategy(
         CompletionPrompter(),
         tokenizer,
-        True,
+        cfg.train_on_inputs,
         cfg.sequence_len,
         max_length=cfg.sequence_len * 64,
     )
@@ -104,15 +90,3 @@ def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
         strat.field = ds_cfg["field"]
 
     return strat
-
-
-def load_json(tokenizer, cfg):
-    strat = CompletionJSONPromptTokenizationStrategy(
-        CompletionPrompter(),
-        tokenizer,
-        True,
-        cfg.sequence_len,
-        max_length=cfg.sequence_len * 64,
-    )
-
-    return strat

src/axolotl/utils/data.py

@@ -698,6 +698,24 @@ def get_dataset_wrapper(
     return dataset_wrapper, dataset_prompter
 
 
+def encode_packed_pretraining(
+    tokenizer: PreTrainedTokenizerBase, max_tokens: int, examples: List[str]
+):
+    # tokenize all the examples
+    # rows get split with stride (overlap)
+    res = tokenizer(
+        examples,
+        truncation=True,
+        max_length=max_tokens,
+        add_special_tokens=True,
+        return_overflowing_tokens=True,
+        stride=256,
+    )
+    # convert to a dataset.from_list
+    # use a dataloader and multipack batch sampler to pack the data
+    pass
+
+
 def encode_pretraining(
     tokenizer: PreTrainedTokenizerBase, max_tokens: int, examples: List[str]
 ) -> Dict[str, List]:
@@ -813,6 +831,7 @@ def load_pretraining_dataset(path, tokenizer, max_tokens=2048, seed=42):
     dataset = dataset.map(
         encode,
         batched=True,
+        batch_size=10_000,
         input_columns="text",
         # remove all the existing columns after mapping since they end up having
         # a different length than the encoded/tokenized column