Récupération contextuelle avec Milvus

La recherche contextuelle d' images est une méthode de recherche avancée proposée par Anthropic pour résoudre le problème de l'isolation sémantique des morceaux, qui se pose dans les solutions actuelles de génération améliorée de recherche (RAG). Dans le paradigme pratique actuel de la RAG, les documents sont divisés en plusieurs morceaux, et une base de données vectorielle est utilisée pour rechercher la requête, en récupérant les morceaux les plus pertinents. Un LLM répond ensuite à la requête en utilisant ces morceaux récupérés. Toutefois, ce processus de découpage peut entraîner la perte d'informations contextuelles, ce qui rend difficile la détermination de la pertinence par l'utilisateur.

La recherche contextuelle améliore les systèmes de recherche traditionnels en ajoutant un contexte pertinent à chaque fragment de document avant l'intégration ou l'indexation, ce qui augmente la précision et réduit les erreurs de recherche. Combinée à des techniques telles que la recherche hybride et le reranking, elle améliore les systèmes RAG (Retrieval-Augmented Generation), en particulier pour les grandes bases de connaissances. En outre, elle offre une solution rentable lorsqu'elle est associée à une mise en cache rapide, réduisant de manière significative la latence et les coûts opérationnels, les morceaux contextualisés coûtant environ 1,02 dollar par million de jetons de document. Il s'agit donc d'une approche évolutive et efficace pour le traitement de grandes bases de connaissances. La solution d'Anthropic présente deux aspects intéressants :

• Document Enhancement: La réécriture des requêtes est une technique cruciale dans la recherche d'information moderne, qui utilise souvent des informations auxiliaires pour rendre la requête plus informative. De même, pour obtenir de meilleures performances dans RAG, le prétraitement des documents avec un LLM (par exemple, le nettoyage de la source de données, le complément des informations perdues, le résumé, etc. En d'autres termes, cette étape de prétraitement aide à rapprocher les documents des requêtes en termes de pertinence.
• Low-Cost Processing by Caching Long Context: Une préoccupation commune lors de l'utilisation des LLM pour traiter les documents est le coût. Le KVCache est une solution populaire qui permet de réutiliser les résultats intermédiaires pour le même contexte précédent. Alors que la plupart des fournisseurs de LLM hébergés rendent cette fonctionnalité transparente pour l'utilisateur, Anthropic donne aux utilisateurs le contrôle sur le processus de mise en cache. Lorsqu'un hit du cache se produit, la plupart des calculs peuvent être sauvegardés (ceci est courant lorsque le contexte long reste le même, mais que l'instruction pour chaque requête change). Pour plus de détails, cliquez ici.

Dans ce carnet, nous démontrerons comment effectuer une recherche contextuelle en utilisant Milvus avec un LLM, en combinant la recherche hybride dense-sparse et un reranker pour créer un système de recherche progressivement plus puissant. Les données et la configuration expérimentale sont basées sur la recherche contextuelle.

Préparation

Installer les dépendances

$ pip install "pymilvus[model]"
$ pip install tqdm
$ pip install anthropic

Vous aurez besoin des clés API de Cohere, Voyage et Anthropic pour exécuter le code.

Télécharger les données

La commande suivante permet de télécharger les données d'exemple utilisées dans la démonstration originale d'Anthropic.

$ wget https://raw.githubusercontent.com/anthropics/anthropic-cookbook/refs/heads/main/skills/contextual-embeddings/data/codebase_chunks.json
$ wget https://raw.githubusercontent.com/anthropics/anthropic-cookbook/refs/heads/main/skills/contextual-embeddings/data/evaluation_set.jsonl

Définir le récupérateur

Cette classe est conçue pour être flexible, vous permettant de choisir entre différents modes de récupération en fonction de vos besoins. En spécifiant des options dans la méthode d'initialisation, vous pouvez déterminer si vous voulez utiliser la recherche contextuelle, la recherche hybride (combinant les méthodes de recherche dense et sparse), ou un reranker pour des résultats améliorés.

from pymilvus.model.dense import VoyageEmbeddingFunction
from pymilvus.model.hybrid import BGEM3EmbeddingFunction
from pymilvus.model.reranker import CohereRerankFunction

from typing import List, Dict, Any
from typing import Callable
from pymilvus import (
    MilvusClient,
    DataType,
    AnnSearchRequest,
    RRFRanker,
)
from tqdm import tqdm
import json
import anthropic


class MilvusContextualRetriever:
    def __init__(
        self,
        uri="milvus.db",
        collection_name="contexual_bgem3",
        dense_embedding_function=None,
        use_sparse=False,
        sparse_embedding_function=None,
        use_contextualize_embedding=False,
        anthropic_client=None,
        use_reranker=False,
        rerank_function=None,
    ):
        self.collection_name = collection_name

        # For Milvus-lite, uri is a local path like "./milvus.db"
        # For Milvus standalone service, uri is like "http://localhost:19530"
        # For Zilliz Clond, please set `uri` and `token`, which correspond to the [Public Endpoint and API key](https://docs.zilliz.com/docs/on-zilliz-cloud-console#cluster-details) in Zilliz Cloud.
        self.client = MilvusClient(uri)

        self.embedding_function = dense_embedding_function

        self.use_sparse = use_sparse
        self.sparse_embedding_function = None

        self.use_contextualize_embedding = use_contextualize_embedding
        self.anthropic_client = anthropic_client

        self.use_reranker = use_reranker
        self.rerank_function = rerank_function

        if use_sparse is True and sparse_embedding_function:
            self.sparse_embedding_function = sparse_embedding_function
        elif sparse_embedding_function is False:
            raise ValueError(
                "Sparse embedding function cannot be None if use_sparse is False"
            )
        else:
            pass

    def build_collection(self):
        schema = self.client.create_schema(
            auto_id=True,
            enable_dynamic_field=True,
        )
        schema.add_field(field_name="pk", datatype=DataType.INT64, is_primary=True)
        schema.add_field(
            field_name="dense_vector",
            datatype=DataType.FLOAT_VECTOR,
            dim=self.embedding_function.dim,
        )
        if self.use_sparse is True:
            schema.add_field(
                field_name="sparse_vector", datatype=DataType.SPARSE_FLOAT_VECTOR
            )

        index_params = self.client.prepare_index_params()
        index_params.add_index(
            field_name="dense_vector", index_type="FLAT", metric_type="IP"
        )
        if self.use_sparse is True:
            index_params.add_index(
                field_name="sparse_vector",
                index_type="SPARSE_INVERTED_INDEX",
                metric_type="IP",
            )

        self.client.create_collection(
            collection_name=self.collection_name,
            schema=schema,
            index_params=index_params,
            enable_dynamic_field=True,
        )

    def insert_data(self, chunk, metadata):
        dense_vec = self.embedding_function([chunk])[0]
        if self.use_sparse is True:
            sparse_result = self.sparse_embedding_function.encode_documents([chunk])
            if type(sparse_result) == dict:
                sparse_vec = sparse_result["sparse"][[0]]
            else:
                sparse_vec = sparse_result[[0]]
            self.client.insert(
                collection_name=self.collection_name,
                data={
                    "dense_vector": dense_vec,
                    "sparse_vector": sparse_vec,
                    **metadata,
                },
            )
        else:
            self.client.insert(
                collection_name=self.collection_name,
                data={"dense_vector": dense_vec, **metadata},
            )

    def insert_contextualized_data(self, doc, chunk, metadata):
        contextualized_text, usage = self.situate_context(doc, chunk)
        metadata["context"] = contextualized_text
        text_to_embed = f"{chunk}\n\n{contextualized_text}"
        dense_vec = self.embedding_function([text_to_embed])[0]
        if self.use_sparse is True:
            sparse_vec = self.sparse_embedding_function.encode_documents(
                [text_to_embed]
            )["sparse"][[0]]
            self.client.insert(
                collection_name=self.collection_name,
                data={
                    "dense_vector": dense_vec,
                    "sparse_vector": sparse_vec,
                    **metadata,
                },
            )
        else:
            self.client.insert(
                collection_name=self.collection_name,
                data={"dense_vector": dense_vec, **metadata},
            )

    def situate_context(self, doc: str, chunk: str):
        DOCUMENT_CONTEXT_PROMPT = """
        <document>
        {doc_content}
        </document>
        """

        CHUNK_CONTEXT_PROMPT = """
        Here is the chunk we want to situate within the whole document
        <chunk>
        {chunk_content}
        </chunk>

        Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk.
        Answer only with the succinct context and nothing else.
        """

        response = self.anthropic_client.beta.prompt_caching.messages.create(
            model="claude-3-haiku-20240307",
            max_tokens=1000,
            temperature=0.0,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": DOCUMENT_CONTEXT_PROMPT.format(doc_content=doc),
                            "cache_control": {
                                "type": "ephemeral"
                            },  # we will make use of prompt caching for the full documents
                        },
                        {
                            "type": "text",
                            "text": CHUNK_CONTEXT_PROMPT.format(chunk_content=chunk),
                        },
                    ],
                },
            ],
            extra_headers={"anthropic-beta": "prompt-caching-2024-07-31"},
        )
        return response.content[0].text, response.usage

    def search(self, query: str, k: int = 20) -> List[Dict[str, Any]]:
        dense_vec = self.embedding_function([query])[0]
        if self.use_sparse is True:
            sparse_vec = self.sparse_embedding_function.encode_queries([query])[
                "sparse"
            ][[0]]

        req_list = []
        if self.use_reranker:
            k = k * 10
        if self.use_sparse is True:
            req_list = []
            dense_search_param = {
                "data": [dense_vec],
                "anns_field": "dense_vector",
                "param": {"metric_type": "IP"},
                "limit": k * 2,
            }
            dense_req = AnnSearchRequest(**dense_search_param)
            req_list.append(dense_req)

            sparse_search_param = {
                "data": [sparse_vec],
                "anns_field": "sparse_vector",
                "param": {"metric_type": "IP"},
                "limit": k * 2,
            }
            sparse_req = AnnSearchRequest(**sparse_search_param)

            req_list.append(sparse_req)

            docs = self.client.hybrid_search(
                self.collection_name,
                req_list,
                RRFRanker(),
                k,
                output_fields=[
                    "content",
                    "original_uuid",
                    "doc_id",
                    "chunk_id",
                    "original_index",
                    "context",
                ],
            )
        else:
            docs = self.client.search(
                self.collection_name,
                data=[dense_vec],
                anns_field="dense_vector",
                limit=k,
                output_fields=[
                    "content",
                    "original_uuid",
                    "doc_id",
                    "chunk_id",
                    "original_index",
                    "context",
                ],
            )
        if self.use_reranker and self.use_contextualize_embedding:
            reranked_texts = []
            reranked_docs = []
            for i in range(k):
                if self.use_contextualize_embedding:
                    reranked_texts.append(
                        f"{docs[0][i]['entity']['content']}\n\n{docs[0][i]['entity']['context']}"
                    )
                else:
                    reranked_texts.append(f"{docs[0][i]['entity']['content']}")
            results = self.rerank_function(query, reranked_texts)
            for result in results:
                reranked_docs.append(docs[0][result.index])
            docs[0] = reranked_docs
        return docs


def evaluate_retrieval(
    queries: List[Dict[str, Any]], retrieval_function: Callable, db, k: int = 20
) -> Dict[str, float]:
    total_score = 0
    total_queries = len(queries)
    for query_item in tqdm(queries, desc="Evaluating retrieval"):
        query = query_item["query"]
        golden_chunk_uuids = query_item["golden_chunk_uuids"]

        # Find all golden chunk contents
        golden_contents = []
        for doc_uuid, chunk_index in golden_chunk_uuids:
            golden_doc = next(
                (
                    doc
                    for doc in query_item["golden_documents"]
                    if doc["uuid"] == doc_uuid
                ),
                None,
            )
            if not golden_doc:
                print(f"Warning: Golden document not found for UUID {doc_uuid}")
                continue

            golden_chunk = next(
                (
                    chunk
                    for chunk in golden_doc["chunks"]
                    if chunk["index"] == chunk_index
                ),
                None,
            )
            if not golden_chunk:
                print(
                    f"Warning: Golden chunk not found for index {chunk_index} in document {doc_uuid}"
                )
                continue

            golden_contents.append(golden_chunk["content"].strip())

        if not golden_contents:
            print(f"Warning: No golden contents found for query: {query}")
            continue

        retrieved_docs = retrieval_function(query, db, k=k)

        # Count how many golden chunks are in the top k retrieved documents
        chunks_found = 0
        for golden_content in golden_contents:
            for doc in retrieved_docs[0][:k]:
                retrieved_content = doc["entity"]["content"].strip()
                if retrieved_content == golden_content:
                    chunks_found += 1
                    break

        query_score = chunks_found / len(golden_contents)
        total_score += query_score

    average_score = total_score / total_queries
    pass_at_n = average_score * 100
    return {
        "pass_at_n": pass_at_n,
        "average_score": average_score,
        "total_queries": total_queries,
    }


def retrieve_base(query: str, db, k: int = 20) -> List[Dict[str, Any]]:
    return db.search(query, k=k)


def load_jsonl(file_path: str) -> List[Dict[str, Any]]:
    """Load JSONL file and return a list of dictionaries."""
    with open(file_path, "r") as file:
        return [json.loads(line) for line in file]


def evaluate_db(db, original_jsonl_path: str, k):
    # Load the original JSONL data for queries and ground truth
    original_data = load_jsonl(original_jsonl_path)

    # Evaluate retrieval
    results = evaluate_retrieval(original_data, retrieve_base, db, k)
    print(f"Pass@{k}: {results['pass_at_n']:.2f}%")
    print(f"Total Score: {results['average_score']}")
    print(f"Total queries: {results['total_queries']}")

Vous devez maintenant initialiser ces modèles pour les expériences suivantes. Vous pouvez facilement passer à d'autres modèles en utilisant la bibliothèque de modèles PyMilvus.

dense_ef = VoyageEmbeddingFunction(api_key="your-voyage-api-key", model_name="voyage-2")
sparse_ef = BGEM3EmbeddingFunction()
cohere_rf = CohereRerankFunction(api_key="your-cohere-api-key")

Fetching 30 files:   0%|          | 0/30 [00:00<?, ?it/s]

path = "codebase_chunks.json"
with open(path, "r") as f:
    dataset = json.load(f)

Expérience I : Récupération standard

La recherche standard utilise uniquement des encastrements denses pour retrouver des documents liés. Dans cette expérience, nous utiliserons Pass@5 pour reproduire les résultats de la base de données originale.

standard_retriever = MilvusContextualRetriever(
    uri="standard.db", collection_name="standard", dense_embedding_function=dense_ef
)

standard_retriever.build_collection()
for doc in dataset:
    doc_content = doc["content"]
    for chunk in doc["chunks"]:
        metadata = {
            "doc_id": doc["doc_id"],
            "original_uuid": doc["original_uuid"],
            "chunk_id": chunk["chunk_id"],
            "original_index": chunk["original_index"],
            "content": chunk["content"],
        }
        chunk_content = chunk["content"]
        standard_retriever.insert_data(chunk_content, metadata)

evaluate_db(standard_retriever, "evaluation_set.jsonl", 5)

Evaluating retrieval: 100%|██████████| 248/248 [01:29<00:00,  2.77it/s]

Pass@5: 80.92%
Total Score: 0.8091877880184332
Total queries: 248

Expérience II : Récupération hybride

Maintenant que nous avons obtenu des résultats prometteurs avec l'intégration de Voyage, nous allons passer à l'extraction hybride en utilisant le modèle BGE-M3 qui génère des intégrations éparses puissantes. Les résultats de l'extraction dense et de l'extraction éparse seront combinés à l'aide de la méthode Reciprocal Rank Fusion (RRF) pour produire un résultat hybride.

hybrid_retriever = MilvusContextualRetriever(
    uri="hybrid.db",
    collection_name="hybrid",
    dense_embedding_function=dense_ef,
    use_sparse=True,
    sparse_embedding_function=sparse_ef,
)

hybrid_retriever.build_collection()
for doc in dataset:
    doc_content = doc["content"]
    for chunk in doc["chunks"]:
        metadata = {
            "doc_id": doc["doc_id"],
            "original_uuid": doc["original_uuid"],
            "chunk_id": chunk["chunk_id"],
            "original_index": chunk["original_index"],
            "content": chunk["content"],
        }
        chunk_content = chunk["content"]
        hybrid_retriever.insert_data(chunk_content, metadata)

evaluate_db(hybrid_retriever, "evaluation_set.jsonl", 5)

Evaluating retrieval: 100%|██████████| 248/248 [02:09<00:00,  1.92it/s]

Pass@5: 84.69%
Total Score: 0.8469182027649771
Total queries: 248

Expérience III : extraction contextuelle

L'extraction hybride montre une amélioration, mais les résultats peuvent être encore améliorés en appliquant une méthode d'extraction contextuelle. Pour ce faire, nous utiliserons le modèle linguistique d'Anthropic pour ajouter le contexte du document entier pour chaque morceau.

anthropic_client = anthropic.Anthropic(
    api_key="your-anthropic-api-key",
)

contextual_retriever = MilvusContextualRetriever(
    uri="contextual.db",
    collection_name="contextual",
    dense_embedding_function=dense_ef,
    use_sparse=True,
    sparse_embedding_function=sparse_ef,
    use_contextualize_embedding=True,
    anthropic_client=anthropic_client,
)

contextual_retriever.build_collection()
for doc in dataset:
    doc_content = doc["content"]
    for chunk in doc["chunks"]:
        metadata = {
            "doc_id": doc["doc_id"],
            "original_uuid": doc["original_uuid"],
            "chunk_id": chunk["chunk_id"],
            "original_index": chunk["original_index"],
            "content": chunk["content"],
        }
        chunk_content = chunk["content"]
        contextual_retriever.insert_contextualized_data(
            doc_content, chunk_content, metadata
        )

evaluate_db(contextual_retriever, "evaluation_set.jsonl", 5)

 Evaluating retrieval: 100%|██████████| 248/248 [01:55<00:00,  2.15it/s]
Pass@5: 87.14%
Total Score: 0.8713517665130568
Total queries: 248 

Expérience IV : Extraction contextuelle avec Reranker

Les résultats peuvent être encore améliorés par l'ajout d'un reranker Cohere. Sans initialiser un nouveau récupérateur avec le reranker séparément, nous pouvons simplement configurer le récupérateur existant pour qu'il utilise le reranker afin d'améliorer les performances.

contextual_retriever.use_reranker = True
contextual_retriever.rerank_function = cohere_rf

evaluate_db(contextual_retriever, "evaluation_set.jsonl", 5)

Evaluating retrieval: 100%|██████████| 248/248 [02:02<00:00,  2.00it/s]
Pass@5: 90.91%
Total Score: 0.9090821812596005
Total queries: 248

Nous avons démontré plusieurs méthodes pour améliorer les performances de récupération. Avec une conception plus ad-hoc adaptée au scénario, la recherche contextuelle montre un potentiel significatif de prétraitement des documents à faible coût, conduisant à un meilleur système RAG.