129 lines
4.3 KiB
Python
129 lines
4.3 KiB
Python
|
'''
|
|||
|
|
支持上传文件进行分析，记录历史上下文，
|
|||
|
|
'''
|
|||
|
|
import time
|
|||
|
|
import tempfile
|
|||
|
|
from dotenv import load_dotenv
|
|||
|
|
|
|||
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|||
|
|
from langchain.prompts import MessagesPlaceholder
|
|||
|
|
from langchain_community.document_loaders import TextLoader, PyPDFLoader, CSVLoader, JSONLoader
|
|||
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|||
|
|
from langchain.chains import create_retrieval_chain
|
|||
|
|
from langchain.chains.combine_documents import create_stuff_documents_chain
|
|||
|
|
from langchain_openai import ChatOpenAI
|
|||
|
|
from langchain_community.chat_models.tongyi import ChatTongyi
|
|||
|
|
# from langchain_community.chat_models.baidu_qianfan_endpoint import QianfanChatEndpoint
|
|||
|
|
from langchain_community.chat_models import QianfanChatEndpoint
|
|||
|
|
from langchain_community.chat_models import ChatZhipuAI
|
|||
|
|
from langchain.retrievers import ContextualCompressionRetriever
|
|||
|
|
from langchain_community.embeddings import HuggingFaceEmbeddings
|
|||
|
|
from langchain_chroma import Chroma
|
|||
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter, RecursiveJsonSplitter
|
|||
|
|
import requests
|
|||
|
|
import os
|
|||
|
|
import re
|
|||
|
|
import pandas as pd
|
|||
|
|
|
|||
|
|
|
|||
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
|||
|
|
device = "cuda"
|
|||
|
|
memory = 6
|
|||
|
|
|
|||
|
|
def init_chain():
|
|||
|
|
try:
|
|||
|
|
# 加载环境变量并初始化模型
|
|||
|
|
|
|||
|
|
# 使用Tongyi
|
|||
|
|
# os.environ["DASHSCOPE_API_KEY"] = 'xxxx'
|
|||
|
|
# llm = ChatTongyi(
|
|||
|
|
# streaming=True,
|
|||
|
|
# model='llama3-70b-instruct'
|
|||
|
|
# )
|
|||
|
|
|
|||
|
|
# 使用Qianfan
|
|||
|
|
# ak = "xxxx"
|
|||
|
|
# sk = "xxxxx"
|
|||
|
|
# llm = QianfanChatEndpoint(
|
|||
|
|
# streaming=True,
|
|||
|
|
# model = 'ERNIE-3.5-8K',
|
|||
|
|
# api_key=ak,
|
|||
|
|
# secret_key=sk
|
|||
|
|
# )
|
|||
|
|
|
|||
|
|
# 使用智谱
|
|||
|
|
os.environ["ZHIPUAI_API_KEY"] = "xxxxx"
|
|||
|
|
llm = ChatZhipuAI(
|
|||
|
|
streaming=True,
|
|||
|
|
model = "glm-4"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
embedding = HuggingFaceEmbeddings(model_name="/home/zhangxj/models/acge_text_embedding", model_kwargs={"device": device})
|
|||
|
|
retriever = Chroma(persist_directory="chromaDB", embedding_function=embedding,)
|
|||
|
|
|
|||
|
|
retrieval_chunks = 20 #前20个相似块
|
|||
|
|
retriever = retriever.as_retriever(search_type="mmr", search_kwargs={"k": 20})
|
|||
|
|
|
|||
|
|
instruct_system_prompt = (
|
|||
|
|
"你是生命周期领域富有经验和知识的专家。"
|
|||
|
|
"使用以下检索到的上下文来回答问题。"
|
|||
|
|
"{context}"
|
|||
|
|
"答案最多使用1句话并保持非常简洁，不能换行。"
|
|||
|
|
)
|
|||
|
|
# instr = "你是生命周期领域富有经验和知识的专家。根据你所掌握的知识只用1句话回答问题。不要列出几点来回答，不需要换行"
|
|||
|
|
instruct_prompt = ChatPromptTemplate.from_messages(
|
|||
|
|
[
|
|||
|
|
("system", instruct_system_prompt),
|
|||
|
|
("human", "{input}"),
|
|||
|
|
]
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Create a chain for passing a list of Documents to a model.
|
|||
|
|
qa_chain = create_stuff_documents_chain(llm, instruct_prompt) #
|
|||
|
|
rag_chain = create_retrieval_chain(retriever, qa_chain)
|
|||
|
|
|
|||
|
|
return rag_chain
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Error in init_chain: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def user_in(uin, rag_chain):
|
|||
|
|
try:
|
|||
|
|
result = rag_chain.invoke({"input": uin})['answer']
|
|||
|
|
# print(rag_chain)
|
|||
|
|
# result = rag_chain(uin)
|
|||
|
|
return result
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Error in user_in: {e}")
|
|||
|
|
return "An error occurred while processing your request."
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
|
|||
|
|
question = []
|
|||
|
|
with open("/home/zhangxj/WorkFile/LCA-GPT/QA/split/question/ques0.txt","r",encoding="utf-8") as file:
|
|||
|
|
for line in file.readlines():
|
|||
|
|
question.append(line.strip())
|
|||
|
|
|
|||
|
|
rag = init_chain()
|
|||
|
|
|
|||
|
|
answers = []
|
|||
|
|
for ques in question:
|
|||
|
|
# print(ques)
|
|||
|
|
response = user_in(ques, rag)
|
|||
|
|
print(response)
|
|||
|
|
# print(len(response))
|
|||
|
|
answers.append(response.strip())
|
|||
|
|
# if len(answers) == 3:
|
|||
|
|
# break
|
|||
|
|
with open("/home/zhangxj/WorkFile/LCA-GPT/QA/eval/RAGpred.txt","w",encoding="utf-8") as file:
|
|||
|
|
for ans in answers:
|
|||
|
|
line = re.sub(r'\s+', '', ans)
|
|||
|
|
file.write(line+'\n')
|
|||
|
|
data = {"ans":answers}
|
|||
|
|
df = pd.DataFrame(data)
|
|||
|
|
df.to_csv("/home/zhangxj/WorkFile/LCA-GPT/QA/eval/rag.csv",index=False)
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|