Skip to content

Rerankings

BaseReranking

Bases: BaseComponent

Source code in libs/kotaemon/kotaemon/rerankings/base.py 
 8
 9
10
11
12
13
class BaseReranking(BaseComponent):
    @abstractmethod
    def run(self, documents: list[Document], query: str) -> list[Document]:
        """Main method to transform list of documents
        (re-ranking, filtering, etc)"""
        ...

run abstractmethod 

1
run(documents, query)

Main method to transform list of documents (re-ranking, filtering, etc)

Source code in libs/kotaemon/kotaemon/rerankings/base.py 
 9
10
11
12
13
@abstractmethod
def run(self, documents: list[Document], query: str) -> list[Document]:
    """Main method to transform list of documents
    (re-ranking, filtering, etc)"""
    ...

CohereReranking

Bases: BaseReranking

Cohere Reranking model

Source code in libs/kotaemon/kotaemon/rerankings/cohere.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class CohereReranking(BaseReranking):
    """Cohere Reranking model"""

    model_name: str = Param(
        "rerank-multilingual-v2.0",
        help=(
            "ID of the model to use. You can go to [Supported Models]"
            "(https://docs.cohere.com/docs/rerank-2) to see the supported models"
        ),
        required=True,
    )
    cohere_api_key: str = Param(
        config("COHERE_API_KEY", ""),
        help="Cohere API key",
        required=True,
    )

    def run(self, documents: list[Document], query: str) -> list[Document]:
        """Use Cohere Reranker model to re-order documents
        with their relevance score"""
        try:
            import cohere
        except ImportError:
            raise ImportError(
                "Please install Cohere " "`pip install cohere` to use Cohere Reranking"
            )

        if not self.cohere_api_key or "COHERE_API_KEY" in self.cohere_api_key:
            print("Cohere API key not found. Skipping rerankings.")
            return documents

        cohere_client = cohere.Client(self.cohere_api_key)
        compressed_docs: list[Document] = []

        if not documents:  # to avoid empty api call
            return compressed_docs

        _docs = [d.content for d in documents]
        response = cohere_client.rerank(
            model=self.model_name, query=query, documents=_docs
        )
        for r in response.results:
            doc = documents[r.index]
            doc.metadata["reranking_score"] = r.relevance_score
            compressed_docs.append(doc)

        return compressed_docs

run 

1
run(documents, query)

Use Cohere Reranker model to re-order documents with their relevance score

Source code in libs/kotaemon/kotaemon/rerankings/cohere.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def run(self, documents: list[Document], query: str) -> list[Document]:
    """Use Cohere Reranker model to re-order documents
    with their relevance score"""
    try:
        import cohere
    except ImportError:
        raise ImportError(
            "Please install Cohere " "`pip install cohere` to use Cohere Reranking"
        )

    if not self.cohere_api_key or "COHERE_API_KEY" in self.cohere_api_key:
        print("Cohere API key not found. Skipping rerankings.")
        return documents

    cohere_client = cohere.Client(self.cohere_api_key)
    compressed_docs: list[Document] = []

    if not documents:  # to avoid empty api call
        return compressed_docs

    _docs = [d.content for d in documents]
    response = cohere_client.rerank(
        model=self.model_name, query=query, documents=_docs
    )
    for r in response.results:
        doc = documents[r.index]
        doc.metadata["reranking_score"] = r.relevance_score
        compressed_docs.append(doc)

    return compressed_docs

TeiFastReranking

Bases: BaseReranking

Text Embeddings Inference (TEI) Reranking model (https://huggingface.co/docs/text-embeddings-inference/en/index)

Source code in libs/kotaemon/kotaemon/rerankings/tei_fast_rerank.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class TeiFastReranking(BaseReranking):
    """Text Embeddings Inference (TEI) Reranking model
    (https://huggingface.co/docs/text-embeddings-inference/en/index)
    """

    endpoint_url: str = Param(
        None, help="TEI Reranking service api base URL", required=True
    )
    model_name: Optional[str] = Param(
        None,
        help=(
            "ID of the model to use. You can go to [Supported Models]"
            "(https://github.com/huggingface"
            "/text-embeddings-inference?tab=readme-ov-file"
            "#supported-models) to see the supported models"
        ),
    )
    is_truncated: Optional[bool] = Param(True, help="Whether to truncate the inputs")

    def client(self, query, texts):
        response = session.post(
            url=self.endpoint_url,
            json={
                "query": query,
                "texts": texts,
                "is_truncated": self.is_truncated,  # default is True
            },
        ).json()
        return response

    def run(self, documents: list[Document], query: str) -> list[Document]:
        """Use the deployed TEI rerankings service to re-order documents
        with their relevance score"""
        if not self.endpoint_url:
            print("TEI API reranking URL not found. Skipping rerankings.")
            return documents

        compressed_docs: list[Document] = []

        if not documents:  # to avoid empty api call
            return compressed_docs

        if isinstance(documents[0], str):
            documents = self.prepare_input(documents)

        batch_size = 6
        num_batch = max(len(documents) // batch_size, 1)
        for i in range(num_batch):
            if i == num_batch - 1:
                mini_batch = documents[batch_size * i :]
            else:
                mini_batch = documents[batch_size * i : batch_size * (i + 1)]

            _docs = [d.content for d in mini_batch]
            rerank_resp = self.client(query, _docs)
            for r in rerank_resp:
                doc = mini_batch[r["index"]]
                doc.metadata["reranking_score"] = r["score"]
                compressed_docs.append(doc)

        compressed_docs = sorted(
            compressed_docs, key=lambda x: x.metadata["reranking_score"], reverse=True
        )
        return compressed_docs

run 

1
run(documents, query)

Use the deployed TEI rerankings service to re-order documents with their relevance score

Source code in libs/kotaemon/kotaemon/rerankings/tei_fast_rerank.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def run(self, documents: list[Document], query: str) -> list[Document]:
    """Use the deployed TEI rerankings service to re-order documents
    with their relevance score"""
    if not self.endpoint_url:
        print("TEI API reranking URL not found. Skipping rerankings.")
        return documents

    compressed_docs: list[Document] = []

    if not documents:  # to avoid empty api call
        return compressed_docs

    if isinstance(documents[0], str):
        documents = self.prepare_input(documents)

    batch_size = 6
    num_batch = max(len(documents) // batch_size, 1)
    for i in range(num_batch):
        if i == num_batch - 1:
            mini_batch = documents[batch_size * i :]
        else:
            mini_batch = documents[batch_size * i : batch_size * (i + 1)]

        _docs = [d.content for d in mini_batch]
        rerank_resp = self.client(query, _docs)
        for r in rerank_resp:
            doc = mini_batch[r["index"]]
            doc.metadata["reranking_score"] = r["score"]
            compressed_docs.append(doc)

    compressed_docs = sorted(
        compressed_docs, key=lambda x: x.metadata["reranking_score"], reverse=True
    )
    return compressed_docs