AutoRAG 0.0.0__py3-none-any.whl

autorag/nodes/retrieval/base.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/nodes/retrieval/run_util.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

autorag/nodes/semanticretrieval/__init__.py

@@ -0,0 +1 @@
+from .vectordb import VectorDB

autorag/nodes/semanticretrieval/run.py

@@ -0,0 +1,148 @@
+import os
+import pathlib
+from typing import List, Dict, Union
+
+import pandas as pd
+
+from autorag.evaluation import evaluate_retrieval
+from autorag.evaluation.retrieval import RETRIEVAL_METRIC_FUNC_DICT
+from autorag.nodes.retrieval.run_util import save_and_summary, find_best
+from autorag.schema.metricinput import MetricInput
+from autorag.strategy import measure_speed
+from autorag.utils.util import apply_recursive, to_list
+
+
+def run_semantic_retrieval_node(
+	modules: List,
+	module_params: List[Dict],
+	previous_result: pd.DataFrame,
+	node_line_dir: str,
+	strategies: Dict,
+) -> pd.DataFrame:
+	"""
+	Run the semantic retrieval node.
+
+	:param modules: Retrieval modules to run.
+	:param module_params: Retrieval module parameters.
+	:param previous_result: Previous result dataframe.
+	    Could be query expansion's best result or qa data.
+	:param node_line_dir: This node line's directory.
+	:param strategies: Strategies for retrieval node.
+	:return: The best result dataframe.
+	    It contains previous result columns and retrieval node's 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()
+		)
+	]
+
+	save_dir = os.path.join(node_line_dir, "semantic_retrieval")
+	if not os.path.exists(save_dir):
+		os.makedirs(save_dir)
+
+	# Run the modules
+	semantic_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),
+		)
+	)
+	semantic_times = list(map(lambda x: x / len(semantic_results[0]), execution_times))
+
+	# run metrics
+	if strategies.get("metrics") is None:
+		raise ValueError("You must at least one metrics for retrieval evaluation.")
+	semantic_results = list(
+		map(
+			lambda x: evaluate_semantic_retrieval_node(
+				x,
+				metric_inputs,
+				strategies.get("metrics"),
+			),
+			semantic_results,
+		)
+	)
+
+	semantic_summary_df = save_and_summary(
+		modules,
+		module_params,
+		semantic_results,
+		semantic_times,
+		0,
+		save_dir,
+		strategies,
+	)
+	semantic_selected_result, semantic_selected_filename = find_best(
+		semantic_results,
+		semantic_times,
+		semantic_summary_df["filename"].tolist(),
+		strategies,
+	)
+	semantic_summary_df["is_best"] = (
+		semantic_summary_df["filename"] == semantic_selected_filename
+	)
+	previous_result.drop(
+		columns=list(RETRIEVAL_METRIC_FUNC_DICT.keys()), inplace=True, errors="ignore"
+	)
+	semantic_selected_result.rename(
+		columns={
+			"retrieved_contents": "retrieved_contents_semantic",
+			"retrieved_ids": "retrieved_ids_semantic",
+			"retrieve_scores": "retrieve_scores_semantic",
+		},
+		inplace=True,
+	)
+	best_result = pd.concat([previous_result, semantic_selected_result], axis=1)
+	best_result.to_parquet(
+		os.path.join(
+			save_dir, f"best_{os.path.splitext(semantic_selected_filename)[0]}.parquet"
+		),
+		index=False,
+	)
+	semantic_summary_df.to_csv(os.path.join(save_dir, "summary.csv"), index=False)
+	return best_result  # The result will be with _semantic suffix
+
+
+def evaluate_semantic_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_semantic', 'retrieved_ids_semantic', 'retrieve_scores_semantic', and metric names.
+	"""
+
+	@evaluate_retrieval(
+		metric_inputs=metric_inputs,
+		metrics=metrics,
+	)
+	def evaluate_this_module(df: pd.DataFrame):
+		return (
+			df["retrieved_contents_semantic"].tolist(),
+			df["retrieved_ids_semantic"].tolist(),
+			df["retrieve_scores_semantic"].tolist(),
+		)
+
+	return evaluate_this_module(result_df)