make again retrieval/base.py and run_util.py

Author: vkehfdl1

autorag/autorag/nodes/retrieval/__init__.py

@@ -0,0 +1,127 @@
+import abc
+import logging
+import os
+from typing import List, Union, Tuple
+
+import pandas as pd
+
+from autorag.schema import BaseModule
+from autorag.support import get_support_modules
+from autorag.utils import fetch_contents, result_to_dataframe, validate_qa_dataset
+from autorag.utils.util import pop_params
+
+logger = logging.getLogger("AutoRAG")
+
+
+class BaseRetrieval(BaseModule, metaclass=abc.ABCMeta):
+	def __init__(self, project_dir: str, *args, **kwargs):
+		logger.info(f"Initialize retrieval node - {self.__class__.__name__}")
+
+		self.resources_dir = os.path.join(project_dir, "resources")
+		data_dir = os.path.join(project_dir, "data")
+		# fetch data from corpus_data
+		self.corpus_df = pd.read_parquet(
+			os.path.join(data_dir, "corpus.parquet"), engine="pyarrow"
+		)
+
+	def __del__(self):
+		logger.info(f"Deleting retrieval node - {self.__class__.__name__} module...")
+
+	def cast_to_run(self, previous_result: pd.DataFrame, *args, **kwargs):
+		logger.info(f"Running retrieval node - {self.__class__.__name__} module...")
+		validate_qa_dataset(previous_result)
+		# find queries columns & type cast queries
+		assert (
+			"query" in previous_result.columns
+		), "previous_result must have query column."
+		if "queries" not in previous_result.columns:
+			previous_result["queries"] = previous_result["query"]
+		previous_result.loc[:, "queries"] = previous_result["queries"].apply(
+			cast_queries
+		)
+		queries = previous_result["queries"].tolist()
+		return queries
+
+
+class HybridRetrieval(BaseRetrieval, metaclass=abc.ABCMeta):
+	def __init__(
+		self, project_dir: str, target_modules, target_module_params, *args, **kwargs
+	):
+		super().__init__(project_dir)
+		self.target_modules = list(
+			map(
+				lambda x, y: get_support_modules(x)(
+					**y,
+					project_dir=project_dir,
+				),
+				target_modules,
+				target_module_params,
+			)
+		)
+		self.target_module_params = target_module_params
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		result_dfs: List[pd.DataFrame] = list(
+			map(
+				lambda x, y: x.pure(
+					**y,
+					previous_result=previous_result,
+				),
+				self.target_modules,
+				self.target_module_params,
+			)
+		)
+		ids = tuple(
+			map(lambda df: df["retrieved_ids"].apply(list).tolist(), result_dfs)
+		)
+		scores = tuple(
+			map(
+				lambda df: df["retrieve_scores"].apply(list).tolist(),
+				result_dfs,
+			)
+		)
+
+		_pure_params = pop_params(self._pure, kwargs)
+		if "ids" in _pure_params or "scores" in _pure_params:
+			raise ValueError(
+				"With specifying ids or scores, you must use HybridRRF.run_evaluator instead."
+			)
+		ids, scores = self._pure(ids=ids, scores=scores, **_pure_params)
+		contents = fetch_contents(self.corpus_df, ids)
+		return contents, ids, scores
+
+
+def cast_queries(queries: Union[str, List[str]]) -> List[str]:
+	if isinstance(queries, str):
+		return [queries]
+	elif isinstance(queries, List):
+		return queries
+	else:
+		raise ValueError(f"queries must be str or list, but got {type(queries)}")
+
+
+def evenly_distribute_passages(
+	ids: List[List[str]], scores: List[List[float]], top_k: int
+) -> Tuple[List[str], List[float]]:
+	assert len(ids) == len(scores), "ids and scores must have same length."
+	query_cnt = len(ids)
+	avg_len = top_k // query_cnt
+	remainder = top_k % query_cnt
+
+	new_ids = []
+	new_scores = []
+	for i in range(query_cnt):
+		if i < remainder:
+			new_ids.extend(ids[i][: avg_len + 1])
+			new_scores.extend(scores[i][: avg_len + 1])
+		else:
+			new_ids.extend(ids[i][:avg_len])
+			new_scores.extend(scores[i][:avg_len])
+
+	return new_ids, new_scores
+
+
+def get_bm25_pkl_name(bm25_tokenizer: str):
+	bm25_tokenizer = bm25_tokenizer.replace("/", "")
+	return f"bm25_{bm25_tokenizer}.pkl"

autorag/autorag/nodes/retrieval/base.py

@@ -0,0 +1,152 @@
+import os
+import pathlib
+from typing import Tuple, List, Union, Dict
+
+import pandas as pd
+
+from autorag.evaluation import evaluate_retrieval
+from autorag.schema.metricinput import MetricInput
+from autorag.strategy import measure_speed, filter_by_threshold, select_best
+
+
+def evaluate_retrieval_node(
+	result_df: pd.DataFrame,
+	metric_inputs: List[MetricInput],
+	metrics: Union[List[str], List[Dict]],
+) -> pd.DataFrame:
+	"""
+	Evaluate retrieval node from retrieval node result dataframe.
+	:param result_df: The result dataframe from a retrieval node.
+	:param metric_inputs: List of metric input schema for AutoRAG.
+	:param metrics: Metric list from input strategies.
+	:return: Return result_df with metrics columns.
+	    The columns will be 'retrieved_contents', 'retrieved_ids', 'retrieve_scores', and metric names.
+	"""
+
+	@evaluate_retrieval(
+		metric_inputs=metric_inputs,
+		metrics=metrics,
+	)
+	def evaluate_this_module(df: pd.DataFrame):
+		return (
+			df["retrieved_contents"].tolist(),
+			df["retrieved_ids"].tolist(),
+			df["retrieve_scores"].tolist(),
+		)
+
+	return evaluate_this_module(result_df)
+
+
+def run(
+	input_modules,
+	input_module_params,
+	project_dir: Union[str, pathlib.Path, pathlib.PurePath],
+	previous_result: pd.DataFrame,
+	strategies,
+	metric_inputs: List[MetricInput],
+) -> Tuple[List[pd.DataFrame], List]:
+	"""
+	Run input modules and parameters.
+	:param input_modules: Input modules
+	:param input_module_params: Input module parameters
+	:param project_dir: Project directory path.
+	:param previous_result: Previous result dataframe.
+	:param strategies: Strategies for retrieval node.
+	:param metric_inputs: List of metric input schema for AutoRAG.
+	:return: First, it returns list of result dataframe.
+	Second, it returns list of execution times.
+	"""
+	result, execution_times = zip(
+		*map(
+			lambda task: measure_speed(
+				task[0].run_evaluator,
+				project_dir=project_dir,
+				previous_result=previous_result,
+				**task[1],
+			),
+			zip(input_modules, input_module_params),
+		)
+	)
+	average_times = list(map(lambda x: x / len(result[0]), execution_times))
+
+	# run metrics before filtering
+	if strategies.get("metrics") is None:
+		raise ValueError("You must at least one metrics for retrieval evaluation.")
+	result = list(
+		map(
+			lambda x: evaluate_retrieval_node(
+				x,
+				metric_inputs,
+				strategies.get("metrics"),
+			),
+			result,
+		)
+	)
+
+	return result, average_times
+
+
+def save_and_summary(
+	input_modules,
+	input_module_params,
+	result_list,
+	execution_time_list,
+	filename_start: int,
+	save_dir: Union[str, pathlib.Path, pathlib.PurePath],
+	strategies,
+):
+	"""
+	Save the result and make summary file
+	:param input_modules: Input modules
+	:param input_module_params: Input module parameters
+	:param result_list: Result list
+	:param execution_time_list: Execution times
+	:param filename_start: The first filename to use
+	:return: First, it returns list of result dataframe.
+	Second, it returns list of execution times.
+	"""
+
+	# save results to folder
+	filepaths = list(
+		map(
+			lambda x: os.path.join(save_dir, f"{x}.parquet"),
+			range(filename_start, filename_start + len(input_modules)),
+		)
+	)
+	list(
+		map(
+			lambda x: x[0].to_parquet(x[1], index=False),
+			zip(result_list, filepaths),
+		)
+	)  # execute save to parquet
+	filename_list = list(map(lambda x: os.path.basename(x), filepaths))
+
+	summary_df = pd.DataFrame(
+		{
+			"filename": filename_list,
+			"module_name": list(map(lambda module: module.__name__, input_modules)),
+			"module_params": input_module_params,
+			"execution_time": execution_time_list,
+			**{
+				metric: list(map(lambda result: result[metric].mean(), result_list))
+				for metric in strategies.get("metrics")
+			},
+		}
+	)
+	summary_df.to_csv(os.path.join(save_dir, "summary.csv"), index=False)
+	return summary_df
+
+
+def find_best(results, average_times, filenames, strategies):
+	# 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"),
+	)
+	return selected_result, selected_filename