Function process_data

finetune/data_preprocess.py:186–379

(
    jsonl: str,
    model_path: str,
    output_dir: str,
    data_name: str = "processd_data",
    use_normalize: bool = True,
)

def process_data(
    jsonl: str,
    model_path: str,
    output_dir: str,
    data_name: str = "processd_data",
    use_normalize: bool = True,
):
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    # Load models
    print("Loading models...")
    tokenizer, spt = load_tokenizer(model_path, SPT_CONFIG_PATH, SPT_CHECKPOINT_PATH)
    spt = spt.to(device)

    # Load the items from the JSONL file
    try:
        with open(jsonl, "r") as f:
            items = [json.loads(line) for line in f.readlines()]
        print(f"Loaded {len(items)} items from {jsonl}")
    except FileNotFoundError:
        print(f"Error: JSONL file '{jsonl}' not found")
        return
    except json.JSONDecodeError as e:
        print(f"Error parsing JSONL file: {e}")
        return

    # Store all processed data and length information
    all_data = []
    offsets = []
    tokens_lengths = []  # Added: store total token length
    tims_lengths = []  # Added: store audio token length

    for idx, item in enumerate(items):
        # Support two JSONL formats
        # Format 1: {"file_path": "path/to/audio.wav", "full_transcript": "speech content..."}
        # Format 2: {"reference_audio": "path1", "reference_text": "text1", "audio": "path2", "text": "text2"}

        if "file_path" in item and "full_transcript" in item:
            # Original format
            file_path = item["file_path"]
            full_text = item["full_transcript"]

            # Check if audio file exists
            if not file_path:
                print(f"Warning: Item {idx} has empty file_path, skipping...")
                continue

            if not os.path.exists(file_path):
                print(
                    f"Warning: Audio file not found: {file_path}, skipping item {idx}..."
                )
                continue

            try: