Upload ./prepare_data.py with huggingface_hub

6808dbe verified 8 months ago

5.84 kB

# coding=utf-8
# Copyright 2025 The Alibaba Qwen team.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Qwen3-TTS-12Hz Audio Codes 提取脚本 (优化版)
修复内容：
1. 使用逐行读取，减少内存占用
2. 添加异常处理，跳过损坏的音频文件
3. 修复最后一个 batch 处理逻辑（合并到主 with 块）
4. 修复进度条更新（最后一个 batch 使用实际数量）
5. 添加 --batch_size 参数动态调整
Usage:
    python prepare_data.py \
        --device cuda:0 \
        --tokenizer_model_path Qwen/Qwen3-TTS-Tokenizer-12Hz \
        --input_jsonl ./processed_data/uyghur_all/train_raw.jsonl \
        --output_jsonl ./processed_data/uyghur_all/train_with_codes.jsonl
"""

import argparse
import json
import os
import sys
from tqdm import tqdm

from qwen_tts import Qwen3TTSTokenizer

BATCH_INFER_NUM = 32
SAVE_INTERVAL = 5000  # Save every N samples

def main():
    parser = argparse.ArgumentParser(description='Extract audio codes (优化版)')
    parser.add_argument("--device", type=str, default="cuda:0")
    parser.add_argument("--tokenizer_model_path", type=str, default="Qwen/Qwen3-TTS-Tokenizer-12Hz")
    parser.add_argument("--input_jsonl", type=str, required=True)
    parser.add_argument("--output_jsonl", type=str, required=True)
    parser.add_argument("--batch_size", type=int, default=BATCH_INFER_NUM,
                        help=f"批次大小 (default: {BATCH_INFER_NUM})")
    args = parser.parse_args()

    print(f"Loading tokenizer from: {args.tokenizer_model_path}")
    tokenizer_12hz = Qwen3TTSTokenizer.from_pretrained(
        args.tokenizer_model_path,
        device_map=args.device,
    )
    print(f"Tokenizer loaded successfully on {args.device}")

    # 修复：使用逐行读取，减少内存占用
    print(f"Reading input file: {args.input_jsonl}")
    total_count = 0
    with open(args.input_jsonl, 'r', encoding='utf-8') as f:
        for _ in f:
            total_count += 1
    print(f"Total samples: {total_count:,}")

    # 创建输出目录
    output_dir = os.path.dirname(args.output_jsonl)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)

    processed_count = 0
    failed_count = 0
    batch_size = args.batch_size

    print(f"Processing with batch_size={batch_size}...")

    # 修复：统一在单个 with 块内处理所有数据
    with open(args.output_jsonl, 'w', encoding='utf-8') as f:
        batch_lines = []
        batch_audios = []

        pbar = tqdm(total=total_count, desc="Extracting audio codes")

        # 逐行读取并处理
        with open(args.input_jsonl, 'r', encoding='utf-8') as infile:
            for line in infile:
                try:
                    line_data = json.loads(line.strip())
                    batch_lines.append(line_data)
                    batch_audios.append(line_data['audio'])
                except Exception as e:
                    print(f"Warning: Failed to parse line: {e}")
                    failed_count += 1
                    pbar.update(1)
                    continue

                # 达到批次大小时处理
                if len(batch_lines) >= batch_size:
                    try:
                        enc_res = tokenizer_12hz.encode(batch_audios)
                        for code, line_data in zip(enc_res.audio_codes, batch_lines):
                            line_data['audio_codes'] = code.cpu().tolist()
                            f.write(json.dumps(line_data, ensure_ascii=False) + '\n')
                            processed_count += 1
                    except Exception as e:
                        print(f"Warning: Failed to encode batch: {e}")
                        failed_count += len(batch_lines)

                    batch_lines.clear()
                    batch_audios.clear()
                    pbar.update(batch_size)

                    # 定期显示进度
                    if processed_count % SAVE_INTERVAL == 0:
                        print(f"  Progress: {processed_count:,} / {total_count:,} ({100*processed_count/total_count:.1f}%)")

        # 修复：处理剩余的最后一个 batch（在主 with 块内）
        if len(batch_audios) > 0:
            try:
                enc_res = tokenizer_12hz.encode(batch_audios)
                for code, line_data in zip(enc_res.audio_codes, batch_lines):
                    line_data['audio_codes'] = code.cpu().tolist()
                    f.write(json.dumps(line_data, ensure_ascii=False) + '\n')
                    processed_count += 1
            except Exception as e:
                print(f"Warning: Failed to encode final batch: {e}")
                failed_count += len(batch_lines)

            # 修复：使用实际数量更新进度条
            pbar.update(len(batch_lines))
            batch_lines.clear()
            batch_audios.clear()

        pbar.close()

    print(f"\n{'='*60}")
    print(f"Done!")
    print(f"  Total samples: {total_count:,}")
    print(f"  Processed: {processed_count:,}")
    print(f"  Failed: {failed_count:,}")
    print(f"  Success rate: {100*processed_count/(processed_count+failed_count):.1f}%")
    print(f"Output saved to: {args.output_jsonl}")
    print(f"{'='*60}")

if __name__ == "__main__":
    main()

	# coding=utf-8
	# Copyright 2025 The Alibaba Qwen team.
	# SPDX-License-Identifier: Apache-2.0
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	"""
	Qwen3-TTS-12Hz Audio Codes 提取脚本 (优化版)

	修复内容：
	1. 使用逐行读取，减少内存占用
	2. 添加异常处理，跳过损坏的音频文件
	3. 修复最后一个 batch 处理逻辑（合并到主 with 块）
	4. 修复进度条更新（最后一个 batch 使用实际数量）
	5. 添加 --batch_size 参数动态调整

	Usage:
	python prepare_data.py \
	--device cuda:0 \
	--tokenizer_model_path Qwen/Qwen3-TTS-Tokenizer-12Hz \
	--input_jsonl ./processed_data/uyghur_all/train_raw.jsonl \
	--output_jsonl ./processed_data/uyghur_all/train_with_codes.jsonl
	"""

	import argparse
	import json
	import os
	import sys
	from tqdm import tqdm

	from qwen_tts import Qwen3TTSTokenizer

	BATCH_INFER_NUM = 32
	SAVE_INTERVAL = 5000 # Save every N samples

	def main():
	parser = argparse.ArgumentParser(description='Extract audio codes (优化版)')
	parser.add_argument("--device", type=str, default="cuda:0")
	parser.add_argument("--tokenizer_model_path", type=str, default="Qwen/Qwen3-TTS-Tokenizer-12Hz")
	parser.add_argument("--input_jsonl", type=str, required=True)
	parser.add_argument("--output_jsonl", type=str, required=True)
	parser.add_argument("--batch_size", type=int, default=BATCH_INFER_NUM,
	help=f"批次大小 (default: {BATCH_INFER_NUM})")
	args = parser.parse_args()

	print(f"Loading tokenizer from: {args.tokenizer_model_path}")
	tokenizer_12hz = Qwen3TTSTokenizer.from_pretrained(
	args.tokenizer_model_path,
	device_map=args.device,
	)
	print(f"Tokenizer loaded successfully on {args.device}")

	# 修复：使用逐行读取，减少内存占用
	print(f"Reading input file: {args.input_jsonl}")
	total_count = 0
	with open(args.input_jsonl, 'r', encoding='utf-8') as f:
	for _ in f:
	total_count += 1
	print(f"Total samples: {total_count:,}")

	# 创建输出目录
	output_dir = os.path.dirname(args.output_jsonl)
	if output_dir and not os.path.exists(output_dir):
	os.makedirs(output_dir, exist_ok=True)

	processed_count = 0
	failed_count = 0
	batch_size = args.batch_size

	print(f"Processing with batch_size={batch_size}...")

	# 修复：统一在单个 with 块内处理所有数据
	with open(args.output_jsonl, 'w', encoding='utf-8') as f:
	batch_lines = []
	batch_audios = []

	pbar = tqdm(total=total_count, desc="Extracting audio codes")

	# 逐行读取并处理
	with open(args.input_jsonl, 'r', encoding='utf-8') as infile:
	for line in infile:
	try:
	line_data = json.loads(line.strip())
	batch_lines.append(line_data)
	batch_audios.append(line_data['audio'])
	except Exception as e:
	print(f"Warning: Failed to parse line: {e}")
	failed_count += 1
	pbar.update(1)
	continue

	# 达到批次大小时处理
	if len(batch_lines) >= batch_size:
	try:
	enc_res = tokenizer_12hz.encode(batch_audios)
	for code, line_data in zip(enc_res.audio_codes, batch_lines):
	line_data['audio_codes'] = code.cpu().tolist()
	f.write(json.dumps(line_data, ensure_ascii=False) + '\n')
	processed_count += 1
	except Exception as e:
	print(f"Warning: Failed to encode batch: {e}")
	failed_count += len(batch_lines)

	batch_lines.clear()
	batch_audios.clear()
	pbar.update(batch_size)

	# 定期显示进度
	if processed_count % SAVE_INTERVAL == 0:
	print(f" Progress: {processed_count:,} / {total_count:,} ({100*processed_count/total_count:.1f}%)")

	# 修复：处理剩余的最后一个 batch（在主 with 块内）
	if len(batch_audios) > 0:
	try:
	enc_res = tokenizer_12hz.encode(batch_audios)
	for code, line_data in zip(enc_res.audio_codes, batch_lines):
	line_data['audio_codes'] = code.cpu().tolist()
	f.write(json.dumps(line_data, ensure_ascii=False) + '\n')
	processed_count += 1
	except Exception as e:
	print(f"Warning: Failed to encode final batch: {e}")
	failed_count += len(batch_lines)

	# 修复：使用实际数量更新进度条
	pbar.update(len(batch_lines))
	batch_lines.clear()
	batch_audios.clear()

	pbar.close()

	print(f"\n{'='*60}")
	print(f"Done!")
	print(f" Total samples: {total_count:,}")
	print(f" Processed: {processed_count:,}")
	print(f" Failed: {failed_count:,}")
	print(f" Success rate: {100*processed_count/(processed_count+failed_count):.1f}%")
	print(f"Output saved to: {args.output_jsonl}")
	print(f"{'='*60}")

	if __name__ == "__main__":
	main()