from llama_index.postprocessor.cohere_rerank import CohereRerank
.....#使用默认的嵌入模型构造向量索引
docs = SimpleDirectoryReader(input_files=["../../data/yiyan.txt"]).load_data()
nodes = SentenceSplitter(chunk_size=100,chunk_overlap=0).get_nodes_from_documents(docs)
vector_index = VectorStoreIndex(nodes)

#检索器retriever =vector_index.as_retriever(similarity_top_k=5)
#不借助cohere rerank，直接检索出结果
nodes = retriever.retrieve("百度文心一言的逻辑推理能力怎么样？")
print('================before rerank================')
print_nodes(nodes)

#借助cohere rerank模型重排结果
cohere_rerank = CohereRerank(model='rerank-multilingual-v3.0',api_key='***', top_n=2)
rerank_nodes = cohere_rerank.postprocess_nodes(nodes,query_str='百度文心一言的逻辑推理能力怎么样？')
print('================after rerank================')
print_nodes(rerank_nodes)

import cohere
co = cohere.Client(api_key="<YOUR API KEY>")
query = "你的问题..."
docs = ["相关文档1...","相关文档2...","相关文档3...",...]
results = co.rerank(model="rerank-multilingual-v3.0", query=query, documents=docs, top_n=5, return_documents=True)

model=BAAI/bge-reranker-large
text-embeddings-router --model-id $model --port 8080

import requests
from typing import List, Optional
from llama_index.core.bridge.pydantic import Field, PrivateAttr
from llama_index.core.postprocessor.types import BaseNodePostprocessor
from llama_index.core.schema import NodeWithScore, QueryBundleclass BgeRerank(BaseNodePostprocessor):
    url: str = Field(description="Rerank server url.")
    top_n: int = Field(description="Top N nodes to return.")

    def __init__(self,top_n: int,url: str):
        super().__init__(url=url, top_n=top_n)

   #调用TEI的rerank模型服务，实现rerank方法
    def rerank(self, query, texts):
        url = f"{self.url}/rerank"
        request_body = {"query": query, "texts": texts, "truncate": False}
        response = requests.post(url, json=request_body)
        if response.status_code != 200:
            raise RuntimeError(f"Failed to rerank, detail: {response}")
        return response.json()

    @classmethod
    def class_name(cls) -> str:
        return "BgeRerank"

   #实现LlamaIndex要求的postprocessor的接口
    def _postprocess_nodes(
        self,
        nodes: List[NodeWithScore],
        query_bundle: Optional[QueryBundle] = None,
    ) -> List[NodeWithScore]:
        if query_bundle is None:
            raise ValueError("Missing query bundle in extra info.")
        if len(nodes) == 0:
            return []

       #调用rerank
        texts = [node.text for node in nodes]
        results = self.rerank(
            query=query_bundle.query_str,
            texts=texts,
        )

       #组装返回nodes
        new_nodes = []
        for result in results[0 : self.top_n]:
            new_node_with_score = NodeWithScore(
                node=nodes[int(result["index"])].node,
                score=result["score"],
            )
            new_nodes.append(new_node_with_score)
        return new_nodes

......
#创建自定义的检索后处理器
customRerank = BgeRerank(url="http://localhost:8080",top_n=2)

#测试处理节点
rerank_nodes = customRerank.postprocess_nodes(nodes,query_str='百度文心一言的逻辑推理能力怎么样?')
......