用 FAISS 语义缓存改进 RAG 系统-未完纪

本笔记本探索一个典型的 RAG 方案：使用开源模型和 Chroma DB 向量数据库，并加入语义缓存。缓存保存用户的查询，判断应从向量数据库还是缓存取回信息，再生成附带上下文的提示词。

语义缓存的目标是识别相同或相近的请求。找到匹配请求时，直接取回对应的缓存信息，减少对原始数据源的访问。比较考虑的是请求的含义，因此文字不必完全一样：不同表述、拼写错误或句子结构差异，仍可能是在询问同一件事。例如，“法国的首都是什么？”“告诉我法国首都的名字”和语序有误的类似问题，都表达了相同意图。

用户可能要求更简短的回答，模型生成的文字因此会不同，但向量数据库需要提供的信息相同。所以，本例把缓存放在用户与向量数据库之间，而不是放在用户与大语言模型之间。

架构图的教学信息是：用户问题先进入语义缓存；命中时复用已检索的文档，未命中时查询 Chroma 并写入缓存；随后把文档与本次问题组合成提示词交给模型。原图位于 Cookbook 图像仓库，图像仓库未注明许可。下图根据流程独立绘制，未复制原图，也不是运行截图。

检索缓存命中与未命中两条路径

许多 RAG 教程针对单个用户和测试环境：在笔记本里与本地向量数据库交互，再调用 API 或本地模型。这种架构面对几十到数千个重复请求时，很快会显得不足。一种改进方法是设置一个或多个语义缓存：新请求到来时，先检查是否有相似的历史请求，有就复用信息。

RAG 中有两个耗时环节：

检索构建增强提示词所需的信息。
调用大语言模型生成回答。

两个位置都可以放缓存，也可以各放一个。如果缓存最终回答，可能减弱用户对回答形式的控制。例如，“用 10 个词解释法国大革命”和“用 100 个词解释法国大革命”可能被视作同一查询，缓存回答就会让用户觉得字数要求没有得到遵守。但两者需要相同的背景信息，因此本例选择缓存检索结果。这是设计取舍：缓存位置取决于请求和回答类型，直接缓存模型输出可能节省更多时间，也可能损失本次指令对回答的影响。

安装和导入依赖

首先安装所需 Python 包。Sentence Transformers 把句子转换为固定长度的向量，也就是嵌入；xformers 提供与 Transformer 模型相关的工具，本例用它配合模型和嵌入；ChromaDB 是保存嵌入的向量数据库；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

加载数据集

可用空间和内存有限，因此使用 MAX_ROWS 限制数据集行数。原例使用 Gemma，访问该模型需要登录 Hugging Face；代码从交互输入读取令牌。

#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 需要唯一标识符。先转成 Pandas 数据，再用索引添加 id 列，查看前 10 行。

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)

编校：数据实际来自医学问答集；原代码注释里的 “News” 是遗留命名，按原样保留。这里只演示检索与缓存，不把样例内容当作医疗建议。

配置向量数据库

本例选择 ChromaDB 保存信息。先导入，再指定持久化目录。

import chromadb

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

写入和查询 ChromaDB

ChromaDB 把数据保存在集合中。原例希望先删除已有同名集合，再创建新集合：

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 需要三类数据：文档内容来自 Answer 列；元数据保存 qtype 主题；标识符由 MAX_ROWS 的整数范围生成。

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)],
)

写入后，查询根据文档与检索词的语义相似性返回结果，不只是匹配原词或完整短语。元数据不直接参与这里的初始语义搜索，但可以用于过滤或进一步缩小结果。定义查询函数：

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

创建语义缓存

本例使用 FAISS 在内存中保存嵌入，与 Chroma 的向量检索用途相近。类 semantic_cache 使用自己的编码器，并提供查询接口。它先查询历史请求的缓存；距离满足阈值时返回缓存文档，否则去 Chroma 检索。缓存数据另外写入 JSON 文件。

!pip install -q faiss-cpu==1.8.0

import faiss
from sentence_transformers import SentenceTransformer
import time
import json

init_cache() 初始化缓存。本例采用 IndexFlatL2：它未必适合所有规模，但适合这个小型数据集和高维向量。不同索引有不同取舍：

FlatL2、FlatIP：适合较小数据集，速度与内存取决于数据规模。
LSH：原文建议用于较小数据集以及不超过 128 维的向量。
HNSW：查询快，但可能需要较多内存。
IVF：原文将它作为较大数据集的选项。

这些是原例的选择说明，具体索引应结合数据与实测决定。进一步说明见 FAISS 索引选择指南。

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

编校：原文一处写成 “FlatLS”，实际代码使用 IndexFlatL2。该索引返回的是平方 L2 距离；距离越小表示越相近，不能把阈值方向写成“相似度越高越命中”。

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)

这些函数由下面的类使用。虽然 ask 较长，其流程很直接：找到最接近的新问题的缓存项，检查距离是否在阈值内；命中则返回对应文档，否则调用 query_database 从 Chroma 获取文档。

作者选择欧氏距离而非余弦相似度，因为此处 FAISS 索引采用该度量。也可以实现余弦检索，但需要调整向量归一化或索引设置。

原例还加入 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}")

编校：原类的 FIFO 路径引用了未定义的 self.max_size，且删除列表数据时没有同步删除 FAISS 向量。读取 JSON 后也未重建索引。保留代码以对应原文，但不能宣称持久化重载、淘汰或生产并发可用。下面只按原文展示全新单会话、默认不启用淘汰的调用过程。

原文中的缓存调用记录

下面的输出来自原作者的笔记本。初始化缓存：

>>> # 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，类随后把它加入缓存。换一个差异较大的问题，预期仍需要访问 Chroma，因为已缓存的问题无法通过距离阈值：

>>> 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

再换成与上一问含义很接近的表述，原例预期无需访问 Chroma，直接命中缓存：

>>> 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

原例记录的距离较小，说明两个问题的向量接近。接着增加“用 20 个词”的形式要求：

>>> 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

原例的距离有所增加，但仍在阈值内，于是复用文档。这也说明为什么缓存的是背景资料：本次字数要求仍可以交给后续模型处理。上面的医疗段落、距离与耗时均为原作者的输出记录。

加载模型

现在使用 Transformers。导入的 AutoTokenizer 用于为预训练模型准备文本输入，AutoModelForCausalLM 用于加载因果语言模型；本例选用 Gemma-2b-it。也可以测试其他文本生成模型，模型各自的许可和访问条件仍需确认。

!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)

编校：虽然设备变量可能取 CPU，后续模型加载和输入仍硬编码为 CUDA；注释中的 Apple Silicon 提示也没有被后续代码实际采用。不能声称这份原样代码覆盖 CPU 或 MPS 环境。

构造增强提示词

提示词由两个部分组成：缓存或数据库返回的相关上下文，以及用户本次提出的问题。把它们组合后提交模型：

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")

然后调用模型生成并解码：

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

原文输出：

<bos>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).<eos>

原例显示模型生成了一段带有用户问题形式要求的回答。这仅是原文记录，不能证明所有输出都会遵守字数或事实要求。