RAG库搭建：从零开始，开启智能问答新世界

从零开始的RAG库的搭建

什么是RAG？RAG是什么？RAG是Retrieval Augmented Generation的缩写，即检索增强生成。RAG是一种信息检索和自然语言处理的结合，它通过将文档和用户问题进行匹配，从而实现信息检索和自然语言处理的结合。

当我们输入一个问题，问题的答案会从文档中检索出来，然后将答案和问题一起输入到大模型中生成答案。

一、基本结构

要实现一个最简单的RAG，我们首先要了解RAG的基础结构。

• 1.向量化结构：能够将文档片段向量化
• 2.文档切分和加载模块：能够将文档切分为多个片段，并加载到内存中
• 3.向量检索模块：能够将向量检索到最相似的向量，并返回对应的文档片段
• 4.向量数据库模块：能够将向量存储到数据库中，并查询到最相似的向量
• 5.问答模块：能够将用户问题向量化，并查询到最相似的向量，并返回对应的答案

二、向量化结构

1.为什么需要向量化？

• 1.高效检索：将文本向量化后，可以利用向量之间的距离（如余弦距离）来衡量文本的相似度，从而快速检索出与用户查询最相关的文档片段
• 2.语义理解：向量化能够捕捉文本的语义信息，使得模型能够理解文本的含义，而不仅仅是基于关键词的匹配。

2.向量化的一般方法

• 1.传统稀疏向量化（Sparse Embedding）：映射成一个高维向量，维度通常与词汇表大小相同。向量的大部分元素为0，非零值表示特定单词在文档中的重要性。典型模型包括TF-IDF和BM25，适合关键词匹配任务。
• 2.密集向量化（Dense Embedding）：映射到一个相对低维的向量，所有维度都非零。典型模型包括基于BERT的模型（如BGE-v1.5）和Sentence Transformers，这些模型能够捕捉语义信息，适用于语义搜索任务

3. 代码实现

from sentence_transformers import SentenceTransformer
'''
SentenceTransformer 是一个非常流行的 Python 库，用于将文本（句子或段落）转换为密集的向量表示（嵌入向量）。它是基于 Hugging Face 的 transformers 库和 PyTorch 的，提供了简单易用的接口来加载预训练的模型并生成文本嵌入。
'''
#下面实现一个基本类，然后继承这个类实现一个向量化的类
class BaseEmbeddings:"""Base class for embeddings"""def __init__(self, path: str, is_api: bool) -> None:self.path = pathself.is_api = is_api# 下面这个函数是得到向量化的函数def get_embedding(self, text: str, model: str) -> List[float]:raise NotImplementedError#下面这个函数大户是计算余弦相似度@classmethoddef cosine_similarity(cls, 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 0return dot_product / magnitude#下面这个类是实现一个向量化的类
class MyEmbedding(BaseEmbeddings):"""class for My embeddings"""def __init__(self, path: str = '', is_api: bool = False) -> None:super().__init__(path, is_api)self.path = pathself.model = SentenceTransformer(path)#这里是加载一个嵌入模型，可以是bert的模型def get_embedding(self, text: str) -> List[float]:embeddings  = self.model.encode(text)return embeddings

三、文档切分和加载模块

1. 为什么需要文档切分？

• 1.提高检索效率：将文档切分为多个片段，可以减少检索的文档数量，提高检索效率。
• 2.提高检索准确性：将文档切分为多个片段，可以提高检索的准确性，因为每个片段都包含一些与用户查询最相关的信息。

2.如何切分

• 1.基于字符长度：将文档切分为固定长度的片段。
• 2.基于句子长度：将文档切分为句子，然后根据句子长度进行切分。
• 3.滑动窗口切分：允许相邻块之间部分重叠，减少信息断裂，缓解上下文不连贯问题，提升检索相关性，但是会增加计算和内存开销
  除此之外还有许多切分方法，这里不在赘述

3. 文档切分和加载模块代码实现

def get_chunk(cls, text: str, max_token_len: int = 600, cover_content: int = 150):chunk_text = []curr_len = 0curr_chunk = ''lines = text.split('\n')  # 假设以换行符分割文本为行for line in lines:line = line.replace(' ', '')line_len = len(enc.encode(line))if line_len > max_token_len:print('warning line_len = ', line_len)if curr_len + line_len <= max_token_len:curr_chunk += linecurr_chunk += '\n'curr_len += line_lencurr_len += 1else:chunk_text.append(curr_chunk)curr_chunk = curr_chunk[-cover_content:]+linecurr_len = line_len + cover_contentif curr_chunk:chunk_text.append(curr_chunk)return chunk_text

四、数据库和向量检索模块

数据库和向量检索模块需要实现下面四个功能：

• persist：数据库持久化，本地保存
• load_vector：从本地加载数据库
• get_vector：获得文档的向量表示
• query：根据问题检索相关的文档片段

1.数据库和向量检索模块代码实现

class VectorStore:def __init__(self, document: List[str] = ['']) -> None:self.document = documentdef get_vector(self, EmbeddingModel: BaseEmbeddings) -> List[List[float]]:self.vectors = []for doc in tqdm(self.document, desc="Calculating embeddings"):self.vectors.append(EmbeddingModel.get_embedding(doc).tolist())return self.vectorsdef persist(self, path: str = 'storage'):if not os.path.exists(path):os.makedirs(path)with open(f"{path}/doecment.json", 'w', encoding='utf-8') as f:json.dump(self.document, f, ensure_ascii=False)if self.vectors:with open(f"{path}/vectors.json", 'w', encoding='utf-8') as f:json.dump(self.vectors, f)def load_vector(self, path: str = 'storage'):with open(f"{path}/vectors.json", 'r', encoding='utf-8') as f:self.vectors = json.load(f)with open(f"{path}/doecment.json", 'r', encoding='utf-8') as f:self.document = json.load(f)def get_similarity(self, vector1: List[float], vector2: List[float]) -> float:return BaseEmbeddings.cosine_similarity(vector1, vector2)def query(self, query: str, EmbeddingModel: BaseEmbeddings, k: int = 1) -> List[str]:query_vector = EmbeddingModel.get_embedding(query)result = np.array([self.get_similarity(query_vector, vector)for vector in self.vectors])return np.array(self.document)[result.argsort()[-k:][::-1]].tolist()

五、问答模块

问答模块其实非常简答，将检索到的内容和用户的问题进行拼接，然后使用LLM进行回答即可。

class BaseModel:def __init__(self, path: str = '') -> None:self.path = pathdef chat(self, prompt: str, history: List[dict], content: str) -> str:passdef load_model(self):passclass InternLMChat(BaseModel):def __init__(self, path: str = '') -> None:super().__init__(path)self.load_model()def chat(self, prompt: str, history: List = [], content: str='') -> str:prompt = PROMPT_TEMPLATE['InternLM_PROMPT_TEMPALTE'].format(question=prompt, context=content)response, history = self.model.chat(self.tokenizer, prompt, history)return responsedef load_model(self):import torchfrom transformers import AutoTokenizer, AutoModelForCausalLMself.tokenizer = AutoTokenizer.from_pretrained(self.path, trust_remote_code=True)self.model = AutoModelForCausalLM.from_pretrained(self.path, torch_dtype=torch.float16, trust_remote_code=True).cuda()