import os import torch import re import hydra import logging from langchain_community.vectorstores import Chroma from BCEmbedding.tools.langchain import BCERerank from langchain_huggingface import HuggingFaceEmbeddings from langchain.prompts import PromptTemplate from langchain.chains import RetrievalQA from langchain.llms.base import LLM from typing import Any, List, Optional, Iterator from langchain.callbacks.manager import CallbackManagerForLLMRun from langchain.retrievers import ContextualCompressionRetriever from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel, TextIteratorStreamer from threading import Thread from modelscope.hub.snapshot_download import snapshot_download import sys sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from rag.chroma_db import get_chroma_db from download_models import download_model # 配置日志 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', filename='ragchat.log', filemode='w') class InternLM(LLM): tokenizer : AutoTokenizer = None model: AutoModelForCausalLM = None llm_system_prompt: str="" def __init__(self, model_path: str, llm_system_prompt: str): super().__init__() logging.info(f"正在从本地:{model_path}加载模型...") try: # 加载分词器和模型 self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.float16).cuda() self.model.eval() # 将模型设置为评估模式 self.llm_system_prompt = llm_system_prompt logging.info("完成本地模型的加载") except Exception as e: logging.error(f"加载模型时发生错误: {e}") raise def _call(self, prompt : str, stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any): # 重写调用函数 system_prompt = self.llm_system_prompt messages = [("system", system_prompt)] response, history = self.model.chat(self.tokenizer, prompt , history=messages) return response def stream(self, prompt: str) -> Iterator[str]: input_ids = self.tokenizer.encode(prompt, return_tensors="pt").cuda() streamer = TextIteratorStreamer(self.tokenizer, skip_special_tokens=True) generation_kwargs = dict( input_ids=input_ids, max_new_tokens=2048, # do_sample=False, # top_k=30, # top_p=0.85, # temperature=0.7, # repetition_penalty=1.1, streamer=streamer, ) thread = Thread(target=self.model.generate, kwargs=generation_kwargs) thread.start() for idx, new_text in enumerate(streamer): # 跳过prompt部分 if idx > 0: yield new_text @property def _llm_type(self) -> str: return "InternLM2" class WuleRAG(): """ 存储检索问答链的对象 """ def __init__(self, data_source_dir, db_persist_directory, base_mode, embeddings_model, reranker_model, rag_prompt_template): # 加载自定义 LLM self.llm = base_mode # 定义 Embeddings ## bce-embedding-base_v1 如果路径不对，则下载默认的模型 if not os.path.exists(embeddings_model): if embeddings_model.endswith("bce-embedding-base_v1"): ## 模型上两级目录是保存路径 save_dir = os.path.dirname(os.path.dirname(embeddings_model)) embeddings_model = snapshot_download("maidalun/bce-embedding-base_v1", cache_dir=save_dir, revision='master') logging.info(f"bce-embedding model not exist, downloading from modelscope \n save to {embeddings_model}") else: raise ValueError(f"{embeddings_model} model not exist, please reset or re-download your model.") self.embeddings = HuggingFaceEmbeddings(model_name=embeddings_model, model_kwargs={"device": "cuda"}, encode_kwargs={"batch_size": 1, "normalize_embeddings": True}) ## bce-reranker-base_v1 如果路径不对，则下载默认的模型 if not os.path.exists(reranker_model): if reranker_model.endswith("bce-reranker-base_v1"): ## 模型上两级目录是保存路径 save_dir = os.path.dirname(os.path.dirname(reranker_model)) reranker_model = snapshot_download("maidalun/bce-reranker-base_v1", cache_dir=save_dir, revision='master') logging.info(f"reranker_model model not exist, downloading from modelscope \n save to {reranker_model}") else: raise ValueError(f"{reranker_model} model not exist, please reset or re-download your model.") reranker_args = {'model': reranker_model, 'top_n': 5, 'device': 'cuda', "use_fp16": True} self.reranker = BCERerank(**reranker_args) vectordb = get_chroma_db(data_source_dir, db_persist_directory, self.embeddings) # 创建基础检索器 # retriever = vectordb.as_retriever(search_type="similarity", search_kwargs={"score_threshold": 0.3, "k": 2}) self.retriever = vectordb.as_retriever(search_kwargs={"k": 3, "score_threshold": 0.6}, search_type="similarity_score_threshold" ) # 创建上下文压缩检索器 self.compression_retriever = ContextualCompressionRetriever( base_compressor=self.reranker, base_retriever=self.retriever ) # 定义包含 system prompt 的模板 self.PROMPT = PromptTemplate( template=rag_prompt_template, input_variables=["context", "question"] ) # 创建 RetrievalQA 链，包含自定义 prompt self.qa_chain = RetrievalQA.from_chain_type( llm=self.llm, chain_type="stuff", # "stuff"、"map_reduce"、"refine"、"map_rerank" retriever=self.compression_retriever, return_source_documents=True, chain_type_kwargs={"prompt": self.PROMPT} ) def query_stream(self, query: str) -> Iterator[str]: docs = self.compression_retriever.get_relevant_documents(query) context = "\n\n".join(doc.page_content for doc in docs) prompt = self.PROMPT.format(context=context, question=query) return self.llm.stream(prompt) def query(self, question): """ 调用问答链进行回答，如果没有找到相关文档，则使用模型自身的回答 #使用示例 question='黑神话悟空发售时间和团队？主要讲了什么故事？' result = self.qa_chain({"query": question}) print(result["result"]) """ if not question: return "请提供个有用的问题。" try: # 使用检索链来获取相关文档 result = self.qa_chain.invoke({"query": question}) # logging.info(f"Get rag res:\n{result}") if 'result' in result: answer = result['result'] final_answer = re.sub(r'^根据提供的信息，\s?', '', answer, flags=re.M).strip() return final_answer else: logging.error("Error: 'result' field not found in the result.") return "悟了悟了目前无法提供答案，请稍后再试。" except Exception as e: # 打印更详细的错误信息，包括traceback import traceback logging.error(f"An error occurred: {e}\n{traceback.format_exc()}") return "悟了悟了遇到了一些技术问题，正在修复中。" @hydra.main(version_base=None, config_path="../configs", config_name="rag_cfg") def main(config): data_source_dir = config.data_source_dir db_persist_directory = config.db_persist_directory llm_model = config.llm_model embeddings_model = config.embeddings_model reranker_model = config.reranker_model llm_system_prompt = config.llm_system_prompt rag_prompt_template = config.rag_prompt_template ## download model from modelscope if not os.path.exists(llm_model): download_model(llm_model_path = llm_model) base_mode = InternLM(model_path=llm_model, llm_system_prompt=llm_system_prompt) # from deploy.lmdeploy_model import LmdeployLM # base_mode = LmdeployLM(model_path=llm_model, llm_system_prompt=llm_system_prompt, cache_max_entry_count=0.2) wulewule_rag = WuleRAG(data_source_dir, db_persist_directory, base_mode, embeddings_model, reranker_model, rag_prompt_template) question="""黑神话悟空发售时间和团队？主要讲了什么故事？""" # 流式显示, used streaming result if config.stream_response: logging.info("Streaming response:") for chunk in wulewule_rag.query_stream(question): print(chunk, end='', flush=True) print("\n") # 一次性显示结果 else: response = wulewule_rag.query(question) logging.info(f"question: {question}\n wulewule answer:\n{response}") if __name__ == "__main__": main() # llm_system_prompt = "你是悟了悟了，由xzyun2011开发的AI助手，专注于回答和《黑神话：悟空》这款游戏相关的问题，你想帮助玩家了解更多这款游戏背后的故事和文化知识。" # rag_prompt_template = """系统: 你是悟了悟了，由xzyun2011开发的AI助手，专注于回答和《黑神话：悟空》这款游戏相关的问题，你想帮助玩家了解更多这款游戏背后的故事和文化知识。 # 人类: {question} # 助手: 我会根据提供的信息来回答。 # 相关上下文: # {context} # 基于以上信息，我的回答是： # """ # data_source_dir = "/root/wulewule/data" # txt数据目录 # db_persist_directory ='/root/wulewule/rag/chroma' # chroma向量库数据目录 # llm_model = "/root/wulewule/models/wulewule_v1_1_8b" # embeddings_model = "/root/share/new_models/maidalun1020/bce-embedding-base_v1" # reranker_model = "/root/share/new_models/maidalun1020/bce-reranker-base_v1"