├── README.md ├── 技术文档 ├── Embeddings.py ├── CustomLLM.py ├── Rag.py └── St.py /README.md: -------------------------------------------------------------------------------- 1 | # LLM_RAG 2 | 利用开源大模型，通过RAG(检索增强生成)技术，实现基于企业内部知识图谱的，可内网运行的大模型智能客服 3 | 使用说明，已分享到我的博客了，https://juejin.cn/user/310524887198384 ，感兴趣的小伙伴可以查看，欢迎一起学习交流。。 4 | -------------------------------------------------------------------------------- /技术文档: -------------------------------------------------------------------------------- 1 | ''' 2 | RAG技术实现文档 3 | ''' 4 | 5 | 1,利用FASTCHAT导入和启动开源本地大模型，并开启API接口。 6 | 2,Embeddings文件，将知识库存储到chroma向量数据库中并本地保存。 7 | 3,CustomLLM文件，重写langchain的llm基类，通过调用API接口自定义生成大模型类。 8 | 4,Rag文件，自定义rag类，传参是用户问题，将用户问题与向量数据库的数据进行相似性查询，返回结果与问题，提示词，一起传给自定义大模型，并返 回 大模型的回答结果。 9 | 5,St文件，利用Streamlit开发前端页面，调用Rag文件将问题传入，前端接收显示返回结果。 10 | -------------------------------------------------------------------------------- /Embeddings.py: -------------------------------------------------------------------------------- 1 | ''' 2 | 嵌入知识库 3 | ''' 4 | from langchain.document_loaders import TextLoader 5 | from langchain.text_splitter import RecursiveCharacterTextSplitter 6 | from langchain_community.embeddings.sentence_transformer import ( 7 | SentenceTransformerEmbeddings, 8 | ) 9 | from langchain.vectorstores import Chroma 10 | 11 | filepath = '/***' #你本地文件及地址 12 | loader = TextLoader(filepath, encoding='utf-8') 13 | text = loader.load() 14 | 15 | text_splitter = RecursiveCharacterTextSplitter( 16 | chunk_size=100, 17 | chunk_overlap=0, 18 | separators=["},\n"], 19 | length_function=len 20 | ) 21 | split_docs = text_splitter.split_documents(text) 22 | 23 | embedding_function = SentenceTransformerEmbeddings(model_name="shibing624/text2vec-base-chinese") 24 | 25 | vectorstore = Chroma.from_documents(split_docs, embedding_function, persist_directory="/") # 本地存储到向量数据库的地址，换成你自己的 26 | vectorstore.persist() 27 | 28 | # 测试是否嵌入成功 29 | query = "你是谁？" #测试问句 30 | doc = vectorstore.similarity_search(query) 31 | print(doc[0].page_content) 32 | -------------------------------------------------------------------------------- /CustomLLM.py: -------------------------------------------------------------------------------- 1 | ''' 2 | 自定义LLM 3 | ''' 4 | from langchain.llms.base import LLM 5 | from typing import Any, List, Optional 6 | from langchain.callbacks.manager import CallbackManagerForLLMRun 7 | import requests 8 | 9 | 10 | class CustomLLM(LLM): 11 | endpoint: str = "***/v1" # 本地大模型的API地址 12 | model: str = "***" #大模型名称 13 | 14 | @property 15 | def _llm_type(self) -> str: 16 | return "***" #大模型名称 17 | 18 | def _call( 19 | self, 20 | prompt: str, 21 | stop: Optional[List[str]] = None, 22 | callbacks: Optional[CallbackManagerForLLMRun] = None, 23 | **kwargs: Any, 24 | ) -> str: 25 | headers = {"Content-Type": "application/json"} 26 | data = {"model": self.model, "messages": [{"role": "user", "content": prompt}]} 27 | response = requests.post(f"{self.endpoint}/chat/completions", headers=headers, json=data) 28 | if response.status_code != 200: 29 | return "error" 30 | 31 | result = response.json() 32 | text = result["choices"][0]["message"]["content"] 33 | return text -------------------------------------------------------------------------------- /Rag.py: -------------------------------------------------------------------------------- 1 | ''' 2 | RAG主函数 3 | ''' 4 | from langchain_community.embeddings.sentence_transformer import ( 5 | SentenceTransformerEmbeddings, 6 | ) 7 | from langchain.prompts import PromptTemplate 8 | from langchain.vectorstores import Chroma 9 | from langchain.chains import LLMChain 10 | from CustomLLM import CustomLLM 11 | 12 | 13 | class Rag: 14 | 15 | def __init__(self, question): 16 | self.q = question 17 | 18 | def __str__(self): 19 | if self.q: 20 | embedding_function = SentenceTransformerEmbeddings(model_name="shibing624/text2vec-base-chinese") 21 | db = Chroma(persist_directory="/", embedding_function=embedding_function) #改成实际的本地向量库地址 22 | doc = db.similarity_search(self.q) 23 | context = doc[0].page_content 24 | 25 | #提示词 26 | prompt_template = """请用以下片段去回答问题. 如果不知道答案, 就说不知道. 不要编造答案. 27 | {context} 28 | Question: {question} 29 | Answer:""" 30 | prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"]) 31 | 32 | llm = CustomLLM() 33 | query_llm = LLMChain(llm=llm, prompt=prompt) 34 | response = query_llm.run({"context": context, "question": self.q}) 35 | return response -------------------------------------------------------------------------------- /St.py: -------------------------------------------------------------------------------- 1 | ''' 2 | 运用Streamlit完成前端页面显示 3 | ''' 4 | from Rag import Rag 5 | import streamlit as st 6 | 7 | with st.sidebar: 8 | job_id = st.text_input("工号", key="job_id") 9 | name = st.text_input("姓名", key="name") 10 | tips = st.checkbox("温馨提示:您与大模型的对话都会被记录,并有可能被用于训练该大模型,勾选即表示您知情并同意", 11 | value=False, key=None) 12 | st.title("💬 RAG大模型") 13 | st.caption("🚀 聊天机器人developed by IT") 14 | if "messages" not in st.session_state: 15 | st.session_state["messages"] = [{"role": "assistant", "content": "您好，有什么可以帮您?"}] 16 | 17 | for msg in st.session_state.messages: 18 | st.chat_message(msg["role"]).write(msg["content"]) 19 | 20 | 21 | def clear_chat_history(): 22 | st.session_state.messages = [{"role": "assistant", "content": "您好，有什么可以帮您?"}] 23 | st.sidebar.button('清空聊天记录', on_click=clear_chat_history) 24 | 25 | 26 | if prompt := st.chat_input(): 27 | if not job_id: 28 | st.info("请输入工号.") 29 | st.stop() 30 | if not name: 31 | st.info("请输入姓名.") 32 | st.stop() 33 | if not tips: 34 | st.info("请先选中知情同意书.") 35 | st.stop() 36 | 37 | st.session_state.messages.append({"role": "user", "content": prompt}) 38 | st.chat_message("user").write(prompt) 39 | # 调用RAG类 40 | msg = Rag(prompt) 41 | st.session_state.messages.append({"role": "assistant", "content": msg}) 42 | st.chat_message("assistant").write(msg) 43 | --------------------------------------------------------------------------------