Update train.py

train.py

@@ -166,8 +166,30 @@ def main():
     model_name = "datalab-to/chandra"
     batch_size = 20
 
+    init_flag = os.environ.get("INIT", "0")
+    is_first_run = init_flag == "0"
+    is_second_run = init_flag == "1"
+
     if not os.path.exists(preprocessed_dataset):
-        run_preprocessing(input_dataset, preprocessed_dataset)
+        print(f"[{'First' if is_first_run else 'Second'} Run] Running preprocessing...")
+        ds_full = datasets.load_dataset(input_dataset, split="train")
+        total_size = len(ds_full)
+        midpoint = total_size // 2
+
+        if is_first_run:
+            ds_to_process = ds_full.select(range(0, midpoint))
+        else:
+            ds_to_process = ds_full.select(range(midpoint, total_size))
+
+        print(
+            f"[{'First' if is_first_run else 'Second'} Run] Saving selected shard to disk..."
+        )
+        ds_to_process.save_to_disk("temp_input_shard")
+
+        run_preprocessing("temp_input_shard", preprocessed_dataset)
+
+        # Clean up temp input shard
+        shutil.rmtree("temp_input_shard")
 
     print("Loading preprocessed dataset...")
     ds = datasets.load_from_disk(preprocessed_dataset)
@@ -220,9 +242,16 @@ def main():
     shards = [cast(Dataset, datasets.load_from_disk(f)) for f in temp_files]
     final_ds = datasets.concatenate_datasets(shards)
 
-    print(f"Final dataset size: {len(final_ds)}")
-    print("Pushing to hub...")
-    final_ds.push_to_hub(output_dataset, create_pr=False)
+    if is_first_run:
+        print("First run: pushing first half to hub...")
+        final_ds.push_to_hub(output_dataset, create_pr=False)
+    else:
+        print("Second run: loading first half and merging...")
+        first_half_ds = datasets.load_dataset(output_dataset, split="train")
+        merged_ds = datasets.concatenate_datasets([first_half_ds, final_ds])
+        print(f"Final merged dataset size: {len(merged_ds)}")
+        print("Pushing full dataset with create_pr=True...")
+        merged_ds.push_to_hub(output_dataset, create_pr=True)
 
     print("Cleaning up temporary files...")
     for f in temp_files: