remove do_tokenization

Author: Jackmin801

src/genesys/data.py

@@ -43,9 +43,7 @@ class DataLoaderGenesys:
     Each dataset that is pass must have a "train" split and the content must be a list of dict with at least a "problem" and a "ground_truth" key.
     """
 
-    def __init__(
-        self, config: DataConfig, tokenizer: AutoTokenizer, prime_metric: PrimeMetric, do_tokenization: bool = False
-    ):
+    def __init__(self, config: DataConfig, tokenizer: AutoTokenizer, prime_metric: PrimeMetric):
         self.config = config
 
         self.paths = list(config.path.split(","))
@@ -74,7 +72,6 @@ def _add_column(dataset, path):
 
         self.total_samples = min(max_samples, total_samples)
 
-        self.do_tokenization = do_tokenization
         self.tokenizer = tokenizer
 
         self.dataset_lengths = [len(dataset) for dataset in self.datasets]
@@ -114,9 +111,7 @@ def _prepare_batch(self, batch: dict, dataset: str) -> tuple:
                 [{"role": "user", "content": b["prompt"]}, {"role": "assistant", "content": "<think>/n"}] for b in batch
             ]
 
-        batch_inputs = self.tokenizer.apply_chat_template(
-            batch_messages, tokenize=self.do_tokenization, continue_final_message=True
-        )
+        batch_inputs = self.tokenizer.apply_chat_template(batch_messages, tokenize=True, continue_final_message=True)
 
         return batch_inputs, batch
 

src/genesys/generate.py

@@ -76,7 +76,7 @@ def main(config: GenerateConfig):
     tokenizer = AutoTokenizer.from_pretrained(config.name_model)
 
     log("[cyan] Loading dataloader...[/]")
-    dataloader = DataLoaderGenesys(config.data, tokenizer=tokenizer, prime_metric=prime_metric, do_tokenization=True)
+    dataloader = DataLoaderGenesys(config.data, tokenizer=tokenizer, prime_metric=prime_metric)
     machine_info = get_machine_info()
 
     log("[bold green]✨ Setup complete! Starting generation...[/]")