add sample_per_file

Author: samsja

src/genesys/generate.py

@@ -1,4 +1,5 @@
 import itertools
+from pydantic import model_validator
 from pydantic_config import BaseConfig, parse_argv
 import sglang as sgl
 from datasets import load_dataset
@@ -12,13 +13,21 @@
 
 class Config(BaseConfig):
     name_model: str = "Qwen/QwQ-32B-Preview"
-    out_file_name: str = "out.jsonl"
     num_responses_per_question: int = 1
     num_gpus: int = 8
     temperature: float = 0.9
-    batch_size: int = 10000
+    batch_size: int = 10_000
     max_samples: int | None = None
-    gcp_bucket: str | None = None
+    gcp_bucket: str | None = None  # optional, if provided, will save the each file with sample_per_file  to GCP
+    sample_per_file: int = 10_000  # how much sample each file contains
+
+    @model_validator(mode="after")
+    def check_batch_size(self):
+        if self.sample_per_file < self.batch_size:
+            raise ValueError("sample_per_file must be greater than or equal to batch_size")
+        if self.max_samples is not None and self.max_samples < self.sample_per_file:
+            raise ValueError("max_samples must be greater than or equal to sample_per_file")
+        return self
 
 
 def main(config: Config):
@@ -33,10 +42,11 @@ def main(config: Config):
 
     sampling_params = dict(temperature=config.temperature, max_new_tokens=8192, stop=["<|eot_id|>"])
 
-    open(config.out_file_name, "w").close()
-
     max_samples = config.max_samples if config.max_samples is not None else len(math_dataset)
 
+    all_results = []
+    file_counter = 0
+
     for i in tqdm(range(0, min(max_samples, len(math_dataset)), config.batch_size), desc="Generating data"):
         batch = math_dataset[i : min(i + config.batch_size, len(math_dataset))]
         batch_ids = list(
@@ -54,7 +64,6 @@ def main(config: Config):
         batch_inputs = tokenizer.apply_chat_template(batch_messages, tokenize=False, add_generation_prompt=True)
         batch_output = llm.generate(batch_inputs, sampling_params)
 
-        all_results = []
         for j, out in enumerate(batch_output):
             result = dict()
             result["prompt"] = batch_messages[j][1]["content"]
@@ -64,7 +73,11 @@ def main(config: Config):
 
             all_results.append(result)
 
-        save_batch_results(all_results, config.out_file_name, gcp_bucket)
+        if len(all_results) >= config.sample_per_file:
+            file_name = f"out_{file_counter}.jsonl"
+            save_batch_results(all_results, file_name, gcp_bucket)
+            all_results = []
+            file_counter += 1
 
 
 if __name__ == "__main__":