contrain instruct datasets

scripts/prepare_contrain_datasets.py

@@ -1,5 +1,6 @@
 from functools import partial
 
+from litgpt.tokenizer import Tokenizer
 from litdata import optimize, TokensLoader, StreamingDataset
 from transformers import AutoTokenizer
 
@@ -16,7 +17,11 @@ chunk_size = block_size * 4000
 output_dir = f'../contrain-data-{i}-{block_size}-{chunk_size}'
 
 outputs = optimize(
-    fn=partial(tokenize_chat_fn, tokenizer=AutoTokenizer.from_pretrained('..', trust_remote_code=True, use_fast=True)),
+    fn=partial(
+        tokenize_chat_fn,
+        hf_tokenizer=AutoTokenizer.from_pretrained('..', trust_remote_code=True, use_fast=True),
+        tokenizer=Tokenizer('..'),
+    ),
     inputs=contrain_datasets,
     output_dir=output_dir,
     chunk_size=chunk_size, # Number of tokens to store by chunks. This is roughly 64MB of tokens per chunk.

scripts/utils.py

@@ -123,9 +123,12 @@ def tokenize_text_fn(dataset_config: list, tokenizer: Tokenizer, min_len: Option
             yield text_ids
 
 
-def tokenize_chat_fn(dataset_config: list, tokenizer: AutoTokenizer, min_len: Optional[int]=None, max_len: Optional[int]=None) -> Iterator[torch.Tensor]:
+def tokenize_chat_fn(dataset_config: list, hf_tokenizer: AutoTokenizer, tokenizer: Tokenizer, min_len: Optional[int]=None, max_len: Optional[int]=None) -> Iterator[torch.Tensor]:
     for messages in batch_chat_iterator(dataset_config):
-        text_ids: torch.Tensor = tokenizer.apply_chat_template(messages, tokenize=True, return_tensors="pt")
+        # text_ids: torch.Tensor = tokenizer.apply_chat_template(messages, tokenize=True, return_tensors='pt')
+        # text_ids = text_ids.to(torch.int)
+        text: str = hf_tokenizer.apply_chat_template(messages, tokenize=False)
+        text_ids: torch.Tensor = tokenizer.encode(text, bos=False, eos=False)
 
         if min_len is None and max_len is None:
             yield text_ids