【深度学习总结】使用PDF构建RAG：结合Langchain和通义千问 - ToB企服应用市场:ToB评测及商务社交产业平台

python-dotenv==1.0.1 # For reading environment variables stored in .env file
langchain==0.2.2
langchain-community==0.2.3
dashscope
unstructured==0.14.4 # Document loading
# onnxruntime==1.17.1 # chromadb dependency: on Mac use `conda install onnxruntime -c conda-forge`
# For Windows users, install Microsoft Visual C++ Build Tools first
# install onnxruntime before installing `chromadb`
chromadb==0.5.0 # Vector storage
tiktoken==0.7.0 # For embeddings

复制代码

from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings.huggingface import HuggingFaceEmbeddings
from dotenv import load_dotenv
import os
import shutil
import dashscope
from dashscope import Generation
from langchain.prompts import ChatPromptTemplate
from http import HTTPStatus

复制代码

# Load environment variables. Assumes that project contains .env file with API keys
load_dotenv()
# 设置镜像，便于下载后面的HuggingFaceEmbeddings
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
# huggingface下载地址
os.environ["HF_HOME"] = "/gemini/code/huggingface"
# huggingface下载地址
os.environ["TRANSFORMERS_CACHE"] = "/gemini/code/huggingface"
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/gemini/code/huggingface/bce-embedding-base_v1"
# 向量存放位置
CHROMA_PATH = "chroma"
# 存放数据
DATA_PATH = "data/books"
dashscope.api_key = "你的通义千问api key"
print(os.getenv('DASHSCOPE_API_KEY'))
def prepare_db():
# 处理读个pdf
pdf_paths = ["data/books/M1.pdf",
"data/books/M2-航空器维修R1.pdf",
"data/books/M3-飞机结构和系统R1.pdf",
"data/books/M4-直升机结构和系统.pdf",
"data/books/M5-航空涡轮发动机R1.pdf",
"data/books/M6-活塞发动机及其维修.pdf",
"data/books/M7-航空器维修基本技能.pdf",
"data/books/M8-航空器维修实践R1.pdf"]
documents = []
count = 0
for pdf_path in pdf_paths:
loader = PyPDFLoader(pdf_path)
doc = loader.load()
documents.extend(doc)
count += 1
print(f"处理第{count}本")
print(len(documents))
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=300,
chunk_overlap=100,
length_function=len,
add_start_index=True,
)
chunks = text_splitter.split_documents(documents)
print(f"Split {len(documents)} documents into {len(chunks)} chunks.")
document = chunks[10]
print(document.page_content)
print(document.metadata)
if os.path.exists(CHROMA_PATH):
shutil.rmtree(CHROMA_PATH)
# 将文本保存为向量存储
model_name = "maidalun1020/bce-embedding-base_v1"
model_kwargs = {'device': 'cuda'}
encode_kwargs = {'normalize_embeddings': False}
embeddings = HuggingFaceEmbeddings(
model_name=model_name,
model_kwargs=model_kwargs,
encode_kwargs=encode_kwargs,
cache_folder="/gemini/code/huggingface/",
)
# Create a new DB from the documents.
db = Chroma.from_documents(
chunks, embeddings, persist_directory=CHROMA_PATH
)
db.persist()
print(f"Saved {len(chunks)} chunks to {CHROMA_PATH}.")

复制代码

def query():
model_name = "maidalun1020/bce-embedding-base_v1"
model_kwargs = {'device': 'cuda'}
encode_kwargs = {'normalize_embeddings': False}
embeddings = HuggingFaceEmbeddings(
model_name=model_name,
model_kwargs=model_kwargs,
encode_kwargs=encode_kwargs,
# 模型缓存路径
cache_folder="/gemini/code/huggingface/",
)
# 改成你的保存的路径
db = Chroma(persist_directory="./chroma/aae4fae7-3477-4094-8a7d-c5df8be2223a", embedding_function=embeddings)
# 提示模板
PROMPT_TEMPLATE = """
仅根据下列文本回答问题：
{context}
"""
while True:
query = input('请输入问题:')
results = db.similarity_search_with_relevance_scores(query, k=5)
if len(results) == 0:
print(f"Unable to find matching results.")
return
# 拼接成输入给大模型的内容
context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
messages = [
{'role': 'system', 'content': PROMPT_TEMPLATE.format(context=context_text)},
{'role': 'user', 'content': f"请回答如下问题:{query}"}
]
print(messages)
responses = Generation.call(Generation.Models.qwen_max,
api_key=os.getenv('DASHSCOPE_API_KEY'),
messages=messages,
result_format='message')
# 如果你不确定responses的结果，可以打印处理
# print(responses.output)
sources = [doc.metadata.get("source", None) for doc, _score in results]
if responses.status_code == HTTPStatus.OK:
whole_message = responses.output["choices"][0]["message"]["content"]
else:
whole_message = "error"
print('Failed request_id: %s, status_code: %s, code: %s, message:%s' %
(responses.request_id, responses.status_code, responses.code,
responses.message))
formatted_response = f"Response: {whole_message}\nSources: {sources}"
print(formatted_response)

复制代码