【记录】LangChain+当地模子的文档问答（webUI） - ToB企服应用市场:ToB评测及商务社交产业平台

# coding: utf-8
# Author: 唐国梁Tommy
# Date: 2023-08-06
import streamlit as st
from langchain_community.embeddings import HuggingFaceInstructEmbeddings
from langchain.chains.conversational_retrieval.base import ConversationalRetrievalChain
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.vectorstores import FAISS, Milvus, Pinecone, Chroma
from langchain.memory import ConversationBufferMemory
from transformers import AutoModelForCausalLM, AutoTokenizer,pipeline
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
import streamlit as st
from PyPDF2 import PdfReader
def main():
# 配置界面
st.set_page_config(page_title="基于PDF文档的 QA ChatBot",
page_icon=":robot:")
st.header("基于LangChain+LLM实现QA ChatBot")
# 参考官网链接：https://github.com/hwchase17/langchain-streamlit-template/blob/master/main.py
# 初始化
# session_state是Streamlit提供的用于存储会话状态的功能
if "conversation" not in st.session_state:
st.session_state.conversation = None
if "chat_history" not in st.session_state:
st.session_state.chat_history = None
# 1. 提供用户输入文本框
user_input = st.text_input("基于上传的PDF文档，请输入你的提问: ")
# 处理用户输入，并返回响应结果
if user_input:
process_user_input(user_input)
with st.sidebar:
# 2. 设置子标题
st.subheader("你的PDF文档")
# 3. 上传文档
files = st.file_uploader("上传PDF文档，然后点击'提交并处理'",
accept_multiple_files=True)
if st.button("提交并处理"):
with st.spinner("请等待，处理中..."):
# 4. 获取PDF文档内容（文本）
texts = extract_text_from_PDF(files)
# 5. 将获取到的文档内容进行切分
content_chunks = split_content_into_chunks(texts)
# st.write(content_chunks)
# 6. 对每个chunk计算embedding，并存入到向量数据库
# 6.1 根据model_type和model_name创建embedding model对象
#embedding_model = get_openaiEmbedding_model()
# embedding_model = get_huggingfaceEmbedding_model(model_name="BAAI/bge-m3")
embedding_model = HuggingFaceInstructEmbeddings(model_name="BAAI/bge-m3")
# 6.2 创建向量数据库对象，并将文本embedding后存入到里面
vector_store = save_chunks_into_vectorstore(content_chunks, embedding_model)
# 7. 创建对话chain
# 官网链接：https://python.langchain.com/docs/modules/memory/types/buffer
st.session_state.conversation = get_chat_chain(vector_store)
def extract_text_from_PDF(files):
# 参考官网链接：https://python.langchain.com/docs/modules/data_connection/document_loaders/pdf
# 加载多个PDF文件
text = ""
for pdf in files:
pdf_reader = PdfReader(pdf)
for page in pdf_reader.pages:
text += page.extract_text()
return text
def split_content_into_chunks(text):
# 参考官网链接：https://python.langchain.com/docs/modules/data_connection/document_transformers/text_splitters/character_text_splitter
text_spliter = CharacterTextSplitter(separator="\n",
chunk_size=500,
chunk_overlap=80,
length_function=len)
chunks = text_spliter.split_text(text)
return chunks
def save_chunks_into_vectorstore(content_chunks, embedding_model):
# 参考官网链接：https://python.langchain.com/docs/modules/data_connection/vectorstores/
# ① FAISS
# pip install faiss-gpu (如果没有GPU，那么 pip install faiss-cpu)
vectorstore = FAISS.from_texts(texts=content_chunks,
embedding=embedding_model)
return vectorstore
def get_chat_chain(vector_store):
# ① 获取 LLM model
#llm = get_openai_model()
# llm = get_huggingfacehub(model_name="google/flan-t5-xxl")
# llm = get_huggingfacehub(model_name="google-bert/bert-base-chinese")
model_path = "baichuan-inc/Baichuan2-7B-Chat"
model = AutoModelForCausalLM.from_pretrained(model_path,trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_path)
print(type(model))
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
print(type(pipe))
llm = HuggingFacePipeline(pipeline=pipe)
print(type(llm))
# ② 存储历史记录
# 参考官网链接：https://python.langchain.com/docs/use_cases/question_answering/how_to/chat_vector_db
# 用于缓存或者保存对话历史记录的对象
memory = ConversationBufferMemory(
memory_key='chat_history', return_messages=True)
# ③ 对话链
conversation_chain = ConversationalRetrievalChain.from_llm(
llm=llm,
retriever=vector_store.as_retriever(
search_type="similarity",
search_kwargs={"k": 5}
),
memory=memory
)
return conversation_chain
def process_user_input(user_input):
print('输入内容 '+user_input)
if st.session_state.conversation is not None:
print('不为空')
# 调用函数st.session_state.conversation，并把用户输入的内容作为一个问题传入，返回响应。
response = st.session_state.conversation({'question': user_input})
print('response '+response)
# session状态是Streamlit中的一个特性，允许在用户的多个请求之间保存数据。
st.session_state.chat_history = response['chat_history']
# 显示聊天记录
# chat_history : 一个包含之前聊天记录的列表
for i, message in enumerate(st.session_state.chat_history):
# 用户输入
if i % 2 == 0:
st.write(user_template.replace(
"{{MSG}}", message.content), unsafe_allow_html=True) # unsafe_allow_html=True表示允许HTML内容被渲染
else:
# 机器人响应
st.write(bot_template.replace(
"{{MSG}}", message.content), unsafe_allow_html=True)
if __name__ == "__main__":
main()

复制代码