76 lines
2.2 KiB
Python
76 lines
2.2 KiB
Python
import os
|
|
# hugging face镜像设置，如果国内环境无法使用启用该设置
|
|
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
|
|
from dotenv import load_dotenv
|
|
from langchain_community.document_loaders import UnstructuredMarkdownLoader
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain_huggingface import HuggingFaceEmbeddings
|
|
from langchain_core.vectorstores import InMemoryVectorStore
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
load_dotenv()
|
|
|
|
markdown_path = "../../data/C1/markdown/easy-rl-chapter1.md"
|
|
|
|
# 加载本地markdown文件
|
|
loader = UnstructuredMarkdownLoader(markdown_path)
|
|
docs = loader.load()
|
|
|
|
# 文本分块
|
|
text_splitter = RecursiveCharacterTextSplitter()
|
|
chunks = text_splitter.split_documents(docs)
|
|
|
|
# 中文嵌入模型
|
|
embeddings = HuggingFaceEmbeddings(
|
|
model_name="BAAI/bge-small-zh-v1.5",
|
|
model_kwargs={'device': 'cpu'},
|
|
encode_kwargs={'normalize_embeddings': True}
|
|
)
|
|
|
|
# 构建向量存储
|
|
vectorstore = InMemoryVectorStore(embeddings)
|
|
vectorstore.add_documents(chunks)
|
|
|
|
# 提示词模板
|
|
prompt = ChatPromptTemplate.from_template("""请根据下面提供的上下文信息来回答问题。
|
|
请确保你的回答完全基于这些上下文。
|
|
如果上下文中没有足够的信息来回答问题，请直接告知：“抱歉，我无法根据提供的上下文找到相关信息来回答此问题。”
|
|
|
|
上下文:
|
|
{context}
|
|
|
|
问题: {question}
|
|
|
|
回答:"""
|
|
)
|
|
|
|
# 配置大语言模型
|
|
|
|
# 使用 AIHubmix
|
|
llm = ChatOpenAI(
|
|
model="glm-4.7-flash-free",
|
|
temperature=0.7,
|
|
max_tokens=4096,
|
|
api_key=os.getenv("DEEPSEEK_API_KEY"),
|
|
base_url="https://aihubmix.com/v1"
|
|
)
|
|
|
|
# llm = ChatOpenAI(
|
|
# model="deepseek-chat",
|
|
# temperature=0.7,
|
|
# max_tokens=4096,
|
|
# api_key=os.getenv("DEEPSEEK_API_KEY"),
|
|
# base_url="https://api.deepseek.com"
|
|
# )
|
|
|
|
# 用户查询
|
|
question = "文中举了哪些例子？"
|
|
|
|
# 在向量存储中查询相关文档
|
|
retrieved_docs = vectorstore.similarity_search(question, k=3)
|
|
docs_content = "\n\n".join(doc.page_content for doc in retrieved_docs)
|
|
|
|
answer = llm.invoke(prompt.format(question=question, context=docs_content))
|
|
print(answer)
|