rag/retriever/rerank/cohere.py


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

import os

import cohere
from loguru import logger as log

from rag.generator.prompt import Prompt
from rag.retriever.rerank.abstract import AbstractReranker


class CohereReranker(metaclass=AbstractReranker):
    def __init__(self) -> None:
        self.client = cohere.Client(os.environ["COHERE_API_KEY"])
        self.top_k = int(os.environ["RERANK_TOP_K"])

    def rank(self, prompt: Prompt) -> Prompt:
        if prompt.documents:
            response = self.client.rerank(
                model="rerank-english-v3.0",
                query=prompt.query,
                documents=[d.text for d in prompt.documents],
                top_n=self.top_k,
            )
            ranking = list(filter(lambda x: x.relevance_score > 0.5, response.results))
            log.debug(
                f"Reranking gave {len(ranking)} relevant documents of {len(prompt.documents)}"
            )
            prompt.documents = [prompt.documents[r.index] for r in ranking]
        return prompt