使用语义缓存改进基于 FAISS 的 RAG 系统。

作者：Pere Martra

在此 notebook 中，我们将探索一个典型的 RAG 解决方案，其中我们将使用开源模型和向量数据库 Chroma DB。**但是，我们将集成一个语义缓存系统，该系统将存储各种用户查询，并决定是生成从向量数据库或缓存中获取信息来丰富 prompt，还是直接从缓存中获取信息。**

语义缓存系统旨在识别相似或相同的用户请求。当找到匹配的请求时，系统会从缓存中检索相应的信息，从而减少从原始来源获取信息的需要。

由于比较考虑了请求的语义含义，因此它们不必完全相同，系统也可以将它们识别为同一个问题。 它们的表述可能不同，或者包含不准确之处（无论是拼写错误还是句子结构错误），但我们仍然可以识别出用户实际上是在请求相同的信息。

例如，像**法国的首都是什么？**、**告诉我法国首都的名称？**和**什么是法国的首都？**这样的查询都表达了相同的意图，应被识别为同一个问题。

虽然模型的响应可能会因第二个示例中对简洁答案的请求而有所不同，但从向量数据库检索的信息应该是相同的。 这就是为什么我将缓存系统放置在用户和向量数据库之间，而不是用户和大型语言模型之间。

大多数指导您创建 RAG 系统的教程都是为单用户使用而设计的，旨在在测试环境中运行。 换句话说，在 notebook 中，与本地向量数据库交互并进行 API 调用或使用本地存储的模型。

当尝试将其中一个模型过渡到生产环境时，这种架构很快就会变得不足，因为在生产环境中，它们可能会遇到来自数十到数千个重复请求。

一种提高性能的方法是通过一个或多个语义缓存。 此缓存保留先前请求的结果，并在解析新请求之前，检查是否之前已收到类似的请求。 如果是这样，则从缓存中检索信息，而不是重新执行该过程。

在 RAG 系统中，有两个耗时的点

• 检索用于构建增强 prompt 的信息
• 调用大型语言模型以获取响应。

在这两个点上，都可以实现语义缓存系统，我们甚至可以为每个点设置两个缓存。

将其放置在模型的响应点可能会导致对获得的响应失去影响。 我们的缓存系统可能会将“用 10 个词解释法国大革命”和“用 100 个词解释法国大革命”视为相同的查询。 如果我们的缓存系统存储模型响应，用户可能会认为他们的指令没有被准确地遵循。

但是，这两个请求都需要相同的信息来丰富 prompt。 这就是我选择将语义缓存系统放置在用户请求和从向量数据库检索信息之间的主要原因。

但是，这是一个设计决策。 根据响应类型和系统请求，它可以放置在一个点或另一个点。 显然，缓存模型响应将节省最多的时间，但正如我已经解释的那样，这是以牺牲用户对响应的影响为代价的。

导入和加载库。

首先，我们需要安装必要的 Python 包。

• sentence transformers。 这个库是必要的，用于将句子转换为固定长度的向量，也称为 Embedding。
• xformers。 这是一个软件包，提供库和实用程序，以方便使用 transformer 模型的工作。 我们需要安装它，以避免在使用模型和 Embedding 时出现错误。
• chromadb。 这是我们的向量数据库。 ChromaDB 易于使用且开源，可能是最常用的用于存储 Embedding 的向量数据库。
• accelerate。 运行 GPU 中的模型所必需。

!pip install -q transformers==4.38.1
!pip install -q accelerate==0.27.2
!pip install -q sentence-transformers==2.5.1
!pip install -q xformers==0.0.24
!pip install -q chromadb==0.4.24
!pip install -q datasets==2.17.1

import numpy as np
import pandas as pd

加载数据集

由于我们是在免费且有限的空间中工作，并且我们只能使用几 GB 的内存，因此我使用变量 MAX_ROWS 限制了从数据集中使用的行数。

#Login to Hugging Face. It is mandatory to use the Gemma Model,
#and recommended to acces public models and Datasets.
from getpass import getpass
if 'hf_key' not in locals():
  hf_key = getpass("Your Hugging Face API Key: ")
!huggingface-cli login --token $hf_key

from datasets import load_dataset

data = load_dataset("keivalya/MedQuad-MedicalQnADataset", split="train")

ChromaDB 要求数据具有唯一的标识符。 我们可以使用此语句来创建它，这将创建一个名为 **Id** 的新列。

data = data.to_pandas()
data["id"] = data.index
data.head(10)

MAX_ROWS = 15000
DOCUMENT = "Answer"
TOPIC = "qtype"

# Because it is just a sample we select a small portion of News.
subset_data = data.head(MAX_ROWS)

导入和配置向量数据库

为了存储信息，我选择使用 ChromaDB，它是最著名和广泛使用的开源向量数据库之一。

首先，我们需要导入 ChromaDB。

import chromadb

现在我们只需要指示向量数据库将存储的路径。

chroma_client = chromadb.PersistentClient(path="/path/to/persist/directory")

填充和查询 ChromaDB 数据库

ChromaDB 中的数据存储在集合中。 如果集合存在，我们需要删除它。

在接下来的几行中，我们通过调用上面创建的 chroma_client 中的 create_collection 函数来创建集合。

collection_name = "news_collection"
if len(chroma_client.list_collections()) > 0 and collection_name in [chroma_client.list_collections()[0].name]:
    chroma_client.delete_collection(name=collection_name)

collection = chroma_client.create_collection(name=collection_name)

现在我们准备使用 add 函数将数据添加到集合中。 此函数需要三个关键信息

• 在 **document** 中，我们存储数据集 **Answer** 列的内容。
• 在 **metadatas** 中，我们可以告知主题列表。 我使用了 qtype 列中的值。
• 在 **id** 中，我们需要为每一行告知一个唯一的标识符。 我正在使用 MAX_ROWS 的范围创建 ID。

collection.add(
    documents=subset_data[DOCUMENT].tolist(),
    metadatas=[{TOPIC: topic} for topic in subset_data[TOPIC].tolist()],
    ids=[f"id{x}" for x in range(MAX_ROWS)],
)

一旦我们在数据库中有了信息，我们就可以查询它，并请求与我们的需求匹配的数据。 搜索是在 document 的内容中完成的，它不会查找确切的单词或短语。 结果将基于搜索词和文档内容之间的相似性。

元数据不直接参与初始搜索过程，它可以用于在检索后过滤或优化结果，从而实现进一步的自定义和精度。

让我们定义一个函数来查询 ChromaDB 数据库。

def query_database(query_text, n_results=10):
    results = collection.query(query_texts=query_text, n_results=n_results)
    return results

创建语义缓存系统

为了实现缓存系统，我们将使用 Faiss，这是一个允许在内存中存储 Embedding 的库。 它与 Chroma 的功能非常相似，但没有持久性。

为此，我们将创建一个名为 semantic_cache 的类，该类将使用自己的编码器，并为用户执行查询提供必要的功能。

在这个类中，我们首先查询使用 Faiss 实现的缓存，其中包含之前的请求，如果返回的结果高于指定的阈值，它将返回缓存的内容。 否则，它将从 Chroma 数据库中获取结果。

缓存存储在 .json 文件中。

!pip install -q faiss-cpu==1.8.0

import faiss
from sentence_transformers import SentenceTransformer
import time
import json

下面的 init_cache() 函数初始化语义缓存。

它采用 FlatLS 索引，这可能不是最快的，但非常适合小型数据集。 根据缓存数据的特性和预期的数据集大小，可以使用另一个索引，例如 HNSW 或 IVF。

我选择此索引是因为它与示例非常吻合。 它可以与高维向量一起使用，消耗最少的内存，并且在小型数据集上表现良好。

我概述了 Faiss 中可用的各种索引的关键特性。

• FlatL2 或 FlatIP。 非常适合小型数据集，可能不是最快的，但其内存消耗并不高。
• LSH。 它在小型数据集上有效工作，建议用于最多 128 维的向量。
• HNSW。 速度非常快，但需要大量的 RAM。
• IVF。 在大型数据集上效果良好，且不会消耗太多内存或影响性能。

有关 Faiss 中可用的不同索引的更多信息，请访问此链接： https://github.com/facebookresearch/faiss/wiki/Guidelines-to-choose-an-index

def init_cache():
    index = faiss.IndexFlatL2(768)
    if index.is_trained:
        print("Index trained")

    # Initialize Sentence Transformer model
    encoder = SentenceTransformer("all-mpnet-base-v2")

    return index, encoder

在 `retrieve_cache` 函数中，.json 文件从磁盘检索，以防需要在会话之间重用缓存。

def retrieve_cache(json_file):
    try:
        with open(json_file, "r") as file:
            cache = json.load(file)
    except FileNotFoundError:
        cache = {"questions": [], "embeddings": [], "answers": [], "response_text": []}

    return cache

store_cache 函数将包含缓存数据的文件保存到磁盘。

def store_cache(json_file, cache):
    with open(json_file, "w") as file:
        json.dump(cache, file)

这些函数将在 `SemanticCache` 类中使用，该类包括搜索函数及其初始化函数。

尽管 `ask` 函数有大量的代码，但其目的是非常直接的。 它在缓存中查找与用户刚刚提出的问题最相似的问题。

之后，检查它是否在指定的阈值内。 如果是，它直接从缓存返回响应； 否则，它调用 query_database 函数从 ChromaDB 检索数据。

我使用了欧几里得距离而不是余弦距离，后者广泛用于向量比较。 这种选择是基于欧几里得距离是 Faiss 使用的默认指标这一事实。 虽然也可以计算余弦距离，但这样做会增加复杂性，而这可能不会对最终结果做出重大贡献。

我在 `semantic_cache` 类中包含了 FIFO 驱逐策略，旨在提高其效率和灵活性。 通过引入驱逐策略，我们为用户提供了控制缓存在达到其最大容量时如何表现的能力。 这对于保持最佳缓存性能以及处理可用内存受限的情况至关重要。

从缓存的结构来看，FIFO 的实现似乎很简单。 每当新的问题-答案对添加到缓存时，它都会附加到列表的末尾。 因此，最旧（先进先出）的项目位于列表的前面。 当缓存达到其最大大小时，并且您需要驱逐一个项目时，您将从每个列表中删除（弹出）第一个项目。 这就是 FIFO 驱逐策略。

另一种驱逐策略是最近最少使用 (LRU) 策略，它更复杂，因为它需要知道缓存中每个项目上次访问的时间。 但是，此策略尚不可用，将在以后实现。

class semantic_cache:
    def __init__(self, json_file="cache_file.json", thresold=0.35, max_response=100, eviction_policy=None):
        """Initializes the semantic cache.

        Args:
        json_file (str): The name of the JSON file where the cache is stored.
        thresold (float): The threshold for the Euclidean distance to determine if a question is similar.
        max_response (int): The maximum number of responses the cache can store.
        eviction_policy (str): The policy for evicting items from the cache.
                                This can be any policy, but 'FIFO' (First In First Out) has been implemented for now.
                                If None, no eviction policy will be applied.
        """

        # Initialize Faiss index with Euclidean distance
        self.index, self.encoder = init_cache()

        # Set Euclidean distance threshold
        # a distance of 0 means identicals sentences
        # We only return from cache sentences under this thresold
        self.euclidean_threshold = thresold

        self.json_file = json_file
        self.cache = retrieve_cache(self.json_file)
        self.max_response = max_response
        self.eviction_policy = eviction_policy

    def evict(self):
        """Evicts an item from the cache based on the eviction policy."""
        if self.eviction_policy and len(self.cache["questions"]) > self.max_size:
            for _ in range((len(self.cache["questions"]) - self.max_response)):
                if self.eviction_policy == "FIFO":
                    self.cache["questions"].pop(0)
                    self.cache["embeddings"].pop(0)
                    self.cache["answers"].pop(0)
                    self.cache["response_text"].pop(0)

    def ask(self, question: str) -> str:
        # Method to retrieve an answer from the cache or generate a new one
        start_time = time.time()
        try:
            # First we obtain the embeddings corresponding to the user question
            embedding = self.encoder.encode([question])

            # Search for the nearest neighbor in the index
            self.index.nprobe = 8
            D, I = self.index.search(embedding, 1)

            if D[0] >= 0:
                if I[0][0] >= 0 and D[0][0] <= self.euclidean_threshold:
                    row_id = int(I[0][0])

                    print("Answer recovered from Cache. ")
                    print(f"{D[0][0]:.3f} smaller than {self.euclidean_threshold}")
                    print(f"Found cache in row: {row_id} with score {D[0][0]:.3f}")
                    print(f"response_text: " + self.cache["response_text"][row_id])

                    end_time = time.time()
                    elapsed_time = end_time - start_time
                    print(f"Time taken: {elapsed_time:.3f} seconds")
                    return self.cache["response_text"][row_id]

            # Handle the case when there are not enough results
            # or Euclidean distance is not met, asking to chromaDB.
            answer = query_database([question], 1)
            response_text = answer["documents"][0][0]

            self.cache["questions"].append(question)
            self.cache["embeddings"].append(embedding[0].tolist())
            self.cache["answers"].append(answer)
            self.cache["response_text"].append(response_text)

            print("Answer recovered from ChromaDB. ")
            print(f"response_text: {response_text}")

            self.index.add(embedding)

            self.evict()

            store_cache(self.json_file, self.cache)

            end_time = time.time()
            elapsed_time = end_time - start_time
            print(f"Time taken: {elapsed_time:.3f} seconds")

            return response_text
        except Exception as e:
            raise RuntimeError(f"Error during 'ask' method: {e}")

测试 semantic_cache 类。

>>> # Initialize the cache.
>>> cache = semantic_cache("4cache.json")

Index trained

>>> results = cache.ask("How do vaccines work?")

Answer recovered from ChromaDB. 
response_text: Summary : Shots may hurt a little, but the diseases they can prevent are a lot worse. Some are even life-threatening. Immunization shots, or vaccinations, are essential. They protect against things like measles, mumps, rubella, hepatitis B, polio, tetanus, diphtheria, and pertussis (whooping cough). Immunizations are important for adults as well as children.    Your immune system helps your body fight germs by producing substances to combat them. Once it does, the immune system "remembers" the germ and can fight it again. Vaccines contain germs that have been killed or weakened. When given to a healthy person, the vaccine triggers the immune system to respond and thus build immunity.     Before vaccines, people became immune only by actually getting a disease and surviving it. Immunizations are an easier and less risky way to become immune.     NIH: National Institute of Allergy and Infectious Diseases
Time taken: 0.057 seconds

正如预期的那样，此响应已从 ChromaDB 获取。 然后，该类将其存储在缓存中。

现在，如果我们发送第二个非常不同的问题，响应也应该从 ChromaDB 检索。 这是因为先前存储的问题非常不同，以至于它将超过欧几里得距离方面的指定阈值。

>>> results = cache.ask("Explain briefly what is a Sydenham chorea")

Answer recovered from ChromaDB. 
response_text: Sydenham chorea (SD) is a neurological disorder of childhood resulting from infection via Group A beta-hemolytic streptococcus (GABHS), the bacterium that causes rheumatic fever. SD is characterized by rapid, irregular, and aimless involuntary movements of the arms and legs, trunk, and facial muscles. It affects girls more often than boys and typically occurs between 5 and 15 years of age. Some children will have a sore throat several weeks before the symptoms begin, but the disorder can also strike up to 6 months after the fever or infection has cleared. Symptoms can appear gradually or all at once, and also may include uncoordinated movements, muscular weakness, stumbling and falling, slurred speech, difficulty concentrating and writing, and emotional instability. The symptoms of SD can vary from a halting gait and slight grimacing to involuntary movements that are frequent and severe enough to be incapacitating. The random, writhing movements of chorea are caused by an auto-immune reaction to the bacterium that interferes with the normal function of a part of the brain (the basal ganglia) that controls motor movements. Due to better sanitary conditions and the use of antibiotics to treat streptococcal infections, rheumatic fever, and consequently SD, are rare in North America and Europe. The disease can still be found in developing nations.
Time taken: 0.082 seconds

完美，语义缓存系统正如预期的那样运行。

让我们继续用一个与我们刚刚提出的问题非常相似的问题来测试它。

在这种情况下，响应应直接来自缓存，而无需访问 ChromaDB 数据库。

>>> results = cache.ask("Briefly explain me what is a Sydenham chorea.")

Answer recovered from Cache. 
0.028 smaller than 0.35
Found cache in row: 1 with score 0.028
response_text: Sydenham chorea (SD) is a neurological disorder of childhood resulting from infection via Group A beta-hemolytic streptococcus (GABHS), the bacterium that causes rheumatic fever. SD is characterized by rapid, irregular, and aimless involuntary movements of the arms and legs, trunk, and facial muscles. It affects girls more often than boys and typically occurs between 5 and 15 years of age. Some children will have a sore throat several weeks before the symptoms begin, but the disorder can also strike up to 6 months after the fever or infection has cleared. Symptoms can appear gradually or all at once, and also may include uncoordinated movements, muscular weakness, stumbling and falling, slurred speech, difficulty concentrating and writing, and emotional instability. The symptoms of SD can vary from a halting gait and slight grimacing to involuntary movements that are frequent and severe enough to be incapacitating. The random, writhing movements of chorea are caused by an auto-immune reaction to the bacterium that interferes with the normal function of a part of the brain (the basal ganglia) that controls motor movements. Due to better sanitary conditions and the use of antibiotics to treat streptococcal infections, rheumatic fever, and consequently SD, are rare in North America and Europe. The disease can still be found in developing nations.
Time taken: 0.019 seconds

这两个问题非常相似，以至于它们的欧几里得距离非常小，几乎就像它们是相同的。

现在，让我们尝试另一个问题，这次稍微不同，并观察系统如何运行。

>>> question_def = "Write in 20 words what is a Sydenham chorea."
>>> results = cache.ask(question_def)

Answer recovered from Cache. 
0.228 smaller than 0.35
Found cache in row: 1 with score 0.228
response_text: Sydenham chorea (SD) is a neurological disorder of childhood resulting from infection via Group A beta-hemolytic streptococcus (GABHS), the bacterium that causes rheumatic fever. SD is characterized by rapid, irregular, and aimless involuntary movements of the arms and legs, trunk, and facial muscles. It affects girls more often than boys and typically occurs between 5 and 15 years of age. Some children will have a sore throat several weeks before the symptoms begin, but the disorder can also strike up to 6 months after the fever or infection has cleared. Symptoms can appear gradually or all at once, and also may include uncoordinated movements, muscular weakness, stumbling and falling, slurred speech, difficulty concentrating and writing, and emotional instability. The symptoms of SD can vary from a halting gait and slight grimacing to involuntary movements that are frequent and severe enough to be incapacitating. The random, writhing movements of chorea are caused by an auto-immune reaction to the bacterium that interferes with the normal function of a part of the brain (the basal ganglia) that controls motor movements. Due to better sanitary conditions and the use of antibiotics to treat streptococcal infections, rheumatic fever, and consequently SD, are rare in North America and Europe. The disease can still be found in developing nations.
Time taken: 0.016 seconds

我们观察到欧几里得距离增加了，但它仍然在指定的阈值内。 因此，它继续直接从缓存返回响应。

加载模型并创建 prompt

现在可以使用库 **transformers** 了，这是 hugging face 最著名的用于处理语言模型的库。

我们正在导入

• Autotokenizer：它是一个实用程序类，用于对与各种预训练语言模型兼容的文本输入进行标记化。
• AutoModelForCausalLM：它为专门为使用因果语言建模（例如，GPT 模型）或本 notebook 中使用的模型 Gemma-2b-it 进行语言生成任务而设计的预训练语言模型提供了一个接口。

请随意测试 不同的模型，您需要搜索针对文本生成训练的 NLP 模型。

!pip install torch

from torch import cuda, torch

# In a MAC Silicon the device must be 'mps'
# device = torch.device('mps') #to use with MAC Silicon
device = f"cuda:{cuda.current_device()}" if cuda.is_available() else "cpu"

from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "google/gemma-2b-it"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16)

创建扩展的 prompt

为了创建 prompt，我们使用从查询 `semantic_cache` 类获得的结果以及用户提出的问题。

prompt 有两个部分，**相关上下文**（即从数据库中恢复的信息）和**用户的问题**。

我们只需要将这两部分放在一起以创建 prompt，然后将其发送到模型。

prompt_template = f"Relevant context: {results}\n\n The user's question: {question_def}"
prompt_template

input_ids = tokenizer(prompt_template, return_tensors="pt").to("cuda")

现在剩下的就是将 prompt 发送到模型并等待其响应！

>>> outputs = model.generate(**input_ids, max_new_tokens=256)
>>> print(tokenizer.decode(outputs[0]))

Relevant context: Sydenham chorea (SD) is a neurological disorder of childhood resulting from infection via Group A beta-hemolytic streptococcus (GABHS), the bacterium that causes rheumatic fever. SD is characterized by rapid, irregular, and aimless involuntary movements of the arms and legs, trunk, and facial muscles. It affects girls more often than boys and typically occurs between 5 and 15 years of age. Some children will have a sore throat several weeks before the symptoms begin, but the disorder can also strike up to 6 months after the fever or infection has cleared. Symptoms can appear gradually or all at once, and also may include uncoordinated movements, muscular weakness, stumbling and falling, slurred speech, difficulty concentrating and writing, and emotional instability. The symptoms of SD can vary from a halting gait and slight grimacing to involuntary movements that are frequent and severe enough to be incapacitating. The random, writhing movements of chorea are caused by an auto-immune reaction to the bacterium that interferes with the normal function of a part of the brain (the basal ganglia) that controls motor movements. Due to better sanitary conditions and the use of antibiotics to treat streptococcal infections, rheumatic fever, and consequently SD, are rare in North America and Europe. The disease can still be found in developing nations.

 The user's question: Write in 20 words what is a Sydenham chorea.

Sure, here is a 20-word answer:

Sydenham chorea is a neurological disorder of childhood resulting from infection via Group A beta-hemolytic streptococcus (GABHS).

结论。

在访问 ChromaDB 和直接访问缓存之间，数据检索时间减少了 50%。 但是，在较大的项目中，这种差异会增加，从而使性能提高 90-95%。

我们的 Chroma 中数据很少，并且只有一个缓存类实例。 通常，缓存系统背后的数据要大得多，可能不仅涉及对向量数据库的查询，还可能来自各个地方。

拥有多个缓存类实例是很常见的，通常基于用户类型，因为问题往往在具有共同特征的用户之间重复更多。

总之，我们创建了一个非常简单的 RAG（检索增强生成）系统，并通过在用户问题和获取创建增强 prompt 所需信息之间添加语义缓存层来对其进行了增强。