import os import re from models.bicodec_tokenizer.spark_tokenizer import SparkTokenizer from models.glm_speech_tokenizer.speech_token_extractor import SpeechTokenExtractor from models.glm_speech_tokenizer.modeling_whisper import WhisperVQEncoder from transformers import PreTrainedTokenizer,AutoTokenizer,WhisperFeatureExtractor import torch import torch.nn.functional as F import logging from typing import List, Dict, Any, Literal, Optional, Union from datasets import load_dataset from torch.utils.data import DataLoader def has_punctuation(text: str) -> bool: # 包含中英文符号 pattern = r"[，。！？；：（）“”‘’、,.!?;:()\[\]{}\"']" return bool(re.search(pattern, text)) ALL_TASKS = ["stt", "tts-a", "vc"] class ark_infer_processor: def __init__( self, glm_tokenizer: SpeechTokenExtractor, bicodec_tokenizer: SparkTokenizer, text_tokenizer: PreTrainedTokenizer, max_length: int = 512, glm_semantic_token_offset: int = 151727, semantic_token_offset: int = 172207, global_token_offset: int = 168111, audio_path_name: str = "audio", device: str = "cpu", ): self.glm_tokenizer = glm_tokenizer self.bicodec_tokenizer = bicodec_tokenizer self.text_tokenizer = text_tokenizer self.max_length = max_length self.glm_semantic_token_offset = glm_semantic_token_offset self.semantic_token_offset = semantic_token_offset self.global_token_offset = global_token_offset self.device = device self.audio_path_name = audio_path_name def _process_example_stt(self, audio_path: str): ##target 音频 with torch.no_grad(): glm_semantic_tokens = self.glm_tokenizer.extract([audio_path]) glm_semantic_tokens = torch.as_tensor(glm_semantic_tokens, device="cpu", dtype=torch.long) semantic_tokens = self.bicodec_tokenizer.tokenize([audio_path])['semantic_tokens'] glm_semantic_tokens_list = ( (glm_semantic_tokens + self.glm_semantic_token_offset).cpu().tolist()[0] ) semantic_tokens_list = ( (semantic_tokens + self.semantic_token_offset).cpu().tolist()[0] ) input_ids = ( self.text_tokenizer.encode("<|start_glm_token|>") + glm_semantic_tokens_list + self.text_tokenizer.encode("<|end_glm_token|>") + self.text_tokenizer.encode("<|start_semantic_token|>") + semantic_tokens_list + self.text_tokenizer.encode("<|end_semantic_token|>") + self.text_tokenizer.encode("<|start_content|>") ) attention_mask = [1] * (len(input_ids)) return input_ids, attention_mask def _process_example_tts_a(self, text: str, ref_audio_path: str): with torch.no_grad(): global_tokens = self.bicodec_tokenizer.tokenize([ref_audio_path])['global_tokens'] all_text = "<|start_content|>" + text + "<|end_content|>" global_tokens_list = ( (global_tokens + self.global_token_offset).cpu().tolist()[0][0] ) text_tokens = self.text_tokenizer( all_text, truncation=True, max_length=self.max_length ) input_ids = ( self.text_tokenizer.encode("<|start_global_token|>") + global_tokens_list + self.text_tokenizer.encode("<|end_global_token|>") + text_tokens["input_ids"] ) attention_mask = [1] * len(input_ids) return input_ids, attention_mask def _process_example_vc(self, audio_path: str, ref_audio_path: str): with torch.no_grad(): semantic_tokens = self.bicodec_tokenizer.tokenize([audio_path])['semantic_tokens'] new_global_tokens = self.bicodec_tokenizer.tokenize([ref_audio_path])['global_tokens'] semantic_tokens_list = ( (semantic_tokens + self.semantic_token_offset).cpu().tolist()[0] ) new_global_tokens_list = ( (new_global_tokens + self.global_token_offset).cpu().tolist()[0][0] ) all_str = ( "<|start_global_token|>" + self.text_tokenizer.decode(new_global_tokens_list) + "<|end_global_token|>" + "<|start_semantic_token|>" + self.text_tokenizer.decode(semantic_tokens_list) + "<|end_semantic_token|>" + "<|end_content|>" ) inputs = self.text_tokenizer(all_str) input_ids = inputs["input_ids"] attention_mask = inputs["attention_mask"] return input_ids, attention_mask def process_input( self, task: Literal["stt", "tts-a", "vc"], audio_path: str | None = None, ref_audio_path: str | None = None, text: str | None = None, ): """加载指定音频、特征并根据任务类型返回 token 化结果。""" if task == "stt": assert audio_path is not None input_ids, attention_mask = self._process_example_stt(audio_path) elif task == "tts-a": assert ref_audio_path is not None and text is not None input_ids, attention_mask = self._process_example_tts_a( text, ref_audio_path ) elif task == "vc": assert audio_path is not None and ref_audio_path is not None input_ids, attention_mask = self._process_example_vc( audio_path, ref_audio_path ) else: raise ValueError( f"Unsupported task: {task}, all supported tasks: {ALL_TASKS}" ) return { "input_ids": input_ids, "attention_mask": attention_mask, } class ark_processor: def __init__(self, glm_tokenizer: SpeechTokenExtractor, bicodec_tokenizer: SparkTokenizer, text_tokenizer:PreTrainedTokenizer, max_length:int = 512, glm_semantic_token_offset:int = 151727, semantic_token_offset: int =172207, global_token_offset: int =168111, audio_path_name:str = "audio", device:str ='cpu'): self.glm_tokenizer = glm_tokenizer self.bicodec_tokenizer = bicodec_tokenizer self.text_tokenizer = text_tokenizer self.max_length = max_length self.glm_semantic_token_offset =glm_semantic_token_offset self.semantic_token_offset=semantic_token_offset self.global_token_offset=global_token_offset self.device = device self.audio_path_name =audio_path_name def process_example(self, example: Dict[str, Any]): """ 这个函数由多个CPU进程并行执行。 它负责加载、重采样和对单个样本进行特征提取/分词。 """ task = example.get("task", "stt") audio_path = example.get(self.audio_path_name, "") ref_audio_path = example.get("ref_audio", "") vc_audio = example.get("vc_audio", "") text = example.get("text", "") if task == "stt": ##target 音频 with torch.no_grad(): glm_semantic_tokens = self.glm_tokenizer.extract([audio_path]) glm_semantic_tokens = torch.as_tensor(glm_semantic_tokens, device="cpu", dtype=torch.long) semantic_tokens = self.bicodec_tokenizer.tokenize([audio_path])['semantic_tokens'] glm_semantic_tokens_list = (glm_semantic_tokens + self.glm_semantic_token_offset).cpu().tolist()[0] semantic_tokens_list = (semantic_tokens + self.semantic_token_offset).cpu().tolist()[0] # print(f"len of semantic is {len(semantic_tokens_list)}") ##对text进行token text_tokens = self.text_tokenizer(text, truncation=True, max_length=self.max_length) input_ids = self.text_tokenizer.encode("<|start_glm_token|>") + glm_semantic_tokens_list + self.text_tokenizer.encode("<|end_glm_token|>") \ + self.text_tokenizer.encode("<|start_semantic_token|>") + semantic_tokens_list + self.text_tokenizer.encode( "<|end_semantic_token|>") \ + self.text_tokenizer.encode("<|start_content|>") + text_tokens["input_ids"] + self.text_tokenizer.encode("<|end_content|>") \ + self.text_tokenizer.encode("<|im_end|>") attention_mask = [1] * (len(input_ids)) labels = [-100] * (len(semantic_tokens_list) + 5 + len(glm_semantic_tokens_list)) + text_tokens["input_ids"] + self.text_tokenizer.encode( "<|end_content|>") + self.text_tokenizer.encode("<|im_end|>") elif task == "tts-a": with torch.no_grad(): semantic_tokens = self.bicodec_tokenizer.tokenize([audio_path])['semantic_tokens'] global_tokens = self.bicodec_tokenizer.tokenize([ref_audio_path])['global_tokens'] all_text = "<|start_content|>" + text + "<|end_content|>" global_tokens_list = (global_tokens + self.global_token_offset).cpu().tolist()[0][0] text_tokens = self.text_tokenizer(all_text, truncation=True, max_length=self.max_length) semantic_tokens_list = (semantic_tokens + self.semantic_token_offset).cpu().tolist()[0] input_ids = self.text_tokenizer.encode("<|start_global_token|>") + global_tokens_list + self.text_tokenizer.encode( "<|end_global_token|>") + text_tokens["input_ids"] + semantic_tokens_list + self.text_tokenizer.encode("<|im_end|>") attention_mask = [1] * len(input_ids) labels = [-100] * (len(text_tokens["input_ids"]) + 2 + len(global_tokens_list)) + semantic_tokens_list + self.text_tokenizer.encode("<|im_end|>") elif task == "vc": with torch.no_grad(): semantic_tokens = self.bicodec_tokenizer.tokenize([audio_path])['semantic_tokens'] global_tokens = self.bicodec_tokenizer.tokenize([audio_path])['global_tokens'] # global_tokens, semantic_tokens=self.bicodec_tokenizer.tokenize(audio_path=audio_path) # new_global_tokens, new_semantic_tokens=self.bicodec_tokenizer.tokenize(vc_audio,ref_audio_path) new_semantic_tokens = self.bicodec_tokenizer.tokenize([vc_audio])['semantic_tokens'] new_global_tokens = self.bicodec_tokenizer.tokenize([ref_audio_path])['global_tokens'] global_tokens_list = (global_tokens + self.global_token_offset).cpu().tolist()[0][0] semantic_tokens_list = (semantic_tokens + self.semantic_token_offset).cpu().tolist()[0] new_global_tokens_list = (new_global_tokens + self.global_token_offset).cpu().tolist()[0][0] new_semantic_tokens_list = (new_semantic_tokens + self.semantic_token_offset).cpu().tolist()[0] all_str = "<|start_global_token|>" + self.text_tokenizer.decode(new_global_tokens_list) + "<|end_global_token|>" + "<|start_semantic_token|>" + self.text_tokenizer.decode( semantic_tokens_list) + "<|end_semantic_token|>" + "<|end_content|>" + self.text_tokenizer.decode(new_semantic_tokens_list) + "<|im_end|>" ##add token and mask inputs = self.text_tokenizer(all_str) input_ids = inputs['input_ids'] attention_mask = inputs['attention_mask'] labels = [-100] * (5 + len(new_global_tokens_list) + len(semantic_tokens_list)) + new_semantic_tokens_list + self.text_tokenizer.encode("<|im_end|>") else: ##默认走stt with torch.no_grad(): glm_semantic_tokens = self.glm_tokenizer.extract([audio_path]) glm_semantic_tokens = torch.as_tensor(glm_semantic_tokens, device="cpu", dtype=torch.long) semantic_tokens = self.bicodec_tokenizer.tokenize([audio_path])['semantic_tokens'] glm_semantic_tokens_list = (glm_semantic_tokens+self.glm_semantic_token_offset).cpu().tolist()[0] semantic_tokens_list = (semantic_tokens+self.semantic_token_offset).cpu().tolist()[0] # print(f"len of semantic is {len(semantic_tokens_list)}") ##对text进行token text_tokens = self.text_tokenizer(text, truncation=True, max_length=self.max_length) input_ids = self.text_tokenizer.encode("<|start_glm_token|>")+ glm_semantic_tokens_list + self.text_tokenizer.encode("<|end_glm_token|>") \ + self.text_tokenizer.encode("<|start_semantic_token|>")+ semantic_tokens_list + self.text_tokenizer.encode("<|end_semantic_token|>") \ + text_tokens["input_ids"] \ + self.text_tokenizer.encode("<|im_end|>") attention_mask = [1]*(len(semantic_tokens_list)+4+len(glm_semantic_tokens_list)) +text_tokens["attention_mask"] +[1] labels = [-100]*(len(semantic_tokens_list)+4+len(glm_semantic_tokens_list))+ text_tokens["input_ids"]+ self.text_tokenizer.encode("<|im_end|>") return { "input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, } def create_tts_collate_fn( pad_token_id: int, processor, # ark_processor max_length: Optional[int]=None,# 传入你想要的截断上限，例如 512 truncation_side: str = "right" # "right" 或 "left"，默认右截断 ): """ 手动填充 + 可选截断的 collate_fn 工厂。 参数： pad_token_id: 用于 input_ids 的 pad 值 processor: 你的 ark_processor，需提供 .process_example() max_length: 若提供，则对每个样本在拼批前先截断到该长度 truncation_side: "right" | "left"，决定从哪侧截断 """ label_pad_value = -100 attention_mask_pad_value = 0 def _truncate_1d(x: torch.Tensor, keep_len: int, side: str) -> torch.Tensor: if x.numel() <= keep_len: return x if side == "right": return x[:keep_len] elif side == "left": return x[-keep_len:] else: raise ValueError(f"Unsupported truncation_side: {side}") def _to_long_tensor(x) -> torch.Tensor: if isinstance(x, torch.Tensor): return x.detach().clone().long() return torch.tensor(x, dtype=torch.long) def collate_fn(examples: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: # 1) 预处理（过滤空样本） proc = [processor.process_example(ex) for ex in examples if ex] proc = [d for d in proc if d and ("input_ids" in d) and ("attention_mask" in d) and ("labels" in d)] if len(proc) == 0: # 返回空批，避免 DataLoader 崩溃 return { "input_ids": torch.empty(0, dtype=torch.long), "attention_mask": torch.empty(0, dtype=torch.long), "labels": torch.empty(0, dtype=torch.long), } # 2) 样本级截断（如果设置了 max_length） if max_length is not None: trimmed = [] for ex in proc: ids = _to_long_tensor(ex["input_ids"]) mask = _to_long_tensor(ex["attention_mask"]) labs = _to_long_tensor(ex["labels"]) keep_len = min(max_length, ids.numel()) ids = _truncate_1d(ids, keep_len, truncation_side) mask = _truncate_1d(mask, keep_len, truncation_side) labs = _truncate_1d(labs, keep_len, truncation_side) trimmed.append({"input_ids": ids, "attention_mask": mask, "labels": labs}) proc = trimmed # 3) 计算本批最大长度（截断后再取最大） max_len_in_batch = max(int(len(ex["input_ids"])) for ex in proc) # 4) 逐样本右侧 pad 到 batch 最大长度 padded_input_ids_list = [] padded_attention_mask_list = [] padded_labels_list = [] for ex in proc: ids = _to_long_tensor(ex["input_ids"]) mask = _to_long_tensor(ex["attention_mask"]) labs = _to_long_tensor(ex["labels"]) need = max_len_in_batch - ids.numel() if need < 0: # 极端情况：有人为 max_length=None 时超长样本溢出 keep_len = max_len_in_batch ids = _truncate_1d(ids, keep_len, "right") mask = _truncate_1d(mask, keep_len, "right") labs = _truncate_1d(labs, keep_len, "right") need = 0 pad_dims = (0, need) ids = F.pad(ids, pad_dims, mode="constant", value=pad_token_id) mask = F.pad(mask, pad_dims, mode="constant", value=attention_mask_pad_value) labs = F.pad(labs, pad_dims, mode="constant", value=label_pad_value) padded_input_ids_list.append(ids) padded_attention_mask_list.append(mask) padded_labels_list.append(labs) # 5) 堆叠成批 batch = { "input_ids": torch.stack(padded_input_ids_list, dim=0), "attention_mask": torch.stack(padded_attention_mask_list, dim=0), "labels": torch.stack(padded_labels_list, dim=0), } return batch return collate_fn if __name__ == "__main__": device = "cuda:0" bicodec_audio_tokenizer_path = "/data/arki_production/model/SparkAudio/Spark-TTS-0___5B/" glm_speech_tokenizer_path = "/data/yumu/model/glm-4-voice-tokenizer" feature_extractor = WhisperFeatureExtractor.from_pretrained(glm_speech_tokenizer_path) audio_model = WhisperVQEncoder.from_pretrained(glm_speech_tokenizer_path).eval().to(device) glm_tokenizer = SpeechTokenExtractor(model=audio_model, feature_extractor=feature_extractor, device=device) text_tokenizer = AutoTokenizer.from_pretrained("/data/yumu/model/ark_audio_v1_0_3_b",trust_remote_code=True) bicodec_tokenizer = SparkTokenizer(model_path=bicodec_audio_tokenizer_path, device=device) # 配置项 DATASET_PATH = "/data/yumu/glm_asr_vllm/test/data/test_meeting.jsonl" MAX_LENGTH = 4096 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" print(f"将使用设备: {DEVICE}") # --- 2. 加载流式数据集 --- print(f"以流式方式加载数据集 '{DATASET_PATH}'...") streaming_dataset = load_dataset("json", data_files=DATASET_PATH, streaming=True)['train'] # --- 4. 构建数据处理流水线 (Pipeline) --- print("正在对数据流进行shuffle，buffer_size=1000...") shuffled_dataset = streaming_dataset.shuffle(buffer_size=10000, seed=42) processor = ark_processor( glm_tokenizer=glm_tokenizer, bicodec_tokenizer=bicodec_tokenizer, text_tokenizer=text_tokenizer, device = DEVICE, audio_path_name="audio") collate_fn = create_tts_collate_fn(text_tokenizer.pad_token_id,processor,max_length=4096) # 创建最终的DataLoader data_loader = DataLoader( shuffled_dataset, batch_size=10, # 根据你的GPU显存和模型大小调整 collate_fn=collate_fn, num_workers=0 # DataLoader的worker，负责从打乱后的流中拉取数据 ) print("\n--- 高性能流式 DataLoader 演示 ---") print("将从DataLoader中获取并展示第一个批次的数据：\n") first_batch = next(iter(data_loader)) print("成功获取第一个批次！数据已在collate_fn中填充。") for key, value in first_batch.items(): if value is not None: # print(f" - {key}: shape={value.shape}, dtype={value.dtype}") print(f" - {key}: shape={value.shape}, dtype={value}")