139 lines
4.6 KiB
Python
139 lines
4.6 KiB
Python
import logging
|
||
|
||
import yaml
|
||
from typing import (
|
||
Tuple,
|
||
List, Any
|
||
)
|
||
from utils.model import AudioItem
|
||
import os
|
||
import zipfile
|
||
import tarfile
|
||
import gzip
|
||
import shutil
|
||
from pathlib import Path
|
||
|
||
|
||
def read_data(dataset_filepath: str) -> Tuple[str, List[AudioItem]]:
|
||
"""
|
||
读取数据文件，返回语言和文本列表
|
||
|
||
参数:
|
||
dataset_filepath (str): 数据文件路径
|
||
|
||
返回:
|
||
Tuple[str, List[str]]:
|
||
- language: 文件中指定的语言字符串
|
||
- datas: 文件中除语言行以外的文本列表（每行为一个元素）
|
||
"""
|
||
try:
|
||
# 认为都是压缩包，先解压数据。
|
||
data_extract_path = "/tmp/datas"
|
||
data_yaml_path = extract_file(dataset_filepath, data_extract_path)
|
||
if not data_yaml_path:
|
||
raise ValueError(f"未找到数据集data.yaml文件。")
|
||
dataset_filepath = str(Path(data_yaml_path).parent.resolve())
|
||
except Exception as e:
|
||
logging.exception(e)
|
||
|
||
|
||
with open(f"{dataset_filepath}/data.yaml") as f:
|
||
datas = yaml.safe_load(f)
|
||
language = datas.get("global", {}).get("lang", "zh")
|
||
query_data = datas.get("query_data", [])
|
||
|
||
audios = []
|
||
"""
|
||
- audio_length: 1.0099999999997635
|
||
duration: 1.0099999999997635
|
||
file: zh/0.wav
|
||
orig_file: ./112801_1/112801_1-631-772.wav
|
||
voice:
|
||
- answer: 好吧。
|
||
end: 1.0099999999997635
|
||
start: 0
|
||
"""
|
||
for item in query_data:
|
||
audio = AudioItem.model_validate(item)
|
||
audio.absolute_path = f"{dataset_filepath}/{audio.file}"
|
||
audios.append(audio)
|
||
|
||
return (
|
||
language,
|
||
audios
|
||
)
|
||
|
||
|
||
def extract_file(filepath: str, output_dir: str = ".") -> None:
|
||
"""
|
||
将数据集解压到指定路径，返回data.yaml文件的路径
|
||
"""
|
||
if not os.path.exists(output_dir):
|
||
os.makedirs(output_dir)
|
||
data_yaml_path = None
|
||
# 硬编码一下 leaderboard_data_samples 数据集没有加拓展名
|
||
# if filepath.endswith(".zip") or filepath.endswith("leaderboard_data_samples"):
|
||
# with zipfile.ZipFile(filepath, 'r') as zf:
|
||
# # 获取所有文件（非目录）
|
||
# all_files = [f for f in zf.namelist() if not f.endswith('/')]
|
||
#
|
||
# if not all_files:
|
||
# raise ValueError(f"数据集文件为空。{filepath}")
|
||
#
|
||
# # 获取公共路径前缀
|
||
# parts_list = [Path(f).parts for f in all_files]
|
||
# common_parts = os.path.commonprefix(parts_list)
|
||
# strip_prefix_len = len(common_parts)
|
||
#
|
||
# for file in all_files:
|
||
# file_parts = Path(file).parts
|
||
# relative_parts = file_parts[strip_prefix_len:]
|
||
# dest_path = Path(output_dir).joinpath(*relative_parts)
|
||
#
|
||
# # 创建父目录
|
||
# dest_path.parent.mkdir(parents=True, exist_ok=True)
|
||
#
|
||
# # 解压写入
|
||
# with zf.open(file) as source, open(dest_path, "wb") as target:
|
||
# shutil.copyfileobj(source, target)
|
||
#
|
||
# # 检查是否是 data.yaml
|
||
# if Path(file).name == "data.yaml":
|
||
# data_yaml_path = str(dest_path.resolve())
|
||
#
|
||
# logging.info(f"数据集解压成功。")
|
||
# else:
|
||
# raise ValueError(f"暂时不支持的压缩格式。{filepath}")
|
||
# TODO 使用的是已有的数据，都不是按照 zip结尾命名的，强制按照zip解压
|
||
with zipfile.ZipFile(filepath, 'r') as zf:
|
||
# 获取所有文件（非目录）
|
||
all_files = [f for f in zf.namelist() if not f.endswith('/')]
|
||
|
||
if not all_files:
|
||
raise ValueError(f"数据集文件为空。{filepath}")
|
||
|
||
# 获取公共路径前缀
|
||
parts_list = [Path(f).parts for f in all_files]
|
||
common_parts = os.path.commonprefix(parts_list)
|
||
strip_prefix_len = len(common_parts)
|
||
|
||
for file in all_files:
|
||
file_parts = Path(file).parts
|
||
relative_parts = file_parts[strip_prefix_len:]
|
||
dest_path = Path(output_dir).joinpath(*relative_parts)
|
||
|
||
# 创建父目录
|
||
dest_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 解压写入
|
||
with zf.open(file) as source, open(dest_path, "wb") as target:
|
||
shutil.copyfileobj(source, target)
|
||
|
||
# 检查是否是 data.yaml
|
||
if Path(file).name == "data.yaml":
|
||
data_yaml_path = str(dest_path.resolve())
|
||
|
||
logging.info(f"数据集解压成功。")
|
||
|
||
return data_yaml_path
|