""" 数据集加载 支持tatoeba和cveto数据集 """ import os import sys import random import psutil from typing import List, Tuple, Optional, Dict, Any from dataclasses import dataclass import torch from torch.utils.data import Dataset, DataLoader from tokenizer import Tokenizer def check_memory(): """检查可用内存""" mem = psutil.virtual_memory() available_gb = mem.available / (1024**3) return available_gb @dataclass class TranslationPair: """翻译句对""" zh: str en: str class TranslationDataset(Dataset): """翻译数据集 - 流式处理，内存友好""" def __init__( self, pairs: List[TranslationPair], zh_tokenizer: Tokenizer, en_tokenizer: Tokenizer, max_len: int = 128, cache_tokenized: bool = True, ): self.pairs = pairs self.zh_tokenizer = zh_tokenizer self.en_tokenizer = en_tokenizer self.max_len = max_len # 小缓存，只缓存最近访问的数据 self._cache: Dict[int, Dict[str, Any]] = {} self._cache_size = min(5000, len(pairs) // 10) # 缓存10%或最多5000条 print(f" 数据集: {len(pairs)} 条 (流式处理)") def __len__(self) -> int: return len(self.pairs) def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: # 检查缓存 if idx in self._cache: return self._cache[idx] # 处理数据 pair = self.pairs[idx] zh_ids = self.zh_tokenizer.encode(pair.zh, add_sos=True, add_eos=True)[:self.max_len] en_ids = self.en_tokenizer.encode(pair.en, add_sos=True, add_eos=True)[:self.max_len] result = { 'zh_ids': torch.tensor(zh_ids, dtype=torch.long), 'en_ids': torch.tensor(en_ids, dtype=torch.long), 'zh_len': len(zh_ids), 'en_len': len(en_ids), 'zh_text': pair.zh, 'en_text': pair.en, } # 添加到缓存 if len(self._cache) < self._cache_size: self._cache[idx] = result return result def collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: """批处理函数，动态padding""" zh_ids_list = [item['zh_ids'] for item in batch] en_ids_list = [item['en_ids'] for item in batch] # 找最大长度 max_zh_len = max(len(ids) for ids in zh_ids_list) max_en_len = max(len(ids) for ids in en_ids_list) # Padding zh_padded = torch.zeros(len(batch), max_zh_len, dtype=torch.long) en_padded = torch.zeros(len(batch), max_en_len, dtype=torch.long) zh_lens = [] en_lens = [] for i, (zh_ids, en_ids) in enumerate(zip(zh_ids_list, en_ids_list)): zh_padded[i, :len(zh_ids)] = zh_ids en_padded[i, :len(en_ids)] = en_ids zh_lens.append(len(zh_ids)) en_lens.append(len(en_ids)) return { 'zh_ids': zh_padded, 'en_ids': en_padded, 'zh_lens': torch.tensor(zh_lens, dtype=torch.long), 'en_lens': torch.tensor(en_lens, dtype=torch.long), 'zh_texts': [item['zh_text'] for item in batch], 'en_texts': [item['en_text'] for item in batch], } def load_tatoeba(path: str, max_samples: Optional[int] = None) -> List[TranslationPair]: """加载tatoeba数据集 格式: 编号\t中文\t编号\t英文 """ pairs = [] seen = set() with open(path, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line: continue parts = line.split('\t') if len(parts) < 4: continue zh = parts[1].strip() en = parts[3].strip() # 去重 key = (zh, en) if key in seen: continue seen.add(key) pairs.append(TranslationPair(zh=zh, en=en)) if max_samples and len(pairs) >= max_samples: break return pairs def load_cveto(zh_path: str, en_path: str, max_samples: Optional[int] = None) -> List[TranslationPair]: """加载cveto数据集 两个文件，行号对应 """ pairs = [] # 先统计总行数 print(" 统计文件行数...", end="", flush=True) with open(zh_path, 'r', encoding='utf-8') as f: total_lines = sum(1 for _ in f) print(f" {total_lines:,} 行") print(" 读取数据...", end="", flush=True) last_print = 0 with open(zh_path, 'r', encoding='utf-8') as zh_f, \ open(en_path, 'r', encoding='utf-8') as en_f: for i, (zh_line, en_line) in enumerate(zip(zh_f, en_f)): zh = zh_line.strip() en = en_line.strip() if zh and en: pairs.append(TranslationPair(zh=zh, en=en)) # 每10万行打印一次进度 if i - last_print >= 100000: print(f".{i//100000}", end="", flush=True) last_print = i if max_samples and len(pairs) >= max_samples: break print(f" 完成") return pairs def load_all_data(config) -> Tuple[List[TranslationPair], List[TranslationPair], List[TranslationPair]]: """加载所有数据，返回训练集、验证集、测试集""" print("加载数据集...") # 加载tatoeba tatoeba_path = config.data.tatoeba_path if os.path.exists(tatoeba_path): print(f" 加载 tatoeba: {tatoeba_path}") tatoeba_pairs = load_tatoeba(tatoeba_path, max_samples=config.data.max_samples) print(f" 句对数: {len(tatoeba_pairs)}") else: tatoeba_pairs = [] print(f" 警告: tatoeba路径不存在: {tatoeba_path}") # 合并所有数据 all_pairs = tatoeba_pairs.copy() # 如果还需要更多数据，加载cveto if config.data.max_samples is None or len(all_pairs) < config.data.max_samples: cveto_zh_path = config.data.cveto_zh_path cveto_en_path = config.data.cveto_en_path if os.path.exists(cveto_zh_path) and os.path.exists(cveto_en_path): print(f" 加载 cveto...") remaining = None if config.data.max_samples: remaining = config.data.max_samples - len(all_pairs) cveto_pairs = load_cveto(cveto_zh_path, cveto_en_path, max_samples=remaining) print(f" 句对数: {len(cveto_pairs)}") all_pairs.extend(cveto_pairs) # 过滤长度 print(f" 过滤数据...", end="", flush=True) filtered_pairs = [] total = len(all_pairs) last_print = 0 for i, pair in enumerate(all_pairs): zh_len = len(pair.zh) en_len = len(pair.en) if config.data.min_len <= zh_len <= config.data.max_len and \ config.data.min_len <= en_len <= config.data.max_len: filtered_pairs.append(pair) # 每10万条打印进度 if i - last_print >= 100000: progress = (i + 1) / total * 100 print(f".{progress:.0f}%", end="", flush=True) last_print = i print(f" 完成") print(f" 过滤后句对数: {len(filtered_pairs)}") # 打乱并分割 random.shuffle(filtered_pairs) n = len(filtered_pairs) # 80% 训练, 10% 验证, 10% 测试 train_end = int(n * 0.8) val_end = int(n * 0.9) train_pairs = filtered_pairs[:train_end] val_pairs = filtered_pairs[train_end:val_end] test_pairs = filtered_pairs[val_end:] print(f" 训练集: {len(train_pairs)}") print(f" 验证集: {len(val_pairs)}") print(f" 测试集: {len(test_pairs)}") return train_pairs, val_pairs, test_pairs def create_dataloaders( train_pairs: List[TranslationPair], val_pairs: List[TranslationPair], zh_tokenizer: Tokenizer, en_tokenizer: Tokenizer, config, ) -> Tuple[DataLoader, DataLoader]: """创建数据加载器""" train_dataset = TranslationDataset( train_pairs, zh_tokenizer, en_tokenizer, max_len=config.model.max_len, ) val_dataset = TranslationDataset( val_pairs, zh_tokenizer, en_tokenizer, max_len=config.model.max_len, ) train_loader = DataLoader( train_dataset, batch_size=config.training.batch_size, shuffle=True, collate_fn=collate_fn, num_workers=0, # CPU环境不用多进程 pin_memory=False, ) val_loader = DataLoader( val_dataset, batch_size=config.training.batch_size, shuffle=False, collate_fn=collate_fn, num_workers=0, pin_memory=False, ) return train_loader, val_loader