Minimind/preprocessing/trex_to_sentences_simple.py

#!/usr/bin/env python3
"""
TREx数据集增强预处理脚本
使用vLLM OpenAI兼容API进行句子后处理和重要性评分

支持两个独立步骤：
1. 句子提取：从TREx数据集提取句子并保存为JSON
2. LLM处理：读取JSON文件进行LLM后处理和重要性评分
"""

import json
import os
import glob
from typing import List, Dict, Any, Union, Set
import re
import asyncio
import time
import logging
from datetime import datetime
import requests
from pydantic import BaseModel, Field
import aiohttp
import concurrent.futures

# 设置日志系统
def setup_logging():
    """设置日志系统"""
    # 确保logs目录存在
    os.makedirs('logs', exist_ok=True)
    
    # 创建日志文件名（包含时间戳）
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    log_file = f'logs/trex_processor_{timestamp}.log'
    
    # 配置日志格式
    log_format = '%(asctime)s - %(levelname)s - [%(funcName)s:%(lineno)d] - %(message)s'
    
    # 配置root logger
    logging.basicConfig(
        level=logging.INFO,
        format=log_format,
        handlers=[
            logging.FileHandler(log_file, encoding='utf-8'),
            logging.StreamHandler()  # 同时输出到控制台
        ]
    )
    
    # 获取logger
    logger = logging.getLogger(__name__)
    logger.info(f"日志系统初始化完成，日志文件: {log_file}")
    return logger

# 全局日志对象
logger = setup_logging()

class ProcessedSentence(BaseModel):
    """处理后的句子结构"""
    corrected_sentence: str = Field(
        ..., 
        description="修正后的句子，只修正语法错误、乱码和不通顺的地方，不进行额外润色"
    )
    importance_score: float = Field(
        ..., 
        description="重要性评分，范围0.0-10.0，以0.1递进。评判这个知识在现实世界中的常用程度和重要度",
        ge=0.0, 
        le=10.0
    )


class EnhancedTRExProcessor:
    def __init__(self, input_dir: str = None, output_file: str = None, max_files: int = None, 
                 sentences_json: str = None, enable_llm_processing: bool = True):
        self.input_dir = input_dir
        
        # 确保output目录存在
        os.makedirs('output', exist_ok=True)
        
        # 确保所有输出文件都在output目录中
        if output_file:
            if not output_file.startswith('output/'):
                self.output_file = os.path.join('output', output_file)
            else:
                self.output_file = output_file
        else:
            self.output_file = None
            
        if sentences_json:
            if not sentences_json.startswith('output/'):
                self.sentences_json = os.path.join('output', sentences_json)
            else:
                self.sentences_json = sentences_json
        else:
            self.sentences_json = "output/extracted_sentences.json"
            
        self.max_files = max_files
        self.enable_llm_processing = enable_llm_processing
        
        # Ollama API配置
        self.model_name = "gemma3:latest"  # Ollama模型名称
        self.ollama_base_url = "http://localhost:11434"  # Ollama服务器地址
        self.batch_size_per_request = 8  # 每个API请求处理的句子数量（Ollama建议较小批次）
        self.max_concurrent_requests = 2  # 最大并发请求数（Ollama建议较低并发）
        self.request_timeout = 180  # 请求超时时间（秒）
        self.retry_attempts = 3  # 重试次数
        
        # 统计信息
        self.total_requests = 0
        self.successful_requests = 0
        self.failed_requests = 0
        
        logger.info(f"处理器初始化完成 - 模型: {self.model_name}, 批次大小: {self.batch_size_per_request}, 并发数: {self.max_concurrent_requests}")
        
        # 扩展的Wikidata属性映射
        self.property_mappings = {
            # 基本关系
            "http://www.wikidata.org/prop/direct/P31": "is a",
            "http://www.wikidata.org/prop/direct/P279": "is a type of",
            
            # 人物相关
            "http://www.wikidata.org/prop/direct/P106": "works as",
            "http://www.wikidata.org/prop/direct/P27": "is a citizen of",
            "http://www.wikidata.org/prop/direct/P19": "was born in",
            "http://www.wikidata.org/prop/direct/P20": "died in",
            "http://www.wikidata.org/prop/direct/P569": "was born on",
            "http://www.wikidata.org/prop/direct/P570": "died on",
            "http://www.wikidata.org/prop/direct/P22": "has father",
            "http://www.wikidata.org/prop/direct/P25": "has mother",
            "http://www.wikidata.org/prop/direct/P26": "is married to",
            
            # 组织相关
            "http://www.wikidata.org/prop/direct/P102": "is a member of",
            "http://www.wikidata.org/prop/direct/P108": "works for",
            "http://www.wikidata.org/prop/direct/P159": "has headquarters in",
            "http://www.wikidata.org/prop/direct/P112": "was founded by",
            "http://www.wikidata.org/prop/direct/P571": "was founded in",
            "http://www.wikidata.org/prop/direct/P169": "has CEO",
            
            # 地理相关
            "http://www.wikidata.org/prop/direct/P17": "is located in",
            "http://www.wikidata.org/prop/direct/P131": "is located in",
            "http://www.wikidata.org/prop/direct/P36": "has capital",
            "http://www.wikidata.org/prop/direct/P47": "borders",
            
            # 其他关系
            "http://www.wikidata.org/prop/direct/P1142": "has ideology",
            "http://www.wikidata.org/prop/direct/P361": "is part of",
            "http://www.wikidata.org/prop/direct/P737": "was influenced by",
            "http://www.wikidata.org/prop/direct/P127": "is owned by",
            "http://www.wikidata.org/prop/direct/P155": "follows",
            "http://www.wikidata.org/prop/direct/P156": "is followed by",
            "http://www.wikidata.org/prop/direct/P138": "is named after"
        }

    def get_system_prompt(self) -> str:
        """获取系统提示"""
        return """You are a professional text processing assistant responsible for correcting errors in sentences and evaluating the importance of knowledge.

### Sentence Correction Rules:
1. Remove Wikipedia-specific markers: such as (disambiguation), (film), (band), etc. in parentheses
2. Ensure grammatical completeness: complete subject+predicate+object structure, avoid dangling 'and is', 'or', etc.
3. Fix obvious grammatical errors: tense consistency, singular/plural consistency, correct preposition usage
4. Clean up garbled text and special characters: such as â, €, ™ and other encoding issues
5. Ensure semantic fluency: if the original sentence cannot be fixed, reorganize the language to make it coherent
6. Do not add information not present in the original text, only correct errors

### Correction Examples:
- Error: 'Argument (disambiguation) is related to philosophy, logic, and is an.'
- Corrected: 'Argument is related to philosophy and logic.'

- Error: 'Beijing is a capital city and are.'
- Corrected: 'Beijing is a capital city.'

Importance scoring criteria (0.0-10.0, in increments of 0.1):

0.0 points - Completely incorrect or meaningless information
Examples: 'Apple is a metal', 'The sun rises from the west', '1+1=3'

0.5 points - Almost worthless information
Examples: 'Color of a fictional character's socks', 'Third line of dialogue from a game NPC', 'What someone had for breakfast yesterday'

1.0 points - Extremely rare, non-practical knowledge
Examples: 'Pet name of a minor novel character', 'Content of the 15th line in movie end credits', 'Nickname of website user ID 123456'

1.5 points - Very niche detailed information
Examples: 'Outfit of a passerby at minute 37 in a movie', 'Duration of background music in a game's hidden level', 'Content of the 3rd dialogue box on page 200 of a manga'

2.0 points - Details in niche professional fields
Examples: 'Color change of rare minerals at specific temperatures', 'Length of an insect's third antenna', 'Molecular formula of chemical reaction byproducts'

2.5 points - Technical details only professionals care about
Examples: 'Release date of specific software library version', 'Time complexity coefficient of an algorithm', 'Thermal expansion coefficient of a material'

3.0 points - Professional knowledge in specific fields
Examples: 'Programming language syntax features', 'Gene sequence of a virus', 'Official system of ancient dynasties'

3.5 points - Professional information with some value
Examples: 'Specific system of historical dynasty', 'Mechanism of action of a drug', 'Development time of a technical standard'

4.0 points - Meaningful knowledge known by few
Examples: 'Unique cultural traditions of a country', 'Important discoveries by a scientist', 'Detailed process of historical events'

4.5 points - Knowledge of interest to some groups
Examples: 'Author's creative background', 'Characteristics of an art movement', 'Detailed rules of a sport'

5.0 points - General knowledge of moderate importance
Examples: 'Famous attractions in cities', 'Development history of a company', 'Living habits of animals'

5.5 points - Fairly useful common sense
Examples: 'Plant growth environment', 'Healthy eating common sense', 'Basic first aid knowledge'

6.0 points - Knowledge most educated people should know
Examples: 'Shakespeare's representative works', 'Basic geometric theorems', 'Major world currencies'

6.5 points - Important cultural or scientific common sense
Examples: 'Basic structure of DNA', 'Newton's three laws', 'Major world religions'

7.0 points - Important foundational knowledge
Examples: 'Time period of World War II', 'Functions of major human organs', 'Basic mathematical operation rules'

7.5 points - Very important common sense
Examples: 'Light speed is the fastest in the universe', 'Earth is round', 'Basic principles of blood circulation'

8.0 points - Core knowledge in basic education
Examples: 'Earth orbits the sun', 'Principle of seasonal formation', 'Basic grammar rules'

8.5 points - Important knowledge everyone should master
Examples: 'Chemical formula of water H2O', 'Basic safety common sense', 'Simple mathematical calculations'

9.0 points - Extremely important basic concepts
Examples: 'Humans need oxygen to survive', 'Fire is hot', 'Basic directional concepts'

9.5 points - Core knowledge everyone must know
Examples: 'A day has 24 hours', 'A year has 12 months', 'Basic number concepts'

10.0 points - Most basic and important common sense
Examples: 'Humans need food and water to survive', 'The sky is blue', 'Stones are heavier than feathers'

When scoring, please consider:
1. Popularity of knowledge - How many people know this knowledge
2. Practical value - How useful this knowledge is in daily life
3. Educational importance - The position of this knowledge in the education system
4. Cultural significance - The importance of this knowledge for understanding world

Please respond with valid JSON in the following format:
{
  "corrected_sentence": "corrected sentence here",
  "importance_score": evaluation score
}"""

    async def process_batch_with_vllm_api(self, sentences: List[str]) -> List[Dict[str, Any]]:
        """使用vLLM OpenAI兼容API处理一批句子"""
        processed_sentences = []
        
        async with aiohttp.ClientSession() as session:
            # 创建并发任务
            semaphore = asyncio.Semaphore(self.max_concurrent_requests)
            tasks = []
            
            # 将句子分成小批次
            for i in range(0, len(sentences), self.batch_size_per_request):
                batch_sentences = sentences[i:i + self.batch_size_per_request]
                task = self.process_single_batch_request(session, semaphore, batch_sentences, i)
                tasks.append(task)
            
            # 等待所有任务完成
            batch_results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # 收集结果
            for result in batch_results:
                if isinstance(result, Exception):
                    logger.error(f"批次处理出错: {result}")
                    continue
                if result:
                    processed_sentences.extend(result)
        
        return processed_sentences

    async def process_single_batch_request(self, session: aiohttp.ClientSession, semaphore: asyncio.Semaphore, 
                                         sentences: List[str], batch_index: int) -> List[Dict[str, Any]]:
        """处理单个批次的API请求"""
        async with semaphore:
            for attempt in range(self.retry_attempts):
                try:
                    # 为每个句子创建单独的消息
                    messages = []
                    for sentence in sentences:
                        messages.append({
                            "role": "user",
                            "content": f"Please correct the errors in the following sentence and evaluate its importance: {sentence}"
                        })
                    
                    # 构建Ollama请求数据
                    request_data = {
                        "model": self.model_name,
                        "messages": [
                            {"role": "system", "content": self.get_system_prompt()}
                        ] + messages,
                        "stream": False,
                        "options": {
                            "temperature": 0.2,
                            "num_predict": 500 * len(sentences)  # 为每个句子分配足够的token
                        },
                        "format": "json"  # Ollama的JSON格式参数
                    }
                    
                    # 发送请求到Ollama
                    async with session.post(
                        f'{self.ollama_base_url}/api/chat',
                        json=request_data,
                        timeout=aiohttp.ClientTimeout(total=self.request_timeout)
                    ) as response:
                        
                        if response.status == 200:
                            result = await response.json()
                            return self.parse_ollama_response(result, sentences, batch_index)
                        else:
                            error_text = await response.text()
                            logger.warning(f"API请求失败 (批次 {batch_index}, 尝试 {attempt + 1}/{self.retry_attempts}): {response.status} - {error_text}")
                            
                            if attempt == self.retry_attempts - 1:  # 最后一次尝试
                                logger.error(f"批次 {batch_index} 在 {self.retry_attempts} 次尝试后仍然失败")
                                self.failed_requests += len(sentences)
                                return self.create_default_responses(sentences)
                            else:
                                # 等待后重试
                                await asyncio.sleep(2 ** attempt)  # 指数退避
                                continue
                                
                except asyncio.TimeoutError:
                    logger.warning(f"批次 {batch_index} 请求超时 (尝试 {attempt + 1}/{self.retry_attempts})")
                    if attempt == self.retry_attempts - 1:
                        logger.error(f"批次 {batch_index} 在 {self.retry_attempts} 次尝试后仍然超时")
                        self.failed_requests += len(sentences)
                        return self.create_default_responses(sentences)
                    else:
                        await asyncio.sleep(2 ** attempt)
                        continue
                        
                except Exception as e:
                    logger.warning(f"处理批次 {batch_index} 时出错 (尝试 {attempt + 1}/{self.retry_attempts}): {e}")
                    if attempt == self.retry_attempts - 1:
                        logger.error(f"批次 {batch_index} 在 {self.retry_attempts} 次尝试后仍然失败")
                        self.failed_requests += len(sentences)
                        return self.create_default_responses(sentences)
                    else:
                        await asyncio.sleep(2 ** attempt)
                        continue
            
            # 如果所有重试都失败了
            return self.create_default_responses(sentences)

    def parse_ollama_response(self, response: Dict[str, Any], original_sentences: List[str], batch_index: int) -> List[Dict[str, Any]]:
        """解析Ollama响应"""
        processed_sentences = []
        
        try:
            # Ollama的响应格式
            message = response.get('message', {})
            content = message.get('content', '')
            
            if not content:
                logger.warning(f"批次 {batch_index} 没有返回内容")
                return self.create_default_responses(original_sentences)
            
            # 尝试解析JSON响应
            try:
                # 如果返回的是单个JSON对象
                if content.strip().startswith('{') and content.strip().endswith('}'):
                    response_data = json.loads(content)
                    processed_sentence = ProcessedSentence(
                        corrected_sentence=response_data.get('corrected_sentence', original_sentences[0] if original_sentences else ""),
                        importance_score=float(response_data.get('importance_score', 5.0))
                    )
                    
                    processed_sentences.append({
                        "original_sentence": original_sentences[0] if original_sentences else "",
                        "corrected_sentence": processed_sentence.corrected_sentence,
                        "importance_score": processed_sentence.importance_score
                    })
                    self.successful_requests += 1
                    
                    # 如果有多个句子但只返回一个结果，为其他句子创建默认响应
                    for i in range(1, len(original_sentences)):
                        processed_sentences.append({
                            "original_sentence": original_sentences[i],
                            "corrected_sentence": original_sentences[i],
                            "importance_score": 5.0
                        })
                        self.failed_requests += 1
                        
                else:
                    # 尝试解析多个JSON对象
                    json_objects = []
                    for line in content.split('\n'):
                        line = line.strip()
                        if line.startswith('{') and line.endswith('}'):
                            try:
                                json_objects.append(json.loads(line))
                            except:
                                continue
                    
                    if json_objects:
                        for i, (sentence, json_obj) in enumerate(zip(original_sentences, json_objects)):
                            try:
                                processed_sentence = ProcessedSentence(
                                    corrected_sentence=json_obj.get('corrected_sentence', sentence),
                                    importance_score=float(json_obj.get('importance_score', 5.0))
                                )
                                
                                processed_sentences.append({
                                    "original_sentence": sentence,
                                    "corrected_sentence": processed_sentence.corrected_sentence,
                                    "importance_score": processed_sentence.importance_score
                                })
                                self.successful_requests += 1
                            except Exception as e:
                                logger.warning(f"解析JSON对象失败: {e}")
                                processed_sentences.append({
                                    "original_sentence": sentence,
                                    "corrected_sentence": sentence,
                                    "importance_score": 5.0
                                })
                                self.failed_requests += 1
                        
                        # 为剩余句子创建默认响应
                        for i in range(len(json_objects), len(original_sentences)):
                            processed_sentences.append({
                                "original_sentence": original_sentences[i],
                                "corrected_sentence": original_sentences[i],
                                "importance_score": 5.0
                            })
                            self.failed_requests += 1
                    else:
                        logger.warning(f"批次 {batch_index} 无法解析JSON响应: {content}")
                        return self.create_default_responses(original_sentences)
                        
            except (json.JSONDecodeError, ValueError) as e:
                logger.warning(f"批次 {batch_index} 解析响应JSON失败: {e}")
                logger.warning(f"原始内容: {content}")
                return self.create_default_responses(original_sentences)
                
        except Exception as e:
            logger.error(f"解析批次 {batch_index} 响应时出错: {e}")
            return self.create_default_responses(original_sentences)
        
        return processed_sentences

    def create_default_responses(self, sentences: List[str]) -> List[Dict[str, Any]]:
        """为失败的请求创建默认响应"""
        default_responses = []
        for sentence in sentences:
            default_responses.append({
                "original_sentence": sentence,
                "corrected_sentence": sentence,
                "importance_score": 5.0
            })
        return default_responses

    async def process_sentences_with_vllm_api(self, sentences: List[str]) -> List[Dict[str, Any]]:
        """使用Ollama API处理句子"""
        logger.info(f"开始使用Ollama API处理 {len(sentences)} 个句子...")
        print(f"开始使用Ollama API处理 {len(sentences)} 个句子...")
        
        start_time = time.time()
        total_sentences = len(sentences)
        total_processed_count = 0
        
        # 检查Ollama服务状态
        if not self.check_ollama_status():
            logger.error("Ollama服务状态异常，无法继续处理")
            print("错误：Ollama服务状态异常，请检查服务是否正常运行")
            return []
        
        # 分大批次处理（用于检查点保存）
        large_batch_size = 1000  # 每1000个句子保存一次检查点
        all_processed_sentences = []
        
        for large_batch_start in range(0, total_sentences, large_batch_size):
            large_batch_end = min(large_batch_start + large_batch_size, total_sentences)
            large_batch_sentences = sentences[large_batch_start:large_batch_end]
            large_batch_number = large_batch_start // large_batch_size + 1
            
            logger.info(f"=== 处理大批次 {large_batch_number} ({large_batch_start + 1}-{large_batch_end}/{total_sentences}) ===")
            print(f"\n=== 处理大批次 {large_batch_number} ({large_batch_start + 1}-{large_batch_end}/{total_sentences}) ===")
            
            large_batch_start_time = time.time()
            
            # 处理当前大批次
            batch_processed = await self.process_batch_with_vllm_api(large_batch_sentences)
            all_processed_sentences.extend(batch_processed)
            total_processed_count += len(batch_processed)
            
            # 保存当前大批次的检查点
            checkpoint_filename = self.save_batch_checkpoint(batch_processed, large_batch_number, total_processed_count)
            
            # 打印进度
            large_batch_time = time.time() - large_batch_start_time
            elapsed_time = time.time() - start_time
            
            logger.info(f"大批次 {large_batch_number} 处理完成！")
            logger.info(f"  - 当前批次：成功 {len(batch_processed)}，用时 {large_batch_time/60:.1f}分钟")
            logger.info(f"  - 总体进度：{total_processed_count}/{total_sentences} ({total_processed_count/total_sentences*100:.1f}%)")
            logger.info(f"  - 已用时间：{elapsed_time/60:.1f}分钟")
            logger.info(f"  - 批次检查点已保存：{checkpoint_filename}")
            
            print(f"大批次 {large_batch_number} 处理完成！")
            print(f"  - 当前批次：成功 {len(batch_processed)}，用时 {large_batch_time/60:.1f}分钟")
            print(f"  - 总体进度：{total_processed_count}/{total_sentences} ({total_processed_count/total_sentences*100:.1f}%)")
            print(f"  - 已用时间：{elapsed_time/60:.1f}分钟")
            print(f"  - 批次检查点已保存：{checkpoint_filename}")
            
            if large_batch_end < total_sentences:
                remaining_sentences = total_sentences - total_processed_count
                avg_time_per_sentence = elapsed_time / total_processed_count
                estimated_remaining_time = avg_time_per_sentence * remaining_sentences
                logger.info(f"  - 预估剩余时间：{estimated_remaining_time/60:.1f}分钟")
                print(f"  - 预估剩余时间：{estimated_remaining_time/60:.1f}分钟")
        
        # 打印最终统计
        total_time = time.time() - start_time
        logger.info(f"=== 全部处理完成！===")
        logger.info(f"  - 总成功：{self.successful_requests}")
        logger.info(f"  - 总失败：{self.failed_requests}")
        logger.info(f"  - 总用时：{total_time/60:.1f}分钟")
        logger.info(f"  - 平均处理速度：{total_processed_count/total_time:.2f}句/秒")
        
        print(f"\n=== 全部处理完成！===")
        print(f"  - 总成功：{self.successful_requests}")
        print(f"  - 总失败：{self.failed_requests}")
        print(f"  - 总用时：{total_time/60:.1f}分钟")
        print(f"  - 平均处理速度：{total_processed_count/total_time:.2f}句/秒")
        
        return all_processed_sentences

    def check_ollama_status(self) -> bool:
        """检查Ollama服务是否正常运行"""
        try:
            # 检查Ollama API是否响应
            response = requests.get(f'{self.ollama_base_url}/api/tags', timeout=10)
            
            if response.status_code == 200:
                models = response.json()
                model_names = [model.get('name', 'unknown') for model in models.get('models', [])]
                logger.info(f"Ollama服务状态正常，可用模型: {model_names}")
                
                # 检查目标模型是否可用
                if self.model_name in model_names:
                    logger.info(f"目标模型 {self.model_name} 可用")
                    return True
                else:
                    logger.warning(f"目标模型 {self.model_name} 不在可用模型列表中: {model_names}")
                    logger.info("尝试拉取模型...")
                    # 尝试拉取模型
                    try:
                        pull_response = requests.post(
                            f'{self.ollama_base_url}/api/pull',
                            json={"name": self.model_name},
                            timeout=300  # 5分钟超时
                        )
                        if pull_response.status_code == 200:
                            logger.info(f"成功拉取模型 {self.model_name}")
                            return True
                        else:
                            logger.error(f"拉取模型失败: {pull_response.status_code}")
                            return False
                    except Exception as e:
                        logger.error(f"拉取模型时出错: {e}")
                        return False
            else:
                logger.error(f"Ollama API响应异常，状态码: {response.status_code}")
                return False
                
        except requests.exceptions.RequestException as e:
            logger.error(f"无法连接到Ollama API: {e}")
            return False
        except Exception as e:
            logger.error(f"检查Ollama状态时出错: {e}")
            return False

    def clean_text(self, text: str) -> str:
        """清理文本，处理特殊字符"""
        if not text:
            return ""
        
        # 处理常见的Unicode字符
        text = text.replace("–", "-")  # en dash
        text = text.replace("—", "-")  # em dash
        text = text.replace("'", "'")  # right single quotation mark
        text = text.replace("'", "'")  # left single quotation mark
        text = text.replace(""", '"')  # left double quotation mark
        text = text.replace(""", '"')  # right double quotation mark
        
        # 处理可能的转义序列
        try:
            text = text.encode('utf-8').decode('utf-8')
        except:
            pass
        
        # 清理多余的空格
        text = re.sub(r'\s+', ' ', text).strip()
        
        # 移除可能的引号
        text = text.strip('"\'')
        
        return text

    def parse_large_json_file(self, file_path: str) -> List[Dict]:
        """解析大型JSON文件，处理可能的格式问题"""
        documents = []
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read().strip()
            
            # 尝试不同的解析方法
            if content.startswith('[') and content.endswith(']'):
                # 标准JSON数组
                documents = json.loads(content)
            else:
                # 可能是连续的JSON对象
                # 尝试在}{"之间分割
                if '}{"' in content:
                    json_strings = content.split('}{')
                    json_strings[0] += '}'  # 第一个对象
                    json_strings[-1] = '{' + json_strings[-1]  # 最后一个对象
                    
                    for i in range(1, len(json_strings) - 1):
                        json_strings[i] = '{' + json_strings[i] + '}'
                    
                    for json_str in json_strings:
                        try:
                            doc = json.loads(json_str)
                            documents.append(doc)
                        except json.JSONDecodeError:
                            continue
                else:
                    # 尝试作为单个JSON对象
                    try:
                        documents = [json.loads(content)]
                    except json.JSONDecodeError:
                        pass
                        
        except Exception as e:
            print(f"Error parsing {file_path}: {e}")
            
        return documents

    def extract_sentences_from_document(self, doc: Dict[str, Any]) -> List[str]:
        """从文档中提取句子"""
        sentences = []
        
        title = self.clean_text(doc.get('title', ''))
        text = self.clean_text(doc.get('text', ''))
        entities = doc.get('entities', [])
        triples = doc.get('triples', [])
        
        # 处理显式三元组
        for triple in triples:
            sentence = self.triple_to_sentence(triple)
            if sentence:
                sentences.append(sentence)
        
        # 从实体和文本中生成基本句子（如果三元组句子不够）
        if title and text and len(sentences) < 5:
            # 基于标题和实体生成句子
            entity_names = []
            for entity in entities[:15]:
                entity_name = self.clean_text(entity.get('surfaceform', ''))
                if entity_name and len(entity_name) > 2:
                    entity_names.append(entity_name)
            
            # 生成简单的描述句子
            if entity_names:
                important_entities = []
                title_lower = title.lower()
                for entity in entity_names:
                    if (entity.lower() != title_lower and 
                        entity not in important_entities and
                        not any(t.lower() in entity.lower() for t in title_lower.split()[:2])):
                        important_entities.append(entity)
                        if len(important_entities) >= 6:
                            break
                
                if important_entities and len(sentences) < 3:
                    entities_str = ', '.join(important_entities[:3])
                    sentences.append(f"{title} is related to {entities_str}.")
        
        return sentences

    def triple_to_sentence(self, triple: Dict[str, Any]) -> str:
        """将三元组转换为自然语言句子"""
        try:
            subject = triple.get('subject', {})
            predicate = triple.get('predicate', {})
            obj = triple.get('object', {})
            
            subject_name = self.clean_text(subject.get('surfaceform', ''))
            object_name = self.clean_text(obj.get('surfaceform', ''))
            predicate_uri = predicate.get('uri', '')
            
            # 检查是否有有效的主语和宾语
            if not subject_name or not object_name:
                return ""
            
            # 检查主语和宾语是否过短或无意义
            if len(subject_name) <= 2 or len(object_name) <= 2:
                return ""
            
            # 获取关系文本
            relation_text = self.property_mappings.get(predicate_uri, "is related to")
            
            # 避免重复的主语宾语
            if subject_name.lower() == object_name.lower():
                return ""
            
            return f"{subject_name} {relation_text} {object_name}."
            
        except Exception as e:
            print(f"Error converting triple to sentence: {e}")
            return ""

    def save_batch_checkpoint(self, processed_sentences: List[Dict[str, Any]], batch_number: int, total_processed_count: int) -> str:
        """保存当前批次的检查点文件"""
        # 生成检查点文件名，确保在output目录中
        base_name = os.path.splitext(os.path.basename(self.output_file))[0]
        checkpoint_filename = os.path.join('output', f"{base_name}_batch_{batch_number}.json")
        
        # 保存检查点
        with open(checkpoint_filename, 'w', encoding='utf-8') as f:
            json.dump({
                "metadata": {
                    "batch_number": batch_number,
                    "batch_size": len(processed_sentences),
                    "total_processed_count": total_processed_count,
                    "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
                },
                "sentences": processed_sentences
            }, f, ensure_ascii=False, indent=2)
        
        return checkpoint_filename

    async def process_files(self) -> List[Dict[str, Any]]:
        """处理所有文件"""
        json_files = glob.glob(os.path.join(self.input_dir, "re-nlg_*.json"))
        
        if not json_files:
            print(f"No JSON files found in {self.input_dir}")
            return []
        
        # 排序文件以确保一致的处理顺序
        json_files.sort()
        
        if self.max_files:
            json_files = json_files[:self.max_files]
        
        print(f"Found {len(json_files)} JSON files to process")
        
        all_sentences = []
        
        for i, file_path in enumerate(json_files):
            print(f"Processing file {i+1}/{len(json_files)}: {os.path.basename(file_path)}")
            
            documents = self.parse_large_json_file(file_path)
            print(f"  Parsed {len(documents)} documents")
            
            for doc in documents:
                sentences = self.extract_sentences_from_document(doc)
                all_sentences.extend(sentences)
            
            print(f"  Generated {len(all_sentences)} total raw sentences so far")
        
        print(f"总共提取了 {len(all_sentences)} 个原始句子")
        
        # 去重
        unique_sentences = []
        seen = set()
        for sentence in all_sentences:
            sentence = sentence.strip()
            if sentence and sentence not in seen and len(sentence) > 10:
                unique_sentences.append(sentence)
                seen.add(sentence)
        
        print(f"去重后剩余 {len(unique_sentences)} 个句子")
        
        # 保存原始句子到JSON文件
        sentences_data = {
            "metadata": {
                "total_sentences": len(unique_sentences),
                "extraction_timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                "source_files": len(json_files),
                "max_files_limit": self.max_files
            },
            "sentences": [{"sentence": sentence, "processed": False} for sentence in unique_sentences]
        }
        
        with open(self.sentences_json, 'w', encoding='utf-8') as f:
            json.dump(sentences_data, f, ensure_ascii=False, indent=2)
        
        print(f"句子提取完成！已保存到: {self.sentences_json}")
        print(f"总计句子数: {len(unique_sentences)}")
        
        return unique_sentences

    def save_sentences(self, processed_sentences: List[Dict[str, Any]]):
        """保存处理后的句子到文件"""
        # 确保输出目录存在
        os.makedirs('output', exist_ok=True)
        
        # 保存为JSON格式，包含完整信息
        json_output_file = self.output_file.replace('.txt', '.json')
        with open(json_output_file, 'w', encoding='utf-8') as f:
            json.dump(processed_sentences, f, ensure_ascii=False, indent=2)
        
        # 保存为简单文本格式（仅修正后的句子）
        with open(self.output_file, 'w', encoding='utf-8') as f:
            for item in processed_sentences:
                f.write(item['corrected_sentence'] + '\n')
        
        # 生成重要性排序文件
        importance_sorted = sorted(processed_sentences, key=lambda x: x['importance_score'], reverse=True)
        importance_file = self.output_file.replace('.txt', '_sorted_by_importance.txt')
        with open(importance_file, 'w', encoding='utf-8') as f:
            for item in importance_sorted:
                f.write(f"[{item['importance_score']:.1f}] {item['corrected_sentence']}\n")
        
        print(f"保存了 {len(processed_sentences)} 个处理后的句子：")
        print(f"  - JSON格式: {json_output_file}")
        print(f"  - 文本格式: {self.output_file}")
        print(f"  - 重要性排序: {importance_file}")
        
        # 统计信息
        scores = [item['importance_score'] for item in processed_sentences]
        avg_score = sum(scores) / len(scores) if scores else 0
        print(f"  - 平均重要性评分: {avg_score:.2f}")
        print(f"  - 最高评分: {max(scores):.1f}")
        print(f"  - 最低评分: {min(scores):.1f}")

    async def run(self):
        """运行处理流程"""
        print("Starting enhanced TREx to sentences conversion...")
        processed_sentences = await self.process_files()
        self.save_sentences(processed_sentences)
        print("Enhanced conversion completed!")

    def find_latest_checkpoint(self) -> Union[tuple, None]:
        """查找最新的检查点文件"""
        base_name = os.path.splitext(os.path.basename(self.output_file))[0]
        pattern = os.path.join('output', f"{base_name}_checkpoint_*.json")
        checkpoint_files = glob.glob(pattern)
        
        if not checkpoint_files:
            return None
        
        # 按检查点编号排序，获取最新的
        latest_file = None
        latest_count = 0
        
        for file in checkpoint_files:
            try:
                # 从文件名中提取数字
                match = re.search(r'checkpoint_(\d+)\.json$', file)
                if match:
                    count = int(match.group(1))
                    if count > latest_count:
                        latest_count = count
                        latest_file = file
            except:
                continue
        
        if latest_file:
            return latest_file, latest_count
        else:
            return None

    def load_checkpoint(self, checkpoint_file: str) -> List[Dict[str, Any]]:
        """从检查点文件加载已处理的句子"""
        try:
            with open(checkpoint_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            if 'sentences' in data:
                return data['sentences']
            else:
                # 旧格式的检查点文件
                return data
        except Exception as e:
            print(f"加载检查点文件失败: {e}")
            return []

    def get_processed_sentences_from_checkpoints(self) -> Set[str]:
        """从检查点文件中获取已处理过的句子集合"""
        if not self.output_file:
            return set()
            
        processed_sentences = set()
        base_name = os.path.splitext(os.path.basename(self.output_file))[0]
        
        # 首先查找新格式的批次文件
        batch_pattern = os.path.join('output', f"{base_name}_batch_*.json")
        batch_files = glob.glob(batch_pattern)
        
        if batch_files:
            print(f"找到 {len(batch_files)} 个批次检查点文件")
            batch_files.sort()  # 确保按顺序处理
            
            for batch_file in batch_files:
                try:
                    with open(batch_file, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                    
                    sentences_data = data.get('sentences', [])
                    for item in sentences_data:
                        original_sentence = item.get('original_sentence', '')
                        if original_sentence:
                            processed_sentences.add(original_sentence)
                    
                    batch_number = data.get('metadata', {}).get('batch_number', 'unknown')
                    print(f"  - 从批次 {batch_number} 加载了 {len(sentences_data)} 个句子")
                    
                except Exception as e:
                    print(f"读取批次文件 {batch_file} 失败: {e}")
                    continue
            
            print(f"从批次文件总计加载了 {len(processed_sentences)} 个已处理的句子")
            logger.info(f"从批次文件总计加载了 {len(processed_sentences)} 个已处理的句子")
            return processed_sentences
        
        # 如果没有批次文件，尝试查找旧格式的检查点文件
        old_pattern = os.path.join('output', f"{base_name}_checkpoint_*.json")
        checkpoint_files = glob.glob(old_pattern)
        
        if not checkpoint_files:
            print("未找到检查点文件，将从头开始处理")
            return set()
        
        # 找到最新的检查点文件
        latest_file = None
        latest_count = 0
        
        for file in checkpoint_files:
            try:
                match = re.search(r'checkpoint_(\d+)\.json$', file)
                if match:
                    count = int(match.group(1))
                    if count > latest_count:
                        latest_count = count
                        latest_file = file
            except:
                continue
        
        if latest_file:
            print(f"找到旧格式检查点: {latest_file} (包含 {latest_count} 条记录)")
            logger.info(f"找到旧格式检查点: {latest_file} (包含 {latest_count} 条记录)")
            try:
                with open(latest_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                
                sentences_data = data.get('sentences', [])
                for item in sentences_data:
                    original_sentence = item.get('original_sentence', '')
                    if original_sentence:
                        processed_sentences.add(original_sentence)
                
                print(f"从旧格式检查点加载了 {len(processed_sentences)} 个已处理的句子")
                logger.info(f"从旧格式检查点加载了 {len(processed_sentences)} 个已处理的句子")
                
            except Exception as e:
                print(f"读取检查点文件失败: {e}")
                return set()
        
        return processed_sentences

    async def process_with_llm(self):
        """步骤2：从JSON文件读取句子并进行vLLM处理（保持兼容性）"""
        await self.process_with_vllm_api()

    async def process_with_vllm_api(self):
        """步骤2：从JSON文件读取句子并进行vLLM处理"""
        if not self.enable_llm_processing:
            print("Error: LLM processing is disabled!")
            return
            
        if not self.output_file:
            print("Error: output_file is required for LLM processing!")
            return
            
        print("=== 步骤2：vLLM处理 ===")
        
        # 读取句子JSON文件
        if not os.path.exists(self.sentences_json):
            print(f"Error: Sentences file {self.sentences_json} not found!")
            print("请先运行步骤1进行句子提取")
            return
        
        print(f"正在读取句子文件: {self.sentences_json}")
        
        try:
            with open(self.sentences_json, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            all_sentences = [item["sentence"] for item in data.get("sentences", [])]
            print(f"从文件中读取了 {len(all_sentences)} 个句子")
            
        except Exception as e:
            print(f"读取句子文件失败: {e}")
            return
        
        # 获取已处理的句子
        processed_sentences_set = self.get_processed_sentences_from_checkpoints()
        
        # 过滤出未处理的句子
        unprocessed_sentences = []
        for sentence in all_sentences:
            if sentence not in processed_sentences_set:
                unprocessed_sentences.append(sentence)
        
        print(f"需要处理的句子数: {len(unprocessed_sentences)} (跳过已处理: {len(processed_sentences_set)})")
        logger.info(f"需要处理的句子数: {len(unprocessed_sentences)} (跳过已处理: {len(processed_sentences_set)})")
        
        if not unprocessed_sentences:
            print("所有句子都已处理完成！")
            
            # 如果有检查点，直接从最新检查点生成最终文件
            if processed_sentences_set:
                latest_checkpoint = self.find_latest_checkpoint()
                if latest_checkpoint:
                    checkpoint_file, _ = latest_checkpoint
                    processed_data = self.load_checkpoint(checkpoint_file)
                    self.save_sentences(processed_data)
                    print("已从检查点生成最终输出文件")
            return
        
        # 处理未处理的句子
        print("开始vLLM处理...")
        
        # 处理新句子（现在返回空列表，数据保存在批次检查点中）
        await self.process_sentences_with_vllm_api(unprocessed_sentences)
        
        # 处理完成后，合并所有批次检查点生成最终文件
        print("合并所有批次检查点生成最终文件...")
        all_processed_sentences = self.merge_all_batch_checkpoints()
        
        if all_processed_sentences:
            # 保存最终结果
            self.save_sentences(all_processed_sentences)
            print("vLLM处理完成！")
        else:
            print("警告：没有找到任何处理结果")

    def merge_all_batch_checkpoints(self) -> List[Dict[str, Any]]:
        """合并所有批次检查点文件"""
        if not self.output_file:
            return []
            
        base_name = os.path.splitext(os.path.basename(self.output_file))[0]
        
        # 查找所有批次检查点文件
        batch_pattern = os.path.join('output', f"{base_name}_batch_*.json")
        batch_files = glob.glob(batch_pattern)
        
        if not batch_files:
            # 如果没有批次文件，尝试查找旧格式的检查点文件
            old_pattern = os.path.join('output', f"{base_name}_checkpoint_*.json")
            old_files = glob.glob(old_pattern)
            if old_files:
                print("找到旧格式检查点文件，尝试读取最新的...")
                latest_checkpoint = self.find_latest_checkpoint()
                if latest_checkpoint:
                    checkpoint_file, _ = latest_checkpoint
                    return self.load_checkpoint(checkpoint_file)
            return []
        
        print(f"找到 {len(batch_files)} 个批次检查点文件")
        
        all_sentences = []
        batch_files.sort()  # 确保按顺序处理
        
        for batch_file in batch_files:
            try:
                with open(batch_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                
                batch_sentences = data.get('sentences', [])
                all_sentences.extend(batch_sentences)
                
                batch_number = data.get('metadata', {}).get('batch_number', 'unknown')
                batch_size = len(batch_sentences)
                print(f"  - 批次 {batch_number}: {batch_size} 个句子")
                
            except Exception as e:
                print(f"读取批次文件 {batch_file} 失败: {e}")
                continue
        
        print(f"总计合并了 {len(all_sentences)} 个句子")
        return all_sentences
    
    def extract_sentences(self):
        """步骤1：从TREx数据集提取句子并保存为JSON"""
        if not self.input_dir:
            print("Error: input_dir is required for sentence extraction!")
            return
            
        print("=== 步骤1：句子提取 ===")
        print("开始从TREx数据集提取句子...")
        
        json_files = glob.glob(os.path.join(self.input_dir, "re-nlg_*.json"))
        
        if not json_files:
            print(f"No JSON files found in {self.input_dir}")
            return
        
        # 排序文件以确保一致的处理顺序
        json_files.sort()
        
        if self.max_files:
            json_files = json_files[:self.max_files]
        
        print(f"Found {len(json_files)} JSON files to process")
        
        all_sentences = []
        
        for i, file_path in enumerate(json_files):
            print(f"Processing file {i+1}/{len(json_files)}: {os.path.basename(file_path)}")
            
            documents = self.parse_large_json_file(file_path)
            print(f"  Parsed {len(documents)} documents")
            
            for doc in documents:
                sentences = self.extract_sentences_from_document(doc)
                all_sentences.extend(sentences)
            
            print(f"  Generated {len(all_sentences)} total raw sentences so far")
        
        print(f"总共提取了 {len(all_sentences)} 个原始句子")
        
        # 去重
        unique_sentences = []
        seen = set()
        for sentence in all_sentences:
            sentence = sentence.strip()
            if sentence and sentence not in seen and len(sentence) > 10:
                unique_sentences.append(sentence)
                seen.add(sentence)
        
        print(f"去重后剩余 {len(unique_sentences)} 个句子")
        
        # 保存原始句子到JSON文件
        sentences_data = {
            "metadata": {
                "total_sentences": len(unique_sentences),
                "extraction_timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                "source_files": len(json_files),
                "max_files_limit": self.max_files
            },
            "sentences": [{"sentence": sentence, "processed": False} for sentence in unique_sentences]
        }
        
        with open(self.sentences_json, 'w', encoding='utf-8') as f:
            json.dump(sentences_data, f, ensure_ascii=False, indent=2)
        
        print(f"句子提取完成！已保存到: {self.sentences_json}")
        print(f"总计句子数: {len(unique_sentences)}")
        
        return unique_sentences


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Convert TREx dataset to enhanced sentences with vLLM processing')
    
    # 选择运行模式
    parser.add_argument('--step', choices=['extract', 'llm', 'all'], default='llm',
                        help='运行步骤: extract=仅提取句子, llm=仅vLLM处理, all=完整流程')
    
    # 文件路径参数
    parser.add_argument('--input_dir', default='dataset/TREx', help='Input directory containing TREx JSON files')
    parser.add_argument('--sentences_json', default='extracted_sentences.json', help='JSON file for extracted sentences (will be saved in output/)')
    parser.add_argument('--output_file', default='trex_sentences_enhanced.txt', help='Output file path (will be saved in output/)')
    
    # 处理参数
    parser.add_argument('--max_files', type=int, help='Maximum number of files to process (for testing)')
    parser.add_argument('--no_llm', action='store_true', help='Disable vLLM processing (basic mode)')
    
    args = parser.parse_args()
    
    # 根据步骤验证参数
    if args.step in ['extract', 'all']:
        if not os.path.exists(args.input_dir):
            print(f"Error: Input directory {args.input_dir} does not exist!")
            return
    
    if args.step in ['llm', 'all']:
        if args.no_llm:
            print("Error: Cannot run vLLM step with --no_llm flag!")
            return
    
    # 创建处理器
    processor = EnhancedTRExProcessor(
        input_dir=args.input_dir,
        sentences_json=args.sentences_json,
        output_file=args.output_file,
        max_files=args.max_files,
        enable_llm_processing=not args.no_llm
    )
    
    # 根据选择的步骤运行
    if args.step == 'extract':
        print("=== 运行模式：仅句子提取 ===")
        processor.extract_sentences()
        
    elif args.step == 'llm':
        print("=== 运行模式：仅vLLM处理 ===")
        asyncio.run(processor.process_with_vllm_api())
        
    elif args.step == 'all':
        print("=== 运行模式：完整流程 ===")
        
        # 步骤1：提取句子
        print("\n--- 开始步骤1：句子提取 ---")
        sentences = processor.extract_sentences()
        
        if not sentences:
            print("句子提取失败，退出")
            return
        
        if args.no_llm:
            print("vLLM处理已禁用，流程结束")
            return
        
        # 步骤2：vLLM处理
        print("\n--- 开始步骤2：vLLM处理 ---")
        asyncio.run(processor.process_with_vllm_api())


if __name__ == "__main__":
    main()