#!/usr/bin/env python3 """ Data Processing Module for ASR Pipeline This module handles data preparation and processing for the ASR pipeline, including downloading datasets, creating data splits, and preparing input files for inference. """ import logging import os import shutil import subprocess import urllib.request import tarfile from pathlib import Path from typing import Dict, List, Tuple data_tag = ["dev"] class DataProcessor: """Base class for data processing""" def __init__(self, config: Dict): self.config = config self.logger = logging.getLogger(__name__) self.data_dir = Path(config.get('data_dir', './data')) def prepare_data(self) -> Dict: """Prepare data for inference - to be implemented by subclasses""" raise NotImplementedError def load_existing_data(self) -> Dict: """Load existing prepared data""" raise NotImplementedError class AISHELLDataProcessor(DataProcessor): """Data processor for AISHELL dataset""" def __init__(self, config: Dict): super().__init__(config) # 使用现有的数据目录结构 self.aishell_dir = Path(config.get('aishell_data_dir', '/data/datasets/0')) self.raw_data_dir = self.aishell_dir / 'data_aishell' self.processed_data_dir = self.data_dir / 'processed' def prepare_data(self) -> Dict: """Prepare AISHELL dataset for inference using existing data""" self.logger.info("Preparing AISHELL dataset from existing data") # 检查现有数据是否存在 if not self._check_data_exists(): raise FileNotFoundError(f"AISHELL dataset not found at {self.raw_data_dir}") self.logger.info("AISHELL dataset found, processing data") # 创建处理后的数据目录 self.processed_data_dir.mkdir(parents=True, exist_ok=True) # 处理数据 data_info = self._process_aishell_data() return data_info def load_existing_data(self) -> Dict: """Load existing prepared AISHELL data""" self.logger.info("Loading existing AISHELL data") # 首先检查是否有处理后的数据 if self._check_processed_data_exists(): return self._load_processed_data() else: # 如果没有处理后的数据，直接从原始数据加载 self.logger.info("No processed data found, loading from raw data") return self.prepare_data() def _check_data_exists(self) -> bool: """Check if AISHELL dataset exists at the specified location""" required_files = [ self.raw_data_dir / 'wav', self.raw_data_dir / 'transcript' / 'aishell_transcript_v0.8.txt' ] return all(f.exists() for f in required_files) def _check_processed_data_exists(self) -> bool: """Check if processed data exists""" test_sets = self.config.get('test_sets', data_tag) required_files = [] for test_set in test_sets: required_files.extend([ self.processed_data_dir / test_set / 'wav.scp', self.processed_data_dir / test_set / 'text' ]) return all(f.exists() for f in required_files) def _process_aishell_data(self) -> Dict: """Process AISHELL data into Kaldi-style format using existing dataset structure""" test_sets = self.config.get('test_sets', data_tag) data_info = {'test_files': {}, 'references': {}} # 加载转录文件 transcript_file = self.raw_data_dir / 'transcript' / 'aishell_transcript_v0.8.txt' if not transcript_file.exists(): raise FileNotFoundError(f"Transcript file not found: {transcript_file}") # 解析转录文件 transcript_dict = self._parse_transcript_file(transcript_file) for test_set in test_sets: self.logger.info(f"Processing {test_set} set") # 创建输出目录 output_dir = self.processed_data_dir / test_set output_dir.mkdir(parents=True, exist_ok=True) # 根据数据集类型设置音频文件路径 audio_dir = self.raw_data_dir / 'wav' / test_set if not audio_dir.exists(): self.logger.warning(f"Audio directory not found: {audio_dir}") continue # 查找音频文件 audio_files = list(audio_dir.rglob('*.wav')) self.logger.info(f"Found {len(audio_files)} audio files for {test_set}") # 创建 wav.scp 和 text 文件 wav_scp_file = output_dir / 'wav.scp' text_file = output_dir / 'text' with open(wav_scp_file, 'w', encoding='utf-8') as wav_f, \ open(text_file, 'w', encoding='utf-8') as text_f: for audio_file in audio_files: # 从文件名提取 utterance ID # AISHELL 的文件名格式如: BAC009S0724W0121.wav utt_id = audio_file.stem # 写入 wav.scp wav_f.write(f"{utt_id} {audio_file.resolve()}\n") # 从转录字典中获取转录文本 transcription = transcript_dict.get(utt_id, "") if not transcription: self.logger.warning(f"No transcription found for {utt_id}") transcription = "[NO_TRANSCRIPTION]" # 写入 text 文件 text_f.write(f"{utt_id} {transcription}\n") # 存储文件信息 data_info['test_files'][utt_id] = str(audio_file.resolve()) data_info['references'][utt_id] = transcription self.logger.info(f"Created {wav_scp_file} and {text_file}") return data_info def _parse_transcript_file(self, transcript_file: Path) -> Dict[str, str]: """Parse AISHELL transcript file into dictionary""" transcript_dict = {} try: with open(transcript_file, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line: continue # AISHELL 转录文件格式: 文件名 转录文本 # 例如: BAC009S0724W0121 而 对 楼市 成交 抑制 最 明显 的 限 购 parts = line.split(maxsplit=1) if len(parts) == 2: utt_id, transcription = parts transcript_dict[utt_id] = transcription else: self.logger.warning(f"Invalid line in transcript file: {line}") self.logger.info(f"Parsed {len(transcript_dict)} transcriptions") except Exception as e: self.logger.error(f"Error parsing transcript file: {e}") raise return transcript_dict def _load_processed_data(self) -> Dict: """Load processed data from files""" test_sets = self.config.get('test_sets', data_tag) data_info = {'test_files': {}, 'references': {}} for test_set in test_sets: wav_scp_file = self.processed_data_dir / test_set / 'wav.scp' text_file = self.processed_data_dir / test_set / 'text' if not wav_scp_file.exists() or not text_file.exists(): self.logger.warning(f"Missing files for {test_set} set") continue # Load wav.scp with open(wav_scp_file, 'r', encoding='utf-8') as f: for line in f: utt_id, audio_path = line.strip().split(maxsplit=1) data_info['test_files'][utt_id] = audio_path # Load text with open(text_file, 'r', encoding='utf-8') as f: for line in f: parts = line.strip().split(maxsplit=1) if len(parts) == 2: utt_id, transcription = parts data_info['references'][utt_id] = transcription self.logger.info(f"Loaded {len(data_info['test_files'])} test files") return data_info def split_data_for_parallel_processing(data_info: Dict, nj: int) -> List[Dict]: """ Split data for parallel processing Args: data_info: Dictionary containing test files and references nj: Number of parallel jobs Returns: List of data_info dictionaries for each job """ if nj <= 1: return [data_info] # Split data evenly test_files = list(data_info['test_files'].items()) chunk_size = len(test_files) // nj splits = [] for i in range(nj): start_idx = i * chunk_size end_idx = start_idx + chunk_size if i < nj - 1 else len(test_files) split_files = dict(test_files[start_idx:end_idx]) split_references = {k: data_info['references'][k] for k in split_files.keys()} splits.append({ 'test_files': split_files, 'references': split_references }) return splits if __name__ == '__main__': # Test the data processor config = { 'data_dir': './data', 'dataset': 'aishell', 'test_sets': data_tag } processor = AISHELLDataProcessor(config) data_info = processor.prepare_data() print(f"Prepared {len(data_info['test_files'])} test files")