From f0e276077f0335c88524b8b2bd999224e59272bc Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 1 May 2025 06:16:35 -0700
Subject: [PATCH] Update synthetic.py

---
 unsloth/dataprep/synthetic.py | 45 ++++++++++++++++++++---------------
 1 file changed, 26 insertions(+), 19 deletions(-)

diff --git a/unsloth/dataprep/synthetic.py b/unsloth/dataprep/synthetic.py
index 2f4a85f143..22065ebc39 100644
--- a/unsloth/dataprep/synthetic.py
+++ b/unsloth/dataprep/synthetic.py
@@ -24,7 +24,7 @@ import gc
 import time
 from unsloth_zoo.vllm_utils import load_vllm
 from transformers import AutoConfig, AutoTokenizer
-import signal
+import numpy as np
 
 from .synthetic_configs import (
     synthetic_qa_config,
@@ -190,34 +190,39 @@ class SyntheticDataKit:
     def __exit__(self, *exc): self.cleanup()
     def __del__(self): self.cleanup()
 
-    def truncate(self, filename = None):
-        # Truncates by summary and max generation
+    def chunk_data(self, filename = None):
+        # Chunks data by max tokens and generation length
         assert(filename is not None)
         assert(os.path.exists(filename))
         assert(hasattr(self, "tokenizer"))
+        if not hasattr(self, "max_seq_length"):
+            raise RuntimeError("Please use SynthetidDataKit.from_pretrained(...) first!")
+        if not hasattr(self, "overlap") or not hasattr(self, "max_generation_tokens"):
+            raise RuntimeError("Please use prepare_qa_generation first!")
 
         with open(filename, "r") as f: text = f.read()
 
         max_tokens = self.max_seq_length - self.max_generation_tokens*2 - 2
-        input_ids = self.tokenizer(text).input_ids
-        length = len(text)
-        original_length = len(text)
-        original_n_tokens = len(input_ids)
+        input_ids = self.tokenizer(text, add_special_tokens = False).input_ids
 
-        if len(input_ids) > max_tokens:
-            # Will fix later, but for now we simply naively truncate by ratios
-            length = original_length
-            while True:
-                input_ids = self.tokenizer(text[:length]).input_ids
-                if len(input_ids) < max_tokens or length == 0: break
-                length = length * (max_tokens/len(input_ids))
-                length = max(int(length), 0)
-            pass
-            print(f"Unsloth: Will truncate your data which has {original_n_tokens} tokens to {len(input_ids)} tokens.")
+        # Get left and right boundaries
+        length = len(input_ids)
+        n_chunks = int(np.ceil(length / (max_tokens - overlap)))
+        boundaries = np.ceil(np.linspace(0, length - overlap, n_chunks)).astype(int)
+        boundaries = np.stack((boundaries[:-1], (boundaries + overlap)[1:])).T
+        boundaries = np.minimum(boundaries, length).tolist()
 
-            with open(filename, "w") as f: f.write(text[:length])
+        # Get extension of filename like .txt
+        filename, extension = os.path.splitext(filename)
+
+        all_filenames = []
+        for i, (left, right) in enumerate(boundaries):
+            chunked_text = self.tokenizer.decode(input_ids[left : right])
+            new_filename = os.path.join(filename + f"_{i}", extension)
+            all_filenames.append(new_filename)
+            with open(new_filename, "w") as f: f.write(chunked_text)
         pass
-        return filename, length
+        return all_filenames
     pass
 
     def prepare_qa_generation(
@@ -258,5 +263,7 @@ class SyntheticDataKit:
             .replace("{cleanup_temperature}", str(cleanup_temperature))
 
         with open("synthetic_data_kit_config.yaml", "w") as f: f.write(config)
+
+        self.overlap = overlap
     pass
 pass