from langchain_community.document_loaders import DirectoryLoader, JSONLoader, UnstructuredMarkdownLoader from langchain.text_splitter import RecursiveCharacterTextSplitter, MarkdownTextSplitter, MarkdownHeaderTextSplitter from langchain_community.embeddings import HuggingFaceEmbeddings from langchain_community.vectorstores import FAISS from pathlib import Path from transformers import AutoModel, AutoTokenizer import torch import torch.nn.functional as F import numpy as np import config as cfg class LocalRAG: def __init__(self, rag_top_k=3, doc_dir="rag/kb/BIGOLIVE及公司介绍/", # 默认加载这个，若选择角色扮演模式，可根据角色选择 vector_db_path="rag/vector_db/", embed_model=cfg.DEFAULT_EMBEDDING_MODEL ): self.rag_top_k = rag_top_k self.doc_dir = doc_dir # 本地知识库的文档目录 self.vector_db_path = vector_db_path # 向量数据库存储路径 self.embed_model = embed_model self.build_vector_db() def build_vector_db(self): # 加载文档（支持PDF、TXT、DOCX） if isinstance(self.doc_dir, list): general_docs = [] json_docs = [] md_docs = [] for doc_dir in self.doc_dir: # 处理一般文件，txt等 loader = DirectoryLoader(doc_dir, glob="**/*.[!json!md]*") # "**/[!.]*" tmp_docs = loader.load() general_docs.extend(tmp_docs) # 额外处理json文件 for json_file in Path(doc_dir).rglob("*.json"): loader = JSONLoader( file_path=str(json_file), jq_schema=".[] | {spk: .spk, text: .text}", text_content=False) data = loader.load() for iidx in range(len(data)): data[iidx].page_content = bytes(data[iidx].page_content, "utf-8").decode("unicode_escape") json_docs.extend(data) # 额外处理md文件 headers_to_split_on = [ ("#", "Header 1"), ("##", "Header 2"), ("###", "Header 3"), ] for md_file in Path(doc_dir).rglob("*.md"): with open(md_file, 'r') as f: content = f.read() # 定义拆分器，拆分markdown内容 markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on) md_header_splits = markdown_splitter.split_text(content) md_docs.extend(md_header_splits) # loader = UnstructuredMarkdownLoader(md_file, mode="elements") # data = loader.load() # docs.extend(data) else: loader = DirectoryLoader(self.doc_dir, glob="**/*.*") docs = loader.load() # 文本分块 if len(general_docs) > 0: text_splitter = RecursiveCharacterTextSplitter( chunk_size=500, chunk_overlap=50 ) chunks = text_splitter.split_documents(docs) else: chunks = json_docs + md_docs # 生成向量并构建FAISS数据库 embeddings = HuggingFaceEmbeddings(model_name=self.embed_model) self.vector_db = FAISS.from_documents(chunks, embeddings) self.vector_db.save_local(self.vector_db_path) def reload_knowledge_base(self, target_doc_dir): self.doc_dir = target_doc_dir self.build_vector_db() # def reset(self): # self.vector_db = None class LocalRAG_new: def __init__(self, rag_top_k=3, doc_dir="rag/kb/BIGOLIVE及公司介绍/", # 默认加载这个，若选择角色扮演模式，可根据角色选择 vector_db_path="rag/vector_db/", embed_model_path="princeton-nlp/sup-simcse-bert-large-uncased", device=torch.device('cuda:2')): self.rag_top_k = rag_top_k self.doc_dir = doc_dir # 本地知识库的文档目录 self.kb_name = '_'.join([Path(doc_dir[i]).name for i in range(len(doc_dir))]) self.embed_model_name = Path(embed_model_path).name self.vector_db_path = vector_db_path # 向量数据库存储路径 self.embed_model = embed_model_path self.device = device # 加载分词器和模型 self.tokenizer = AutoTokenizer.from_pretrained(self.embed_model) self.embed_model = AutoModel.from_pretrained(self.embed_model).to(device) self.vector_db = None self._vector_db = None self.build_vector_db() class VectorDB: def __init__(self, rag): self._data = rag._vector_db self.rag = rag def similarity_search(self, query, k): # 可能的输入预处理，暂无 # query = input_optimize(query) # 计算query的embedding并与库中比较 with torch.inference_mode(): query_token = self.rag.tokenizer(query, padding=True, truncation=False, return_tensors="pt").to(self.rag.device) query_embed = self.rag.embed_model(**query_token)['last_hidden_state'].mean(dim=1) sim_query = F.cosine_similarity(query_embed.repeat(len(self._data['embeds']), 1), self._data['embeds'], dim=1, eps=1e-8) max_ids_query = torch.argsort(sim_query, descending=True)[:self.rag.rag_top_k].cpu().detach().numpy() return list(zip(np.array(self._data['chunks'])[max_ids_query], sim_query[max_ids_query])) def build_vector_db(self): # 加载文档（支持PDF、TXT、DOCX） if isinstance(self.doc_dir, list): docs = [] for doc_dir in self.doc_dir: loader = DirectoryLoader(doc_dir, glob="**/*.[!json!md]*") # "**/[!.]*" tmp_docs = loader.load() docs.extend(tmp_docs) # # 额外处理json文件 # for json_file in Path(doc_dir).rglob("*.json"): # loader = JSONLoader( # file_path=str(json_file), # jq_schema='.messages[].content', # text_content=False) # data = loader.load() # 额外处理md文件 headers_to_split_on = [ ("#", "Header 1"), ("##", "Header 2"), ("###", "Header 3"), ] for md_file in Path(doc_dir).rglob("*.md"): with open(md_file, 'r') as f: content = f.read() # 定义拆分器，拆分markdown内容 markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on) md_header_splits = markdown_splitter.split_text(content) docs.extend(md_header_splits) # loader = UnstructuredMarkdownLoader(md_file, mode="elements") # data = loader.load() # docs.extend(data) else: loader = DirectoryLoader(self.doc_dir, glob="**/*.*") docs = loader.load() # 文本分块 text_splitter = RecursiveCharacterTextSplitter( chunk_size=500, chunk_overlap=50 ) chunks = text_splitter.split_documents(docs) with torch.inference_mode(): chunk_and_embed = [] for chunk in chunks: chunk_token = self.tokenizer(chunk.page_content, padding=True, truncation=False, return_tensors="pt").to(self.device) chunk_embed = self.embed_model(**chunk_token)['last_hidden_state'].mean(dim=1) chunk_and_embed.append((chunk, chunk_embed)) all_chunks, all_embeds = list(zip(*chunk_and_embed)) all_chunks, all_embeds = list(all_chunks), list(all_embeds) all_embeds = torch.cat(all_embeds, dim=0) self._vector_db = {'chunks': all_chunks, 'embeds': all_embeds} self.vector_db = self.VectorDB(self) torch.save(self.vector_db, str(Path(self.vector_db_path) / f'{self.kb_name}_{self.embed_model_name}.pt')) def reload_knowledge_base(self, target_doc_dir): self.doc_dir = target_doc_dir self.build_vector_db() # def reset(self): # self.vector_db = None class CosPlayer: def __init__(self, description_file): self.update(description_file) def update(self, description_file): self.description_file = description_file with open(description_file, 'r') as f: all_lines = f.readlines() self.core_setting = ''.join(all_lines) self.characters_dir = Path(description_file).parent self.prologue_file = self.description_file.replace('/characters/', '/prologues/') if not Path(self.prologue_file).exists(): self.prologue_file = None def get_all_characters(self): return [str(i) for i in list(self.characters_dir.rglob('*.txt'))] def get_core_setting(self): return self.core_setting def get_prologue(self): if self.prologue_file: with open(self.prologue_file, 'r') as f: all_lines = f.readlines() return ''.join(all_lines) else: return None if __name__ == "__main__": rag = LocalRAG() # # rag.build_vector_db() # doc_dir = "rag/debug" # loader = DirectoryLoader(doc_dir, glob="**/*.*") # docs = loader.load() # # 文本分块 # text_splitter = RecursiveCharacterTextSplitter( # chunk_size=500, # chunk_overlap=50 # ) # chunks = text_splitter.split_documents(docs) # pass