AutoRAG 0.0.0__py3-none-any.whl

autorag/nodes/passagecompressor/run.py

@@ -0,0 +1,186 @@
+import os.path
+import pathlib
+from typing import List, Dict
+
+import pandas as pd
+
+from autorag.evaluation.metric import (
+	retrieval_token_recall,
+	retrieval_token_precision,
+	retrieval_token_f1,
+)
+from autorag.schema.metricinput import MetricInput
+from autorag.strategy import measure_speed, filter_by_threshold, select_best
+from autorag.utils.util import fetch_contents
+
+
+def run_passage_compressor_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 compressor modules.
+
+	:param modules: Passage compressor modules to run.
+	:param module_params: Passage compressor 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 compressor node.
+	    In this node, we use
+	    You can skip evaluation when you use only one module and a module parameter.
+	:return: The best result dataframe with previous result columns.
+	    This node will replace 'retrieved_contents' to compressed passages, so its length will be one.
+	"""
+	if not os.path.exists(node_line_dir):
+		os.makedirs(node_line_dir)
+	project_dir = pathlib.PurePath(node_line_dir).parent.parent
+	data_dir = os.path.join(project_dir, "data")
+	save_dir = os.path.join(node_line_dir, "passage_compressor")
+	if not os.path.exists(save_dir):
+		os.makedirs(save_dir)
+
+	# make retrieval contents gt
+	qa_data = pd.read_parquet(os.path.join(data_dir, "qa.parquet"), engine="pyarrow")
+	corpus_data = pd.read_parquet(
+		os.path.join(data_dir, "corpus.parquet"), engine="pyarrow"
+	)
+	# check qa_data have retrieval_gt
+	assert all(len(x[0]) > 0 for x in qa_data["retrieval_gt"].tolist()), (
+		"Can't use passage compressor if you don't have retrieval gt values in QA dataset."
+	)
+
+	# run modules
+	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),
+		)
+	)
+	results = list(results)
+	average_times = list(map(lambda x: x / len(results[0]), execution_times))
+
+	retrieval_gt_contents = list(
+		map(lambda x: fetch_contents(corpus_data, x), qa_data["retrieval_gt"].tolist())
+	)
+
+	metric_inputs = [
+		MetricInput(retrieval_gt_contents=ret_cont_gt)
+		for ret_cont_gt in retrieval_gt_contents
+	]
+
+	# run metrics before filtering
+	if strategies.get("metrics") is None:
+		raise ValueError(
+			"You must at least one metrics for retrieval contents evaluation."
+			"It can be 'retrieval_token_f1', 'retrieval_token_precision', 'retrieval_token_recall'."
+		)
+	results = list(
+		map(
+			lambda x: evaluate_passage_compressor_node(
+				x, metric_inputs, strategies.get("metrics")
+			),
+			results,
+		)
+	)
+
+	# save results to folder
+	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))
+
+	# make summary file
+	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_compressor_{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"),
+	)
+	new_retrieved_contents = selected_result["retrieved_contents"]
+	previous_result["retrieved_contents"] = new_retrieved_contents
+	selected_result = selected_result.drop(columns=["retrieved_contents"])
+	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
+
+	# add prefix 'passage_compressor' to best_result columns
+	best_result = best_result.rename(
+		columns={
+			metric_name: f"passage_compressor_{metric_name}"
+			for metric_name in strategies.get("metrics")
+		}
+	)
+
+	# save the result files
+	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_passage_compressor_node(
+	result_df: pd.DataFrame, metric_inputs: List[MetricInput], metrics: List[str]
+):
+	metric_funcs = {
+		retrieval_token_recall.__name__: retrieval_token_recall,
+		retrieval_token_precision.__name__: retrieval_token_precision,
+		retrieval_token_f1.__name__: retrieval_token_f1,
+	}
+	for metric_input, generated_text in zip(
+		metric_inputs, result_df["retrieved_contents"].tolist()
+	):
+		metric_input.retrieved_contents = generated_text
+	metrics = list(filter(lambda x: x in metric_funcs.keys(), metrics))
+	if len(metrics) <= 0:
+		raise ValueError(f"metrics must be one of {metric_funcs.keys()}")
+	metrics_scores = dict(
+		map(
+			lambda metric: (
+				metric,
+				metric_funcs[metric](
+					metric_inputs=metric_inputs,
+				),
+			),
+			metrics,
+		)
+	)
+	result_df = pd.concat([result_df, pd.DataFrame(metrics_scores)], axis=1)
+	return result_df

autorag/nodes/passagecompressor/tree_summarize.py

@@ -0,0 +1,56 @@
+from typing import List, Optional
+
+from llama_index.core import PromptTemplate
+from llama_index.core.prompts import PromptType
+from llama_index.core.prompts.utils import is_chat_model
+from llama_index.core.response_synthesizers import TreeSummarize as ts
+
+from autorag.nodes.passagecompressor.base import LlamaIndexCompressor
+from autorag.utils.util import get_event_loop, process_batch
+
+
+class TreeSummarize(LlamaIndexCompressor):
+	def _pure(
+		self,
+		queries: List[str],
+		contents: List[List[str]],
+		prompt: Optional[str] = None,
+		chat_prompt: Optional[str] = None,
+		batch: int = 16,
+	) -> List[str]:
+		"""
+		Recursively merge retrieved texts and summarizes them in a bottom-up fashion.
+		This function is a wrapper for llama_index.response_synthesizers.TreeSummarize.
+		For more information, visit https://docs.llamaindex.ai/en/latest/examples/response_synthesizers/tree_summarize.html.
+
+		:param queries: The queries for retrieved passages.
+		:param contents: The contents of retrieved passages.
+		:param prompt: The prompt template for summarization.
+		    If you want to use chat prompt, you should pass chat_prompt instead.
+		    At prompt, you must specify where to put 'context_str' and 'query_str'.
+		    Default is None. When it is None, it will use llama index default prompt.
+		:param chat_prompt: The chat prompt template for summarization.
+		    If you want to use normal prompt, you should pass prompt instead.
+		    At prompt, you must specify where to put 'context_str' and 'query_str'.
+		    Default is None. When it is None, it will use llama index default chat prompt.
+		:param batch: The batch size for llm.
+		    Set low if you face some errors.
+		    Default is 16.
+		:return: The list of compressed texts.
+		"""
+		if prompt is not None and not is_chat_model(self.llm):
+			summary_template = PromptTemplate(prompt, prompt_type=PromptType.SUMMARY)
+		elif chat_prompt is not None and is_chat_model(self.llm):
+			summary_template = PromptTemplate(
+				chat_prompt, prompt_type=PromptType.SUMMARY
+			)
+		else:
+			summary_template = None
+		summarizer = ts(llm=self.llm, summary_template=summary_template, use_async=True)
+		tasks = [
+			summarizer.aget_response(query, content)
+			for query, content in zip(queries, contents)
+		]
+		loop = get_event_loop()
+		results = loop.run_until_complete(process_batch(tasks, batch_size=batch))
+		return results

autorag/nodes/passagefilter/__init__.py

@@ -0,0 +1,6 @@
+from .pass_passage_filter import PassPassageFilter
+from .percentile_cutoff import PercentileCutoff
+from .recency import RecencyFilter
+from .similarity_percentile_cutoff import SimilarityPercentileCutoff
+from .similarity_threshold_cutoff import SimilarityThresholdCutoff
+from .threshold_cutoff import ThresholdCutoff

autorag/nodes/passagefilter/base.py

@@ -0,0 +1,40 @@
+import abc
+import logging
+from pathlib import Path
+from typing import Union
+
+import pandas as pd
+
+from autorag.schema.base import BaseModule
+from autorag.utils import validate_qa_dataset
+from autorag.utils.cast import cast_retrieve_infos
+
+logger = logging.getLogger("AutoRAG")
+
+
+class BasePassageFilter(BaseModule, metaclass=abc.ABCMeta):
+	def __init__(self, project_dir: Union[str, Path], *args, **kwargs):
+		logger.info(f"Initialize passage filter node - {self.__class__.__name__}")
+
+	def __del__(self):
+		logger.info(f"Prompt maker node - {self.__class__.__name__} module is deleted.")
+
+	def cast_to_run(self, previous_result: pd.DataFrame, *args, **kwargs):
+		logger.info(
+			f"Running passage filter 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/passagefilter/pass_passage_filter.py

@@ -0,0 +1,14 @@
+import pandas as pd
+
+from autorag.nodes.passagefilter.base import BasePassageFilter
+from autorag.utils import result_to_dataframe
+
+
+class PassPassageFilter(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 contents, ids, scores
+
+	def _pure(self, *args, **kwargs):
+		pass

autorag/nodes/passagefilter/percentile_cutoff.py

@@ -0,0 +1,58 @@
+from typing import List, Tuple
+
+import pandas as pd
+
+from autorag.nodes.passagefilter.base import BasePassageFilter
+from autorag.utils.util import sort_by_scores, select_top_k, result_to_dataframe
+
+
+class PercentileCutoff(BasePassageFilter):
+	@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)
+		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]],
+		percentile: float,
+		reverse: bool = False,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Filter out the contents that are below the content's length times percentile.
+		If This is a filter and does not override scores.
+		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 reverse: If True, the lower the score, the better
+		    Default is False.
+		:return: Tuple of lists containing the filtered contents, ids, and scores
+		"""
+		num_top_k = max(1, int(len(scores_list[0]) * percentile))
+
+		df = pd.DataFrame(
+			{
+				"contents": contents_list,
+				"ids": ids_list,
+				"scores": scores_list,
+			}
+		)
+
+		reverse = not reverse
+		df[["contents", "ids", "scores"]] = df.apply(
+			sort_by_scores, axis=1, result_type="expand", reverse=reverse
+		)
+		results = select_top_k(df, ["contents", "ids", "scores"], num_top_k)
+
+		return (
+			results["contents"].tolist(),
+			results["ids"].tolist(),
+			results["scores"].tolist(),
+		)

autorag/nodes/passagefilter/recency.py

@@ -0,0 +1,105 @@
+import logging
+import os
+from datetime import datetime, date
+from pathlib import Path
+from typing import List, Tuple, Union
+
+import pandas as pd
+
+from autorag.nodes.passagefilter.base import BasePassageFilter
+from autorag.utils import fetch_contents, result_to_dataframe
+
+logger = logging.getLogger("AutoRAG")
+
+
+class RecencyFilter(BasePassageFilter):
+	def __init__(self, project_dir: Union[str, Path], *args, **kwargs):
+		super().__init__(project_dir, *args, **kwargs)
+		self.corpus_df = pd.read_parquet(
+			os.path.join(project_dir, "data", "corpus.parquet"), engine="pyarrow"
+		)
+
+	@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, *args, **kwargs)
+		metadatas = fetch_contents(self.corpus_df, ids, column_name="metadata")
+		times = [
+			[time["last_modified_datetime"] for time in time_list]
+			for time_list in metadatas
+		]
+		return self._pure(contents, scores, ids, times, *args, **kwargs)
+
+	def _pure(
+		self,
+		contents_list: List[List[str]],
+		scores_list: List[List[float]],
+		ids_list: List[List[str]],
+		time_list: List[List[datetime]],
+		threshold_datetime: Union[datetime, date],
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Filter out the contents that are below the threshold datetime.
+		If all contents are filtered, keep the only one recency content.
+		If the threshold date format is incorrect, return the original contents.
+
+		: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 time_list: The list of lists of datetime retrieved
+		:param threshold_datetime: The threshold to cut off.
+			In recency filter, you have to use the datetime.datetime object or datetime.date object.
+			All you need to do is to set the date at your YAML file.
+			For example, you can write "2010-09-09 3:45:06" or "2010-09-09" in the YAML file.
+		:return: Tuple of lists containing the filtered contents, ids, and scores
+		"""
+		if not (
+			isinstance(threshold_datetime, datetime)
+			or isinstance(threshold_datetime, date)
+		):
+			raise ValueError(
+				f"Threshold should be a datetime object, but got {type(threshold_datetime)}"
+			)
+
+		if not isinstance(threshold_datetime, datetime):
+			threshold_datetime = datetime.combine(
+				threshold_datetime, datetime.min.time()
+			)
+
+		time_list = [
+			list(
+				map(
+					lambda t: datetime.combine(t, datetime.min.time())
+					if not isinstance(t, datetime)
+					else t,
+					time,
+				)
+			)
+			for time in time_list
+		]
+
+		def sort_row(contents, scores, ids, time, _datetime_threshold):
+			combined = list(zip(contents, scores, ids, time))
+			combined_filtered = [
+				item for item in combined if item[3] >= _datetime_threshold
+			]
+
+			if combined_filtered:
+				remain_contents, remain_scores, remain_ids, _ = zip(*combined_filtered)
+			else:
+				combined.sort(key=lambda x: x[3], reverse=True)
+				remain_contents, remain_scores, remain_ids, _ = zip(*combined[:1])
+
+			return list(remain_contents), list(remain_ids), list(remain_scores)
+
+		remain_contents_list, remain_ids_list, remain_scores_list = zip(
+			*map(
+				sort_row,
+				contents_list,
+				scores_list,
+				ids_list,
+				time_list,
+				[threshold_datetime] * len(contents_list),
+			)
+		)
+
+		return remain_contents_list, remain_ids_list, remain_scores_list

autorag/nodes/passagefilter/run.py

@@ -0,0 +1,138 @@
+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 to_list, apply_recursive
+
+
+def run_passage_filter_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 filter node results.
+
+	:param modules: Passage filter modules to run.
+	:param module_params: Passage filter module parameters.
+	:param previous_result: Previous result dataframe.
+	    Could be retrieval, reranker, passage filter 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 filter 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_filter 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_filter")  # 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_filter_{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"),
+	)
+	selected_result = selected_result.rename(
+		columns={
+			metric_name: f"passage_filter_{metric_name}"
+			for metric_name in strategies["metrics"]
+		}
+	)
+	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