AutoRAG 0.0.0__py3-none-any.whl

autorag/nodes/passagereranker/openvino.py

@@ -0,0 +1,191 @@
+from pathlib import Path
+from typing import Any, List, Tuple
+
+import numpy as np
+import pandas as pd
+
+from autorag.nodes.passagereranker.base import BasePassageReranker
+
+
+from autorag.utils.util import (
+	make_batch,
+	sort_by_scores,
+	flatten_apply,
+	select_top_k,
+	result_to_dataframe,
+	pop_params,
+	empty_cuda_cache,
+)
+
+
+class OpenVINOReranker(BasePassageReranker):
+	def __init__(
+		self,
+		project_dir: str,
+		model: str = "BAAI/bge-reranker-large",
+		*args,
+		**kwargs,
+	):
+		super().__init__(project_dir)
+
+		try:
+			from huggingface_hub import HfApi
+			from transformers import AutoTokenizer
+		except ImportError as e:
+			raise ValueError(
+				"Could not import huggingface_hub python package. "
+				"Please install it with: "
+				"`pip install -U huggingface_hub`."
+			) from e
+
+		def require_model_export(
+			model_id: str, revision: Any = None, subfolder: Any = None
+		) -> bool:
+			model_dir = Path(model_id)
+			if subfolder is not None:
+				model_dir = model_dir / subfolder
+			if model_dir.is_dir():
+				return (
+					not (model_dir / "openvino_model.xml").exists()
+					or not (model_dir / "openvino_model.bin").exists()
+				)
+			hf_api = HfApi()
+			try:
+				model_info = hf_api.model_info(model_id, revision=revision or "main")
+				normalized_subfolder = (
+					None if subfolder is None else Path(subfolder).as_posix()
+				)
+				model_files = [
+					file.rfilename
+					for file in model_info.siblings
+					if normalized_subfolder is None
+					or file.rfilename.startswith(normalized_subfolder)
+				]
+				ov_model_path = (
+					"openvino_model.xml"
+					if subfolder is None
+					else f"{normalized_subfolder}/openvino_model.xml"
+				)
+				return (
+					ov_model_path not in model_files
+					or ov_model_path.replace(".xml", ".bin") not in model_files
+				)
+			except Exception:
+				return True
+
+		try:
+			from optimum.intel.openvino import OVModelForSequenceClassification
+		except ImportError:
+			raise ImportError(
+				"Please install optimum package to use OpenVINOReranker"
+				"pip install 'optimum[openvino,nncf]'"
+			)
+
+		model_kwargs = pop_params(
+			OVModelForSequenceClassification.from_pretrained, kwargs
+		)
+
+		if require_model_export(model):
+			# use remote model
+			self.model = OVModelForSequenceClassification.from_pretrained(
+				model, export=True, **model_kwargs
+			)
+		else:
+			# use local model
+			self.model = OVModelForSequenceClassification.from_pretrained(
+				model, **model_kwargs
+			)
+
+		self.tokenizer = AutoTokenizer.from_pretrained(model)
+
+	def __del__(self):
+		del self.model
+		del self.tokenizer
+		empty_cuda_cache()
+		super().__del__()
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		queries, contents, _, ids = self.cast_to_run(previous_result)
+		top_k = kwargs.get("top_k", 3)
+		batch = kwargs.get("batch", 64)
+		return self._pure(queries, contents, ids, top_k, batch)
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents_list: List[List[str]],
+		ids_list: List[List[str]],
+		top_k: int,
+		batch: int = 64,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Rerank a list of contents based on their relevance to a query using MonoT5.
+
+		:param queries: The list of queries to use for reranking
+		:param contents_list: The list of lists of contents to rerank
+		: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
+		:return: tuple of lists containing the reranked contents, ids, and scores
+		"""
+		# Retrieve the tokens used by the model to represent false and true predictions
+
+		nested_list = [
+			list(map(lambda x: [query, x], content_list))
+			for query, content_list in zip(queries, contents_list)
+		]
+
+		rerank_scores = flatten_apply(
+			openvino_run_model,
+			nested_list,
+			model=self.model,
+			batch_size=batch,
+			tokenizer=self.tokenizer,
+		)
+
+		df = pd.DataFrame(
+			{
+				"contents": contents_list,
+				"ids": ids_list,
+				"scores": rerank_scores,
+			}
+		)
+		df[["contents", "ids", "scores"]] = df.apply(
+			sort_by_scores, axis=1, result_type="expand"
+		)
+		results = select_top_k(df, ["contents", "ids", "scores"], top_k)
+
+		return (
+			results["contents"].tolist(),
+			results["ids"].tolist(),
+			results["scores"].tolist(),
+		)
+
+
+def openvino_run_model(
+	input_texts,
+	model,
+	batch_size: int,
+	tokenizer,
+):
+	batch_input_texts = make_batch(input_texts, batch_size)
+	results = []
+	for batch_texts in batch_input_texts:
+		input_tensors = tokenizer(
+			batch_texts,
+			padding=True,
+			truncation=True,
+			return_tensors="pt",
+		)
+
+		outputs = model(**input_tensors, return_dict=True)
+		if outputs[0].shape[1] > 1:
+			scores = outputs[0][:, 1]
+		else:
+			scores = outputs[0].flatten()
+
+		scores = list(map(float, (1 / (1 + np.exp(-np.array(scores))))))
+		results.extend(scores)
+	return results

autorag/nodes/passagereranker/pass_reranker.py

@@ -0,0 +1,31 @@
+from typing import List
+
+import pandas as pd
+
+from autorag.nodes.passagereranker.base import BasePassageReranker
+from autorag.utils import result_to_dataframe
+
+
+class PassReranker(BasePassageReranker):
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		top_k = kwargs.pop("top_k")
+
+		_, contents_list, scores_list, ids_list = self.cast_to_run(previous_result)
+		return self._pure(contents_list, scores_list, ids_list, top_k)
+
+	def _pure(
+		self,
+		contents_list: List[List[str]],
+		scores_list: List[List[float]],
+		ids_list: List[List[str]],
+		top_k: int,
+	):
+		"""
+		Do not perform reranking.
+		Return the given top-k passages as is.
+		"""
+		contents_list = list(map(lambda x: x[:top_k], contents_list))
+		scores_list = list(map(lambda x: x[:top_k], scores_list))
+		ids_list = list(map(lambda x: x[:top_k], ids_list))
+		return contents_list, ids_list, scores_list

autorag/nodes/passagereranker/rankgpt.py

@@ -0,0 +1,170 @@
+from typing import List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import pandas as pd
+from llama_index.core.base.llms.types import ChatMessage, ChatResponse
+from llama_index.core.llms import LLM
+from llama_index.core.postprocessor.rankGPT_rerank import RankGPTRerank
+from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode
+from llama_index.core.utils import print_text
+from llama_index.llms.openai import OpenAI
+
+from autorag import generator_models
+from autorag.nodes.passagereranker.base import BasePassageReranker
+from autorag.utils.util import (
+	get_event_loop,
+	process_batch,
+	pop_params,
+	result_to_dataframe,
+	empty_cuda_cache,
+)
+
+
+class RankGPT(BasePassageReranker):
+	def __init__(
+		self, project_dir: str, llm: Optional[Union[str, LLM]] = None, **kwargs
+	):
+		"""
+		Initialize the RankGPT reranker.
+
+		:param project_dir: The project directory
+		:param llm: The LLM model to use for RankGPT rerank.
+			It is a llama index model.
+			Default is the OpenAI model with gpt-4o-mini.
+		:param kwargs: The keyword arguments for the LLM model.
+		"""
+		super().__init__(project_dir)
+		if llm is None:
+			self.llm = OpenAI(model="gpt-4o-mini")
+		else:
+			if not isinstance(llm, LLM):
+				llm_class = generator_models[llm]
+				llm_param = pop_params(llm_class.__init__, kwargs)
+				self.llm = llm_class(**llm_param)
+			else:
+				self.llm = llm
+
+	def __del__(self):
+		del self.llm
+		empty_cuda_cache()
+		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.get("top_k", 1)
+		verbose = kwargs.get("verbose", False)
+		rankgpt_rerank_prompt = kwargs.get("rankgpt_rerank_prompt", None)
+		batch = kwargs.get("batch", 16)
+		return self._pure(
+			queries=queries,
+			contents_list=contents,
+			scores_list=scores,
+			ids_list=ids,
+			top_k=top_k,
+			verbose=verbose,
+			rankgpt_rerank_prompt=rankgpt_rerank_prompt,
+			batch=batch,
+		)
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents_list: List[List[str]],
+		scores_list: List[List[float]],
+		ids_list: List[List[str]],
+		top_k: int,
+		verbose: bool = False,
+		rankgpt_rerank_prompt: Optional[str] = None,
+		batch: int = 16,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Rerank given context paragraphs using RankGPT.
+		Return pseudo scores, since the actual scores are not available on RankGPT.
+
+		: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 verbose: Whether to print intermediate steps.
+		:param rankgpt_rerank_prompt: The prompt template for RankGPT rerank.
+			Default is RankGPT's default prompt.
+		:param batch: The number of queries to be processed in a batch.
+		:return: Tuple of lists containing the reranked contents, ids, and scores
+		"""
+		query_bundles = list(map(lambda query: QueryBundle(query_str=query), queries))
+		nodes_list = [
+			list(
+				map(
+					lambda x: NodeWithScore(node=TextNode(text=x[0]), score=x[1]),
+					zip(content_list, score_list),
+				)
+			)
+			for content_list, score_list in zip(contents_list, scores_list)
+		]
+
+		reranker = AsyncRankGPTRerank(
+			top_n=top_k,
+			llm=self.llm,
+			verbose=verbose,
+			rankgpt_rerank_prompt=rankgpt_rerank_prompt,
+		)
+
+		tasks = [
+			reranker.async_postprocess_nodes(nodes, query, ids)
+			for nodes, query, ids in zip(nodes_list, query_bundles, ids_list)
+		]
+		loop = get_event_loop()
+		rerank_result = loop.run_until_complete(process_batch(tasks, batch_size=batch))
+		content_result = [
+			list(map(lambda x: x.node.text, res[0])) for res in rerank_result
+		]
+		score_result = [
+			np.linspace(1.0, 0.0, len(res[0])).tolist() for res in rerank_result
+		]
+		id_result = [res[1] for res in rerank_result]
+
+		del reranker
+
+		return content_result, id_result, score_result
+
+
+class AsyncRankGPTRerank(RankGPTRerank):
+	async def async_run_llm(self, messages: Sequence[ChatMessage]) -> ChatResponse:
+		return await self.llm.achat(messages)
+
+	async def async_postprocess_nodes(
+		self,
+		nodes: List[NodeWithScore],
+		query_bundle: QueryBundle,
+		ids: Optional[List[str]] = None,
+	) -> Tuple[List[NodeWithScore], List[str]]:
+		if ids is None:
+			ids = [str(i) for i in range(len(nodes))]
+
+		items = {
+			"query": query_bundle.query_str,
+			"hits": [{"content": node.get_content()} for node in nodes],
+		}
+
+		messages = self.create_permutation_instruction(item=items)
+		permutation = await self.async_run_llm(messages=messages)
+		if permutation.message is not None and permutation.message.content is not None:
+			rerank_ranks = self._receive_permutation(
+				items, str(permutation.message.content)
+			)
+			if self.verbose:
+				print_text(f"After Reranking, new rank list for nodes: {rerank_ranks}")
+
+			initial_results: List[NodeWithScore] = []
+			id_results = []
+
+			for idx in rerank_ranks:
+				initial_results.append(
+					NodeWithScore(node=nodes[idx].node, score=nodes[idx].score)
+				)
+				id_results.append(ids[idx])
+			return initial_results[: self.top_n], id_results[: self.top_n]
+		else:
+			return nodes[: self.top_n], ids[: self.top_n]

autorag/nodes/passagereranker/run.py

@@ -0,0 +1,145 @@
+import logging
+import os
+import pathlib
+from typing import List, Dict
+
+import pandas as pd
+
+from autorag.nodes.retrieval.run_util import evaluate_retrieval_node
+from autorag.schema.metricinput import MetricInput
+from autorag.strategy import measure_speed, filter_by_threshold, select_best
+from autorag.utils.util import apply_recursive, to_list
+
+logger = logging.getLogger("AutoRAG")
+
+
+def run_passage_reranker_node(
+	modules: List,
+	module_params: List[Dict],
+	previous_result: pd.DataFrame,
+	node_line_dir: str,
+	strategies: Dict,
+) -> pd.DataFrame:
+	"""
+	Run evaluation and select the best module among passage reranker node results.
+
+	:param modules: Passage reranker modules to run.
+	:param module_params: Passage reranker module parameters.
+	:param previous_result: Previous result dataframe.
+	    Could be retrieval, reranker modules result.
+	    It means it must contain 'query', 'retrieved_contents', 'retrieved_ids', 'retrieve_scores' columns.
+	:param node_line_dir: This node line's directory.
+	:param strategies: Strategies for passage reranker node.
+	    In this node, we use 'retrieval_f1', 'retrieval_recall' and 'retrieval_precision'.
+	    You can skip evaluation when you use only one module and a module parameter.
+	:return: The best result dataframe with previous result columns.
+	"""
+	if not os.path.exists(node_line_dir):
+		os.makedirs(node_line_dir)
+	project_dir = pathlib.PurePath(node_line_dir).parent.parent
+	qa_df = pd.read_parquet(
+		os.path.join(project_dir, "data", "qa.parquet"), engine="pyarrow"
+	)
+	retrieval_gt = qa_df["retrieval_gt"].tolist()
+	retrieval_gt = apply_recursive(lambda x: str(x), to_list(retrieval_gt))
+
+	# make rows to metric_inputs
+	metric_inputs = [
+		MetricInput(retrieval_gt=ret_gt, query=query, generation_gt=gen_gt)
+		for ret_gt, query, gen_gt in zip(
+			retrieval_gt, qa_df["query"].tolist(), qa_df["generation_gt"].tolist()
+		)
+	]
+
+	results, execution_times = zip(
+		*map(
+			lambda task: measure_speed(
+				task[0].run_evaluator,
+				project_dir=project_dir,
+				previous_result=previous_result,
+				**task[1],
+			),
+			zip(modules, module_params),
+		)
+	)
+	average_times = list(map(lambda x: x / len(results[0]), execution_times))
+
+	# run metrics before filtering
+	if strategies.get("metrics") is None:
+		raise ValueError(
+			"You must at least one metrics for passage_reranker evaluation."
+		)
+	results = list(
+		map(
+			lambda x: evaluate_retrieval_node(
+				x,
+				metric_inputs,
+				strategies.get("metrics"),
+			),
+			results,
+		)
+	)
+
+	# save results to folder
+	save_dir = os.path.join(node_line_dir, "passage_reranker")  # node name
+	if not os.path.exists(save_dir):
+		os.makedirs(save_dir)
+	filepaths = list(
+		map(lambda x: os.path.join(save_dir, f"{x}.parquet"), range(len(modules)))
+	)
+	list(
+		map(lambda x: x[0].to_parquet(x[1], index=False), zip(results, filepaths))
+	)  # execute save to parquet
+	filenames = list(map(lambda x: os.path.basename(x), filepaths))
+
+	summary_df = pd.DataFrame(
+		{
+			"filename": filenames,
+			"module_name": list(map(lambda module: module.__name__, modules)),
+			"module_params": module_params,
+			"execution_time": average_times,
+			**{
+				f"passage_reranker_{metric}": list(
+					map(lambda result: result[metric].mean(), results)
+				)
+				for metric in strategies.get("metrics")
+			},
+		}
+	)
+
+	# filter by strategies
+	if strategies.get("speed_threshold") is not None:
+		results, filenames = filter_by_threshold(
+			results, average_times, strategies["speed_threshold"], filenames
+		)
+	selected_result, selected_filename = select_best(
+		results,
+		strategies.get("metrics"),
+		filenames,
+		strategies.get("strategy", "mean"),
+	)
+	# change metric name columns to passage_reranker_metric_name
+	selected_result = selected_result.rename(
+		columns={
+			metric_name: f"passage_reranker_{metric_name}"
+			for metric_name in strategies["metrics"]
+		}
+	)
+	# drop retrieval result columns in previous_result
+	previous_result = previous_result.drop(
+		columns=["retrieved_contents", "retrieved_ids", "retrieve_scores"]
+	)
+	best_result = pd.concat([previous_result, selected_result], axis=1)
+
+	# add 'is_best' column to summary file
+	summary_df["is_best"] = summary_df["filename"] == selected_filename
+
+	# save files
+	summary_df.to_csv(os.path.join(save_dir, "summary.csv"), index=False)
+	best_result.to_parquet(
+		os.path.join(
+			save_dir, f"best_{os.path.splitext(selected_filename)[0]}.parquet"
+		),
+		index=False,
+	)
+	return best_result

autorag/nodes/passagereranker/sentence_transformer.py

@@ -0,0 +1,129 @@
+from typing import List, Tuple
+
+import pandas as pd
+
+from autorag.nodes.passagereranker.base import BasePassageReranker
+from autorag.utils.util import (
+	flatten_apply,
+	make_batch,
+	select_top_k,
+	sort_by_scores,
+	pop_params,
+	result_to_dataframe,
+	empty_cuda_cache,
+)
+
+
+class SentenceTransformerReranker(BasePassageReranker):
+	def __init__(
+		self,
+		project_dir: str,
+		model_name: str = "cross-encoder/ms-marco-MiniLM-L-2-v2",
+		*args,
+		**kwargs,
+	):
+		"""
+		Initialize the Sentence Transformer reranker node.
+
+		:param project_dir: The project directory
+		:param model_name: The name of the Sentence Transformer model to use for reranking
+		Default is "cross-encoder/ms-marco-MiniLM-L-2-v2"
+		:param kwargs: The CrossEncoder parameters
+		"""
+		super().__init__(project_dir, *args, **kwargs)
+		try:
+			import torch
+			from sentence_transformers import CrossEncoder
+		except ImportError:
+			raise ImportError(
+				"You have to install AutoRAG[gpu] to use SentenceTransformerReranker"
+			)
+		self.device = "cuda" if torch.cuda.is_available() else "cpu"
+		model_params = pop_params(CrossEncoder.__init__, kwargs)
+		self.model = CrossEncoder(model_name, device=self.device, **model_params)
+
+	def __del__(self):
+		del self.model
+		empty_cuda_cache()
+		super().__del__()
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		"""
+		Rerank a list of contents based on their relevance to a query using a Sentence Transformer model.
+
+		:param previous_result: The previous result
+		:param top_k: The number of passages to be retrieved
+		:param batch: The number of queries to be processed in a batch
+		:return: pd DataFrame containing the reranked contents, ids, and scores
+		"""
+		queries, contents_list, scores_list, ids_list = self.cast_to_run(
+			previous_result
+		)
+		top_k = kwargs.get("top_k", 1)
+		batch = kwargs.get("batch", 64)
+		return self._pure(queries, contents_list, ids_list, top_k, batch)
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents_list: List[List[str]],
+		ids_list: List[List[str]],
+		top_k: int,
+		batch: int = 64,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Rerank a list of contents based on their relevance to a query using a Sentence Transformer model.
+
+		:param queries: The list of queries to use for reranking
+		:param contents_list: The list of lists of contents to rerank
+		: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
+
+		:return: tuple of lists containing the reranked contents, ids, and scores
+		"""
+		nested_list = [
+			list(map(lambda x: [query, x], content_list))
+			for query, content_list in zip(queries, contents_list)
+		]
+		rerank_scores = flatten_apply(
+			sentence_transformer_run_model,
+			nested_list,
+			model=self.model,
+			batch_size=batch,
+		)
+
+		df = pd.DataFrame(
+			{
+				"contents": contents_list,
+				"ids": ids_list,
+				"scores": rerank_scores,
+			}
+		)
+		df[["contents", "ids", "scores"]] = df.apply(
+			sort_by_scores, axis=1, result_type="expand"
+		)
+		results = select_top_k(df, ["contents", "ids", "scores"], top_k)
+
+		return (
+			results["contents"].tolist(),
+			results["ids"].tolist(),
+			results["scores"].tolist(),
+		)
+
+
+def sentence_transformer_run_model(input_texts, model, batch_size: int):
+	try:
+		import torch
+	except ImportError:
+		raise ImportError(
+			"You have to install AutoRAG[gpu] to use SentenceTransformerReranker"
+		)
+	batch_input_texts = make_batch(input_texts, batch_size)
+	results = []
+	for batch_texts in batch_input_texts:
+		with torch.no_grad():
+			pred_scores = model.predict(batch_texts)
+		results.extend(pred_scores.tolist())
+	return results

autorag/nodes/passagereranker/tart/__init__.py

@@ -0,0 +1 @@
+from .tart import Tart