Task3案例：AI科研助手
输入“/”快速插入内容
🌊Task3案例：AI科研助手
Task3案例：AI科研助手
2024年7月6日创建
4463
5661
2
1
学习赛提交要求
1.
项目背景
随着学术研究领域的快速发展，研究人员需要处理大量的文献资料。传统的文献综述方法往往耗时且效率低下，而且很难及时跟进最新的研究成果。近年来，自然语言处理技术（NLP）的进步，使得机器能够理解和处理人类语言成为可能。因此，开发一款能够自动分析、理解并回答关于论文内容的问题的AI科研助手变得十分必要。
2.
产品功能
•
论文概括：AI科研助手能够快速阅读并理解一篇论文的主要内容，并生成简洁明了的摘要。
•
内容问答：用户可以通过自然语言提出与论文相关的问题，系统能够给出准确的答案。
•
关键词提取：从论文中提取关键概念和术语，帮助用户快速了解研究重点。
•
论文对比：比较两篇或者多篇论文的内容，总结他们之间的异同点。
•
相关工作推荐：根据用户的兴趣领域推荐相关的最新研究成果。
3.
应用价值
•
提高工作效率：通过自动化文献分析减少研究人员的时间消耗。
•
促进知识传播：使更多人能够迅速掌握某一领域的核心发现。
•
增强学术合作：通过协作平台加强跨学科的研究合作。
•
辅助教学：在教育场景下，帮助学生更好地理解复杂的研究主题。
4.
技术方案
方案架构图
本项目基于源大模型RAG技术来解决用户的问题。
具体来说，项目主要包含一个Streamlit开发的客户端，以及一个部署好的浪潮源大模型的服务端。
客户端接收到用户上传的PDF后，发送到服务端。服务端首先完成PDF内容解析，然后拼接摘要Prompt并输入源大模型，得到模型输出结果后，返回给客户端并展示给用户。
如果用户接下来进行提问，客户端将用户请求发送到服务端，服务端进行Embedding和Faiss检索，然后将检索到的chunks与用户请求拼接成Prompt并输入到源大模型，得到模型输出结果后，返回给客户端进行结构化，然后展示给用户。
画板
核心代码
依赖安装
代码块
pip install pypdf faiss-gpu langchain langchain_community langchain_huggingface streamlit==1.24.0
启动脚本
代码块
streamlit run Task\ 3\ 案例：AI科研助手.py --server.address 127.0.0.1 --server.port 6006
源代码
代码块
# 导入所需的库
import torch
import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
from langchain.prompts import PromptTemplate
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.chains import LLMChain
from langchain.chains.question_answering import load_qa_chain
from langchain.llms.base import LLM
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.text_splitter import RecursiveCharacterTextSplitter
from typing import Any, List, Optional
# 向量模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('AI-ModelScope/bge-small-en-v1.5', cache_dir='./')
# 源大模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('IEITYuan/Yuan2-2B-Mars-hf', cache_dir='./')
# 定义模型路径
model_path = './IEITYuan/Yuan2-2B-Mars-hf'
# 定义向量模型路径
embedding_model_path = './AI-ModelScope/bge-small-en-v1___5'
# 定义模型数据类型
torch_dtype = torch.bfloat16 # A10
# torch_dtype = torch.float16 # P100
# 定义源大模型类
class Yuan2_LLM(LLM):
"""
class for Yuan2_LLM
"""
tokenizer: AutoTokenizer = None
model: AutoModelForCausalLM = None
def __init__(self, mode_path :str):
super().__init__()
# 加载预训练的分词器和模型
print("Creat tokenizer...")
self.tokenizer = AutoTokenizer.from_pretrained(mode_path, add_eos_token=False, add_bos_token=False, eos_token='<eod>')
self.tokenizer.add_tokens(['<sep>', '<pad>', '<mask>', '<predict>', '<FIM_SUFFIX>', '<FIM_PREFIX>', '<FIM_MIDDLE>','<commit_before>','<commit_msg>','<commit_after>','<jupyter_start>','<jupyter_text>','<jupyter_code>','<jupyter_output>','<empty_output>'], special_tokens=True)
print("Creat model...")
self.model = AutoModelForCausalLM.from_pretrained(mode_path, torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()
def _call(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
prompt = prompt.strip()
prompt += "<sep>"
inputs = self.tokenizer(prompt, return_tensors="pt")["input_ids"].cuda()
outputs = self.model.generate(inputs,do_sample=False,max_length=4096)
output = self.tokenizer.decode(outputs[0])
response = output.split("<sep>")[-1].split("<eod>")[0]
return response
@property
def _llm_type(self) -> str:
return "Yuan2_LLM"
# 定义一个函数，用于获取llm和embeddings
@st.cache_resource
def get_models():
llm = Yuan2_LLM(model_path)
model_kwargs = {'device': 'cuda'}
encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
embeddings = HuggingFaceEmbeddings(
model_name=embedding_model_path,
model_kwargs=model_kwargs,
encode_kwargs=encode_kwargs,
)
return llm, embeddings
summarizer_template = """
假设你是一个AI科研助手，请用一段话概括下面文章的主要内容，200字左右。
{text}
"""
# 定义Summarizer类
class Summarizer:
"""
class for Summarizer.
"""
def __init__(self, llm):
self.llm = llm
self.prompt = PromptTemplate(
input_variables=["text"],
template=summarizer_template
)
self.chain = LLMChain(llm=self.llm, prompt=self.prompt)
def summarize(self, docs):
# 从第一页中获取摘要
content = docs[0].page_content.split('ABSTRACT')[1].split('KEY WORDS')[0]
summary = self.chain.run(content)
return summary
chatbot_template = '''
假设你是一个AI科研助手，请基于背景，简要回答问题。
背景：
{context}
问题：
{question}
'''.strip()
# 定义ChatBot类
class ChatBot:
"""
class for ChatBot.
"""
def __init__(self, llm, embeddings):
self.prompt = PromptTemplate(
input_variables=["text"],
template=chatbot_template
)
self.chain = load_qa_chain(llm=llm, chain_type="stuff", prompt=self.prompt)
self.embeddings = embeddings
# 加载 text_splitter
self.text_splitter = RecursiveCharacterTextSplitter(
chunk_size=450,
chunk_overlap=10,
length_function=len
)
def run(self, docs, query):
# 读取所有内容
text = ''.join([doc.page_content for doc in docs])
# 切分成chunks
all_chunks = self.text_splitter.split_text(text=text)
# 转成向量并存储
VectorStore = FAISS.from_texts(all_chunks, embedding=self.embeddings)
# 检索相似的chunks
chunks = VectorStore.similarity_search(query=query, k=1)
# 生成回复
response = self.chain.run(input_documents=chunks, question=query)
return chunks, response
def main():
# 创建一个标题
st.title('💬 Yuan2.0 AI科研助手')
# 获取llm和embeddings
llm, embeddings = get_models()
# 初始化summarizer
summarizer = Summarizer(llm)
# 初始化ChatBot
chatbot = ChatBot(llm, embeddings)
# 上传pdf
uploaded_file = st.file_uploader("Upload your PDF", type='pdf')
if uploaded_file:
# 加载上传PDF的内容
file_content = uploaded_file.read()
# 写入临时文件
temp_file_path = "temp.pdf"
with open(temp_file_path, "wb") as temp_file:
temp_file.write(file_content)
# 加载临时文件中的内容
loader = PyPDFLoader(temp_file_path)
docs = loader.load()
st.chat_message("assistant").write(f"正在生成论文概括，请稍候...")
# 生成概括
summary = summarizer.summarize(docs)
# 在聊天界面上显示模型的输出
st.chat_message("assistant").write(summary)
# 接收用户问题
if query := st.text_input("Ask questions about your PDF file"):
# 检索 + 生成回复
chunks, response = chatbot.run(docs, query)
# 在聊天界面上显示模型的输出
st.chat_message("assistant").write(f"正在检索相关信息，请稍候...")
st.chat_message("assistant").write(chunks)
st.chat_message("assistant").write(f"正在生成回复，请稍候...")
st.chat_message("assistant").write(response)
if __name__ == '__main__':
main()