概述:本笔记用于对作者第一次RAG实战的记录以及思考
一、在魔塔平台创建实例并安装配置环境(魔塔创建实例的步骤参考task01笔记)
git lfs install
git clone https://www.modelscope.cn/datasets/Datawhale/AICamp_yuan_baseline.git
cp AICamp_yuan_baseline/Task\ 3:源大模型RAG实战/* .
二、终端代码执行结束以后,依次执行代码块代码即可完成
三、模型下载
在RAG实战中,构建一个向量模型。
向量模型通常采用BERT架构,它是一个Transformer Encoder。
# 向量模型下载
from modelscope import snapshot_download
model_dir = snapshot_download("AI-ModelScope/bge-small-zh-v1.5", cache_dir='.')
# 源大模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('IEITYuan/Yuan2-2B-Mars-hf', cache_dir='.')
四、构造索引
封装一个向量模型类EmbeddingModel
# 定义向量模型类
class EmbeddingModel:
"""
class for EmbeddingModel
"""
def __init__(self, path: str) -> None:
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = AutoModel.from_pretrained(path).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')
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()
通过传入模型路径,新建一个 EmbeddingModel
对象 embed_model
。
初始化时自动加载向量模型的tokenizer和模型参数。
#新建一个 EmbeddingModel 对象 embed_model
print("> Create embedding model...")
embed_model_path = './AI-ModelScope/bge-small-zh-v1___5'
embed_model = EmbeddingModel(embed_model_path)
五、实现向量检索
为了实现向量检索,我们定义了一个向量库索引类 VectorStoreIndex
# 定义向量库索引类
class VectorStoreIndex:
"""
class for VectorStoreIndex
"""
def __init__(self, doecment_path: str, embed_model: EmbeddingModel) -> None:
self.documents = []
for line in open(doecment_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 {doecment_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()
通过传入知识库文件路径,新建一个 VectorStoreIndex
对象 index
。
初始化时会自动读取知识库的内容,然后传入向量模型,获得向量表示。
#新建一个 VectorStoreIndex 对象 index
print("> Create index...")
doecment_path = './knowledge.txt'
index = VectorStoreIndex(doecment_path, embed_model)
上文提到 get_embeddings()
函数支持一次性传入多条文本,但由于GPU的显存有限,输入的文本不宜太多。(这就是为什么市面上的AI对话机器人会有输入文本限制)如果知识库很大,需要将知识库切分成多个batch,然后分批次送入向量模型。
六、基于RAG的生成
为了实现基于RAG的生成,我们还需要定义一个大语言模型类 LLM
# 定义大语言模型类
class LLM:
"""
class for Yuan2.0 LLM
"""
def __init__(self, model_path: str) -> None:
print("Creat 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("Creat model...")
self.model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, trust_remote_code=True).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"].cuda()
outputs = self.model.generate(inputs, do_sample=False, max_length=1024)
output = self.tokenizer.decode(outputs[0])
print(output.split("<sep>")[-1])
通过传入模型路径新建一个LLM对象llm。
初始化时自动加载源大模型的tokenizer和模型参数。
#新建一个 LLM 对象 llm
print("> Create Yuan2.0 LLM...")
model_path = './IEITYuan/Yuan2-2B-Mars-hf'
llm = LLM(model_path)
LLM
类的入口是生成函数 generate()
,它有两个参数:
1.question
: 用户提问,是一个str
2.context
: 检索到的上下文信息,是一个List,默认是[],代表没有使用RAG
最后运行以下代码即可体验有RAG和无RAG的大模型回答效果
print('> Without RAG:')
llm.generate(question, [])
print('> With RAG:')
llm.generate(question, context)
如下所示,如果模型依赖底层知识进行问题回答,很有可能出现经典的计算机幻觉。生成错误的内容。
使用RAG之后,模型能够结果知识库中的知识,准确回答用户的提问。