AutoRAG 0.0.0__py3-none-any.whl

autorag/nodes/hybridretrieval/hybrid_cc.py

@@ -0,0 +1,227 @@
+from pathlib import Path
+from typing import Tuple, List, Union
+
+import numpy as np
+import pandas as pd
+
+from autorag.nodes.hybridretrieval.base import HybridRetrieval
+from autorag.nodes.hybridretrieval.run import evaluate_retrieval_node
+from autorag.strategy import select_best
+
+
+def normalize_mm(scores: List[str], fixed_min_value: float = 0):
+	arr = np.array(scores)
+	max_value = np.max(arr)
+	min_value = np.min(arr)
+	norm_score = (arr - min_value) / (max_value - min_value)
+	return norm_score
+
+
+def normalize_tmm(scores: List[str], fixed_min_value: float):
+	arr = np.array(scores)
+	max_value = np.max(arr)
+	norm_score = (arr - fixed_min_value) / (max_value - fixed_min_value)
+	return norm_score
+
+
+def normalize_z(scores: List[str], fixed_min_value: float = 0):
+	arr = np.array(scores)
+	mean_value = np.mean(arr)
+	std_value = np.std(arr)
+	norm_score = (arr - mean_value) / std_value
+	return norm_score
+
+
+def normalize_dbsf(scores: List[str], fixed_min_value: float = 0):
+	arr = np.array(scores)
+	mean_value = np.mean(arr)
+	std_value = np.std(arr)
+	min_value = mean_value - 3 * std_value
+	max_value = mean_value + 3 * std_value
+	norm_score = (arr - min_value) / (max_value - min_value)
+	return norm_score
+
+
+normalize_method_dict = {
+	"mm": normalize_mm,
+	"tmm": normalize_tmm,
+	"z": normalize_z,
+	"dbsf": normalize_dbsf,
+}
+
+
+class HybridCC(HybridRetrieval):
+	def _pure(
+		self,
+		info: dict,
+		top_k: int,
+		weight: float,
+		normalize_method: str = "mm",
+		semantic_theoretical_min_value: float = -1.0,
+		lexical_theoretical_min_value: float = 0.0,
+	):
+		return hybrid_cc(
+			(info["retrieved_ids_semantic"], info["retrieved_ids_lexical"]),
+			(info["retrieve_scores_semantic"], info["retrieve_scores_lexical"]),
+			top_k,
+			weight,
+			normalize_method,
+			semantic_theoretical_min_value,
+			lexical_theoretical_min_value,
+		)
+
+	@classmethod
+	def run_evaluator(
+		cls,
+		project_dir: Union[str, Path],
+		previous_result: pd.DataFrame,
+		*args,
+		**kwargs,
+	):
+		assert "strategy" in kwargs, "You must specify the strategy to use."
+		assert "input_metrics" in kwargs, (
+			"You must specify the input metrics to use, which is list of MetricInput."
+		)
+		strategies = kwargs.pop("strategy")
+		input_metrics = kwargs.pop("input_metrics")
+		weight_range = kwargs.pop("weight_range", (0.0, 1.0))
+		test_weight_size = kwargs.pop("test_weight_size", 101)
+		weight_candidates = np.linspace(
+			weight_range[0], weight_range[1], test_weight_size
+		).tolist()
+
+		result_list = []
+		instance = cls(project_dir, *args, **kwargs)
+		for weight_value in weight_candidates:
+			result_df = instance.pure(previous_result, weight=weight_value, **kwargs)
+			result_list.append(result_df)
+
+		if strategies.get("metrics") is None:
+			raise ValueError("You must at least one metrics for retrieval evaluation.")
+		result_list = list(
+			map(
+				lambda x: evaluate_retrieval_node(
+					x,
+					input_metrics,
+					strategies.get("metrics"),
+				),
+				result_list,
+			)
+		)
+
+		# select best result
+		best_result_df, best_weight = select_best(
+			result_list,
+			strategies.get("metrics"),
+			metadatas=weight_candidates,
+			strategy_name=strategies.get("strategy", "normalize_mean"),
+		)
+		return {
+			"best_result": best_result_df,
+			"best_weight": best_weight,
+		}
+
+
+def hybrid_cc(
+	ids: Tuple,
+	scores: Tuple,
+	top_k: int,
+	weight: float,
+	normalize_method: str = "mm",
+	semantic_theoretical_min_value: float = -1.0,
+	lexical_theoretical_min_value: float = 0.0,
+) -> Tuple[List[List[str]], List[List[float]]]:
+	"""
+	Hybrid CC function.
+	CC (convex combination) is a method to fuse lexical and semantic retrieval results.
+	It is a method that first normalizes the scores of each retrieval result,
+	and then combines them with the given weights.
+	It is uniquer than other retrieval modules, because it does not really execute retrieval,
+	but just fuse the results of other retrieval functions.
+	So you have to run more than two retrieval modules before running this function.
+	And collect ids and scores result from each retrieval module.
+	Make it as tuple and input it to this function.
+
+	:param ids: The tuple of ids that you want to fuse.
+	    The length of this must be the same as the length of scores.
+	    The semantic retrieval ids must be the first index.
+	:param scores: The retrieve scores that you want to fuse.
+	    The length of this must be the same as the length of ids.
+	    The semantic retrieval scores must be the first index.
+	:param top_k: The number of passages to be retrieved.
+	:param normalize_method: The normalization method to use.
+	  There are some normalization method that you can use at the hybrid cc method.
+	  AutoRAG support following.
+	    - `mm`: Min-max scaling
+	    - `tmm`: Theoretical min-max scaling
+	    - `z`: z-score normalization
+	    - `dbsf`: 3-sigma normalization
+	:param weight: The weight value. If the weight is 1.0, it means the
+	  weight to the semantic module will be 1.0 and weight to the lexical module will be 0.0.
+	:param semantic_theoretical_min_value: This value used by `tmm` normalization method. You can set the
+	    theoretical minimum value by yourself. Default is -1.
+	:param lexical_theoretical_min_value: This value used by `tmm` normalization method. You can set the
+	    theoretical minimum value by yourself. Default is 0.
+	:return: The tuple of ids and fused scores that fused by CC. Plus, the third element is selected weight value.
+	"""
+	assert len(ids) == len(scores), "The length of ids and scores must be the same."
+	assert len(ids) > 1, "You must input more than one retrieval results."
+	assert top_k > 0, "top_k must be greater than 0."
+	assert weight >= 0, "The weight must be greater than 0."
+	assert weight <= 1, "The weight must be less than 1."
+
+	df = pd.DataFrame(
+		{
+			"semantic_ids": ids[0],
+			"lexical_ids": ids[1],
+			"semantic_score": scores[0],
+			"lexical_score": scores[1],
+		}
+	)
+
+	def cc_pure_apply(row):
+		return fuse_per_query(
+			row["semantic_ids"],
+			row["lexical_ids"],
+			row["semantic_score"],
+			row["lexical_score"],
+			normalize_method=normalize_method,
+			weight=weight,
+			top_k=top_k,
+			semantic_theoretical_min_value=semantic_theoretical_min_value,
+			lexical_theoretical_min_value=lexical_theoretical_min_value,
+		)
+
+	# fixed weight
+	df[["cc_id", "cc_score"]] = df.apply(
+		lambda row: cc_pure_apply(row), axis=1, result_type="expand"
+	)
+	return df["cc_id"].tolist(), df["cc_score"].tolist()
+
+
+def fuse_per_query(
+	semantic_ids: List[str],
+	lexical_ids: List[str],
+	semantic_scores: List[float],
+	lexical_scores: List[float],
+	normalize_method: str,
+	weight: float,
+	top_k: int,
+	semantic_theoretical_min_value: float,
+	lexical_theoretical_min_value: float,
+):
+	normalize_func = normalize_method_dict[normalize_method]
+	norm_semantic_scores = normalize_func(
+		semantic_scores, semantic_theoretical_min_value
+	)
+	norm_lexical_scores = normalize_func(lexical_scores, lexical_theoretical_min_value)
+	ids = [semantic_ids, lexical_ids]
+	scores = [norm_semantic_scores, norm_lexical_scores]
+	df = pd.concat(
+		[pd.Series(dict(zip(_id, score))) for _id, score in zip(ids, scores)], axis=1
+	)
+	df.columns = ["semantic", "lexical"]
+	df = df.fillna(0)
+	df["weighted_sum"] = df.mul((weight, 1.0 - weight)).sum(axis=1)
+	df = df.sort_values(by="weighted_sum", ascending=False)
+	return df.index.tolist()[:top_k], df["weighted_sum"][:top_k].tolist()

autorag/nodes/hybridretrieval/hybrid_rrf.py

@@ -0,0 +1,149 @@
+from pathlib import Path
+from typing import List, Tuple, Union
+
+import numpy as np
+import pandas as pd
+
+from autorag.nodes.hybridretrieval.base import HybridRetrieval
+from autorag.nodes.hybridretrieval.run import evaluate_retrieval_node
+from autorag.strategy import select_best
+
+
+class HybridRRF(HybridRetrieval):
+	def _pure(self, info: dict, top_k: int, weight: int = 60, rrf_k: int = -1):
+		return hybrid_rrf(
+			(info["retrieved_ids_semantic"], info["retrieved_ids_lexical"]),
+			(info["retrieve_scores_semantic"], info["retrieve_scores_lexical"]),
+			top_k,
+			weight,
+			rrf_k,
+		)
+
+	@classmethod
+	def run_evaluator(
+		cls,
+		project_dir: Union[str, Path],
+		previous_result: pd.DataFrame,
+		*args,
+		**kwargs,
+	):
+		assert "strategy" in kwargs, "You must specify the strategy to use."
+		assert "input_metrics" in kwargs, (
+			"You must specify the input metrics to use, which is list of MetricInput."
+		)
+		strategies = kwargs.pop("strategy")
+		input_metrics = kwargs.pop("input_metrics")
+		weight_range = kwargs.pop("weight_range", (4, 80))
+		test_weight_size = weight_range[1] - weight_range[0] + 1
+		weight_candidates = np.linspace(
+			weight_range[0], weight_range[1], test_weight_size
+		).tolist()
+
+		result_list = []
+		instance = cls(project_dir, *args, **kwargs)
+		for weight_value in weight_candidates:
+			result_df = instance.pure(previous_result, weight=weight_value, **kwargs)
+			result_list.append(result_df)
+
+		if strategies.get("metrics") is None:
+			raise ValueError("You must at least one metrics for retrieval evaluation.")
+		result_list = list(
+			map(
+				lambda x: evaluate_retrieval_node(
+					x,
+					input_metrics,
+					strategies.get("metrics"),
+				),
+				result_list,
+			)
+		)
+
+		# select best result
+		best_result_df, best_weight = select_best(
+			result_list,
+			strategies.get("metrics"),
+			metadatas=weight_candidates,
+			strategy_name=strategies.get("strategy", "normalize_mean"),
+		)
+		return {
+			"best_result": best_result_df,
+			"best_weight": best_weight,
+		}
+
+
+def hybrid_rrf(
+	ids: Tuple,
+	scores: Tuple,
+	top_k: int,
+	weight: int = 60,
+	rrf_k: int = -1,
+) -> Tuple[List[List[str]], List[List[float]]]:
+	"""
+	Hybrid RRF function.
+	RRF (Rank Reciprocal Fusion) is a method to fuse multiple retrieval results.
+	It is common to fuse dense retrieval and sparse retrieval results using RRF.
+	To use this function, you must input ids and scores as tuple.
+	It is more unique than other retrieval modules because it does not really execute retrieval but just fuses
+	the results of other retrieval functions.
+	So you have to run more than two retrieval modules before running this function.
+	And collect ids and scores result from each retrieval module.
+	Make it as a tuple and input it to this function.
+
+	:param ids: The tuple of ids that you want to fuse.
+	    The length of this must be the same as the length of scores.
+	:param scores: The retrieve scores that you want to fuse.
+	    The length of this must be the same as the length of ids.
+	:param top_k: The number of passages to be retrieved.
+	:param weight: Hyperparameter for RRF.
+	    It was originally rrf_k value.
+	    Default is 60.
+	    For more information, please visit our documentation.
+	:param rrf_k: (Deprecated) Hyperparameter for RRF.
+	    It was originally rrf_k value. Will remove at a further version.
+	:return: The tuple of ids and fused scores that are fused by RRF.
+	"""
+	assert len(ids) == len(scores), "The length of ids and scores must be the same."
+	assert len(ids) > 1, "You must input more than one retrieval results."
+	assert top_k > 0, "top_k must be greater than 0."
+	assert weight > 0, "rrf_k must be greater than 0."
+
+	if rrf_k != -1:
+		weight = int(rrf_k)
+	else:
+		weight = int(weight)
+
+	id_df = pd.DataFrame({f"id_{i}": id_list for i, id_list in enumerate(ids)})
+	score_df = pd.DataFrame(
+		{f"score_{i}": score_list for i, score_list in enumerate(scores)}
+	)
+	df = pd.concat([id_df, score_df], axis=1)
+
+	def rrf_pure_apply(row):
+		ids_tuple = tuple(row[[f"id_{i}" for i in range(len(ids))]].values)
+		scores_tuple = tuple(row[[f"score_{i}" for i in range(len(scores))]].values)
+		return pd.Series(rrf_pure(ids_tuple, scores_tuple, weight, top_k))
+
+	df[["rrf_id", "rrf_score"]] = df.apply(rrf_pure_apply, axis=1)
+	return df["rrf_id"].tolist(), df["rrf_score"].tolist()
+
+
+def rrf_pure(
+	ids: Tuple, scores: Tuple, rrf_k: int, top_k: int
+) -> Tuple[List[str], List[float]]:
+	df = pd.concat(
+		[pd.Series(dict(zip(_id, score))) for _id, score in zip(ids, scores)], axis=1
+	)
+	rank_df = df.rank(ascending=False, method="min")
+	rank_df = rank_df.fillna(0)
+	rank_df["rrf"] = rank_df.apply(lambda row: rrf_calculate(row, rrf_k), axis=1)
+	rank_df = rank_df.sort_values(by="rrf", ascending=False)
+	return rank_df.index.tolist()[:top_k], rank_df["rrf"].tolist()[:top_k]
+
+
+def rrf_calculate(row, rrf_k):
+	result = 0
+	for r in row:
+		if r == 0:
+			continue
+		result += 1 / (r + rrf_k)
+	return result

autorag/nodes/hybridretrieval/run.py

@@ -0,0 +1,137 @@
+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_hybrid_retrieval_node(
+	modules: List,
+	module_params: List[Dict],
+	previous_result: pd.DataFrame,
+	node_line_dir: str,
+	strategies: Dict,
+) -> pd.DataFrame:
+	"""
+	Run the hybrid retrieval node with the given modules and parameters.
+
+	:param modules: List of hybrid retrieval modules to run.
+	:param module_params: List of parameters for each module.
+	:param previous_result: DataFrame containing the previous results.
+	:param node_line_dir: Directory where the results will be saved.
+	:param strategies: Dictionary containing strategies for filtering and evaluation.
+	:return: DataFrame containing the results of the hybrid retrieval.
+	"""
+	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, "hybrid_retrieval")  # node name
+	if not os.path.exists(save_dir):
+		os.makedirs(save_dir)
+
+	results, execution_times = [], []
+	for module, module_param in zip(modules, module_params):
+		result, speed = measure_speed(
+			module.run_evaluator,
+			project_dir=project_dir,
+			previous_result=previous_result,
+			strategy=strategies,
+			input_metrics=metric_inputs,
+			**module_param,
+		)
+		if "best_weight" in result.keys():
+			module_param["weight"] = result["best_weight"]
+		results.append(result["best_result"])
+		execution_times.append(speed)
+
+	hybrid_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 retrieval evaluation.")
+	results = list(
+		map(
+			lambda x: evaluate_retrieval_node(
+				x,
+				metric_inputs,
+				strategies.get("metrics"),
+			),
+			results,
+		)
+	)
+
+	summary_df = save_and_summary(
+		modules,
+		module_params,
+		results,
+		hybrid_times,
+		0,
+		save_dir,
+		strategies,
+	)
+	selected_result, selected_filename = find_best(
+		results, hybrid_times, summary_df["filename"].tolist(), strategies
+	)
+	summary_df["is_best"] = summary_df["filename"] == selected_filename
+	previous_result.drop(
+		columns=list(RETRIEVAL_METRIC_FUNC_DICT.keys()), inplace=True, errors="ignore"
+	)
+	best_result = pd.concat([previous_result, selected_result], axis=1)
+	best_result.to_parquet(
+		os.path.join(
+			save_dir, f"best_{os.path.splitext(selected_filename)[0]}.parquet"
+		),
+		index=False,
+	)
+	summary_df.to_csv(os.path.join(save_dir, "summary.csv"), index=False)
+	return best_result
+
+
+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)

autorag/nodes/lexicalretrieval/__init__.py

@@ -0,0 +1 @@
+from .bm25 import BM25