FunASR/funasr/datasets/audio_datasets/index_ds.py
zhifu gao 0a4a1d5257
Dev gzf exp (#1645)
* sensevoice finetune

* sensevoice finetune

* sensevoice finetune

* sensevoice finetune

* sensevoice finetune

* sensevoice finetune

* sensevoice finetune

* sensevoice finetune

* sensevoice finetune

* sensevoice finetune

* bugfix

* update with main (#1631)

* update seaco finetune

* v1.0.24

---------

Co-authored-by: 维石 <shixian.shi@alibaba-inc.com>

* sensevoice

* sensevoice

* sensevoice

* update with main (#1638)

* update seaco finetune

* v1.0.24

* update rwkv template

---------

Co-authored-by: 维石 <shixian.shi@alibaba-inc.com>

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sensevoice

* sense voice

* sense voice

* sense voice

* sense voice

---------

Co-authored-by: 维石 <shixian.shi@alibaba-inc.com>
2024-04-23 12:48:52 +08:00

228 lines
9.1 KiB
Python

import os
import json
import torch
import logging
import concurrent.futures
import librosa
import torch.distributed as dist
from funasr.register import tables
# @tables.register("index_ds_classes", "IndexDSJsonlRankSplit")
# class IndexDSJsonlRankSplit(torch.utils.data.Dataset):
#
# def __init__(self, path):
# super().__init__()
#
# contents = []
# with open(path, encoding='utf-8') as fin:
# for line in fin:
# data = json.loads(line.strip())
# if "text" in data: # for sft
# self.contents.append(data['text'])
# if "source" in data: # for speech lab pretrain
# prompt = data["prompt"]
# source = data["source"]
# target = data["target"]
# source_len = data["source_len"]
# target_len = data["target_len"]
#
# contents.append({"source": source,
# "prompt": prompt,
# "target": target,
# "source_len": source_len,
# "target_len": target_len,
# }
# )
#
# self.contents = []
# total_num = len(contents)
# try:
# rank = dist.get_rank()
# world_size = dist.get_world_size()
# except:
# rank = 0
# world_size = 1
# logging.warning("distributed is not initialized, only single shard")
# num_per_rank = total_num // world_size
#
# # rank = 0
# # import ipdb; ipdb.set_trace()
# self.contents = contents[rank * num_per_rank:(rank + 1) * num_per_rank]
#
# logging.info("in rank: {}, num of samplers: {}, total_num of samplers across ranks: {}".format(rank, len(self.contents), len(contents)))
#
# def __len__(self):
# return len(self.contents)
#
# def __getitem__(self, index):
# try:
# data = self.contents[index]
# except:
# print(index)
# return data
#
# def get_source_len(self, data_dict):
# return data_dict["source_len"]
#
# def get_target_len(self, data_dict):
#
# return data_dict["target_len"] if "target_len" in data_dict else 0
@tables.register("index_ds_classes", "IndexDSJsonl")
@tables.register("index_ds_classes", "IndexDSJsonlRankFull")
class IndexDSJsonlRankFull(torch.utils.data.Dataset):
def __init__(self, path: str, **kwargs):
super().__init__()
self.max_source_length = kwargs.get("max_source_length", 2048)
self.min_source_length = kwargs.get("min_source_length", 0)
self.max_target_length = kwargs.get("max_target_length", 2048)
self.min_target_length = kwargs.get("min_target_length", 0)
if isinstance(path, (list, tuple)): # wav.scp, text.txt/text.trans
from funasr.datasets.audio_datasets.scp2jsonl import gen_jsonl_from_wav_text_list
jsonl_outdir = os.path.dirname(path[0])
jsonl_name = "datalist_train.jsonl" if kwargs.get("is_training", True) else "datalist_val.jsonl"
jsonl_file_out = os.path.join(jsonl_outdir, jsonl_name)
if not os.path.exists(jsonl_file_out):
print(f"datalist is: {path}, generate jsonl from it")
gen_jsonl_from_wav_text_list(path, jsonl_file_out=jsonl_file_out, **kwargs)
path = jsonl_file_out
contents = []
with open(path, encoding='utf-8') as fin:
for line in fin:
data = json.loads(line.strip())
if "text" in data: # for sft
contents.append(data['text'])
if "source" in data: # for speech lab pretrain
prompt = data.get("prompt", "<ASR>")
source = data["source"]
target = data["target"]
source_len = data.get("source_len", 1)
target_len = data.get("target_len", 0)
if "aishell" in source:
target = target.replace(" ", "")
if source_len < self.min_source_length or source_len > self.max_source_length:
continue
if target_len < self.min_target_length or target_len > self.max_target_length:
continue
contents_i = {"source": source,
"prompt": prompt,
"target": target,
"source_len": source_len,
"target_len": target_len,
}
text_language = data.get("text_language", None)
if text_language is not None:
contents_i["text_language"] = text_language
audio_language = data.get("audio_language", None)
if audio_language is not None:
contents_i["audio_language"] = audio_language
contents.append(contents_i)
self.contents = contents
logging.info(
"total_num of samplers across ranks: {}".format(len(self.contents)))
def __len__(self):
return len(self.contents)
def __getitem__(self, index):
try:
data = self.contents[index]
except:
print(index)
return data
def get_source_len(self, data_dict):
return data_dict.get("source_len", 1)
def get_target_len(self, data_dict):
return data_dict.get("target_len", 0)
@tables.register("index_ds_classes", "IndexDSJsonlRankSplit")
class IndexDSJsonlRankSplit(torch.utils.data.Dataset):
def __init__(self, path: str, **kwargs):
super().__init__()
self.max_source_length = kwargs.get("max_source_length", 2048)
self.min_source_length = kwargs.get("min_source_length", 0)
self.max_target_length = kwargs.get("max_target_length", 2048)
self.min_target_length = kwargs.get("min_target_length", 0)
with open(path, encoding='utf-8') as fin:
file_list = fin.readlines()
total_num = len(file_list)
try:
rank = dist.get_rank()
world_size = dist.get_world_size()
except:
rank = 0
world_size = 1
logging.warning("distributed is not initialized, only single shard")
num_per_rank = total_num // world_size
if num_per_rank * world_size < total_num:
logging.warning(f"Warning, jsonl file:{total_num} could not be divided by world_size: {world_size}, {path}")
file_list_rank = file_list[rank * num_per_rank:(rank + 1) * num_per_rank]
contents = []
for file_json in file_list_rank:
with open(file_json.strip(), encoding='utf-8') as fin:
for line in fin:
data = json.loads(line.strip())
if "text" in data: # for sft
contents.append(data['text'])
if "source" in data: # for speech lab pretrain
prompt = data.get("prompt", "<ASR>")
source = data["source"].replace("/cpfs01", "/cpfs_speech/data")
target = data["target"]
source_len = data.get("source_len", 1)
target_len = data.get("target_len", 0)
if source_len < self.min_source_length or source_len > self.max_source_length:
continue
if target_len < self.min_target_length or target_len > self.max_target_length:
continue
contents_i = {"source": source,
"prompt": prompt,
"target": target,
"source_len": source_len,
"target_len": target_len,
}
text_language = data.get("text_language", None)
if text_language is not None:
contents_i["text_language"] = text_language
audio_language = data.get("audio_language", None)
if audio_language is not None:
contents_i["audio_language"] = audio_language
contents.append(contents_i)
self.contents = contents
logging.info(f"total_num: {len(self.contents)} of samplers in ranks: {rank}")
def __len__(self):
return len(self.contents)
def __getitem__(self, index):
try:
data = self.contents[index]
except:
print(index)
return data
def get_source_len(self, data_dict):
return data_dict.get("source_len", 1)
def get_target_len(self, data_dict):
return data_dict.get("target_len", 0)