| import itertools
|
| import os
|
| import re
|
| from collections import defaultdict
|
| from functools import partial
|
| from multiprocessing import Pool
|
| from pathlib import Path
|
|
|
| import click
|
| import numpy as np
|
| from loguru import logger
|
| from tqdm import tqdm
|
|
|
| from fish_speech.datasets.protos.text_data_pb2 import Semantics, Sentence, TextData
|
| from fish_speech.datasets.protos.text_data_stream import pack_pb_stream
|
| from tools.file import load_filelist
|
|
|
|
|
| os.environ["MKL_NUM_THREADS"] = "1"
|
| os.environ["OMP_NUM_THREADS"] = "1"
|
|
|
|
|
| def task_generator_folder(root: Path, text_extension: str):
|
| files = list(tqdm(Path(root).rglob("*.npy"), desc=f"Loading {root}"))
|
| files = sorted(files)
|
|
|
| grouped_files = defaultdict(list)
|
| for file in tqdm(files, desc=f"Grouping {root}"):
|
| p = str(file.parent)
|
| speaker = file.parent.name
|
|
|
| try:
|
| if isinstance(text_extension, str):
|
| texts = [file.with_suffix(text_extension).read_text(encoding="utf-8")]
|
| else:
|
| texts = [
|
| file.with_suffix(ext).read_text(encoding="utf-8")
|
| for ext in text_extension
|
| ]
|
| except Exception as e:
|
| logger.error(f"Failed to read text {file}: {e}")
|
| continue
|
|
|
| grouped_files[p].append((speaker, file, texts))
|
|
|
| logger.info(
|
| f"Found {len(grouped_files)} groups in {root}, {list(grouped_files.keys())[:5]}..."
|
| )
|
|
|
| for i in grouped_files.values():
|
| subset = [(f, t) for _, f, t in i]
|
| yield i[0][0], subset, "folder"
|
|
|
|
|
| def task_generator_filelist(filelist):
|
| grouped_files = defaultdict(list)
|
| for filename, speaker, _, text in load_filelist(filelist):
|
| grouped_files[speaker].append((Path(filename), [text]))
|
|
|
| logger.info(f"Found {len(grouped_files)} groups in {filelist}")
|
| for speaker, values in grouped_files.items():
|
| yield speaker, values, "filelist"
|
|
|
|
|
| def run_task(task):
|
| name, subset, source = task
|
|
|
|
|
| sentences = []
|
| for file, texts in subset:
|
| np_file = file.with_suffix(".npy")
|
| if np_file.exists() is False:
|
| logger.warning(f"Can't find {np_file}")
|
| continue
|
|
|
| new_texts = []
|
|
|
| for text in texts:
|
|
|
| text = re.sub(r"\{.*?\}", " ", text)
|
| text = re.sub(r"<.*?>", " ", text)
|
| text = re.sub(r"\s+", " ", text)
|
| new_texts.append(text)
|
|
|
| try:
|
| semantics = np.load(np_file)
|
| except Exception as e:
|
| logger.error(f"Failed to parse {file}: {e}")
|
| continue
|
|
|
| if isinstance(semantics, np.ndarray):
|
| semantics = semantics.tolist()
|
|
|
| sentences.append(
|
| Sentence(
|
| texts=new_texts,
|
| semantics=[Semantics(values=s) for s in semantics],
|
| )
|
| )
|
|
|
|
|
| return pack_pb_stream(
|
| TextData(
|
| source=source,
|
| name=name,
|
| sentences=sentences,
|
| )
|
| )
|
|
|
|
|
| @click.command()
|
| @click.option(
|
| "--input",
|
| type=click.Path(path_type=Path),
|
| required=True,
|
| help="A folder containing the dataset or a filelist",
|
| multiple=True,
|
| )
|
| @click.option(
|
| "--output", type=click.Path(path_type=Path), default="data/quantized-dataset-ft"
|
| )
|
| @click.option("--num-workers", type=int, default=16)
|
| @click.option("--text-extension", type=str, default=[".txt"], multiple=True)
|
| @click.option(
|
| "--shard-size", type=int, default=10, help="The maximum size of each shard in mb"
|
| )
|
| def main(input, output, num_workers, text_extension, shard_size):
|
| generator_fns = []
|
|
|
| for f in input:
|
| assert f.exists(), f"{f} not found"
|
|
|
| if f.is_dir():
|
| generator_fn = task_generator_folder(f, text_extension)
|
| else:
|
| generator_fn = task_generator_filelist(f)
|
|
|
| generator_fns.append(generator_fn)
|
|
|
| generator_fn = itertools.chain(*generator_fns)
|
| output.mkdir(parents=True, exist_ok=True)
|
|
|
| dataset_fp = None
|
| tar_idx = 0
|
| written_size = 0
|
|
|
| with Pool(num_workers) as p:
|
| for result in tqdm(p.imap_unordered(run_task, generator_fn)):
|
| if dataset_fp is None:
|
| dataset_fp = open(Path(output) / f"{tar_idx:08d}.protos", "wb")
|
|
|
| dataset_fp.write(result)
|
| written_size += len(result)
|
|
|
| if written_size > shard_size * 1024 * 1024:
|
| logger.info(f"Finished writing {tar_idx} shards to {output}")
|
| dataset_fp.close()
|
| dataset_fp = None
|
| written_size = 0
|
| tar_idx += 1
|
|
|
| if dataset_fp is not None:
|
| dataset_fp.close()
|
|
|
| logger.info(f"Finished writing {tar_idx + 1} shards to {output}")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
