Skip to content

Cohere

CohereReranking

Bases: BaseReranking

Source code in libs/kotaemon/kotaemon/indices/rankings/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
57
58
59
60
class CohereReranking(BaseReranking):
    model_name: str = "rerank-multilingual-v2.0"
    cohere_api_key: str = config("COHERE_API_KEY", "")
    use_key_from_ktem: bool = False

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

        # try to get COHERE_API_KEY from embeddings
        if not self.cohere_api_key and self.use_key_from_ktem:
            try:
                from ktem.embeddings.manager import (
                    embedding_models_manager as embeddings,
                )

                cohere_model = embeddings.get("cohere")
                ktem_cohere_api_key = cohere_model._kwargs.get(  # type: ignore
                    "cohere_api_key"
                )
                if ktem_cohere_api_key != "your-key":
                    self.cohere_api_key = ktem_cohere_api_key
            except Exception as e:
                print("Cannot get Cohere API key from `ktem`", e)

        if not 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/indices/rankings/cohere.py 
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
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"
        )

    # try to get COHERE_API_KEY from embeddings
    if not self.cohere_api_key and self.use_key_from_ktem:
        try:
            from ktem.embeddings.manager import (
                embedding_models_manager as embeddings,
            )

            cohere_model = embeddings.get("cohere")
            ktem_cohere_api_key = cohere_model._kwargs.get(  # type: ignore
                "cohere_api_key"
            )
            if ktem_cohere_api_key != "your-key":
                self.cohere_api_key = ktem_cohere_api_key
        except Exception as e:
            print("Cannot get Cohere API key from `ktem`", e)

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