data/utils/llm_dataset.py

import datasets
import os
import json
import torch
import random
import time
random.seed(time.time())
import logging
from tqdm import tqdm
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def verify_jsonl_files(data_files):
    """检查每个 jsonl 文件的有效性"""
    invalid_files = []
    
    for file_path in tqdm(data_files, desc="验证文件"):
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                for i, line in enumerate(f):
                    try:
                        json.loads(line)
                    except json.JSONDecodeError:
                        invalid_files.append((file_path, i+1))
                        logging.error(f"文件 {file_path} 在第 {i+1} 行有无效的 JSON")
                        break
        except Exception as e:
            invalid_files.append((file_path, f"读取错误: {str(e)}"))
            logging.error(f"无法读取文件 {file_path}: {str(e)}")
    
    return invalid_files
def load_jsonl_dataset(directory,tokenizer):
    '''
    load jsonl files in a directory recursively
    '''
    data_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.jsonl'):
                data_files.append(os.path.join(root, file))
    
    logging.info(f"找到 {len(data_files)} 个 JSONL 文件")
    # 验证文件
    invalid_files = verify_jsonl_files(data_files)
    if invalid_files:
        logging.error(f"发现 {len(invalid_files)} 个无效文件:")
        for file_info in invalid_files:
            if isinstance(file_info[1], int):
                logging.error(f"  - {file_info[0]} (错误在第 {file_info[1]} 行)")
            else:
                logging.error(f"  - {file_info[0]} ({file_info[1]})")
        
        # 移除无效文件
        valid_files = [f for f in data_files if f not in [info[0] for info in invalid_files]]
        logging.info(f"继续处理剩余的 {len(valid_files)} 个有效文件")
        data_files = valid_files
    # 手动收集所有样本，确保特征一致性
    all_samples = []
    
    for file_path in tqdm(data_files, desc="加载数据集"):
        try:
            # 手动解析JSONL文件，避免datasets加载时的类型推断问题
            with open(file_path, 'r', encoding='utf-8') as f:
                for line in f:
                    try:
                        data = json.loads(line)
                        # 确保所有字段存在且类型一致
                        llm_prompt_speech_token = data.get('llm_prompt_speech_token', [])
                        tts_speech_tokens = data.get('tts_speech_tokens', [])
                        text = str(data.get('text', ""))
                        prompt_text = str(data.get('prompt_text', ""))
                        
                        # 确保列表类型
                        if not isinstance(llm_prompt_speech_token, list):
                            llm_prompt_speech_token = []
                        if not isinstance(tts_speech_tokens, list):
                            tts_speech_tokens = []
                            
                        # 添加处理后的样本
                        all_samples.append({
                            'llm_prompt_speech_token': llm_prompt_speech_token,
                            'tts_speech_tokens': tts_speech_tokens,
                            'text': text,
                            'prompt_text': prompt_text
                        })
                    except json.JSONDecodeError:
                        continue  # 跳过无效的JSON行
                    except Exception as e:
                        logging.error(f"处理样本时出错: {str(e)}")
        except Exception as e:
            logging.error(f"打开文件 {file_path} 时出错: {str(e)}")
    
    if not all_samples:
        raise ValueError("没有成功加载任何样本")
    
    # 创建数据集
    logging.info(f"手动创建数据集，包含 {len(all_samples)} 个样本")
    dataset = datasets.Dataset.from_list(all_samples)
    
    logging.info(f"成功加载 {len(dataset)} 个样本")
    
    #1. concatenate llm_prompt_speech_token and tts_speech_tokens (list of int)
    #delay the concatenation to collate_fn since sometimes we want to drop the prompt
    # dataset = dataset.map(lambda x: {'speech_token': x['llm_prompt_speech_token'] + x['tts_speech_tokens']},remove_columns=['tts_speech_tokens','llm_prompt_speech_token'])
    #2. Filter the data either :
    #   1. the length of the speech_token is less than 1
    #   2. the length of the speech_token is greater than 1000
    #   3. the length of the text is greater than 500
    #   4. the length of the prompt_text is greater than 500
    #   5. the length of the text_token is less than 1
    #   6. the length of the prompt_text_token is less than 1
    dataset = dataset.filter(lambda x:len(x['llm_prompt_speech_token']) < 2048 and len(x['tts_speech_tokens']) < 2048
                             and len(tokenizer.encode(x['text'])) < 2048 and len(tokenizer.encode(x['prompt_text'])) < 2048 )
    logging.info(f"过滤后剩余 {len(dataset)} 个样本")
    #2. tokenize the text to text_tokens and prompt_text to prompt_text_tokens
    # dataset = dataset.map(lambda x: {'text_tokens': tokenizer.encode(x['text']), 'prompt_text_tokens': tokenizer.encode(x['prompt_text'])},remove_columns=['text','prompt_text'])
    return dataset

def collate_fn(batch, tokenizer, pad_to_max_length=True, max_length=2048, drop_prompt_audio_rate=-0.1):
    '''
    convert the data to torch tensors
    1. call tokenizer.encode('text') and tokenizer.encode('prompt_text'), concatenate them to get the text_token, record each sample's length to text_token_len
    2. convert the text_tokens and text_token_len to torch tensor
    3. record each sample's speech_token length to speech_token_len
    4. convert the speech_token and speech_token_len to torch tensor
    5. We will drop prompt with drop_prompt_audio_rate to ask model to learn generate audio without guaidance
    By default we won't drop anything
    '''
    all_text_tokens = []
    all_speech_tokens = []
    speech_token_len = []
    text_token_len = []
    my_max_length = 0
    is_drop_prompt = random.random() < drop_prompt_audio_rate
    
    for sample in batch:
        tts_speech_tokens = sample['tts_speech_tokens']
        llm_prompt_speech_token = sample['llm_prompt_speech_token']
        
        if is_drop_prompt:
            # 只使用文本部分，不使用提示
            text_tokens = tokenizer.encode(sample['text'])
            all_text_tokens.append(torch.tensor(text_tokens, dtype=torch.int32))
            text_token_len.append(len(text_tokens))
            
            # 只使用语音部分，不使用提示语音
            current_speech_tokens = tts_speech_tokens
            all_speech_tokens.append(torch.tensor(current_speech_tokens, dtype=torch.int32))
            speech_token_len.append(len(current_speech_tokens))
            
            total_length = len(text_tokens) + len(current_speech_tokens)
        else:
            # 使用提示+文本
            text_tokens = tokenizer.encode(sample['text'])
            prompt_tokens = tokenizer.encode(sample['prompt_text'])
            combined_text_tokens = prompt_tokens + text_tokens
            all_text_tokens.append(torch.tensor(combined_text_tokens, dtype=torch.int32))
            text_token_len.append(len(combined_text_tokens))
            
            # 使用提示语音+语音
            current_speech_tokens = llm_prompt_speech_token + tts_speech_tokens
            all_speech_tokens.append(torch.tensor(current_speech_tokens, dtype=torch.int32))
            speech_token_len.append(len(current_speech_tokens))
            
            total_length = len(combined_text_tokens) + len(current_speech_tokens)
            
        if total_length > my_max_length:
            my_max_length = total_length
    
    # 检查长度是否超出最大长度
    skip = my_max_length > max_length
    
    # 将列表转换为填充后的张量
    all_text_tokens = torch.nn.utils.rnn.pad_sequence(all_text_tokens, batch_first=True, padding_value=0)
    all_speech_tokens = torch.nn.utils.rnn.pad_sequence(all_speech_tokens, batch_first=True, padding_value=0)
    
    # 如果需要填充到最大长度
    if pad_to_max_length and not skip:
        pad_length = max_length - my_max_length
        if pad_length > 0:
            all_speech_tokens = torch.nn.functional.pad(all_speech_tokens, (0, pad_length), value=0)
    
    return {
        'text_token': all_text_tokens, 
        'text_token_len': torch.tensor(text_token_len, dtype=torch.int32), 
        'speech_token': all_speech_tokens,  # 确保命名一致
        'speech_token_len': torch.tensor(speech_token_len, dtype=torch.int32),
        'skip': skip
    }
        
        
if __name__ == '__main__':
    from transformers import AutoTokenizer
    model_path = "/external_data/models/rwkv7-2.9B-world"
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    directory = '/external_data/yueyudata/speech_corpus'
    dataset = load_jsonl_dataset(directory,tokenizer)
    print(dataset)
    print(dataset[0])
    from functools import partial
    collate_fn = partial(collate_fn,tokenizer=tokenizer,pad_to_max_length=False)
    dataloader = torch.utils.data.DataLoader(dataset,batch_size=1,collate_fn=collate_fn)
    for data in dataloader:
        print(data)
        print(data['speech_token'].shape)
        print(data['text_token'].shape)
        break