@lzhzzzzwill/cofos 1.0.1 → 1.0.2

package/runtime/src/retrieval/build_bm25_index.py

@@ -0,0 +1,382 @@
+#!/usr/bin/env python3
+"""
+COF-ROS-Reasoner: BM25 索引构建脚本
+
+目标：
+- 构建 BM25 关键词索引用于精确匹配检索
+- 支持来自 JSON evidence、PDF chunks、WOS abstracts 的文本
+"""
+
+import os
+import json
+import yaml
+import logging
+import pickle
+from pathlib import Path
+from typing import Dict, List, Any, Optional
+from datetime import datetime
+
+# 添加 src 路径
+
+try:
+    from rank_bm25 import BM25Okapi
+except ImportError:
+    BM25Okapi = None
+
+from data_processing.utils import setup_logging, ensure_directory
+from retrieval.text_search import tokenize_search_text
+
+
+class BM25IndexBuilder:
+    """构建 BM25 索引"""
+
+    def __init__(self, config: Dict[str, Any]):
+        self.config = config
+        self.output_dir = Path(config['paths']['bm25_store'])
+
+        # 设置日志
+        setup_logging(config)
+
+        # BM25 索引
+        self.bm25 = None
+        self.metadata: List[Dict[str, Any]] = []
+
+    def tokenize(self, text: str) -> List[str]:
+        """Use the same tokenizer as query-time retrieval."""
+        return tokenize_search_text(text)
+
+    def build_index(self, documents: List[str], metadata: List[Dict[str, Any]]):
+        """构建 BM25 索引"""
+        if BM25Okapi is None:
+            logging.error("rank-bm25 not installed. Run: pip install rank-bm25")
+            raise ImportError("rank-bm25 required")
+
+        logging.info("Building BM25 index...")
+
+        # 分词
+        tokenized_docs = [self.tokenize(doc) for doc in documents]
+
+        # 创建 BM25 索引
+        self.bm25 = BM25Okapi(tokenized_docs)
+        self.metadata = metadata
+
+        logging.info(f"BM25 index built with {len(documents)} documents")
+
+    def get_schema_record(self, record: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Return Gold schema record or validated Silver extracted payload."""
+        extracted = record.get('extracted_json')
+        if isinstance(extracted, dict):
+            if extracted.get('status') == 'insufficient_evidence' or extracted.get('error'):
+                return None
+            payload = dict(extracted)
+            payload['source_chunk_id'] = record.get('chunk_id', '')
+            payload['evidence_level'] = record.get('evidence_level', 'low')
+            payload['support_status'] = record.get('support_status', 'unsupported')
+            return payload
+
+        if 'material_context' in record or 'dominant_ros_or_product' in record:
+            return record
+
+        return None
+
+    def json_record_to_text(self, record: Dict[str, Any]) -> str:
+        """构造 JSON 记录的 BM25 文本。"""
+        text_parts = []
+
+        material = record.get('material_context', {})
+        for key in ['material_type', 'active_site', 'structural_feature']:
+            if material.get(key):
+                text_parts.append(material[key])
+
+        reaction_env = record.get('reaction_environment', {})
+        for key in ['phase', 'oxidant_source', 'driving_force', 'notes']:
+            if reaction_env.get(key):
+                text_parts.append(reaction_env[key])
+
+        for field in ['dominant_ros_or_product', 'secondary_or_intermediate_ros', 'mechanism_intermediates']:
+            for item in record.get(field, []):
+                if not isinstance(item, dict):
+                    continue
+                for key in ['species', 'category', 'role', 'basis', 'evidence', 'notation_in_paper']:
+                    if item.get(key):
+                        text_parts.append(item[key])
+
+        for step in record.get('mechanism', []):
+            if step:
+                text_parts.append(step)
+
+        return ' '.join(text_parts)
+
+    def save_index(self):
+        """保存 BM25 索引"""
+        ensure_directory(self.output_dir)
+
+        # 保存元数据
+        metadata_path = self.output_dir / 'bm25_metadata.jsonl'
+        with open(metadata_path, 'w', encoding='utf-8') as f:
+            for meta in self.metadata:
+                f.write(json.dumps(meta, ensure_ascii=False) + '\n')
+        logging.info(f"BM25 metadata saved to {metadata_path}")
+
+        # 保存文本，用于检索阶段重建 BM25Okapi
+        documents_path = self.output_dir / 'bm25_documents.jsonl'
+        with open(documents_path, 'w', encoding='utf-8') as f:
+            for meta in self.metadata:
+                f.write(json.dumps({'text': meta.get('text', '')}, ensure_ascii=False) + '\n')
+        logging.info(f"BM25 documents saved to {documents_path}")
+
+        index_path = self.output_dir / 'bm25_index.pkl'
+        with open(index_path, 'wb') as f:
+            pickle.dump(self.bm25, f, protocol=pickle.HIGHEST_PROTOCOL)
+        logging.info(f"BM25 index saved to {index_path}")
+
+        # 保存索引信息
+        index_info_path = self.output_dir / 'bm25_info.json'
+        index_info = {
+            'document_count': len(self.metadata),
+            'created_at': datetime.now().isoformat(),
+            'metadata_path': str(metadata_path),
+            'documents_path': str(documents_path),
+            'index_path': str(index_path),
+        }
+        with open(index_info_path, 'w', encoding='utf-8') as f:
+            json.dump(index_info, f, ensure_ascii=False, indent=2)
+        logging.info(f"BM25 info saved to {index_info_path}")
+
+    def build_from_json_records(self, jsonl_path: Path):
+        """从 JSON 记录构建索引"""
+        logging.info(f"Building BM25 index from {jsonl_path}")
+
+        documents = []
+        metadata = []
+
+        if not jsonl_path.exists():
+            logging.warning(f"File not found: {jsonl_path}")
+            return
+
+        with open(jsonl_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                if line.strip():
+                    record = json.loads(line)
+
+                    # 构造文档文本
+                    text_parts = []
+
+                    material = record.get('material_context', {})
+                    material_type = material.get('material_type', '')
+                    active_site = material.get('active_site', '')
+
+                    if material_type:
+                        text_parts.append(material_type)
+                    if active_site:
+                        text_parts.append(active_site)
+
+                    reaction_env = record.get('reaction_environment', {})
+                    phase = reaction_env.get('phase', '')
+                    oxidant = reaction_env.get('oxidant_source', '')
+                    driving = reaction_env.get('driving_force', '')
+
+                    if phase:
+                        text_parts.append(phase)
+                    if oxidant:
+                        text_parts.append(oxidant)
+                    if driving:
+                        text_parts.append(driving)
+
+                    dominant = record.get('dominant_ros_or_product', [])
+                    for item in dominant:
+                        species = item.get('species', '')
+                        evidence = item.get('evidence', '')
+                        if species:
+                            text_parts.append(species)
+                        if evidence:
+                            text_parts.append(evidence)
+
+                    secondary = record.get('secondary_or_intermediate_ros', [])
+                    for item in secondary:
+                        species = item.get('species', '')
+                        evidence = item.get('evidence', '')
+                        if species:
+                            text_parts.append(species)
+                        if evidence:
+                            text_parts.append(evidence)
+
+                    mechanism = record.get('mechanism', [])
+                    for step in mechanism:
+                        if step:
+                            text_parts.append(step)
+
+                    text = ' '.join(text_parts)
+                    if not text.strip():
+                        continue
+
+                    documents.append(text)
+
+                    meta = {
+                        'doc_id': record.get('source_chunk_id', ''),
+                        'source_type': 'gold_silver_json',
+                        'evidence_level': record.get('evidence_level', 'low'),
+                        'text': text
+                    }
+                    metadata.append(meta)
+
+        logging.info(f"Loaded {len(documents)} records from {jsonl_path}")
+        self.build_index(documents, metadata)
+
+    def build_from_chunks(self, chunks_path: Path, source_type: str = 'pdf'):
+        """从 chunks 文件构建索引"""
+        logging.info(f"Building BM25 index from {chunks_path}")
+
+        documents = []
+        metadata = []
+
+        if not chunks_path.exists():
+            logging.warning(f"File not found: {chunks_path}")
+            return
+
+        with open(chunks_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                if line.strip():
+                    record = json.loads(line)
+
+                    text = record.get('text', '')
+                    if not text.strip():
+                        continue
+
+                    # 如果只添加 ROS 相关的 chunks
+                    if source_type == 'pdf' and not record.get('is_ros_related', False):
+                        continue
+
+                    documents.append(text)
+
+                    meta = {
+                        'doc_id': record.get('chunk_id', ''),
+                        'source_type': source_type,
+                        'source': record.get('source_type', 'pdf_fulltext'),
+                        'title': record.get('title', ''),
+                        'doi': record.get('doi', ''),
+                        'page': record.get('page', 0),
+                        'is_ros_related': record.get('is_ros_related', False),
+                        'text': text
+                    }
+                    metadata.append(meta)
+
+        logging.info(f"Loaded {len(documents)} {source_type} chunks")
+        self.build_index(documents, metadata)
+
+    def run(self):
+        """运行索引构建流程"""
+        logging.info("Starting BM25 index construction")
+
+        documents = []
+        metadata = []
+
+        # 从 Gold/Silver JSON 记录构建
+        processed_dir = Path(self.config['paths']['processed_dir'])
+        json_paths = [
+            processed_dir / 'json_records_cleaned.jsonl',
+            processed_dir / 'silver_json_validated.jsonl',
+        ]
+        for json_path in json_paths:
+            if not json_path.exists():
+                logging.warning(f"Input file not found, skipping: {json_path}")
+                continue
+
+            json_docs, json_meta = [], []
+            with open(json_path, 'r', encoding='utf-8') as f:
+                for line in f:
+                    if line.strip():
+                        raw_record = json.loads(line)
+                        record = self.get_schema_record(raw_record)
+                        if not record:
+                            continue
+
+                        text = self.json_record_to_text(record)
+                        if text.strip():
+                            json_docs.append(text)
+                            json_meta.append({
+                                'doc_id': record.get('source_chunk_id', raw_record.get('chunk_id', '')),
+                                'source_type': 'gold_silver_json',
+                                'evidence_level': record.get('evidence_level', 'low'),
+                                'text': text
+                            })
+            documents.extend(json_docs)
+            metadata.extend(json_meta)
+            logging.info(f"Added {len(json_docs)} records from {json_path}")
+
+        # 从 PDF chunks 构建
+        pdf_path = Path(self.config['paths']['processed_dir']) / 'pdf_chunks.jsonl'
+        if pdf_path.exists():
+            pdf_docs, pdf_meta = [], []
+            with open(pdf_path, 'r', encoding='utf-8') as f:
+                for line in f:
+                    if line.strip():
+                        record = json.loads(line)
+                        text = record.get('text', '')
+                        if text.strip() and record.get('is_ros_related', False):
+                            pdf_docs.append(text)
+                            pdf_meta.append({
+                                'doc_id': record.get('chunk_id', ''),
+                                'source_type': 'pdf',
+                                'title': record.get('title', ''),
+                                'doi': record.get('doi', ''),
+                                'is_ros_related': record.get('is_ros_related', False),
+                                'text': text
+                            })
+            documents.extend(pdf_docs)
+            metadata.extend(pdf_meta)
+            logging.info(f"Added {len(pdf_docs)} PDF chunks")
+
+        # 从 WoS abstracts 构建
+        wos_path = Path(self.config['paths']['processed_dir']) / 'wos_abstract_chunks.jsonl'
+        if wos_path.exists():
+            wos_docs, wos_meta = [], []
+            with open(wos_path, 'r', encoding='utf-8') as f:
+                for line in f:
+                    if line.strip():
+                        record = json.loads(line)
+                        text = record.get('text', '')
+                        if text.strip():
+                            wos_docs.append(text)
+                            wos_meta.append({
+                                'doc_id': record.get('doc_id', ''),
+                                'source_type': 'wos',
+                                'title': record.get('title', ''),
+                                'doi': record.get('doi', ''),
+                                'year': record.get('year'),
+                                'text': text
+                            })
+            documents.extend(wos_docs)
+            metadata.extend(wos_meta)
+            logging.info(f"Added {len(wos_docs)} WOS abstracts")
+
+        # 构建索引
+        self.build_index(documents, metadata)
+
+        # 保存索引
+        self.save_index()
+
+        logging.info(f"BM25 index construction completed. Total documents: {len(self.metadata)}")
+
+
+def main():
+    """主函数"""
+    import argparse
+
+    parser = argparse.ArgumentParser(description='Build BM25 index for COF-ROS retrieval')
+    parser.add_argument('--config', type=str, default='config/config.yaml',
+                        help='Path to config file')
+    args = parser.parse_args()
+
+    # 加载配置
+    config_path = Path(args.config)
+    with open(config_path, 'r', encoding='utf-8') as f:
+        config = yaml.safe_load(f)
+
+    # 运行构建
+    builder = BM25IndexBuilder(config)
+    builder.run()
+
+
+if __name__ == '__main__':
+    main()