FinRAG-Agent/FinRAG-Pro
1
1"""2金融投研RAG系统 - 知识库构建模块 v5 (精准制导版)3完全适配 reports/annual 和 reports/industry 目录结构4"""5 6import os7import re8import json9import pickle10import numpy as np11 12BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))13# 精准定位你的两个子文件夹14ANNUAL_DIR = os.path.join(BASE_DIR, "reports", "annual")15INDUSTRY_DIR = os.path.join(BASE_DIR, "reports", "industry")16 17PROCESSED_DIR = os.path.join(BASE_DIR, "data", "processed")18INDEX_DIR = os.path.join(BASE_DIR, "data", "index")19 20os.makedirs(PROCESSED_DIR, exist_ok=True)21os.makedirs(INDEX_DIR, exist_ok=True)22 23def load_specific_directory(directory: str, source_type: str) -> list[dict]:24 """精准读取指定目录下的 TXT 和 PDF"""25 documents = []26 27 if not os.path.exists(directory):28 print(f"[WARN] 找不到文件夹: {directory}")29 return documents30 31 print(f"\n📁 正在精准扫描目录: {directory}")32 files = os.listdir(directory)33 34 for filename in files:35 file_path = os.path.join(directory, filename)36 if not os.path.isfile(file_path):37 continue38 39 name_no_ext, ext = os.path.splitext(filename)40 ext = ext.lower() # 无视大小写41 42 if ext not in ['.txt', '.pdf']:43 continue44 45 # --- 读取 TXT 文件 ---46 if ext == '.txt':47 try:48 with open(file_path, "r", encoding="utf-8") as f:49 content = f.read()50 documents.append({51 "id": name_no_ext, "filename": filename, "title": name_no_ext,52 "content": content, "source": filename, 53 "source_type": source_type, "source_dir": os.path.basename(directory)54 })55 print(f" ✅ [TXT] 成功读取: {filename}")56 except Exception as e:57 print(f" ❌ [TXT] 读取失败 {filename}: {e}")58 59 # --- 读取 PDF 文件 ---60 elif ext == '.pdf':61 try:62 import fitz # PyMuPDF63 pdf_doc = fitz.open(file_path)64 content_parts = []65 for page in pdf_doc:66 page_text = page.get_text()67 if page_text.strip():68 content_parts.append(page_text)69 content = "\n\n".join(content_parts)70 pdf_doc.close()71 72 documents.append({73 "id": name_no_ext, "filename": filename, "title": name_no_ext,74 "content": content, "source": filename, 75 "source_type": source_type, "source_dir": os.path.basename(directory)76 })77 print(f" ✅ [PDF] 成功读取: {filename}")78 except Exception as e:79 print(f" ❌ [PDF] 读取失败 {filename}: {e}")80 81 return documents82 83def clean_text(text: str) -> str:84 text = re.sub(r'={5,}.*?={5,}', '', text)85 text = re.sub(r'\n{3,}', '\n\n', text)86 lines = [line.strip() for line in text.split('\n')]87 text = '\n'.join(lines)88 text = re.sub(r'[ \t]+', ' ', text)89 text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)90 return text.strip()91 92def chunk_text(text: str) -> list[str]:93 from langchain_text_splitters import RecursiveCharacterTextSplitter94 splitter = RecursiveCharacterTextSplitter(95 chunk_size=800, chunk_overlap=150,96 separators=["\n\n", "\n", "。", "！", "？", "，", " "],97 length_function=len98 )99 chunks = splitter.split_text(text)100 return [c.strip() for c in chunks if c.strip()]101 102def build_index():103 print("=" * 60)104 print(" 金融投研RAG - 知识库构建 v5（精准制导版）")105 print("=" * 60)106 107 # 1. 加载文档 (分两批精准加载)108 print("\n[步骤1] 按分类加载文档...")109 docs_annual = load_specific_directory(ANNUAL_DIR, "官方年报")110 docs_industry = load_specific_directory(INDUSTRY_DIR, "第三方研报")111 112 all_documents = docs_annual + docs_industry113 114 if not all_documents:115 print("[ERROR] 苍天啊！文件依然没读到，请检查路径！")116 return False117 118 print(f"\n 🎉 共成功加载 {len(all_documents)} 份文档！")119 120 # 2. 清洗121 print("\n[步骤2] 清洗文本...")122 for doc in all_documents:123 doc["content"] = clean_text(doc["content"])124 125 # 3. 切分126 print("\n[步骤3] 切分文本chunks...")127 all_chunks = []128 chunk_metadata = []129 130 for doc in all_documents:131 chunks = chunk_text(doc["content"])132 for i, chunk in enumerate(chunks):133 chunk_id = f"{doc['id']}_chunk_{i:04d}"134 all_chunks.append(chunk)135 chunk_metadata.append({136 "chunk_id": chunk_id, "doc_id": doc["id"], "doc_title": doc["title"],137 "source": doc["source"], "source_type": doc["source_type"], 138 "source_dir": doc["source_dir"], "chunk_index": i,139 })140 141 # 保存Chunks142 chunks_path = os.path.join(PROCESSED_DIR, "chunks.json")143 meta_path = os.path.join(PROCESSED_DIR, "chunk_metadata.json")144 with open(chunks_path, "w", encoding="utf-8") as f:145 json.dump(all_chunks, f, ensure_ascii=False, indent=2)146 with open(meta_path, "w", encoding="utf-8") as f:147 json.dump(chunk_metadata, f, ensure_ascii=False, indent=2)148 149 # 4. 生成向量150 print("\n[步骤4] 生成向量 (TF-IDF)...")151 from sklearn.feature_extraction.text import TfidfVectorizer152 vectorizer = TfidfVectorizer(max_features=5000, analyzer='char', ngram_range=(1, 3), sublinear_tf=True)153 embeddings = vectorizer.fit_transform(all_chunks).toarray().astype(np.float32)154 155 v_path = os.path.join(INDEX_DIR, "tfidf_vectorizer.pkl")156 with open(v_path, "wb") as f:157 pickle.dump(vectorizer, f)158 159 # 5. 保存FAISS160 print("\n[步骤5] 保存FAISS索引...")161 import faiss162 dimension = embeddings.shape[1]163 index = faiss.IndexFlatIP(dimension)164 faiss.normalize_L2(embeddings)165 index.add(embeddings.astype(np.float32))166 index_path = os.path.join(INDEX_DIR, "faiss_index.bin")167 faiss.write_index(index, index_path)168 169 summary = {"total_documents": len(all_documents), "total_chunks": len(all_chunks)}170 with open(os.path.join(PROCESSED_DIR, "summary.json"), "w", encoding="utf-8") as f:171 json.dump(summary, f, ensure_ascii=False, indent=2)172 173 print("\n✅ 奇迹发生！知识库构建彻底完成！")174 return True175 176if __name__ == "__main__":177 build_index()