def main():
    """Main function to orchestrate data partitioning and processing."""
    args = get_args()
    print("Arguments received:")
    for arg, value in sorted(vars(args).items()):
        print(f"  {arg}: {value}")

    # Setup NLTK path if splitting sentences
    if args.split_sentences:
        if nltk_available:
             # Ensure NLTK data path is configured (e.g., via environment variable NLTK_DATA)
             nltk_data_path = os.environ.get("NLTK_DATA")
             if nltk_data_path:
                 print(f"Using NLTK data path: {nltk_data_path}")
                 # Check if 'punkt' is available, download if not (be cautious with auto-downloads)
                 try:
                      nltk.data.find('tokenizers/punkt')
                 except nltk.downloader.DownloadError:
                      print("NLTK 'punkt' tokenizer not found. Attempting download...")
                      try:
                           nltk.download("punkt", quiet=False, download_dir=nltk_data_path)
                      except Exception as e:
                           raise RuntimeError(f"Failed to download NLTK 'punkt' tokenizer. Please install it manually or check NLTK_DATA path. Error: {e}")
             else:
                  print("Warning: NLTK_DATA environment variable not set. NLTK will use default paths.")
                  try:
                       nltk.data.find('tokenizers/punkt')
                  except Exception:
                       raise RuntimeError("NLTK 'punkt' tokenizer not found. Please install it or set NLTK_DATA.")

        else:
            raise ImportError("nltk library required for --split-sentences is not available.")


    # --- File Handling and Partitioning ---
    in_ss_out_names = [] # List to store dictionaries of file names for each partition
    input_pattern = args.input
    output_prefix = args.output_prefix
    num_partitions = args.partitions

    if num_partitions == 1:
        # Check if input is a single file or a pattern that resolves to one file
        in_file_names = glob.glob(input_pattern)
        if not in_file_names:
            print(f"Error: No input files found matching pattern: {input_pattern}")
            sys.exit(1)
        if len(in_file_names) > 1:
            print(f"Warning: Input pattern '{input_pattern}' matched multiple files, but partitions=1. Using first file: {in_file_names[0]}")
        input_file = in_file_names[0]
        file_name_base, extension = os.path.splitext(os.path.basename(input_file))
        # Use output_prefix directly for the single partition's output files
        partition_output_prefix = output_prefix
        sentence_split_file = f"{partition_output_prefix}_ss{extension}" # Suffix added to output prefix

        file_names = {
            'partition': input_file, # Original input file
            'sentence_split': sentence_split_file, # Temporary file if splitting sentences
            'output_prefix': partition_output_prefix # Prefix for final .bin/.idx files
        }
        in_ss_out_names.append(file_names)
    else:
        # Handle multiple partitions
        in_file_names = glob.glob(input_pattern)
        if not in_file_names:
            print(f"Error: No input files found matching pattern: {input_pattern} for {num_partitions} partitions.")
            sys.exit(1)
        print(f"Found {len(in_file_names)} input files matching pattern.")

        # Create partition file names (temporary files)
        partition_base = os.path.join(os.path.dirname(output_prefix), "partition_files")
        os.makedirs(partition_base, exist_ok=True) # Ensure directory exists

        for idx in range(num_partitions):
             # Generate names for temporary partition files and their outputs
             # Use a consistent naming scheme based on the overall output prefix and partition index
             partition_file_path = os.path.join(partition_base, f"input_part_{idx:03d}.jsonl")
             partition_output_prefix = f"{output_prefix}_part_{idx:03d}" # Partition-specific output prefix
             sentence_split_file = f"{partition_output_prefix}_ss.jsonl" # Temp sentence split file for this partition

             in_ss_out_name = {
                 'partition': partition_file_path, # Path to temp input partition file
                 'sentence_split': sentence_split_file, # Path to temp sentence split file
                 'output_prefix': partition_output_prefix # Prefix for this partition's .bin/.idx
                 }
             in_ss_out_names.append(in_ss_out_name)

        # Check if temporary partition files already exist
        partitions_present = check_files_exist(in_ss_out_names, 'partition', num_partitions)
        split_sentences_present = check_files_exist(in_ss_out_names, 'sentence_split', num_partitions) if args.split_sentences else True # Assume present if not splitting

        if not partitions_present:
            print("Creating temporary partition files...")
            # Distribute lines from input files into partition files
            partitioned_input_files = [open(name['partition'], 'w', encoding='utf-8') for name in in_ss_out_names]

            line_count = 0
            processed_files_count = 0
            try:
                for in_file_name in sorted(in_file_names): # Sort for determinism if needed
                    print(f"Reading input file: {in_file_name}")
                    # Handle gzip
                    if in_file_name.endswith(".gz"):
                        fin = gzip.open(in_file_name, 'rt', encoding='utf-8')
                    else:
                        fin = open(in_file_name, 'r', encoding='utf-8')

                    with fin:
                         for line in fin:
                              # Distribute lines round-robin or sequentially based on args.keep_sequential_samples
                              # Simple round-robin distribution:
                              target_partition_index = line_count % num_partitions
                              partitioned_input_files[target_partition_index].write(line)
                              line_count += 1
                    processed_files_count += 1
                    print(f"Finished reading {in_file_name}. Total lines distributed so far: {line_count}")

            except Exception as e:
                 print(f"Error distributing lines to partitions: {e}")
                 # Clean up open files before exiting
                 for f in partitioned_input_files: f.close()
                 sys.exit(1)
            finally:
                for f in partitioned_input_files: f.close()
            print(f"Finished creating {num_partitions} partition files. Total lines processed: {line_count} from {processed_files_count} input files.")
        else:
            print("Temporary partition files already exist. Skipping creation.")


    # --- Worker Setup ---
    if args.workers <= 0:
        print("Error: Number of workers must be positive.")
        sys.exit(1)
    if args.workers % num_partitions != 0:
        print(f"Warning: Number of workers ({args.workers}) is not divisible by the number of partitions ({num_partitions}). This might lead to uneven load.")
        workers_per_partition = args.workers // num_partitions
        if workers_per_partition == 0: workers_per_partition = 1 # Ensure at least one worker per partition
        print(f"Assigning approximately {workers_per_partition} workers per partition.")
    else:
        workers_per_partition = args.workers // num_partitions
        print(f"Assigning {workers_per_partition} workers per partition.")

    partition_handler = Partition(args, workers_per_partition)


    # --- Optional: Sentence Splitting of Partition Files ---
    input_key_for_processing = 'partition' # Default: process the initial partition files
    if args.split_sentences:
        split_sentences_present = check_files_exist(in_ss_out_names, 'sentence_split', num_partitions)
        if not split_sentences_present:
            print("Splitting sentences in partition files...")
            processes = []
            for name in in_ss_out_names:
                try:
                    p = multiprocessing.Process(target=partition_handler.split_sentences,
                                                args=((name['partition'], name['sentence_split']),))
                    p.start()
                    processes.append(p)
                except Exception as e:
                    print(f"Error starting sentence splitting process for {name['partition']}: {e}")
                    # Handle error: maybe terminate already started processes?

            # Wait for sentence splitting processes to complete
            for i, p in enumerate(processes):
                try:
                    p.join()
                    if p.exitcode != 0:
                         print(f"Warning: Sentence splitting process for partition {i} exited with code {p.exitcode}")
                except Exception as e:
                    print(f"Error joining sentence splitting process for partition {i}: {e}")
            print("Sentence splitting complete.")
            input_key_for_processing = 'sentence_split' # Process the sentence-split files next
        else:
             print("Sentence-split files already exist. Skipping splitting.")
             input_key_for_processing = 'sentence_split'


    # --- Main Encoding Process ---
    print(f"Starting encoding process using input key: '{input_key_for_processing}'...")
    processes = []
    for name in in_ss_out_names:
        try:
            p = multiprocessing.Process(target=partition_handler.process_json_file,
                                        args=((name[input_key_for_processing], name['output_prefix']),))
            p.start()
            processes.append(p)
        except Exception as e:
            print(f"Error starting encoding process for {name[input_key_for_processing]}: {e}")


    # Wait for encoding processes to complete
    for i, p in enumerate(processes):
        try:
            p.join()
            if p.exitcode != 0:
                 print(f"Warning: Encoding process for partition {i} exited with code {p.exitcode}")
        except Exception as e:
            print(f"Error joining encoding process for partition {i}: {e}")

    print("All encoding processes finished.")

    if num_partitions == 1:
        print("Processing complete for single partition.")
        return # Nothing more to do

    # --- Merge Bin/Idx Partitions ---
    print("Merging partition results...")
    level = "sentence" if args.split_sentences else "document"
    final_output_bin_files = {}
    final_output_idx_files = {}
    final_builders = {}

    try:
        # Re-initialize tokenizer to get vocab size for final builders
        tokenizer = _MMSentencePieceTokenizer(args.tokenizer_model, vocab_extra_ids=args.vocab_extra_ids)
        dtype = indexed_dataset.DType.optimal_dtype(tokenizer.vocab_size)
    except Exception as e:
        print(f"Error re-initializing tokenizer for merging: {e}")
        sys.exit(1)


    for key in args.json_keys:
        # Skip 'codec' if 'text' is present (as it was handled during processing)
        if key == 'codec' and "text" in args.json_keys:
            continue

        print(f"Merging results for key: '{key}'")
        final_output_bin_files[key] = f"{args.output_prefix}_{key}_{level}.bin"
        final_output_idx_files[key] = f"{args.output_prefix}_{key}_{level}.idx"

        try:
            final_builders[key] = indexed_dataset.MMapIndexedDatasetBuilder(
                final_output_bin_files[key], dtype=dtype
            )

            # Add indices from each partition's output
            num_indices_added = 0
            for name in in_ss_out_names:
                partition_output_prefix = name['output_prefix'] # e.g., "output/data_part_000"
                partition_index_prefix = f"{partition_output_prefix}_{key}_{level}" # e.g., "output/data_part_000_text_document"

                # Check if the partition's bin and idx files exist before adding
                partition_bin_file = f"{partition_index_prefix}.bin"
                partition_idx_file = f"{partition_index_prefix}.idx"
                if os.path.exists(partition_bin_file) and os.path.exists(partition_idx_file):
                    try:
                         final_builders[key].add_index(partition_index_prefix)
                         num_indices_added += 1
                         if DEBUG: print(f"Added index: {partition_index_prefix}")
                    except Exception as e:
                         print(f"Error adding index {partition_index_prefix} for key '{key}': {e}")
                else:
                    print(f"Warning: Index files not found for partition {partition_index_prefix}. Skipping.")


            # Finalize the merged index
            if num_indices_added > 0:
                 print(f"Finalizing merged index for key '{key}' with {num_indices_added} partitions...")
                 final_builders[key].finalize(final_output_idx_files[key])
                 print(f"Finalized merged index: {final_output_idx_files[key]}")
            else:
                 print(f"Warning: No partition indices were added for key '{key}'. Cannot finalize merged index.")
                 # Clean up potentially created empty bin file
                 if os.path.exists(final_output_bin_files[key]):
                      try: os.remove(final_output_bin_files[key])
                      except OSError: pass


        except Exception as e:
            print(f"Error merging or finalizing index for key '{key}': {e}")

    print("Finished merging partitions.")
    # Optional: Clean up temporary partition files here if desired
Main function to orchestrate data partitioning and processing.
finetune/core/preprocess_data_conditional_xcodec_segment.py