#!/usr/bin/env python # -*- coding: utf-8 -*- ''' PDF文档转RAG知识库处理模块 支持将世界观文档(如COC.pdf)转换为RAG检索格式 ''' import json import os import re from typing import List, Dict, Tuple import hashlib from datetime import datetime try: import PyPDF2 import fitz # pymupdf PYMUPDF_AVAILABLE = True except ImportError: PYMUPDF_AVAILABLE = False print("Warning: pymupdf not available, using PyPDF2 only") try: from sentence_transformers import SentenceTransformer import faiss import numpy as np EMBEDDING_AVAILABLE = True except ImportError: EMBEDDING_AVAILABLE = False print("Warning: sentence-transformers or faiss not available, using simple text matching") class PDFToRAGProcessor: def __init__(self, embedding_model: str = "./sentence-transformers/all-MiniLM-L6-v2"): """ 初始化PDF处理器 Args: embedding_model: 向量化模型名称 """ global EMBEDDING_AVAILABLE self.chunks = [] self.embeddings = None self.index = None if EMBEDDING_AVAILABLE: try: self.embedding_model = SentenceTransformer(embedding_model) print(f"✓ 向量模型加载成功: {embedding_model}") except Exception as e: print(f"✗ 向量模型加载失败: {e}") self.embedding_model = None EMBEDDING_AVAILABLE = False else: self.embedding_model = None def extract_text_from_pdf(self, pdf_path: str) -> str: """从PDF提取文本""" if not os.path.exists(pdf_path): raise FileNotFoundError(f"PDF文件不存在: {pdf_path}") text = "" # 优先使用pymupdf，效果更好 if PYMUPDF_AVAILABLE: try: doc = fitz.open(pdf_path) for page_num in range(len(doc)): page = doc.load_page(page_num) text += page.get_text() doc.close() print(f"✓ 使用pymupdf提取文本成功") return text except Exception as e: print(f"✗ pymupdf提取失败: {e}, 尝试PyPDF2") # 备用PyPDF2 try: with open(pdf_path, 'rb') as file: pdf_reader = PyPDF2.PdfReader(file) for page in pdf_reader.pages: text += page.extract_text() print(f"✓ 使用PyPDF2提取文本成功") except Exception as e: raise Exception(f"PDF文本提取失败: {e}") return text def clean_text(self, text: str) -> str: """清理文本""" # 移除多余空行 text = re.sub(r'\n\s*\n', '\n\n', text) # 移除页码等 text = re.sub(r'第\s*\d+\s*页', '', text) text = re.sub(r'Page\s*\d+', '', text) # 统一标点符号 text = text.replace('　', ' ') # 全角空格转半角 text = re.sub(r'\s+', ' ', text) # 多个空格合并 return text.strip() def chunk_text_by_semantic(self, text: str, max_chunk_size: int = 500) -> List[Dict]: """按语义分块文本""" # 按段落分割 paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()] chunks = [] current_chunk = "" current_size = 0 for para in paragraphs: para_size = len(para) # 如果当前段落太长，需要进一步分割 if para_size > max_chunk_size: # 保存当前chunk if current_chunk: chunks.append({ "content": current_chunk.strip(), "size": current_size, "type": "paragraph" }) current_chunk = "" current_size = 0 # 按句子分割长段落 sentences = re.split(r'[。！？；]\s*', para) for sentence in sentences: if sentence.strip(): sentence += '。' # 添加标点 if current_size + len(sentence) > max_chunk_size and current_chunk: chunks.append({ "content": current_chunk.strip(), "size": current_size, "type": "sentence" }) current_chunk = sentence current_size = len(sentence) else: current_chunk += sentence current_size += len(sentence) else: # 正常段落处理 if current_size + para_size > max_chunk_size and current_chunk: chunks.append({ "content": current_chunk.strip(), "size": current_size, "type": "paragraph" }) current_chunk = para current_size = para_size else: current_chunk += "\n" + para if current_chunk else para current_size += para_size # 保存最后的chunk if current_chunk: chunks.append({ "content": current_chunk.strip(), "size": current_size, "type": "paragraph" }) # 为每个chunk添加唯一ID和元数据 for i, chunk in enumerate(chunks): chunk["id"] = f"chunk_{i:04d}" chunk["hash"] = hashlib.md5(chunk["content"].encode()).hexdigest()[:8] return chunks def extract_key_concepts(self, chunks: List[Dict]) -> List[Dict]: """提取关键概念和术语""" concepts = [] # 简单的关键词提取规则 concept_patterns = [ r'(?:技能|属性|规则)[:：]\s*([^。\n]+)', r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)', # 英文专有名词 r'【([^】]+)】', # 中文术语标记 r'"([^"]+)"', # 引号内容 ] for chunk in chunks: content = chunk["content"] chunk_concepts = [] for pattern in concept_patterns: matches = re.findall(pattern, content) chunk_concepts.extend(matches) if chunk_concepts: concepts.append({ "chunk_id": chunk["id"], "concepts": list(set(chunk_concepts)), # 去重 "content_preview": content[:100] + "..." if len(content) > 100 else content }) return concepts def build_vector_index(self, chunks: List[Dict]) -> bool: """构建向量索引""" if not EMBEDDING_AVAILABLE or not self.embedding_model: print("✗ 向量化功能不可用，将使用文本匹配") return False try: # 提取文本内容 texts = [chunk["content"] for chunk in chunks] # 生成向量 print("生成文档向量...") embeddings = self.embedding_model.encode(texts, show_progress_bar=True) # 构建FAISS索引 dimension = embeddings.shape[1] self.index = faiss.IndexFlatL2(dimension) self.index.add(embeddings.astype(np.float32)) self.embeddings = embeddings print(f"✓ 向量索引构建完成，维度: {dimension}, 文档数: {len(chunks)}") return True except Exception as e: print(f"✗ 向量索引构建失败: {e}") return False def process_pdf_to_rag(self, pdf_path: str, output_dir: str = "./rag_knowledge") -> Dict: """ 完整处理PDF到RAG知识库 Args: pdf_path: PDF文件路径 output_dir: 输出目录 Returns: 处理结果统计 """ print(f"开始处理PDF: {pdf_path}") # 创建输出目录 os.makedirs(output_dir, exist_ok=True) # 1. 提取文本 raw_text = self.extract_text_from_pdf(pdf_path) print(f"✓ 提取文本长度: {len(raw_text)} 字符") # 2. 清理文本 clean_text = self.clean_text(raw_text) # 保存清理后的文本 text_output_path = os.path.join(output_dir, "extracted_text.txt") with open(text_output_path, 'w', encoding='utf-8') as f: f.write(clean_text) print(f"✓ 清理后文本保存至: {text_output_path}") # 3. 语义分块 chunks = self.chunk_text_by_semantic(clean_text) self.chunks = chunks print(f"✓ 文档分块完成: {len(chunks)} 个块") # 4. 提取关键概念 concepts = self.extract_key_concepts(chunks) print(f"✓ 提取关键概念: {len(concepts)} 个") # 5. 构建向量索引 vector_success = self.build_vector_index(chunks) # 6. 保存知识库文件 knowledge_base = { "metadata": { "source_file": os.path.basename(pdf_path), "processed_time": datetime.now().isoformat(), "total_chunks": len(chunks), "total_concepts": len(concepts), "vector_enabled": vector_success }, "chunks": chunks, "concepts": concepts } # 保存JSON知识库 kb_output_path = os.path.join(output_dir, "knowledge_base.json") with open(kb_output_path, 'w', encoding='utf-8') as f: json.dump(knowledge_base, f, ensure_ascii=False, indent=2) print(f"✓ 知识库保存至: {kb_output_path}") # 保存向量索引 if vector_success: index_path = os.path.join(output_dir, "vector_index.faiss") faiss.write_index(self.index, index_path) embeddings_path = os.path.join(output_dir, "embeddings.npy") np.save(embeddings_path, self.embeddings) print(f"✓ 向量索引保存至: {index_path}") return { "status": "success", "chunks_count": len(chunks), "concepts_count": len(concepts), "vector_enabled": vector_success, "output_dir": output_dir } def load_knowledge_base(self, knowledge_dir: str) -> bool: """加载已有知识库""" try: # 加载JSON知识库 kb_path = os.path.join(knowledge_dir, "knowledge_base.json") with open(kb_path, 'r', encoding='utf-8') as f: knowledge_base = json.load(f) self.chunks = knowledge_base["chunks"] # 加载向量索引 if EMBEDDING_AVAILABLE: index_path = os.path.join(knowledge_dir, "vector_index.faiss") embeddings_path = os.path.join(knowledge_dir, "embeddings.npy") if os.path.exists(index_path) and os.path.exists(embeddings_path): self.index = faiss.read_index(index_path) self.embeddings = np.load(embeddings_path) print(f"✓ 向量索引加载成功") print(f"✓ 知识库加载成功: {len(self.chunks)} 个文档块") return True except Exception as e: print(f"✗ 知识库加载失败: {e}") return False def search_relevant_content(self, query: str, top_k: int = 3) -> List[Dict]: """搜索相关内容""" if not self.chunks: return [] # 向量搜索 if EMBEDDING_AVAILABLE and self.embedding_model and self.index: try: query_vector = self.embedding_model.encode([query]) distances, indices = self.index.search(query_vector.astype(np.float32), top_k) results = [] for i, (distance, idx) in enumerate(zip(distances[0], indices[0])): if idx < len(self.chunks): result = self.chunks[idx].copy() result["relevance_score"] = float(1 / (1 + distance)) # 转换为相似度分数 result["rank"] = i + 1 results.append(result) return results except Exception as e: print(f"向量搜索失败: {e}, 使用文本匹配") # 文本匹配搜索 query_lower = query.lower() scored_chunks = [] for chunk in self.chunks: content_lower = chunk["content"].lower() # 简单的相关性评分 score = 0 query_words = query_lower.split() for word in query_words: if word in content_lower: score += content_lower.count(word) if score > 0: result = chunk.copy() result["relevance_score"] = score scored_chunks.append(result) # 按分数排序 scored_chunks.sort(key=lambda x: x["relevance_score"], reverse=True) # 添加排名 for i, chunk in enumerate(scored_chunks[:top_k]): chunk["rank"] = i + 1 return scored_chunks[:top_k] def main(): """测试PDF处理功能""" processor = PDFToRAGProcessor() # 示例：处理COC规则书 pdf_path = input("请输入PDF文件路径 (如: ./coc.pdf): ").strip() if not os.path.exists(pdf_path): print(f"文件不存在: {pdf_path}") return try: result = processor.process_pdf_to_rag(pdf_path) print(f"\n{'='*50}") print(f"PDF处理完成!") print(f"状态: {result['status']}") print(f"文档块数量: {result['chunks_count']}") print(f"关键概念数量: {result['concepts_count']}") print(f"向量索引: {'启用' if result['vector_enabled'] else '未启用'}") print(f"输出目录: {result['output_dir']}") # 测试搜索 print(f"\n{'='*50}") print("测试知识库搜索:") while True: query = input("\n请输入搜索关键词 (输入'quit'退出): ").strip() if query.lower() == 'quit': break results = processor.search_relevant_content(query, top_k=3) if results: print(f"\n找到 {len(results)} 个相关结果:") for result in results: print(f"\n排名 {result['rank']} (相关度: {result['relevance_score']:.3f}):") content = result['content'] preview = content[:200] + "..." if len(content) > 200 else content print(f"{preview}") print("-" * 40) else: print("未找到相关内容") except Exception as e: print(f"处理失败: {e}") if __name__ == '__main__': main()