基于向量检索的RAG原理和实战_rag与向量搜索

什么是RAG

RAG（Retrieval Augmented Generation）检索增强生成
通常用于在大模型交互过程中通过检索的方法来增强生成模型的能力。

由于大模型存在的局限性（大模型的知识不是实时的，大模型可能不知道某些私有领域的知识），可以使用检索的方式给大模型提供相应的知识，再参考回答问题

一张图说明RAG的流程

1. 将知识灌入数据库
2. 用户提问
3. 先去数据库检索相关知识
4. 将检索出的内容结合prompt送给大模型
5. 大模型给出回答

什么是向量检索

提到检索，我们通常会使用elastic search，但是es是基于关键字的检索，有其局限性，同一个语义，用词不同，就会导致无法检索出想要的数据
举例：
es中存储 “小明喜欢吃汉堡”，如果检索关键字“小明喜欢吃”，那么可以找到“汉堡”，但是如果检索“小明爱吃”，那么就找不到这个数据了

为了解决这种问题，所以需要使用向量检索

文本向量

1. 将文本转成一组浮点数：每个下标 $i$，对应一个维度
2. 整个数组对应一个 $n$ 维空间的一个点，即文本向量又叫 Embeddings
3. 向量之间可以计算距离，距离远近对应语义相似度大小

文本向量是怎么得到的

1. 构建相关（正立）与不相关（负例）的句子对照样本
2. 训练双塔式模型，让正例间的距离小，负例间的距离大

向量间的相似度计算
两种计算方式
欧氏距离
余弦距离

def cos_sim(a, b):
    '''余弦距离 -- 越大越相似'''
    return dot(a, b)/(norm(a)*norm(b))


def l2(a, b):
    '''欧式距离 -- 越小越相似'''
    x = np.asarray(a)-np.asarray(b)
    return norm(x)
1
2
3
4
5
6
7
8
9

向量数据库
向量数据库，是专门为向量检索设计的中间件
向量数据库本身不生成向量，向量是由 Embedding 模型产生的
向量数据库与传统的关系型数据库是互补的，不是替代关系，在实际应用中根据实际需求经常同时使用。

主流的向量数据库

实现一个基于文档向量检索的RAG

流程如下：

1. 文档加载

def extract_text_from_pdf(filename,page_numbers=None,min_line_length=10):
    """从 PDF 文件中（按指定页码）提取文字"""
    paragraphs = []
    buffer = ''
    full_text = ''
    # 提取全部文本
    for i, page_layout in enumerate(extract_pages(filename)):
        # 如果指定了页码范围，跳过范围外的页
        if page_numbers is not None and i not in page_numbers:
            continue
        for element in page_layout:
            if isinstance(element, LTTextContainer):
                full_text += element.get_text() + '\n'
    # 按空行分隔，将文本重新组织成段落
    lines = full_text.split('\n')
    for text in lines:
        if len(text) >= min_line_length:
            buffer += (' '+text) if not text.endswith('-') else text.strip('-')
        elif buffer:
            paragraphs.append(buffer)
            buffer = ''
    if buffer:
        paragraphs.append(buffer)
    return paragraphs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

2. 文档切割（交叠切割防止问题的答案跨两个片段，使上下文更完整）

def split_text(paragraphs,chunk_size=300,overlap_size=100):
    """按指定 chunk_size 和 overlap_size 交叠割文本"""
    sentences = [s.strip() for p in paragraphs for s in sent_tokenize(p)]
    chunks = []
    i= 0
    while i < len(sentences):
        chunk = sentences[i]
        overlap = ''
        prev_len = 0
        prev = i - 1
        # 向前计算重叠部分
        while prev >= 0 and len(sentences[prev])+len(overlap) <= overlap_size:
            overlap = sentences[prev] + ' ' + overlap
            prev -= 1
        chunk = overlap+chunk
        next = i + 1
        # 向后计算当前chunk
        while next < len(sentences) and len(sentences[next])+len(chunk) <= chunk_size:
            chunk = chunk + ' ' + sentences[next]
            next += 1
        chunks.append(chunk)
        i = next
    return chunks
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

3. 向量化（这里使用openai的向量化模型）

def get_embedding(text, model="text-embedding-ada-002"):
    """封装 OpenAI 的 Embedding 模型接口"""
    return openai.Embedding.create(input=[text], model=model)['data'][0]['embedding']
1
2
3

4. 灌入向量库（使用chromadb）

def __init__(self, name="demo"):
        self.chroma_client = chromadb.Client(Settings(allow_reset=True))
        self.chroma_client.reset()
        self.name = name
        self.collection = self.chroma_client.get_or_create_collection(name=name)

def add_documents(self, documents):
        self.collection.add(
            embeddings=[get_embedding(doc) for doc in documents],
            documents=documents,
            metadatas=[{"source": self.name} for _ in documents],
            ids=[f"id_{i}" for i in range(len(documents))]
        )
1
2
3
4
5
6
7
8
9
10
11
12
13

5. 检索向量数据库

def search(self, query, top_n):
        """检索向量数据库"""
        results = self.collection.query(
            query_embeddings=[get_embedding(query)],
            n_results=top_n
        )
        return results['documents'][0]
1
2
3
4
5
6
7

6. 将检索数据带入提示词

def build_prompt(template=prompt_template, **kwargs):
    """将 Prompt 模板赋值"""
    prompt = template
    for k, v in kwargs.items():
        if isinstance(v, str):
            val = v
        elif isinstance(v, list) and all(isinstance(elem, str) for elem in v):
            val = '\n'.join(v)
        else:
            val = str(v)
        prompt = prompt.replace(f"__{k.upper()}__", val)
    return prompt
1
2
3
4
5
6
7
8
9
10
11
12

7. 调用大模型（使用gpt-3.5-turbo）

def get_completion(prompt, context, model="gpt-3.5-turbo"):
    """封装 openai 接口"""
    messages = context + [{"role": "user", "content": prompt}]
    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=0,  # 模型输出的随机性，0 表示随机性最小
    )
    return response.choices[0].message["content"]
1
2
3
4
5
6
7
8
9