AutoRAG 0.0.0__py3-none-any.whl

autorag/nodes/passagefilter/similarity_percentile_cutoff.py

@@ -0,0 +1,134 @@
+from pathlib import Path
+from typing import List, Tuple, Union
+
+import numpy as np
+import pandas as pd
+
+from autorag.embedding.base import EmbeddingModel
+from autorag.evaluation.metric.util import calculate_cosine_similarity
+from autorag.nodes.passagefilter.base import BasePassageFilter
+from autorag.nodes.passagefilter.similarity_threshold_cutoff import (
+	embedding_query_content,
+)
+from autorag.utils import result_to_dataframe
+from autorag.utils.util import empty_cuda_cache, pop_params
+
+
+class SimilarityPercentileCutoff(BasePassageFilter):
+	def __init__(self, project_dir: Union[str, Path], *args, **kwargs):
+		"""
+		Initialize the SimilarityPercentileCutoff module
+
+		:param project_dir: The project directory to use for initializing the module
+		:param embedding_model: The embedding model string to use for calculating similarity
+		        Default is "openai" which is OpenAI text-embedding-ada-002 embedding model.
+		"""
+		super().__init__(project_dir, *args, **kwargs)
+		embedding_model = kwargs.pop("embedding_model", "openai")
+		self.embedding_model = EmbeddingModel.load(embedding_model)()
+
+	def __del__(self):
+		super().__del__()
+		del self.embedding_model
+
+		empty_cuda_cache()
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, **kwargs):
+		queries, contents, scores, ids = self.cast_to_run(previous_result)
+		kwargs = pop_params(self._pure, kwargs)
+		return self._pure(queries, contents, scores, ids, **kwargs)
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents_list: List[List[str]],
+		scores_list: List[List[float]],
+		ids_list: List[List[str]],
+		percentile: float,
+		batch: int = 128,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Re-calculate each content's similarity with the query and filter out the contents that are below the content's
+		length times percentile. If This is a filter and does not override scores. The output of scores is not coming from
+		query-content similarity.
+		If the value of content's length times percentile is less than 1, keep the only one highest similarity content.
+
+		:param queries: The list of queries to use for filtering
+		:param contents_list: The list of lists of contents to filter
+		:param scores_list: The list of lists of scores retrieved
+		:param ids_list: The list of lists of ids retrieved
+		:param percentile: The percentile to cut off
+		:param batch: The number of queries to be processed in a batch
+		    Default is 128.
+		:return: Tuple of lists containing the filtered contents, ids, and scores
+		"""
+		query_embeddings, content_embeddings = embedding_query_content(
+			queries, contents_list, self.embedding_model, batch
+		)
+
+		results = list(
+			map(
+				lambda x: self.__row_pure(x[0], x[1], x[2], x[3], x[4], percentile),
+				zip(
+					query_embeddings,
+					content_embeddings,
+					contents_list,
+					ids_list,
+					scores_list,
+				),
+			)
+		)
+
+		remain_content_list = list(map(lambda x: x[0], results))
+		remain_ids_list = list(map(lambda x: x[1], results))
+		remain_scores_list = list(map(lambda x: x[2], results))
+
+		return remain_content_list, remain_ids_list, remain_scores_list
+
+	@staticmethod
+	def __row_pure(
+		query_embedding: str,
+		content_embeddings: List[List[float]],
+		content_list: List[str],
+		ids_list: List[str],
+		scores_list: List[float],
+		percentile: float,
+	) -> Tuple[List[str], List[str], List[float]]:
+		"""
+		Return tuple of lists containing the filtered contents, ids, and scores
+
+		:param query_embedding: Query embedding
+		:param content_embeddings: Each content embedding
+		:param content_list: Each content
+		:param ids_list: Each id
+		:param scores_list: Each score
+		:param percentile: The percentile to cut off
+		:return: Tuple of lists containing the filtered contents, ids, and scores
+		"""
+		num_top_k = int(len(content_embeddings) * percentile)
+
+		if num_top_k == 0:
+			num_top_k = 1
+
+		similarities = np.array(
+			list(
+				map(
+					lambda x: calculate_cosine_similarity(query_embedding, x),
+					content_embeddings,
+				)
+			)
+		).tolist()
+
+		content_id_score_similarity = list(
+			zip(ids_list, content_list, scores_list, similarities)
+		)
+
+		sorted_content_id_score_similarity = sorted(
+			content_id_score_similarity, key=lambda x: x[3], reverse=True
+		)[:num_top_k]
+
+		content_result, id_result, score_result, _ = zip(
+			*sorted_content_id_score_similarity
+		)
+		return list(content_result), list(id_result), list(score_result)

autorag/nodes/passagefilter/similarity_threshold_cutoff.py

@@ -0,0 +1,112 @@
+from typing import List, Tuple
+
+import numpy as np
+import pandas as pd
+
+from autorag.embedding.base import EmbeddingModel
+from autorag.evaluation.metric.util import calculate_cosine_similarity
+from autorag.nodes.passagefilter.base import BasePassageFilter
+from autorag.utils.util import (
+	embedding_query_content,
+	empty_cuda_cache,
+	result_to_dataframe,
+	pop_params,
+)
+
+
+class SimilarityThresholdCutoff(BasePassageFilter):
+	def __init__(self, project_dir: str, *args, **kwargs):
+		"""
+		Initialize the SimilarityThresholdCutoff module
+
+		:param project_dir: The project directory to use for initializing the module
+		:param embedding_model: The embedding model string to use for calculating similarity
+		        Default is "openai" which is OpenAI text-embedding-ada-002 embedding model.
+		"""
+		super().__init__(project_dir, *args, **kwargs)
+		embedding_model = kwargs.get("embedding_model", "openai")
+		self.embedding_model = EmbeddingModel.load(embedding_model)()
+
+	def __del__(self):
+		del self.embedding_model
+		empty_cuda_cache()
+		super().__del__()
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		kwargs = pop_params(self._pure, kwargs)
+		queries, contents, scores, ids = self.cast_to_run(previous_result)
+		return self._pure(queries, contents, scores, ids, *args, **kwargs)
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents_list: List[List[str]],
+		scores_list: List[List[float]],
+		ids_list: List[List[str]],
+		threshold: float,
+		batch: int = 128,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Re-calculate each content's similarity with the query and filter out the contents that are below the threshold.
+		If all contents are filtered, keep the only one highest similarity content.
+		This is a filter and does not override scores.
+		The output of scores is not coming from query-content similarity.
+
+		:param queries: The list of queries to use for filtering
+		:param contents_list: The list of lists of contents to filter
+		:param scores_list: The list of lists of scores retrieved
+		:param ids_list: The list of lists of ids retrieved
+		:param threshold: The threshold to cut off
+		:param batch: The number of queries to be processed in a batch
+		    Default is 128.
+		:return: Tuple of lists containing the filtered contents, ids, and scores
+		"""
+		query_embeddings, content_embeddings = embedding_query_content(
+			queries, contents_list, self.embedding_model, batch
+		)
+
+		remain_indices = list(
+			map(
+				lambda x: self.__row_pure(x[0], x[1], threshold),
+				zip(query_embeddings, content_embeddings),
+			)
+		)
+
+		remain_content_list = list(
+			map(lambda c, idx: [c[i] for i in idx], contents_list, remain_indices)
+		)
+		remain_scores_list = list(
+			map(lambda s, idx: [s[i] for i in idx], scores_list, remain_indices)
+		)
+		remain_ids_list = list(
+			map(lambda _id, idx: [_id[i] for i in idx], ids_list, remain_indices)
+		)
+		return remain_content_list, remain_ids_list, remain_scores_list
+
+	@staticmethod
+	def __row_pure(
+		query_embedding: str, content_embeddings: List[List[float]], threshold: float
+	) -> List[int]:
+		"""
+		Return indices that have to remain.
+		Return at least one index if there is nothing to remain.
+
+		:param query_embedding: Query embedding
+		:param content_embeddings: Each content embedding
+		:param threshold: The threshold to cut off
+		:return: Indices to remain at the contents
+		"""
+
+		similarities = np.array(
+			list(
+				map(
+					lambda x: calculate_cosine_similarity(query_embedding, x),
+					content_embeddings,
+				)
+			)
+		)
+		result = np.where(similarities >= threshold)[0].tolist()
+		if len(result) > 0:
+			return result
+		return [np.argmax(similarities)]

autorag/nodes/passagefilter/threshold_cutoff.py

@@ -0,0 +1,78 @@
+from typing import List, Tuple
+
+import pandas as pd
+
+from autorag.nodes.passagefilter.base import BasePassageFilter
+from autorag.utils.util import convert_inputs_to_list, result_to_dataframe
+
+
+class ThresholdCutoff(BasePassageFilter):
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		_, contents, scores, ids = self.cast_to_run(previous_result)
+		return self._pure(contents, scores, ids, *args, **kwargs)
+
+	def _pure(
+		self,
+		contents_list: List[List[str]],
+		scores_list: List[List[float]],
+		ids_list: List[List[str]],
+		threshold: float,
+		reverse: bool = False,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Filters the contents, scores, and ids based on a previous result's score.
+		Keeps at least one item per query if all scores are below the threshold.
+
+		:param contents_list: List of content strings for each query.
+		:param scores_list: List of scores for each content.
+		:param ids_list: List of ids for each content.
+		:param threshold: The minimum score to keep an item.
+		:param reverse: If True, the lower the score, the better.
+		    Default is False.
+		:return: Filtered lists of contents, ids, and scores.
+		"""
+		remain_indices = list(
+			map(lambda x: self.__row_pure(x, threshold, reverse), scores_list)
+		)
+
+		remain_content_list = list(
+			map(lambda c, idx: [c[i] for i in idx], contents_list, remain_indices)
+		)
+		remain_scores_list = list(
+			map(lambda s, idx: [s[i] for i in idx], scores_list, remain_indices)
+		)
+		remain_ids_list = list(
+			map(lambda _id, idx: [_id[i] for i in idx], ids_list, remain_indices)
+		)
+
+		return remain_content_list, remain_ids_list, remain_scores_list
+
+	@convert_inputs_to_list
+	def __row_pure(
+		self, scores_list: List[float], threshold: float, reverse: bool = False
+	) -> List[int]:
+		"""
+		Return indices that have to remain.
+		Return at least one index if there is nothing to remain.
+
+		:param scores_list: Each score
+		:param threshold: The threshold to cut off
+		:param reverse: If True, the lower the score, the better
+			Default is False.
+		:return: Indices to remain at the contents
+		"""
+		assert isinstance(scores_list, list), "scores_list must be a list."
+
+		if reverse:
+			remain_indices = [
+				i for i, score in enumerate(scores_list) if score <= threshold
+			]
+			default_index = scores_list.index(min(scores_list))
+		else:
+			remain_indices = [
+				i for i, score in enumerate(scores_list) if score >= threshold
+			]
+			default_index = scores_list.index(max(scores_list))
+
+		return remain_indices if remain_indices else [default_index]

autorag/nodes/passagereranker/__init__.py

@@ -0,0 +1,16 @@
+from .cohere import CohereReranker
+from .colbert import ColbertReranker
+from .flag_embedding import FlagEmbeddingReranker
+from .flag_embedding_llm import FlagEmbeddingLLMReranker
+from .jina import JinaReranker
+from .koreranker import KoReranker
+from .monot5 import MonoT5
+from .pass_reranker import PassReranker
+from .rankgpt import RankGPT
+from .sentence_transformer import SentenceTransformerReranker
+from .time_reranker import TimeReranker
+from .upr import Upr
+from .openvino import OpenVINOReranker
+from .voyageai import VoyageAIReranker
+from .mixedbreadai import MixedbreadAIReranker
+from .flashrank import FlashRankReranker

autorag/nodes/passagereranker/base.py

@@ -0,0 +1,44 @@
+import abc
+import logging
+from pathlib import Path
+from typing import Union
+
+import pandas as pd
+
+from autorag.schema import BaseModule
+from autorag.utils import validate_qa_dataset
+from autorag.utils.cast import cast_retrieve_infos
+
+logger = logging.getLogger("AutoRAG")
+
+
+class BasePassageReranker(BaseModule, metaclass=abc.ABCMeta):
+	def __init__(self, project_dir: Union[str, Path], *args, **kwargs):
+		logger.info(
+			f"Initialize passage reranker node - {self.__class__.__name__} module..."
+		)
+
+	def __del__(self):
+		logger.info(
+			f"Deleting passage reranker node - {self.__class__.__name__} module..."
+		)
+
+	def cast_to_run(self, previous_result: pd.DataFrame, *args, **kwargs):
+		logger.info(
+			f"Running passage reranker node - {self.__class__.__name__} module..."
+		)
+		validate_qa_dataset(previous_result)
+
+		# find queries columns
+		assert "query" in previous_result.columns, (
+			"previous_result must have query column."
+		)
+		queries = previous_result["query"].tolist()
+
+		retrieve_infos = cast_retrieve_infos(previous_result)
+		return (
+			queries,
+			retrieve_infos["retrieved_contents"],
+			retrieve_infos["retrieve_scores"],
+			retrieve_infos["retrieved_ids"],
+		)

autorag/nodes/passagereranker/cohere.py

@@ -0,0 +1,118 @@
+import os
+from typing import List, Tuple
+
+import cohere
+import pandas as pd
+from cohere import RerankResponseResultsItem
+
+from autorag.nodes.passagereranker.base import BasePassageReranker
+from autorag.utils.util import get_event_loop, process_batch, result_to_dataframe
+
+
+class CohereReranker(BasePassageReranker):
+	def __init__(self, project_dir: str, *args, **kwargs):
+		"""
+		Initialize Cohere rerank node.
+
+		:param project_dir: The project directory path.
+		:param api_key: The API key for Cohere rerank.
+		    You can set it in the environment variable COHERE_API_KEY.
+		    Or, you can directly set it on the config YAML file using this parameter.
+		    Default is env variable "COHERE_API_KEY".
+		:param kwargs: Extra arguments that are not affected
+		"""
+		super().__init__(project_dir)
+		api_key = kwargs.pop("api_key", None)
+		api_key = os.getenv("COHERE_API_KEY", None) if api_key is None else api_key
+		if api_key is None:
+			api_key = os.getenv("CO_API_KEY", None)
+		if api_key is None:
+			raise KeyError(
+				"Please set the API key for Cohere rerank in the environment variable COHERE_API_KEY "
+				"or directly set it on the config YAML file."
+			)
+
+		self.cohere_client = cohere.AsyncClientV2(api_key=api_key)
+
+	def __del__(self):
+		del self.cohere_client
+		super().__del__()
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		queries, contents, scores, ids = self.cast_to_run(previous_result)
+		top_k = kwargs.pop("top_k")
+		batch = kwargs.pop("batch", 64)
+		model = kwargs.pop("model", "rerank-v3.5")
+		return self._pure(queries, contents, scores, ids, top_k, batch, model)
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents_list: List[List[str]],
+		scores_list: List[List[float]],
+		ids_list: List[List[str]],
+		top_k: int,
+		batch: int = 64,
+		model: str = "rerank-v3.5",
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Rerank a list of contents with Cohere rerank models.
+		You can get the API key from https://cohere.com/rerank and set it in the environment variable COHERE_API_KEY.
+
+		:param queries: The list of queries to use for reranking
+		:param contents_list: The list of lists of contents to rerank
+		:param scores_list: The list of lists of scores retrieved from the initial ranking
+		:param ids_list: The list of lists of ids retrieved from the initial ranking
+		:param top_k: The number of passages to be retrieved
+		:param batch: The number of queries to be processed in a batch
+		:param model: The model name for Cohere rerank.
+		    You can choose between "rerank-v3.5", "rerank-english-v3.0", and "rerank-multilingual-v3.0".
+		    Default is "rerank-v3.5".
+		:return: Tuple of lists containing the reranked contents, ids, and scores
+		"""
+		# Run async cohere_rerank_pure function
+		tasks = [
+			cohere_rerank_pure(self.cohere_client, model, query, document, ids, top_k)
+			for query, document, ids in zip(queries, contents_list, ids_list)
+		]
+		loop = get_event_loop()
+		results = loop.run_until_complete(process_batch(tasks, batch_size=batch))
+		content_result = list(map(lambda x: x[0], results))
+		id_result = list(map(lambda x: x[1], results))
+		score_result = list(map(lambda x: x[2], results))
+
+		return content_result, id_result, score_result
+
+
+async def cohere_rerank_pure(
+	cohere_client: cohere.AsyncClient,
+	model: str,
+	query: str,
+	documents: List[str],
+	ids: List[str],
+	top_k: int,
+) -> Tuple[List[str], List[str], List[float]]:
+	"""
+	Rerank a list of contents with Cohere rerank models.
+
+	:param cohere_client: The Cohere AsyncClient to use for reranking
+	:param model: The model name for Cohere rerank
+	:param query: The query to use for reranking
+	:param documents: The list of contents to rerank
+	:param ids: The list of ids corresponding to the documents
+	:param top_k: The number of passages to be retrieved
+	:return: Tuple of lists containing the reranked contents, ids, and scores
+	"""
+	rerank_results = await cohere_client.rerank(
+		model=model,
+		query=query,
+		documents=documents,
+		top_n=top_k,
+	)
+	results: List[RerankResponseResultsItem] = rerank_results.results
+	reranked_scores: List[float] = list(map(lambda x: x.relevance_score, results))
+	indices = list(map(lambda x: x.index, results))
+	reranked_contents: List[str] = list(map(lambda i: documents[i], indices))
+	reranked_ids: List[str] = list(map(lambda i: ids[i], indices))
+	return reranked_contents, reranked_ids, reranked_scores