122 `TransformRetriever` 类详解 llamaindex.core.retrievers.transform

本文链接：https://blog.csdn.net/xycxycooo/article/details/141672677

`TransformRetriever` 类详解

TransformRetriever 类是 llamaindex.core.retrievers 模块中的一个重要组件，用于在执行检索操作之前对查询进行转换。本文将详细解析该类的实现和使用方法。

类定义与初始化

class TransformRetriever(BaseRetriever):
    """Transform Retriever.

    Takes in an existing retriever and a query transform and runs the query transform
    before running the retriever.
    """

    def __init__(
        self,
        retriever: BaseRetriever,
        query_transform: BaseQueryTransform,
        transform_metadata: Optional[dict] = None,
        callback_manager: Optional[CallbackManager] = None,
        object_map: Optional[dict] = None,
        verbose: bool = False,
    ) -> None:
        self._retriever = retriever
        self._query_transform = query_transform
        self._transform_metadata = transform_metadata
        super().__init__(
            callback_manager=callback_manager, object_map=object_map, verbose=verbose
        )

参数说明

retriever: 一个现有的检索器实例，用于执行实际的检索操作。
query_transform: 一个查询转换器实例，用于在检索之前对查询进行转换。
transform_metadata: 可选参数，包含查询转换所需的元数据。
callback_manager: 可选参数，用于管理和调度回调函数。
object_map: 可选参数，对象映射字典。
verbose: 是否输出详细信息，默认为False。

获取提示模块

def _get_prompt_modules(self) -> PromptMixinType:
    """Get prompt sub-modules."""
    # NOTE: don't include tools for now
    return {"query_transform": self._query_transform}

该方法返回一个包含查询转换器的提示模块字典。

检索方法

def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
    query_bundle = self._query_transform.run(
        query_bundle, metadata=self._transform_metadata
    )
    return self._retriever.retrieve(query_bundle)

方法解析

功能：该方法对输入的查询包进行转换，然后使用现有的检索器执行检索操作。
参数：query_bundle，一个QueryBundle实例，表示要处理的查询。
返回值：一个包含NodeWithScore实例的列表，表示检索结果。

处理流程

查询转换：

query_bundle = self._query_transform.run(
    query_bundle, metadata=self._transform_metadata
)

调用查询转换器的run方法，对查询包进行转换。

执行检索：
```
return self._retriever.retrieve(query_bundle)
```
使用现有的检索器执行转换后的查询，并返回检索结果。

实际应用示例

假设我们有一个现有的检索器和一个查询转换器，需要对查询进行转换后再执行检索：

from some_module import BaseRetriever, BaseQueryTransform, QueryBundle, TransformRetriever

# 假设retriever和query_transform已经初始化
retriever = BaseRetriever(...)
query_transform = BaseQueryTransform(...)

# 初始化TransformRetriever实例
transform_retriever = TransformRetriever(
    retriever=retriever,
    query_transform=query_transform,
    verbose=True,
)

# 定义查询
query_bundle = QueryBundle("example query")

# 执行检索
results = transform_retriever._retrieve(query_bundle)

# 输出生成的查询
for node_with_score in results:
    print(f"Node: {node_with_score.node}, Score: {node_with_score.score}")

总结

通过本文的详细解析，我们深入理解了TransformRetriever类的实现原理和应用方法。该类通过在执行检索操作之前对查询进行转换，提供了灵活的查询处理能力，从而有效地提升了检索系统的性能和准确性。希望本文能为您的编程实践提供有益的参考和帮助。

from typing import List, Optional

from llama_index.core.base.base_retriever import BaseRetriever
from llama_index.core.callbacks.base import CallbackManager
from llama_index.core.indices.query.query_transform.base import BaseQueryTransform
from llama_index.core.prompts.mixin import PromptMixinType
from llama_index.core.schema import NodeWithScore, QueryBundle


class TransformRetriever(BaseRetriever):
    """Transform Retriever.

    Takes in an existing retriever and a query transform and runs the query transform
    before running the retriever.

    """

    def __init__(
        self,
        retriever: BaseRetriever,
        query_transform: BaseQueryTransform,
        transform_metadata: Optional[dict] = None,
        callback_manager: Optional[CallbackManager] = None,
        object_map: Optional[dict] = None,
        verbose: bool = False,
    ) -> None:
        self._retriever = retriever
        self._query_transform = query_transform
        self._transform_metadata = transform_metadata
        super().__init__(
            callback_manager=callback_manager, object_map=object_map, verbose=verbose
        )

    def _get_prompt_modules(self) -> PromptMixinType:
        """Get prompt sub-modules."""
        # NOTE: don't include tools for now
        return {"query_transform": self._query_transform}

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        query_bundle = self._query_transform.run(
            query_bundle, metadata=self._transform_metadata
        )
        return self._retriever.retrieve(query_bundle)