""" Sentence-BERT训练数据准备脚本 从QA数据集构建语义相似度训练数据 """ import json import random from pathlib import Path from typing import List, Dict, Tuple from datasets import load_from_disk import numpy as np class SBERTDataPreparator: """SBERT训练数据准备器""" def __init__(self, qa_dataset_path: str, output_dir: str): """ Args: qa_dataset_path: QA数据集路径 output_dir: 输出目录 """ self.qa_dataset_path = qa_dataset_path self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) # 加载QA数据集 print(f"加载QA数据集: {qa_dataset_path}") self.qa_dataset = load_from_disk(qa_dataset_path) print(f"数据集大小: {len(self.qa_dataset['train'])}") def prepare_training_data( self, num_negatives: int = 5, hard_negative_ratio: float = 0.3, train_ratio: float = 0.7, val_ratio: float = 0.15 ) -> Tuple[List[Dict], List[Dict], List[Dict]]: """ 准备训练数据 Args: num_negatives: 每个正样本的负样本数量 hard_negative_ratio: 困难负样本的比例 train_ratio: 训练集比例 val_ratio: 验证集比例 Returns: (train_data, val_data, test_data) """ print("\n准备训练数据...") # 转换为列表便于处理 qa_list = list(self.qa_dataset['train']) # 构建正样本对 positive_pairs = self._create_positive_pairs(qa_list) # 构建负样本 all_pairs = self._add_negatives( positive_pairs, qa_list, num_negatives=num_negatives, hard_negative_ratio=hard_negative_ratio ) # 打乱数据 random.shuffle(all_pairs) # 划分数据集 total = len(all_pairs) train_end = int(total * train_ratio) val_end = int(total * (train_ratio + val_ratio)) train_data = all_pairs[:train_end] val_data = all_pairs[train_end:val_end] test_data = all_pairs[val_end:] print(f"\n数据集划分:") print(f" 训练集: {len(train_data)} 样本") print(f" 验证集: {len(val_data)} 样本") print(f" 测试集: {len(test_data)} 样本") return train_data, val_data, test_data def _create_positive_pairs(self, qa_list: List[Dict]) -> List[Dict]: """创建正样本对 (question, answer_context)""" positive_pairs = [] for qa in qa_list: question = qa.get('question', '').strip() answer_context = qa.get('answer_context', '').strip() if question and answer_context: positive_pairs.append({ 'anchor': question, 'positive': answer_context, 'label': 1 # 相似 }) print(f"创建正样本对: {len(positive_pairs)}") return positive_pairs def _add_negatives( self, positive_pairs: List[Dict], qa_list: List[Dict], num_negatives: int = 5, hard_negative_ratio: float = 0.3 ) -> List[Dict]: """ 添加负样本 Args: positive_pairs: 正样本对 qa_list: 所有QA数据 num_negatives: 负样本数量 hard_negative_ratio: 困难负样本比例 """ print(f"\n添加负样本 (每样本 {num_negatives} 个负样本)...") all_answers = [qa.get('answer_context', '').strip() for qa in qa_list] all_answers = [a for a in all_answers if a] extended_pairs = [] for pair in positive_pairs: anchor = pair['anchor'] positive = pair['positive'] # 添加原始正样本 extended_pairs.append(pair) # 生成负样本 hard_neg_count = int(num_negatives * hard_negative_ratio) random_neg_count = num_negatives - hard_neg_count # 困难负样本: 同领域但不同的答案 hard_negatives = self._sample_hard_negatives( anchor, all_answers, n=hard_neg_count, exclude=positive ) # 随机负样本 random_negatives = self._sample_random_negatives( all_answers, n=random_neg_count, exclude=positive ) # 添加负样本对 for neg in hard_negatives + random_negatives: extended_pairs.append({ 'anchor': anchor, 'positive': neg, # 在SBERT训练中作为负样本 'label': 0 # 不相似 }) print(f"总样本数: {len(extended_pairs)}") return extended_pairs def _sample_hard_negatives( self, anchor: str, all_answers: List[str], n: int, exclude: str ) -> List[str]: """采样困难负样本（简单实现：随机采样后可改进）""" candidates = [a for a in all_answers if a != exclude] if len(candidates) <= n: return candidates return random.sample(candidates, n) def _sample_random_negatives( self, all_answers: List[str], n: int, exclude: str ) -> List[str]: """采样随机负样本""" candidates = [a for a in all_answers if a != exclude] if len(candidates) <= n: return candidates return random.sample(candidates, n) def save_data( self, train_data: List[Dict], val_data: List[Dict], test_data: List[Dict], format: str = 'jsonl' ): """保存数据到文件""" print(f"\n保存数据到 {self.output_dir}...") if format == 'jsonl': # JSONL格式 (适合sentence-transformers) self._save_jsonl(train_data, 'train.jsonl') self._save_jsonl(val_data, 'val.jsonl') self._save_jsonl(test_data, 'test.jsonl') elif format == 'csv': # CSV格式 import pandas as pd pd.DataFrame(train_data).to_csv( self.output_dir / 'train.csv', index=False ) pd.DataFrame(val_data).to_csv( self.output_dir / 'val.csv', index=False ) pd.DataFrame(test_data).to_csv( self.output_dir / 'test.csv', index=False ) print("✓ 数据保存完成") def _save_jsonl(self, data: List[Dict], filename: str): """保存为JSONL格式""" filepath = self.output_dir / filename with open(filepath, 'w', encoding='utf-8') as f: for item in data: f.write(json.dumps(item, ensure_ascii=False) + '\n') print(f" 保存: {filepath} ({len(data)} 样本)") def print_statistics(self, train_data: List[Dict], val_data: List[Dict], test_data: List[Dict]): """打印数据统计信息""" print("\n=== 数据统计 ===") # 正负样本比例 for name, data in [('训练集', train_data), ('验证集', val_data), ('测试集', test_data)]: pos_count = sum(1 for item in data if item.get('label') == 1) neg_count = len(data) - pos_count print(f"\n{name}:") print(f" 总样本: {len(data)}") print(f" 正样本: {pos_count} ({pos_count/len(data)*100:.1f}%)") print(f" 负样本: {neg_count} ({neg_count/len(data)*100:.1f}%)") # 文本长度统计 all_anchors = [item['anchor'] for item in train_data] all_positives = [item['positive'] for item in train_data] anchor_lengths = [len(a.split()) for a in all_anchors] positive_lengths = [len(p.split()) for p in all_positives] print(f"\n文本长度统计 (训练集):") print(f" Anchor: 平均 {np.mean(anchor_lengths):.1f} 词, " f"最大 {max(anchor_lengths)}, 最小 {min(anchor_lengths)}") print(f" Positive: 平均 {np.mean(positive_lengths):.1f} 词, " f"最大 {max(positive_lengths)}, 最小 {min(positive_lengths)}") def main(): """主函数""" import argparse parser = argparse.ArgumentParser(description='准备SBERT训练数据') parser.add_argument( '--qa_dataset', type=str, default='hr-multiwoz-dataset/qa_dataset', help='QA数据集路径' ) parser.add_argument( '--output_dir', type=str, default='data/processed/sbert', help='输出目录' ) parser.add_argument( '--num_negatives', type=int, default=5, help='每个正样本的负样本数量' ) parser.add_argument( '--hard_negative_ratio', type=float, default=0.3, help='困难负样本比例' ) parser.add_argument( '--format', type=str, default='jsonl', choices=['jsonl', 'csv'], help='输出格式' ) args = parser.parse_args() # 设置随机种子 random.seed(42) np.random.seed(42) # 创建数据准备器 preparator = SBERTDataPreparator( qa_dataset_path=args.qa_dataset, output_dir=args.output_dir ) # 准备数据 train_data, val_data, test_data = preparator.prepare_training_data( num_negatives=args.num_negatives, hard_negative_ratio=args.hard_negative_ratio ) # 保存数据 preparator.save_data(train_data, val_data, test_data, format=args.format) # 打印统计 preparator.print_statistics(train_data, val_data, test_data) print("\n✓ 数据准备完成!") print(f"\n输出目录: {args.output_dir}") print(f"下一步: 运行训练脚本") print(f" python scripts/train_sbert.py --train_data {args.output_dir}/train.jsonl") if __name__ == '__main__': main()