zhifu gao
2024-04-24 861147c7308b91068ffa02724fdf74ee623a909e
funasr/datasets/audio_datasets/index_ds.py
@@ -31,17 +31,19 @@
            if not is_training:
                data_split_num = 1
                data_split_i = 0
            with open(path, encoding='utf-8') as fin:
            with open(path, encoding="utf-8") as fin:
                file_list_all = fin.readlines()
    
                num_per_slice = (len(file_list_all)-1) // data_split_num + 1
                file_list = file_list_all[data_split_i * num_per_slice:(data_split_i + 1) * num_per_slice]
                file_list = file_list_all[
                    data_split_i * num_per_slice : (data_split_i + 1) * num_per_slice
                ]
                logging.info(
                    f"is_training: {is_training}, data_split_num: {data_split_num}, data_split_i: {data_split_i}, \nfile_list: {file_list}, \nfile_list_all: {file_list_all}")
                    f"is_training: {is_training}, data_split_num: {data_split_num}, data_split_i: {data_split_i}, \nfile_list: {file_list}, \nfile_list_all: {file_list_all}"
                )
        
        else:
            file_list = [path]
        # total_num = len(file_list)
        # try:
@@ -76,24 +78,33 @@
        # for file_json in file_list_rank:
        contents = []
        for file_json in file_list:
            with open(file_json.strip(), encoding='utf-8') as fin:
            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'])
                        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") # only use in alibaba gpu group: .replace("/cpfs01", "/cpfs_speech/data")
                        source = data["source"].replace(
                            "/cpfs01", "/cpfs_speech/data"
                        )  # only use in alibaba gpu group: .replace("/cpfs01", "/cpfs_speech/data")
                        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:
                        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:
                        if (
                            target_len < self.min_target_length
                            or target_len > self.max_target_length
                        ):
                            continue
                        contents_i = {"source": source,
                        contents_i = {
                            "source": source,
                                     "prompt": prompt,
                                     "target": target,
                                     "source_len": source_len,
@@ -109,8 +120,7 @@
        self.contents = contents
        
        logging.info(
            "total_num of samplers: {}, {}".format(len(self.contents), path))
        logging.info("total_num of samplers: {}, {}".format(len(self.contents), path))
    
    def __len__(self):
        return len(self.contents)
@@ -127,4 +137,3 @@
    def get_target_len(self, data_dict):
        
        return data_dict.get("target_len", 0)