ByteBrain:信息时代您的计算机科学智能知识助手
GitHub项目链接（https://github.com/Stars-niu/ByteBrain）
文章目录
baseline一键部署体验
1、个人云账号授权实例：可以开通阿里云PAI-DSW试用，时长三个月
2、魔塔平台免费实例：注册并绑定阿里云账号试用GPU，时长100h
# JupyterLab->Other->Terminal->Ctrl+V
git clone https://github.com/Stars-niu/ByteBrain.git
cd ByteBrain
pip install --upgrade pip setuptools
pip install -r requirements.txt
python finetune_model.py
pip install streamlit
pip install tf-keras
streamlit run app.py --server.address 127.0.0.1 --server.port 1005
问题注意:
- app.py根据使用的版本自行更换
- 重复打开应用时，可以更换监听端口的四个数字(即1001)，如果出现某端口已占用的情况。补充:127.0.0.1:表示服务器只监听本地回环地址，也就是说，只有在本机上的浏览器才能访问这个应用。(如果你想让其他设备也能访问，可以设置为 0.0.0.0) 1001:这个选项指定了Streamlit服务器监听的端口。这里设置为 1001，意味着应用程序将在 http://127.0.0.1:1001 地址上运行。
程序源码
# 导入所需的库
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import streamlit as st
# 创建一个标题和一个副标题
st.title("✨ ByteBrain")
# 源大模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('IEITYuan/Yuan2-2B-Mars-hf', cache_dir='./')
# model_dir = snapshot_download('IEITYuan/Yuan2-2B-July-hf', cache_dir='./')
# 定义模型路径
path = './IEITYuan/Yuan2-2B-Mars-hf'
# path = './IEITYuan/Yuan2-2B-July-hf'
# 定义模型数据类型
torch_dtype = torch.bfloat16 # A10
# torch_dtype = torch.float16 # P100
# 定义一个函数，用于获取模型和tokenizer
@st.cache_resource
def get_model():
print("创建tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(path, add_eos_token=False, add_bos_token=False, eos_token='<eod>')
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("创建模型...")
model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch_dtype, trust_remote_code=True).cuda()
print("完成.")
return tokenizer, model
# 加载model和tokenizer
tokenizer, model = get_model()
# 初次运行时，session_state中没有"messages"，需要创建一个空列表
if "messages" not in st.session_state:
st.session_state["messages"] = []
# 每次对话时，都需要遍历session_state中的所有消息，并显示在聊天界面上
for msg in st.session_state.messages:
st.chat_message(msg["role"]).write(msg["content"])
# 如果用户在聊天输入框中输入了内容，则执行以下操作
if prompt := st.chat_input("请输入您的问题:"):
# 将用户的输入添加到session_state中的messages列表中
st.session_state.messages.append({"role": "user", "content": prompt})
# 在聊天界面上显示用户的输入
st.chat_message("user").write(prompt)
# 调用模型
prompt = "<n>".join(msg["content"] for msg in st.session_state.messages) + "<sep>" # 拼接对话历史
inputs = tokenizer(prompt, return_tensors="pt")["input_ids"].cuda()
outputs = model.generate(inputs, do_sample=False, max_length=1024) # 设置解码方式和最大生成长度
output = tokenizer.decode(outputs[0])
response = output.split("<sep>")[-1].replace("<eod>", '')
# 将模型的输出添加到session_state中的messages列表中
st.session_state.messages.append({"role": "assistant", "content": response})
# 在聊天界面上显示模型的输出
st.chat_message("assistant").write(response)
2.0版本（引入RAG方法）
RAG（Retrieval-Augmented Generation）：这个名字听起来可能有点复杂，但实际上它就是一个帮助人工智能更好地理解和回答问题的技术。让我们来简单地了解一下RAG是什么以及它是怎么工作的。
RAG 是什么
RAG 就像是一个人工智能助手的超级记忆功能。通常情况下，AI在回答问题时，会依赖于它之前学习过的大量知识。但是有时候这些知识可能不够全面或者不够新。这时候RAG就派上用场了——它可以让AI在回答问题的时候去查找最新的信息，就像我们人类在回答问题前会去查阅资料一样。
RAG 怎么工作
想象一下，如果你要写一篇关于恐龙的文章，你会怎么做？你可能会先回忆自己知道的一些基本事实，然后去图书馆或者上网找一些最新的研究资料来丰富你的文章。RAG的工作原理和这个很相似：
- 理解问题：首先，AI需要理解用户提出的问题是什么意思。
- 搜索信息：接下来，AI会在数据库中查找与问题相关的最新信息。这就像你去图书馆或者上网查资料。
- 整合信息：找到相关信息后，AI会把这些信息和它已有的知识结合起来，形成一个更完整的答案。
- 生成回答：最后，AI会根据整合好的信息来生成一个回答，这样就能提供准确且最新的答案给用户了。
为什么需要 RAG
有时候，传统的AI模型可能不知道最新的数据或者事件，比如新的科学研究发现、新闻报道等。有了RAG的帮助，AI就可以实时地获取这些信息，并利用它们来生成更准确的回答。
举个例子来说，如果有人问：“谁是当前世界上最富有的人？”没有RAG的AI可能会给出一个几年前的答案，而有RAG的AI则会去查找最新的财富排行榜来给出最新的名字。
总的来说，RAG就像是给AI装上了“即时更新”的功能，让它们能够更好地适应不断变化的信息环境，从而提供更加准确和有用的答案。
部署体验2.0版本
# cmd
git clone https://github.com/Stars-niu/ByteBrain.git
cd ByteBrain
pip install --upgrade pip setuptools
pip install -r requirements.txt
python finetune_model.py
pip install streamlit
pip install tf-keras
streamlit run appRAG.py --server.address 127.0.0.1 --server.port 1003
程序源码
# 导入所需的库
import streamlit as st
from modelscope import snapshot_download
from typing import List
import numpy as np
import torch
from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM
import os
# 设置CUDA环境变量
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
# 设置页面配置
st.set_page_config(
page_title="ByteBrain",
page_icon="🪐",
layout="wide",
initial_sidebar_state="auto",
menu_items={
"Get Help": "https://bytebrain.com/help",
"Report a bug": "https://bytebrain.com/bug-report",
"About": "https://bytebrain.com/about"
}
)
# 添加自定义CSS样式
st.markdown(
"""
<style>
.main {
background-image: url('Background.png');
background-size: cover;
background-position: center;
padding: 10px;
}
.stButton>button {
background-color: #4CAF50;
color: blue;
border: none;
padding: 10px 24px;
text-align: center;
text-decoration: none;
display: inline-block;
font-size: 16px;
margin: 4px 2px;
cursor: pointer;
border-radius: 12px;
}
.stTextInput>div>div>input {
border: 2px solid #4CAF50;
border-radius: 12px;
padding: 10px;
}
.fixed-right {
position: fixed;
top: 20px;
right: 20px;
width: 30%;
}
.chat-container {
width: 100%;
max-height: 70vh;
overflow-y: auto;
padding-right: 20px;
}
.stChatMessage {
width: 100%;
}
</style>
""",
unsafe_allow_html=True,
)
# 创建一个标题和一个副标题
st.markdown(
'<span style="font-size: 60px">✨ ByteBrain</span> <span style="font-size: 24px">——计算机科学智能知识助手</span>',
unsafe_allow_html=True
)
# 向量模型下载
embed_model_dir = snapshot_download("AI-ModelScope/bge-small-zh-v1.5", cache_dir='.')
# 源大模型下载
llm_model_dir = snapshot_download('IEITYuan/Yuan2-2B-Mars-hf', cache_dir='.')
# 定义向量模型类
class EmbeddingModel:
"""
class for EmbeddingModel
"""
def __init__(self, path: str) -> None:
# 加载预训练的分词器和模型
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = AutoModel.from_pretrained(path)
if torch.cuda.is_available():
self.model = self.model.cuda()
print(f'Loading EmbeddingModel from {path}.')
def get_embeddings(self, texts: List) -> List[float]:
"""
calculate embedding for text list
"""
# 对输入文本进行编码
encoded_input = self.tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
if torch.cuda.is_available():
encoded_input = {k: v.cuda() for k, v in encoded_input.items()}
with torch.no_grad():
# 获取模型输出
model_output = self.model(**encoded_input)
sentence_embeddings = model_output[0][:, 0]
# 对嵌入向量进行归一化
sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
return sentence_embeddings.tolist()
print("> Create embedding model...")
embed_model = EmbeddingModel(embed_model_dir)
# 定义向量库索引类
class VectorStoreIndex:
"""
class for VectorStoreIndex
"""
def __init__(self, document_path: str, embed_model: EmbeddingModel) -> None:
# 加载文档
self.documents = []
for line in open(document_path, 'r', encoding='utf-8'):
line = line.strip()
self.documents.append(line)
self.embed_model = embed_model
# 获取文档的嵌入向量
self.vectors = self.embed_model.get_embeddings(self.documents)
print(f'Loading {len(self.documents)} documents for {document_path}.')
def get_similarity(self, vector1: List[float], vector2: List[float]) -> float:
"""
calculate cosine similarity between two vectors
"""
dot_product = np.dot(vector1, vector2)
magnitude = np.linalg.norm(vector1) * np.linalg.norm(vector2)
if not magnitude:
return 0
return dot_product / magnitude
def query(self, question: str, k: int = 1) -> List[str]:
# 获取问题的嵌入向量
question_vector = self.embed_model.get_embeddings([question])[0]
# 计算问题向量与文档向量的相似度
result = np.array([self.get_similarity(question_vector, vector) for vector in self.vectors])
# 返回相似度最高的文档
return np.array(self.documents)[result.argsort()[-k:][::-1]].tolist()
print("> Create index...")
document_path = './knowledge.txt'
index = VectorStoreIndex(document_path, embed_model)
# 定义大语言模型类
class LLM:
"""
class for Yuan2.0 LLM
"""
def __init__(self, model_path: str) -> None:
print("Create tokenizer...")
# 加载预训练的分词器
self.tokenizer = AutoTokenizer.from_pretrained(model_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("Create model...")
# 加载预训练的语言模型
self.model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, trust_remote_code=True)
if torch.cuda.is_available():
self.model = self.model.cuda()
print(f'Loading Yuan2.0 model from {model_path}.')
def generate(self, question: str, context: List):
# 构建提示词
if context:
prompt = f'背景：{context}\n问题：{question}\n请基于背景，回答问题。'
else:
prompt = question
prompt += "<sep>"
inputs = self.tokenizer(prompt, return_tensors="pt")["input_ids"]
if torch.cuda.is_available():
inputs = inputs.cuda()
# 截断输入以确保长度不超过最大长度
max_length = 2048 # 增大最大长度限制
if inputs.shape[1] > max_length:
inputs = inputs[:, -max_length:]
# 生成输出
outputs = self.model.generate(inputs, do_sample=False, max_new_tokens=512) # 增大最大生成长度
output = self.tokenizer.decode(outputs[0])
# 移除不需要的字符
output = output.split("<sep>")[-1].replace("<eod>", "").strip()
return output
print("> Create Yuan2.0 LLM...")
llm = LLM(llm_model_dir)
# 初次运行时，session_state中没有"messages"，需要创建一个空列表
if "messages" not in st.session_state:
st.session_state["messages"] = []
# 使用分栏布局
col1, col2 = st.columns([2, 1])
with col1:
st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
# 每次对话时，都需要遍历session_state中的所有消息，并显示在聊天界面上
for msg in st.session_state.messages:
st.chat_message(msg["role"]).write(msg["content"])
st.markdown("</div>", unsafe_allow_html=True)
with col2:
st.markdown("<div class='fixed-right'>", unsafe_allow_html=True)
st.image("logo.png", caption="ByteBrain Logo", width=150)
st.markdown("ByteBrain——一个智能知识助手，旨在帮助用户快速获取信息和解决问题。")
st.markdown("### 联系我们")
st.markdown("如果您有任何问题或建议，请通过以下方式联系我们：")
st.markdown("- 邮箱: support@bytebrain.com")
st.markdown("- 电话: 520-1314")
st.markdown("</div>", unsafe_allow_html=True)
# 聊天输入框放在col1之外
prompt = st.chat_input("请输入您的问题:")
if prompt:
# 将用户的输入添加到session_state中的messages列表中
st.session_state.messages.append({"role": "user", "content": prompt})
# 在聊天界面上显示用户的输入
st.chat_message("user").write(prompt)
# 检索相关知识
context = index.query(prompt)
# 调用模型
response = llm.generate(prompt, context)
# 将模型的输出添加到session_state中的messages列表中
st.session_state.messages.append({"role": "assistant", "content": response})
# 在聊天界面上显示模型的输出
st.chat_message("assistant").write(response)
该方法需要自行准备知识库
#knowledge.txt
广州大学（Guangzhou University），简称广大（GU），是由广东省广州市人民政府举办的全日制普通高等学校，实行省市共建、以市为主的办学体制，是国家“111计划”建设高校、广东省和广州市高水平大学重点建设高校。广州大学的办学历史可以追溯到1927年创办的私立广州大学；1951年并入华南联合大学；1983年筹备复办，1984年定名为广州大学；2000年7月，经教育部批准，与广州教育学院（1953年创办）、广州师范学院（1958年创办）、华南建设学院西院（1984年创办）、广州高等师范专科学校（1985年创办）合并组建成立新的广州大学。
郑州机械研究所有限公司（以下简称郑机所）的前身机械科学研究院1956年始建于北京，是原机械工业部直属一类综合研究院所，现隶属于国资委中国机械科学研究总院集团有限公司。郑机所伴随着共和国的成长一路走来，应运而生于首都，碧玉年华献中原。多次搬迁，驻地从北京经漯河再到郑州；数易其名，由机械科学研究院到漯河机械研究所再到郑州机械研究所，现为郑州机械研究所有限公司。1956～1958年应运而生：依据全国人大一届二次会议的提议和第一机械工业部的决策，1956年3月6日，第一机械工业部发文《（56）机技研究第66号》，通知“机械科学实验研究院”（后改名为“机械科学研究院”）在北京成立。1959～1968年首次创业：承担国家重大科研项目与开发任务，以及行业发展规划以及标准制定等工作，如“九大设备”的若干关键技术等。1969～1972年搬迁河南：1969年按照“战备疏散”的要求，机械科学研究院主体迁建河南漯河，成立“漯河机械研究所”；1972年因发展需要，改迁河南郑州，成立郑州机械研究所。1973～1998年二次创业：先后隶属于国家机械工业委员会、机械电子工业部、机械工业部；1981年4月罗干由铸造室主任升任副所长，同年经国务院批准具备硕士学位授予权；1985年“葛洲坝二、三江工程及其水电机组项目”荣获国家科技进步特等奖。1999～2016年发展壮大：1999年转企改制，隶属于国资委中国机械科学研究总院；2008年被河南省首批认定为“高新技术企业”；2011年获批组建新型钎焊材料与技术国家重点实验室；2014年被工信部认定为“国家技术创新示范企业”；历经十多年开发出填补国内外空白的大型齿轮齿条试验装备，完成了对三峡升船机齿条42.2万次应力循环次数的疲劳寿命试验测试；营业收入从几千万发展到近6亿；2017年至今协同发展：2017年经公司制改制，更名为郑州机械研究所有限公司，一以贯之地坚持党对国有企业的领导，充分发挥党委把方向、管大局、保落实的领导作用；一以贯之地建立现代企业制度，持续推进改革改制，努力实现以高质量党建引领郑机所高质量发展。
非洲野犬，属于食肉目犬科非洲野犬属哺乳动物。 又称四趾猎狗或非洲猎犬； 其腿长身短、体形细长；身上有鲜艳的黑棕色、黄色和白色斑块；吻通常黑色，头部中间有一黑带，颈背有一块浅黄色斑；尾基呈浅黄色，中段呈黑色，末端为白色，因此又有“杂色狼”之称。 非洲野犬分布于非洲东部、中部、南部和西南部一带。 栖息于开阔的热带疏林草原或稠密的森林附近，有时也到高山地区活动。其结群生活，没有固定的地盘，一般在一个较大的范围内逗留时间较长。非洲野犬性情凶猛，以各种羚羊、斑马、啮齿类等为食。奔跑速度仅次于猎； 雌犬妊娠期为69-73天，一窝十只仔，哺乳期持续6-12个星期。 其寿命11年。 非洲野犬正处在灭绝边缘，自然界中仅存两三千只。 非洲野犬被列入《世界自然保护联盟濒危物种红色名录》中，为濒危（EN）保护等级。 ",非洲野犬共有42颗牙齿（具体分布为：i=3/3；c=1/1；p=4/4；m=2/3x2），前臼齿比相对比其他犬科动物要大，因此可以磨碎大量的骨头，这一点很像鬣狗。 主要生活在非洲的干燥草原和半荒漠地带，活跃于草原、稀树草原和开阔的干燥灌木丛，甚至包括撒哈拉沙漠南部一些多山的地带。非洲野犬从来不到密林中活动。
3.0版本（引入大模型微调）
大模型微调（Fine tuning）:想象一下，你有一个非常聪明的助手，它已经学会了很多基本技能，比如理解和回答问题、翻译语言、识别图片上的东西等等。这个助手就像是一个大模型，它通过学习大量的信息来掌握这些技能。
但是，这个助手虽然很聪明，它学到的东西可能并不完全适合你的具体需求。比如，你可能需要它特别擅长理解医学问题或者法律文件。这时候，我们就需要对助手进行一些特别的训练，让它在某些方面变得更加擅长。这个过程就叫做“微调”。
大模型微调的步骤
选择一个大模型(这个模型已经通过学习大量的数据，具备了广泛的知识和技能):
- 准备特定领域的数据：这些数据是专门为你的需要准备的，比如医学问题和答案的集合。
- 微调过程：将这个大模型和你准备的数据一起训练一段时间。在这个过程中，模型会学习到如何更好地处理和理解你的特定领域数据。
- 调整模型参数：在训练过程中，模型的一些内部参数会被调整，以更好地适应新的数据。
- 评估和测试：训练完成后，我们会测试模型的表现，看看它是否已经足够擅长处理特定领域的任务。
- 部署应用：如果测试结果令人满意，这个经过微调的模型就可以被用来解决实际问题了。
为什么需要微调
提高准确性：微调可以帮助模型更准确地理解和处理特定类型的数据。
适应特定需求：每个领域都有其独特的术语和语境，微调可以让模型更好地适应这些需求。
提升效率：相比于从头开始训练一个模型，微调可以在较短的时间内提升模型在特定任务上的表现。
微调的好处
灵活性：可以针对不同的需求快速调整模型。灵活性：可以针对不同的需求快速调整模型。
成本效益：相比于全面重新训练，微调通常需要较少的资源和时间。成本效益：相比于全面重新训练，微调通常需要较少的资源和时间。
持续学习：随着时间的推移，可以不断地对模型进行微调，以适应新的情况和数据。持续学习：随着时间的推移，可以不断地对模型进行微调，以适应新的情况和数据。
通过微调，我们可以让一个通用的智能助手变得更加专业和高效，更好地服务于特定的任务和需求。
部署体验3.0
git clone https://github.com/Stars-niu/ByteBrain.git
cd ByteBrain
pip install --upgrade pip setuptools
pip install -r requirements.txt
python finetune_model.py
pip install streamlit
pip install tf-keras
streamlit run appFineTuning.py --server.address 127.0.0.1 --server.port 1005
程序源码
import streamlit as st
from modelscope import snapshot_download
from typing import List, Dict
import numpy as np
import torch
from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
import os
import json
from datasets import load_dataset, Dataset
# 设置页面配置
st.set_page_config(
page_title="ByteBrain",
page_icon="🪐",
layout="wide",
initial_sidebar_state="auto",
menu_items={
"Get Help": "https://bytebrain.com/help",
"Report a bug": "https://bytebrain.com/bug-report",
"About": "https://bytebrain.com/about"
}
)
# 添加自定义 CSS 样式，美化页面
st.markdown(
"""
<style>
body {
background-color: #282c34; /* 深灰色背景 */
color: #fff; /* 白色文字 */
font-family: Arial, sans-serif; /* 字体 */
}
.main {
padding: 20px;
}
.stButton>button {
background-color: #61dafb; /* 浅蓝色按钮背景 */
color: #000; /* 黑色文字 */
border: none;
padding: 10px 24px;
text-align: center;
text-decoration: none;
display: inline-block;
font-size: 16px;
margin: 4px 2px;
cursor: pointer;
border-radius: 12px;
}
.stTextInput>div>div>input {
border: 2px solid #61dafb; /* 浅蓝色边框 */
border-radius: 12px;
padding: 10px;
background-color: #fff; /* 白色输入框背景 */
color: #000; /* 黑色文字 */
}
.fixed-right {
position: fixed;
top: 20px;
right: 20px;
width: 30%;
}
.chat-container {
width: 100%;
max-height: 70vh;
overflow-y: auto;
padding-right: 20px;
}
.stChatMessage {
width: 100%;
}
.answer-box {
background-color: rgba(0, 0, 0, 0.5);
padding: 15px;
border-radius: 12px;
margin-top: 20px;
}
h1 {
color: #61dafb; /* 浅蓝色标题 */
}
h2 {
color: #98c379; /* 浅绿色副标题 */
}
</style>
""",
unsafe_allow_html=True,
)
# 向量模型下载
try:
embed_model_dir = snapshot_download("AI-ModelScope/bge-small-zh-v1.5", cache_dir='.')
print(f'Successfully loaded embedding model from cache.')
except Exception as e:
st.error(f'Error loading embedding model: {e}')
embed_model_dir = None
# 源大模型下载
try:
llm_model_dir = snapshot_download('IEITYuan/Yuan2-2B-Mars-hf', cache_dir='.')
print(f'Successfully loaded LLM model from cache.')
except Exception as e:
st.error(f'Error loading LLM model: {e}')
llm_model_dir = None
# 定义向量模型类
class EmbeddingModel:
"""
class for EmbeddingModel
"""
def __init__(self, path: str) -> None:
try:
# 加载预训练的分词器和模型
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = AutoModel.from_pretrained(path)
# 确保不使用 GPU
self.model = self.model.to('cpu')
print(f'Loading EmbeddingModel from {path}.')
except Exception as e:
print(f'Error initializing EmbeddingModel: {e}')
def get_embeddings(self, texts: List[str]) -> List[float]:
"""
calculate embedding for text list
"""
try:
# 对输入文本进行编码
encoded_input = self.tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
# 确保不使用 GPU
encoded_input = {k: v.to('cpu') for k, v in encoded_input.items()}
with torch.no_grad():
# 获取模型输出
model_output = self.model(**encoded_input)
sentence_embeddings = model_output[0][:, 0]
# 对嵌入向量进行归一化
sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
return sentence_embeddings.tolist()
except Exception as e:
print(f'Error in get_embeddings: {e}')
return []
print("> Create embedding model...")
if embed_model_dir:
embed_model = EmbeddingModel(embed_model_dir)
else:
st.error('Embedding model not loaded. Cannot continue.')
embed_model = None
# 定义向量库索引类
class VectorStoreIndex:
"""
class for VectorStoreIndex
"""
def __init__(self, document_path: str, embed_model: EmbeddingModel) -> None:
try:
# 加载文档
self.documents = []
for line in open(document_path, 'r', encoding='utf-8'):
line = line.strip()
self.documents.append(line)
self.embed_model = embed_model
# 获取文档的嵌入向量
self.vectors = self.embed_model.get_embeddings(self.documents)
print(f'Loading {len(self.documents)} documents for {document_path}.')
except Exception as e:
print(f'Error initializing VectorStoreIndex: {e}')
def get_similarity(self, vector1: List[float], vector2: List[float]) -> float:
"""
calculate cosine similarity between two vectors
"""
try:
dot_product = np.dot(vector1, vector2)
magnitude = np.linalg.norm(vector1) * np.linalg.norm(vector2)
if not magnitude:
return 0
return dot_product / magnitude
except Exception as e:
print(f'Error in get_similarity: {e}')
return 0
def query(self, question: str, k: int = 1) -> List[str]:
"""
query the vector store for similar documents
"""
try:
# 获取问题的嵌入向量
question_vector = self.embed_model.get_embeddings([question])[0]
# 计算问题向量与文档向量的相似度
result = np.array([self.get_similarity(question_vector, vector) for vector in self.vectors])
# 返回相似度最高的文档
return np.array(self.documents)[result.argsort()[-k:][::-1]].tolist()
except Exception as e:
print(f'Error in query: {e}')
return []
print("> Create index...")
document_path = './knowledge.txt'
if embed_model:
index = VectorStoreIndex(document_path, embed_model)
else:
st.error('Embedding model not available. Cannot create index.')
index = None
# 定义大语言模型类（包含微调功能）
class LLM:
"""
class for Yuan2.0 LLM
"""
def __init__(self, model_path: str) -> None:
try:
print("Create tokenizer...")
# 加载预训练的分词器
self.tokenizer = AutoTokenizer.from_pretrained(model_path, add_eos_token=False, add_bos_token=False, eos_token='<eod>', legacy=False)
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("Create model...")
# 加载预训练的语言模型，添加 trust_remote_code=True
self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
# 确保不使用 GPU
self.model = self.model.to('cpu')
print(f'Loading Yuan2.0 model from {model_path}.')
except Exception as e:
print(f'Error initializing LLM: {e}')
def generate(self, question: str, context: List[str]) -> str:
"""
generate answer using the language model
"""
try:
# 构建提示词
if context:
prompt = f'背景：{context}\n问题：{question}\n请基于背景，回答问题。'
else:
prompt = question
prompt += "<sep>"
inputs = self.tokenizer(prompt, return_tensors="pt")["input_ids"]
# 确保不使用 GPU
inputs = inputs.to('cpu')
# 截断输入以确保长度不超过最大长度
max_length = 2048
if inputs.shape[1] > max_length:
inputs = inputs[:, -max_length:]
# 确保输入数据类型正确
if inputs.dtype == torch.bfloat16:
inputs = inputs.to(torch.long)
elif inputs.dtype!= torch.long:
st.error(f'Unexpected input data type: {inputs.dtype}. Expected torch.long.')
return 'Sorry, an error occurred while generating the answer.'
# 创建注意力掩码
attention_mask = torch.ones_like(inputs)
# 生成输出时确保模型和输入数据的数据类型一致
outputs = self.model.generate(inputs, attention_mask=attention_mask, do_sample=False, max_new_tokens=512)
output = self.tokenizer.decode(outputs[0])
# 移除不需要的字符
output = output.split("<sep>")[-1].replace("<eod>", "").strip()
return output
except Exception as e:
print(f'Error in generate: {e}')
return 'Sorry, an error occurred while generating the answer.'
def fine_tune(self, data_path: str) -> None:
"""
fine-tune the language model
"""
try:
# 加载微调数据
dataset = load_dataset('json', data_files=data_path, split='train')
def preprocess_function(examples):
inputs = examples['input_text']
targets = examples['output_text']
model_inputs = self.tokenizer(inputs, max_length=512, truncation=True)
# Setup the tokenizer for targets
with self.tokenizer.as_target_tokenizer():
labels = self.tokenizer(targets, max_length=512, truncation=True)
model_inputs['labels'] = labels['input_ids']
return model_inputs
tokenized_dataset = dataset.map(preprocess_function, batched=True)
data_collator = DataCollatorForLanguageModeling(tokenizer=self.tokenizer, mlm=False)
training_args = TrainingArguments(
output_dir='./results',
overwrite_output_dir=True,
num_train_epochs=3,
per_device_train_batch_size=4,
save_steps=10_000,
save_total_limit=2,
logging_dir='./logs',
logging_steps=200,
)
trainer = Trainer(
model=self.model,
args=training_args,
data_collator=data_collator,
train_dataset=tokenized_dataset,
)
trainer.train()
except Exception as e:
print(f'Error in fine_tune: {e}')
print("> Create LLM model...")
if llm_model_dir:
llm_model = LLM(llm_model_dir)
else:
st.error('LLM model not loaded. Cannot continue.')
llm_model = None
# Streamlit 用户界面逻辑
st.sidebar.title("ByteBrain 侧边栏")
user_input = st.sidebar.text_area("请输入你的问题：", "", key="unique_text_area_key")
if st.sidebar.button("查询"):
if user_input:
st.sidebar.markdown("### 查询结果")
# 查询 VectorStoreIndex
try:
context = index.query(user_input)
st.sidebar.markdown("#### 上下文")
for doc in context:
st.sidebar.markdown(f"- {doc}")
except Exception as e:
st.sidebar.error(f'查询时出现错误：{e}')
# 使用大语言模型生成回答
if llm_model:
try:
answer = llm_model.generate(user_input, context)
st.sidebar.markdown("#### 回答")
st.sidebar.markdown(f"{answer}")
except Exception as e:
st.sidebar.error(f'生成回答时出现错误：{e}')
else:
st.sidebar.error("语言模型加载失败，无法生成回答。")
else:
st.sidebar.warning("请输入你的问题！")
Python环境的配置文件
# requirements.txt
streamlit==1.24.0
packaging
torch==2.0.1
torchvision==0.15.2
pandas
datasets
transformers==4.35.0
peft
modelscope
tb-nightly
sentencepiece
flash_attn
einops
tf-keras
requests
beautifulsoup4
tensorboard<2.17,>=2.16
tensorflow==2.16.1
timm==0.5.4
huggingface_hub==0.14.1 # Add the compatible version
最后的最后
过程很难，一遍遍的修改、一遍遍的提交；但却真真实实很畅快，由零到一，从整体构思，到功能实现；从logo设计，到PPT制作；收获颇多。
尽管很粗糙，尽管是和Ai携手，但完成重于完美，迈出第一步，才有后续的无数步。
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
