@lzhzzzzwill/cofos 1.0.1 → 1.0.4

package/runtime/src/retrieval/build_faiss_index.py

@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+"""
+COF-ROS-Reasoner: FAISS 向量索引构建脚本
+
+目标：
+- 构建 FAISS 向量索引用于相似性检索
+- 支持来自 JSON evidence、PDF chunks、WOS abstracts 的文本
+"""
+
+import os
+import json
+import yaml
+import logging
+from pathlib import Path
+from typing import Dict, List, Any, Optional
+from datetime import datetime
+
+import numpy as np
+
+# 添加 src 路径
+
+try:
+    from sentence_transformers import SentenceTransformer
+except ImportError:
+    SentenceTransformer = None
+
+try:
+    import faiss
+except ImportError:
+    faiss = None
+
+from data_processing.utils import setup_logging, ensure_directory
+
+
+class FAISSIndexBuilder:
+    """构建 FAISS 向量索引"""
+
+    def __init__(self, config: Dict[str, Any]):
+        self.config = config
+        self.output_dir = Path(config['paths']['retrieval_store'])
+
+        # 设置日志
+        setup_logging(config)
+
+        # 嵌入模型配置
+        self.embedding_model_name = config.get('models', {}).get('embedding_model', '')
+
+        # 加载嵌入模型
+        self._load_embedding_model()
+
+        # 索引参数
+        self.dimension = 1024  # BGE-M3 的维度
+        self.index = None
+
+        # 存储元数据
+        self.metadata: List[Dict[str, Any]] = []
+
+    def _load_embedding_model(self):
+        """加载嵌入模型"""
+        if not self.embedding_model_name:
+            raise ValueError(
+                "FAISS is enabled but models.embedding_model is empty. "
+                "Set it to a local embedding model directory, or keep retrieval.enable_faiss=false."
+            )
+
+        if SentenceTransformer is None:
+            logging.error("sentence-transformers not installed. Run: pip install sentence-transformers")
+            raise ImportError("sentence-transformers required")
+
+        logging.info(f"Loading embedding model: {self.embedding_model_name}")
+        try:
+            self.model = SentenceTransformer(self.embedding_model_name)
+        except Exception as exc:
+            message = (
+                f"Failed to load embedding model '{self.embedding_model_name}'. "
+                "If this machine is offline, download the model on a machine with network access "
+                "and set models.embedding_model in config/config.yaml to the local model directory."
+            )
+            logging.error(message)
+            raise RuntimeError(message) from exc
+
+        # 获取模型维度
+        try:
+            self.dimension = self.model.get_sentence_embedding_dimension()
+        except Exception:
+            self.dimension = 1024
+
+        logging.info(f"Model dimension: {self.dimension}")
+
+    def create_index(self):
+        """创建 FAISS 索引"""
+        if faiss is None:
+            logging.error("faiss not installed. Run: pip install faiss-cpu")
+            raise ImportError("faiss required")
+
+        # 使用 IVFFlat 索引
+        nlist = 100  # 簇的数量
+        self.index = faiss.index_factory(self.dimension, f"IVF{nlist},Flat", faiss.METRIC_INNER_PRODUCT)
+
+        # 训练索引（需要一些样本）
+        logging.info("Training FAISS index...")
+        # 使用随机样本训练
+        samples = np.random.randn(1000, self.dimension).astype('float32')
+        self.index.train(samples)
+
+        logging.info("FAISS index created")
+
+    def embed_texts(self, texts: List[str]) -> np.ndarray:
+        """嵌入文本列表"""
+        embeddings = self.model.encode(texts, normalize_embeddings=True)
+        return embeddings.astype('float32')
+
+    def add_texts(self, texts: List[str], metadata: List[Dict[str, Any]]):
+        """添加文本到索引"""
+        if len(texts) == 0:
+            return
+
+        logging.info(f"Adding {len(texts)} texts to index...")
+
+        # 嵌入文本
+        embeddings = self.embed_texts(texts)
+
+        # 添加到 FAISS 索引
+        self.index.add(embeddings)
+
+        # 保存元数据
+        self.metadata.extend(metadata)
+
+        logging.info(f"Added {len(texts)} vectors. Total: {self.index.ntotal}")
+
+    def save_index(self):
+        """保存 FAISS 索引和元数据"""
+        ensure_directory(self.output_dir)
+
+        # 保存 FAISS 索引
+        index_path = self.output_dir / 'faiss_index.bin'
+        faiss.write_index(self.index, str(index_path))
+        logging.info(f"FAISS index saved to {index_path}")
+
+        # 保存元数据
+        metadata_path = self.output_dir / 'faiss_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"Metadata saved to {metadata_path}")
+
+    def build_from_json_records(self, jsonl_path: Path):
+        """从 JSON 记录构建索引"""
+        logging.info(f"Building index from {jsonl_path}")
+
+        texts = []
+        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)
+                    payload = record.get('extracted_json')
+                    if isinstance(payload, dict) and payload.get('status') != 'insufficient_evidence':
+                        source_record = payload
+                    else:
+                        source_record = record
+
+                    # 构造文本
+                    text_parts = []
+
+                    # 来自 JSON evidence
+                    dominant = source_record.get('dominant_ros_or_product', [])
+                    for item in dominant:
+                        species = item.get('species', '')
+                        evidence = item.get('evidence', '')
+                        if species:
+                            text_parts.append(f"Species: {species}")
+                        if evidence:
+                            text_parts.append(f"Evidence: {evidence}")
+
+                    secondary = source_record.get('secondary_or_intermediate_ros', [])
+                    for item in secondary:
+                        species = item.get('species', '')
+                        evidence = item.get('evidence', '')
+                        if species:
+                            text_parts.append(f"Secondary: {species}")
+                        if evidence:
+                            text_parts.append(f"Evidence: {evidence}")
+
+                    material = source_record.get('material_context', {})
+                    material_type = material.get('material_type', '')
+                    active_site = material.get('active_site', '')
+
+                    if material_type:
+                        text_parts.insert(0, f"Material: {material_type}")
+                    if active_site:
+                        text_parts.insert(1, f"Active site: {active_site}")
+
+                    reaction_env = source_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(f"Phase: {phase}")
+                    if oxidant:
+                        text_parts.append(f"Oxidant: {oxidant}")
+                    if driving:
+                        text_parts.append(f"Driving force: {driving}")
+
+                    mechanism = source_record.get('mechanism', [])
+                    for step in mechanism:
+                        if step:
+                            text_parts.append(f"Mechanism: {step}")
+
+                    # 合并文本
+                    text = ' '.join(text_parts)
+                    if not text.strip():
+                        continue
+
+                    texts.append(text)
+
+                    # 构造元数据
+                    meta = {
+                        'doc_id': record.get('source_chunk_id', record.get('chunk_id', '')),
+                        'source_type': 'gold_silver_json',
+                        'evidence_level': record.get('evidence_level', 'low'),
+                        'text': text
+                    }
+                    metadata.append(meta)
+
+        logging.info(f"Loaded {len(texts)} records from {jsonl_path}")
+        self.add_texts(texts, metadata)
+
+    def build_from_chunks(self, chunks_path: Path, source_type: str = 'pdf'):
+        """从 chunks 文件构建索引"""
+        logging.info(f"Building index from {chunks_path}")
+
+        texts = []
+        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
+
+                    texts.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(texts)} {source_type} chunks")
+        self.add_texts(texts, metadata)
+
+    def run(self):
+        """运行索引构建流程"""
+        logging.info("Starting FAISS index construction")
+
+        # 创建索引
+        self.create_index()
+
+        # 从 JSON 记录构建
+        gold_json_path = Path(self.config['paths']['processed_dir']) / 'json_records_cleaned.jsonl'
+        if gold_json_path.exists():
+            self.build_from_json_records(gold_json_path)
+
+        json_path = Path(self.config['paths']['processed_dir']) / 'silver_json_validated.jsonl'
+        if json_path.exists():
+            self.build_from_json_records(json_path)
+
+        # 从 PDF chunks 构建
+        pdf_path = Path(self.config['paths']['processed_dir']) / 'pdf_chunks.jsonl'
+        if pdf_path.exists():
+            self.build_from_chunks(pdf_path, 'pdf')
+
+        # 从 WoS abstracts 构建
+        wos_path = Path(self.config['paths']['processed_dir']) / 'wos_abstract_chunks.jsonl'
+        if wos_path.exists():
+            self.build_from_chunks(wos_path, 'wos')
+
+        # 保存索引
+        self.save_index()
+
+        logging.info(f"FAISS index construction completed. Total vectors: {self.index.ntotal}")
+
+
+def main():
+    """主函数"""
+    import argparse
+
+    parser = argparse.ArgumentParser(description='Build FAISS 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 = FAISSIndexBuilder(config)
+    builder.run()
+
+
+if __name__ == '__main__':
+    main()