summaryrefslogtreecommitdiff
path: root/rag/retriever
diff options
context:
space:
mode:
Diffstat (limited to 'rag/retriever')
-rw-r--r--rag/retriever/rerank.py77
-rw-r--r--rag/retriever/rerank/__init__.py15
-rw-r--r--rag/retriever/rerank/abstract.py17
-rw-r--r--rag/retriever/rerank/cohere.py28
-rw-r--r--rag/retriever/rerank/local.py30
5 files changed, 90 insertions, 77 deletions
diff --git a/rag/retriever/rerank.py b/rag/retriever/rerank.py
deleted file mode 100644
index 08a9a27..0000000
--- a/rag/retriever/rerank.py
+++ /dev/null
@@ -1,77 +0,0 @@
-import os
-from abc import abstractmethod
-from typing import Type
-
-import cohere
-from loguru import logger as log
-from sentence_transformers import CrossEncoder
-
-from rag.generator.prompt import Prompt
-
-
-class AbstractReranker(type):
- _instances = {}
-
- def __call__(cls, *args, **kwargs):
- if cls not in cls._instances:
- instance = super().__call__(*args, **kwargs)
- cls._instances[cls] = instance
- return cls._instances[cls]
-
- @abstractmethod
- def rank(self, prompt: Prompt) -> Prompt:
- return prompt
-
-
-class Reranker(metaclass=AbstractReranker):
- def __init__(self) -> None:
- self.model = CrossEncoder(os.environ["RERANK_MODEL"])
- self.top_k = int(os.environ["RERANK_TOP_K"])
-
- def rank(self, prompt: Prompt) -> Prompt:
- if prompt.documents:
- results = self.model.rank(
- query=prompt.query,
- documents=[d.text for d in prompt.documents],
- return_documents=False,
- top_k=self.top_k,
- )
- ranking = list(filter(lambda x: x.get("score", 0.0) > 0.5, results))
- log.debug(
- f"Reranking gave {len(ranking)} relevant documents of {len(prompt.documents)}"
- )
- prompt.documents = [
- prompt.documents[r.get("corpus_id", 0)] for r in ranking
- ]
- return prompt
-
-
-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
-
-
-def get_reranker(model: str) -> Type[AbstractReranker]:
- match model:
- case "local":
- return Reranker()
- case "cohere":
- return CohereReranker()
- case _:
- exit(1)
diff --git a/rag/retriever/rerank/__init__.py b/rag/retriever/rerank/__init__.py
new file mode 100644
index 0000000..16b2fac
--- /dev/null
+++ b/rag/retriever/rerank/__init__.py
@@ -0,0 +1,15 @@
+from typing import Type
+
+from rag.retriever.rerank.abstract import AbstractReranker
+from rag.retriever.rerank.cohere import CohereReranker
+from rag.retriever.rerank.local import Reranker
+
+
+def get_reranker(model: str) -> Type[AbstractReranker]:
+ match model:
+ case "local":
+ return Reranker()
+ case "cohere":
+ return CohereReranker()
+ case _:
+ exit(1)
diff --git a/rag/retriever/rerank/abstract.py b/rag/retriever/rerank/abstract.py
new file mode 100644
index 0000000..b96b70a
--- /dev/null
+++ b/rag/retriever/rerank/abstract.py
@@ -0,0 +1,17 @@
+from abc import abstractmethod
+
+from rag.generator.prompt import Prompt
+
+
+class AbstractReranker(type):
+ _instances = {}
+
+ def __call__(cls, *args, **kwargs):
+ if cls not in cls._instances:
+ instance = super().__call__(*args, **kwargs)
+ cls._instances[cls] = instance
+ return cls._instances[cls]
+
+ @abstractmethod
+ def rank(self, prompt: Prompt) -> Prompt:
+ return prompt
diff --git a/rag/retriever/rerank/cohere.py b/rag/retriever/rerank/cohere.py
new file mode 100644
index 0000000..dac9ab5
--- /dev/null
+++ b/rag/retriever/rerank/cohere.py
@@ -0,0 +1,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
diff --git a/rag/retriever/rerank/local.py b/rag/retriever/rerank/local.py
new file mode 100644
index 0000000..758c5dc
--- /dev/null
+++ b/rag/retriever/rerank/local.py
@@ -0,0 +1,30 @@
+import os
+
+from loguru import logger as log
+from sentence_transformers import CrossEncoder
+
+from rag.generator.prompt import Prompt
+from rag.retriever.rerank.abstract import AbstractReranker
+
+
+class Reranker(metaclass=AbstractReranker):
+ def __init__(self) -> None:
+ self.model = CrossEncoder(os.environ["RERANK_MODEL"])
+ self.top_k = int(os.environ["RERANK_TOP_K"])
+
+ def rank(self, prompt: Prompt) -> Prompt:
+ if prompt.documents:
+ results = self.model.rank(
+ query=prompt.query,
+ documents=[d.text for d in prompt.documents],
+ return_documents=False,
+ top_k=self.top_k,
+ )
+ ranking = list(filter(lambda x: x.get("score", 0.0) > 0.5, results))
+ log.debug(
+ f"Reranking gave {len(ranking)} relevant documents of {len(prompt.documents)}"
+ )
+ prompt.documents = [
+ prompt.documents[r.get("corpus_id", 0)] for r in ranking
+ ]
+ return prompt