数据初始化使用了缓存

.gitignore

@@ -5,3 +5,5 @@ wandb/
 **/*.log
 models/sentence_transformers/
 models/sentence_transformers_cache/
+**/*.pyc
+qwen2-1.7B/

preprocessing/README_trex_processor.md

@@ -4,10 +4,35 @@
 1. **句子提取**：从 TREx 数据集提取三元组并转换为自然语言句子
 2. **LLM 处理**：使用 ollama qwen3:4b 模型进行句子修正和重要性评分
 
+## 🆕 防卡死机制
+
+为了解决LLM处理时可能出现的卡死问题，新增了以下功能：
+
+### 超时和重试机制
+- **超时时间**：每个LLM请求60秒超时
+- **重试机制**：失败后最多重试2次，采用指数退避策略
+- **并发控制**：降低并发数至4个，减少服务器压力
+
+### 心跳监控系统
+- **实时监控**：每30秒检查一次LLM响应状态
+- **异常警告**：超过30秒无成功响应时发出警告
+- **服务检测**：自动检查ollama服务状态
+- **详细统计**：实时显示成功率、超时率等统计信息
+
+### 日志系统
+- **详细日志**：所有操作都记录在 `logs/` 目录下
+- **双重输出**：同时输出到日志文件和控制台
+- **时间戳标记**：日志文件包含启动时间戳
+
+### 改进的错误处理
+- **异常恢复**：LLM处理失败时使用原句子和默认评分
+- **状态监控**：处理前检查ollama服务状态
+- **批次间休息**：批次之间休息5秒，避免过度压力
+
 ## 安装依赖
 
 ```bash
-pip install agno asyncio pydantic
+pip install agno asyncio pydantic requests
 ```
 
 确保已安装并启动 ollama，并下载 qwen3:4b 模型：
@@ -50,24 +75,52 @@ python trex_to_sentences_simple.py --step llm --sentences_json my_sentences.json
 
 ## 输出文件
 
-**注意：所有输出文件都会自动保存在 `./output/` 目录中**
+**注意：所有输出文件都会自动保存在相应目录中**
 
-### 步骤1输出
+### 句子提取输出
 - `output/extracted_sentences.json`: 提取的原始句子，包含元数据
 
-### 步骤2输出  
+### LLM处理输出  
 - `output/{output_file}.txt`: 修正后的句子文本文件
 - `output/{output_file}.json`: 完整的处理结果（包含原句、修正句、评分）
 - `output/{output_file}_sorted_by_importance.txt`: 按重要性评分排序的句子
 
 ### 检查点文件
-- `output/{output_file}_checkpoint_{数量}.json`: 每2000条句子自动保存的检查点
+- `output/{output_file}_checkpoint_{数量}.json`: 每1000条句子自动保存的检查点
+
+### 日志文件
+- `logs/trex_processor_{时间戳}.log`: 详细的处理日志
+
+## 🆕 故障诊断
+
+### 如果遇到卡死问题：
+
+1. **检查日志文件**：查看 `logs/` 目录下的最新日志
+2. **观察心跳监控**：注意控制台的心跳警告信息
+3. **检查ollama服务**：
+   ```bash
+   ps aux | grep ollama
+   curl http://localhost:11434/api/tags
+   ```
+4. **重启ollama服务**（如果需要）：
+   ```bash
+   pkill ollama
+   ollama serve &
+   ```
+
+### 常见警告信息：
+
+- `⚠️ 心跳检测`: 30秒无成功响应（正常情况下会自动恢复）
+- `❌ 严重警告`: 90秒无成功响应（可能需要检查服务）
+- `💀 Ollama服务异常`: ollama服务可能已停止
+- `💀 致命错误`: 连续多次警告（建议重启程序）
 
 ## 检查点恢复机制
 
 - 步骤2会自动检测已有的检查点文件（在 `output/` 目录中）
 - 只处理尚未处理的句子，避免重复工作
 - 如果所有句子都已处理，会直接生成最终输出文件
+- 中断后重新运行会自动从最新检查点继续
 
 ## 示例工作流
 
@@ -84,10 +137,11 @@ python trex_to_sentences_simple.py --step llm
 
 ## 性能特点
 
-- **并发处理**: 最大54个并发LLM请求
-- **检查点保存**: 每2000条句子自动保存，支持断点续传
-- **进度显示**: 详细的处理进度和时间预估
-- **错误处理**: LLM请求失败时使用原句子和默认评分
+- **保守的并发**: 最大4个并发LLM请求（降低卡死风险）
+- **检查点保存**: 每1000条句子自动保存，支持断点续传
+- **智能监控**: 详细的处理进度和时间预估
+- **健壮的错误处理**: LLM请求失败时使用原句子和默认评分
+- **服务监控**: 自动检测ollama服务状态
 
 ## 注意事项
 
@@ -95,3 +149,6 @@ python trex_to_sentences_simple.py --step llm
 2. 检查点文件会占用额外磁盘空间（每个都包含所有已处理数据）  
 3. LLM处理速度取决于模型性能和网络状况
 4. 建议先用`--max_files`参数测试小批量数据
+5. **新增**：如果遇到卡死，查看日志文件和心跳监控信息
+6. **新增**：程序会自动检测并报告ollama服务状态
+7. **新增**：所有处理过程都有详细日志记录，便于问题诊断 

preprocessing/merge_output_json.py

@@ -0,0 +1,225 @@
+#!/usr/bin/env python3
+"""
+JSON文件合并脚本
+读取多个JSON文件并合并为一个JSON文件
+"""
+
+import json
+import os
+from typing import Dict, List, Any, Union
+
+# 需要合并的JSON文件列表
+JSON_FILES_TO_MERGE = [
+    "output/trex_sentences_enhanced_checkpoint_360000.json"
+]
+for i in range(1, 1010):
+    JSON_FILES_TO_MERGE.append(f"output/trex_sentences_enhanced_batch_{i}.json")
+
+def load_json_file(file_path: str) -> Union[Dict, List, None]:
+    """加载JSON文件"""
+    if not os.path.exists(file_path):
+        print(f"警告: 文件 {file_path} 不存在")
+        return None
+    
+    try:
+        with open(file_path, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        print(f"成功加载: {file_path}")
+        return data
+    except json.JSONDecodeError as e:
+        print(f"错误: 无法解析JSON文件 {file_path} - {e}")
+        return None
+    except Exception as e:
+        print(f"错误: 读取文件 {file_path} 失败 - {e}")
+        return None
+
+def merge_json_data(data1: Union[Dict, List], data2: Union[Dict, List]) -> Union[Dict, List]:
+    """合并两个JSON数据结构"""
+    
+    # 如果两个都是列表，直接合并
+    if isinstance(data1, list) and isinstance(data2, list):
+        print(f"合并两个列表: {len(data1)} + {len(data2)} = {len(data1) + len(data2)} 项")
+        return data1 + data2
+    
+    # 如果两个都是字典
+    elif isinstance(data1, dict) and isinstance(data2, dict):
+        print("合并两个字典结构")
+        merged = data1.copy()
+        
+        # 特殊处理：如果都有'sentences'字段且为列表，合并sentences
+        if 'sentences' in data1 and 'sentences' in data2:
+            if isinstance(data1['sentences'], list) and isinstance(data2['sentences'], list):
+                print(f"合并sentences字段: {len(data1['sentences'])} + {len(data2['sentences'])} = {len(data1['sentences']) + len(data2['sentences'])} 项")
+                merged['sentences'] = data1['sentences'] + data2['sentences']
+                
+                # 更新metadata if exists
+                if 'metadata' in merged:
+                    if isinstance(merged['metadata'], dict):
+                        merged['metadata']['total_sentences'] = len(merged['sentences'])
+                        merged['metadata']['merged_from'] = [os.path.basename(f) for f in JSON_FILES_TO_MERGE if os.path.exists(f)]
+                
+                # 合并其他字段
+                for key, value in data2.items():
+                    if key != 'sentences' and key not in merged:
+                        merged[key] = value
+                        
+                return merged
+        
+        # 普通字典合并
+        for key, value in data2.items():
+            if key in merged:
+                # 如果key重复且都是列表，合并列表
+                if isinstance(merged[key], list) and isinstance(value, list):
+                    merged[key] = merged[key] + value
+                # 如果key重复且都是字典，递归合并
+                elif isinstance(merged[key], dict) and isinstance(value, dict):
+                    merged[key] = merge_json_data(merged[key], value)
+                else:
+                    # 其他情况保留第二个文件的值
+                    merged[key] = value
+                    print(f"字段 '{key}' 被覆盖")
+            else:
+                merged[key] = value
+        
+        return merged
+    
+    # 类型不匹配的情况，创建一个包含两者的新结构
+    else:
+        print("数据类型不匹配，创建包含两者的新结构")
+        return {
+            "data_from_save.json": data1,
+            "data_from_save2.json": data2,
+            "merged_at": "test.py"
+        }
+
+def save_merged_json(data: Union[Dict, List], output_path: str):
+    """保存合并后的JSON数据"""
+    try:
+        # 确保输出目录存在
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        
+        with open(output_path, 'w', encoding='utf-8') as f:
+            json.dump(data, f, ensure_ascii=False, indent=2)
+        
+        print(f"合并结果已保存到: {output_path}")
+        
+        # 显示统计信息
+        if isinstance(data, dict):
+            if 'sentences' in data and isinstance(data['sentences'], list):
+                print(f"总计句子数: {len(data['sentences'])}")
+            print(f"总计字段数: {len(data)}")
+        elif isinstance(data, list):
+            print(f"总计列表项数: {len(data)}")
+            
+    except Exception as e:
+        print(f"错误: 保存文件失败 - {e}")
+
+def remove_duplicates_from_sentences(data: Union[Dict, List]) -> Union[Dict, List]:
+    """从合并结果中移除重复的句子（基于句子内容）"""
+    if isinstance(data, dict) and 'sentences' in data:
+        if isinstance(data['sentences'], list):
+            original_count = len(data['sentences'])
+            seen_sentences = set()
+            unique_sentences = []
+            
+            for item in data['sentences']:
+                if isinstance(item, dict):
+                    # 如果是字典，使用sentence字段或corrected_sentence字段作为唯一标识
+                    sentence_key = item.get('sentence') or item.get('corrected_sentence') or item.get('original_sentence')
+                elif isinstance(item, str):
+                    sentence_key = item
+                else:
+                    sentence_key = str(item)
+                
+                if sentence_key and sentence_key not in seen_sentences:
+                    seen_sentences.add(sentence_key)
+                    unique_sentences.append(item)
+            
+            data['sentences'] = unique_sentences
+            
+            # 更新metadata
+            if 'metadata' in data and isinstance(data['metadata'], dict):
+                data['metadata']['total_sentences'] = len(unique_sentences)
+                data['metadata']['duplicates_removed'] = original_count - len(unique_sentences)
+            
+            print(f"去重完成: {original_count} -> {len(unique_sentences)} (移除了 {original_count - len(unique_sentences)} 个重复项)")
+    
+    return data
+
+def merge_multiple_json_data(data_list: List[Union[Dict, List]]) -> Union[Dict, List]:
+    """合并多个JSON数据结构"""
+    if not data_list:
+        return {}
+    
+    if len(data_list) == 1:
+        return data_list[0]
+    
+    print(f"准备合并 {len(data_list)} 个JSON数据结构")
+    
+    # 从第一个数据开始，逐步合并其他数据
+    merged_data = data_list[0]
+    
+    for i, data in enumerate(data_list[1:], 1):
+        print(f"正在合并第 {i+1} 个数据结构...")
+        merged_data = merge_json_data(merged_data, data)
+    
+    return merged_data
+
+def main():
+    """主函数"""
+    print("=== JSON文件合并脚本 ===")
+    
+    # 输出路径
+    output_path = "output/merged.json"
+    
+    print(f"准备合并以下文件:")
+    for i, file_path in enumerate(JSON_FILES_TO_MERGE, 1):
+        print(f"  {i}. {file_path}")
+    print(f"输出文件: {output_path}")
+    print()
+    
+    # 加载所有文件
+    loaded_data = []
+    successfully_loaded = []
+    
+    for file_path in JSON_FILES_TO_MERGE:
+        data = load_json_file(file_path)
+        if data is not None:
+            loaded_data.append(data)
+            successfully_loaded.append(file_path)
+    
+    # 检查是否至少有一个文件加载成功
+    if not loaded_data:
+        print("错误: 没有文件能够成功加载，退出")
+        return
+    
+    print(f"成功加载了 {len(loaded_data)} 个文件:")
+    for file_path in successfully_loaded:
+        print(f"  ✓ {file_path}")
+    
+    if len(loaded_data) < len(JSON_FILES_TO_MERGE):
+        failed_count = len(JSON_FILES_TO_MERGE) - len(loaded_data)
+        print(f"警告: {failed_count} 个文件加载失败")
+    print()
+    
+    # 合并所有数据
+    if len(loaded_data) == 1:
+        print("只有一个文件可用，直接使用...")
+        merged_data = loaded_data[0]
+    else:
+        print("开始合并所有文件...")
+        merged_data = merge_multiple_json_data(loaded_data)
+    
+    # 去重处理
+    print("\n检查并去除重复项...")
+    merged_data = remove_duplicates_from_sentences(merged_data)
+    
+    # 保存合并结果
+    print("\n保存合并结果...")
+    save_merged_json(merged_data, output_path)
+    
+    print("\n=== 合并完成 ===")
+    print(f"合并了 {len(successfully_loaded)} 个文件的数据")
+
+if __name__ == "__main__":
+    main()

train_pretrain_accelerate.py

@@ -92,6 +92,47 @@ def init_model(lm_config, pretrained_embedding_path=None, database_init_path=Non
         from sentence_transformers import SentenceTransformer
         import os
         
+        # 聚类参数（需要提前定义用于缓存检查）
+        knowledge_num = args.knowledge_num
+        knowledge_length = args.knowledge_length
+        
+        # 检查是否使用缓存（提前检查，避免不必要的数据处理）
+        cache_dir = os.path.dirname(args.cluster_cache_path)
+        if cache_dir:
+            os.makedirs(cache_dir, exist_ok=True)
+        
+        clustered_tensor = None
+        
+        # 尝试加载缓存的聚类结果
+        if not args.recompute_clusters and os.path.exists(args.cluster_cache_path):
+            try:
+                Logger(f"Loading cached cluster results from {args.cluster_cache_path}")
+                clustered_tensor = torch.load(args.cluster_cache_path)
+                
+                # 验证缓存文件的形状是否可用
+                cached_knowledge_num, cached_knowledge_length = clustered_tensor.shape
+                
+                if cached_knowledge_length == knowledge_length:
+                    if cached_knowledge_num >= knowledge_num:
+                        # 缓存足够大，可以截取使用
+                        clustered_tensor = clustered_tensor[:knowledge_num, :]
+                        Logger(f"Successfully loaded cached clusters with shape {clustered_tensor.shape}")
+                        Logger(f"Truncated from cached shape ({cached_knowledge_num}, {cached_knowledge_length}) to required shape ({knowledge_num}, {knowledge_length})")
+                        Logger("Skipping database initialization and clustering - using cached results")
+                    else:
+                        # 缓存太小，需要重新计算
+                        Logger(f"Cached knowledge_num ({cached_knowledge_num}) < required knowledge_num ({knowledge_num}), recomputing...")
+                        clustered_tensor = None
+                else:
+                    # knowledge_length不匹配，需要重新计算
+                    Logger(f"Cached knowledge_length ({cached_knowledge_length}) != required knowledge_length ({knowledge_length}), recomputing...")
+                    clustered_tensor = None
+            except Exception as e:
+                Logger(f"Failed to load cached clusters: {e}, recomputing...")
+                clustered_tensor = None
+        
+        # 只有在没有有效缓存时才进行数据库初始化和聚类计算
+        if clustered_tensor is None:
             Logger(f"Loading database initialization data from {database_init_path}")
             
             # 1. 加载JSON文件并转换为字典
@@ -174,19 +215,6 @@ def init_model(lm_config, pretrained_embedding_path=None, database_init_path=Non
                     'original_index': i
                 })
             
-        # # Create a JSON-serializable version for saving
-        # json_serializable_sentences = []
-        # for sentence in processed_sentences:
-        #     json_sentence = sentence.copy()
-        #     # Convert embedding to list if it's a numpy array
-        #     if hasattr(json_sentence['embedding'], 'tolist'):
-        #         json_sentence['embedding'] = json_sentence['embedding'].tolist()
-        #     json_serializable_sentences.append(json_sentence)
-        
-        # json.dump(json_serializable_sentences, open('processed_sentences.json', 'w', encoding='utf-8'))
-        
-        # processed_sentences = json.load(open('processed_sentences.json', 'r', encoding='utf-8'))
-        
             # 转换为numpy数组以便后续处理
             embeddings_array = np.array(embeddings_list)
             token_lengths_array = np.array(token_lengths)
@@ -197,12 +225,7 @@ def init_model(lm_config, pretrained_embedding_path=None, database_init_path=Non
             Logger(f"  - Average token length: {np.mean(token_lengths_array):.2f}")
             Logger(f"  - Token length range: {np.min(token_lengths_array)} - {np.max(token_lengths_array)}")
             
-        # 2. 聚类处理 - 优化版本
-        Logger("Starting optimized clustering process...")
-        
-        # 聚类参数
-        knowledge_num = args.knowledge_num
-        knowledge_length = args.knowledge_length
+            # 聚类参数定义
             min_tokens = int(0.85 * knowledge_length)
             max_tokens = int(0.95 * knowledge_length)
             
@@ -265,7 +288,7 @@ def init_model(lm_config, pretrained_embedding_path=None, database_init_path=Non
                         all_remaining = high_importance + medium_importance + low_importance
                         if all_remaining:
                             # 随机采样候选句子（而不是计算所有相似度）
-                        sample_size = min(100, len(all_remaining))
+                            sample_size = min(2000, len(all_remaining))
                             candidates = random.sample(all_remaining, sample_size)
                             
                             # 简单按token长度和重要性选择
@@ -403,6 +426,13 @@ def init_model(lm_config, pretrained_embedding_path=None, database_init_path=Non
             Logger(f"  - Cluster shape: {clustered_tensor.shape}")
             Logger(f"  - Expected shape: ({knowledge_num}, {knowledge_length})")
             
+            # 保存聚类结果到缓存文件
+            try:
+                torch.save(clustered_tensor, args.cluster_cache_path)
+                Logger(f"Cluster results saved to {args.cluster_cache_path}")
+            except Exception as e:
+                Logger(f"Failed to save cluster results: {e}")
+        
         # 3. 初始化模型的weight_down_embed
         if hasattr(model, 'extract_db') and hasattr(model.extract_db, 'weight_down_embed'):
             model.extract_db.weight_down_embed.data.copy_(clustered_tensor)
@@ -651,10 +681,12 @@ def main():
     parser.add_argument("--profile", action="store_true", default=True, help="启用性能分析")
     parser.add_argument("--profile_interval", type=int, default=10, help="性能分析打印间隔（步数）")
     parser.add_argument("--use_flash_attn", action="store_true", default=True, help="启用FlashAttention")
-    parser.add_argument("--knowledge_num", type=int, default=64*64,help="知识库的数据数目")
+    parser.add_argument("--knowledge_num", type=int, default=65536,help="知识库的数据数目")
     parser.add_argument("--knowledge_length", type=int, default=64,help="知识库的句子长度")
     parser.add_argument("--database_init_path", type=str, default="./dataset/database_init.json", help="数据库初始化路径")
     parser.add_argument("--fast_clustering", action="store_true", default=True, help="使用快速近似聚类算法（适用于大数据集）")
+    parser.add_argument("--cluster_cache_path", type=str, default="./cache/cluster_tokens.pt", help="聚类结果缓存文件路径")
+    parser.add_argument("--recompute_clusters", action="store_true", default=False, help="强制重新计算聚类，忽略缓存文件")
     args = parser.parse_args()
 
     #########################################################