Function process_batch

generation_utils.py:445–761

Process a batch of data items and generate audio, return audio data and metadata

(
    batch_items,
    tokenizer,
    model,
    spt,
    device,
    system_prompt,
    start_idx,
    use_normalize=False,
    silence_duration=0,
)

def process_batch(
    batch_items,
    tokenizer,
    model,
    spt,
    device,
    system_prompt,
    start_idx,
    use_normalize=False,
    silence_duration=0,
):
    """Process a batch of data items and generate audio, return audio data and metadata"""
    try:
        # Prepare batch data
        batch_size = len(batch_items)
        texts = []
        prompts = [system_prompt] * batch_size
        prompt_audios = []
        actual_texts_data = []  # Store actual text data used

        print(f"Processing {batch_size} samples starting from index {start_idx}...")

        # Extract text and audio from each sample
        for i, item in enumerate(batch_items):
            # Use new processing function
            processed_item = process_jsonl_item(item)

            text = processed_item["text"]
            prompt_text = processed_item["prompt_text"]

            # Merge text, if prompt_text is empty, full_text is just text
            full_text = prompt_text + text if prompt_text else text
            original_full_text = full_text  # Save original text

            # Apply text normalization based on parameter
            if use_normalize:
                full_text = normalize_text(full_text)

            # Replace speaker tags
            final_text = full_text.replace("[S1]", "<speaker1>").replace(
                "[S2]", "<speaker2>"
            )
            texts.append(final_text)

            # Save actual text information used
            actual_texts_data.append(
                {
                    "index": start_idx + i,
                    "original_text": original_full_text,
                    "normalized_text": (
                        normalize_text(original_full_text) if use_normalize else None
                    ),
                    "final_text": final_text,
                    "use_normalize": use_normalize,
                }
            )

            # Get reference audio