AutoRAG 0.0.0__py3-none-any.whl

autorag/nodes/passagereranker/voyageai.py

@@ -0,0 +1,109 @@
+import os
+from typing import List, Tuple
+import pandas as pd
+import voyageai
+
+from autorag.nodes.passagereranker.base import BasePassageReranker
+from autorag.utils.util import result_to_dataframe, get_event_loop, process_batch
+
+
+class VoyageAIReranker(BasePassageReranker):
+	def __init__(self, project_dir: str, *args, **kwargs):
+		super().__init__(project_dir)
+		api_key = kwargs.pop("api_key", None)
+		api_key = os.getenv("VOYAGE_API_KEY", None) if api_key is None else api_key
+		if api_key is None:
+			raise KeyError(
+				"Please set the API key for VoyageAI rerank in the environment variable VOYAGE_API_KEY "
+				"or directly set it on the config YAML file."
+			)
+
+		self.voyage_client = voyageai.AsyncClient(api_key=api_key)
+
+	def __del__(self):
+		del self.voyage_client
+		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.pop("top_k")
+		batch = kwargs.pop("batch", 8)
+		model = kwargs.pop("model", "rerank-2")
+		truncation = kwargs.pop("truncation", True)
+		return self._pure(queries, contents, ids, top_k, model, batch, truncation)
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents_list: List[List[str]],
+		ids_list: List[List[str]],
+		top_k: int,
+		model: str = "rerank-2",
+		batch: int = 8,
+		truncation: bool = True,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Rerank a list of contents with VoyageAI rerank models.
+		You can get the API key from https://docs.voyageai.com/docs/api-key-and-installation and set it in the environment variable VOYAGE_API_KEY.
+
+		: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 model: The model name for VoyageAI rerank.
+		    You can choose between "rerank-2" and "rerank-2-lite".
+		    Default is "rerank-2".
+		:param batch: The number of queries to be processed in a batch
+		:param truncation: Whether to truncate the input to satisfy the 'context length limit' on the query and the documents.
+		:return: Tuple of lists containing the reranked contents, ids, and scores
+		"""
+		tasks = [
+			voyageai_rerank_pure(
+				self.voyage_client, model, query, contents, ids, top_k, truncation
+			)
+			for query, contents, ids in zip(queries, contents_list, ids_list)
+		]
+		loop = get_event_loop()
+		results = loop.run_until_complete(process_batch(tasks, batch))
+
+		content_result, id_result, score_result = zip(*results)
+
+		return list(content_result), list(id_result), list(score_result)
+
+
+async def voyageai_rerank_pure(
+	voyage_client: voyageai.AsyncClient,
+	model: str,
+	query: str,
+	documents: List[str],
+	ids: List[str],
+	top_k: int,
+	truncation: bool = True,
+) -> Tuple[List[str], List[str], List[float]]:
+	"""
+	Rerank a list of contents with VoyageAI rerank models.
+
+	:param voyage_client: The Voyage Client to use for reranking
+	:param model: The model name for VoyageAI rerank
+	:param query: The query to use for reranking
+	:param documents: The list of contents to rerank
+	:param ids: The list of ids corresponding to the documents
+	:param top_k: The number of passages to be retrieved
+	:param truncation: Whether to truncate the input to satisfy the 'context length limit' on the query and the documents.
+	:return: Tuple of lists containing the reranked contents, ids, and scores
+	"""
+	rerank_results = await voyage_client.rerank(
+		model=model,
+		query=query,
+		documents=documents,
+		top_k=top_k,
+		truncation=truncation,
+	)
+	reranked_scores: List[float] = list(
+		map(lambda x: x.relevance_score, rerank_results.results)
+	)
+	indices = list(map(lambda x: x.index, rerank_results.results))
+	reranked_contents: List[str] = list(map(lambda i: documents[i], indices))
+	reranked_ids: List[str] = list(map(lambda i: ids[i], indices))
+	return reranked_contents, reranked_ids, reranked_scores

autorag/nodes/promptmaker/__init__.py

@@ -0,0 +1,12 @@
+from .long_context_reorder import LongContextReorder
+from .window_replacement import WindowReplacement
+from .fstring import Fstring
+from .chat_fstring import ChatFstring
+
+
+__all__ = [
+	"Fstring",
+	"ChatFstring",
+	"LongContextReorder",
+	"WindowReplacement",
+]

autorag/nodes/promptmaker/base.py

@@ -0,0 +1,32 @@
+import logging
+from abc import ABCMeta
+from pathlib import Path
+from typing import Union
+
+import pandas as pd
+
+from autorag.schema.base import BaseModule
+from autorag.utils.cast import cast_retrieved_contents
+
+logger = logging.getLogger("AutoRAG")
+
+
+class BasePromptMaker(BaseModule, metaclass=ABCMeta):
+	def __init__(self, project_dir: Union[str, Path], *args, **kwargs):
+		logger.info(
+			f"Initialize prompt maker node - {self.__class__.__name__} module..."
+		)
+
+	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 prompt maker node - {self.__class__.__name__} module...")
+		# get query and retrieved contents from previous_result
+		assert "query" in previous_result.columns, (
+			"previous_result must have query column."
+		)
+
+		query = previous_result["query"].tolist()
+		prompt = kwargs.pop("prompt")
+		return query, cast_retrieved_contents(previous_result), prompt

autorag/nodes/promptmaker/chat_fstring.py

@@ -0,0 +1,73 @@
+import copy
+from typing import List
+
+import pandas as pd
+
+from autorag.nodes.promptmaker.base import BasePromptMaker
+from autorag.utils import result_to_dataframe
+
+
+class ChatFstring(BasePromptMaker):
+	@result_to_dataframe(["prompts"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		query, retrieved_contents, prompt = self.cast_to_run(
+			previous_result, *args, **kwargs
+		)
+		return self._pure(prompt, query, retrieved_contents)
+
+	def _pure(
+		self,
+		prompt: List[dict[str, str]],
+		queries: List[str],
+		retrieved_contents: List[List[str]],
+	) -> List[List[dict[str, str]]]:
+		"""
+		Make a prompt using f-string from a query and retrieved_contents.
+		You must type a prompt or prompt list at a config YAML file like this:
+
+		.. Code:: yaml
+		nodes:
+		- node_type: prompt_maker
+		  modules:
+		  - module_type: chatfstring
+			prompt:
+			- - role: system
+			  content: You are a helpful assistant that helps people find information.
+			  - role: user
+			    content: |
+			    Answer this question: {query}
+			    {retrieved_contents}
+			- - role: system
+			    content: You are helpful.
+			  - role: user
+				content: |
+			    Read the passages carefully and answer this question: {query}
+
+			    Passages: {retrieved_contents}
+
+		:param prompt: A prompt string.
+		:param queries: List of query strings.
+		:param retrieved_contents: List of retrieved contents.
+		:return: Prompts that are made by chat f-string.
+			It is the list of OpenAI chat format prompts.
+		"""
+
+		def fstring_row(
+			_prompt: List[dict[str, str]], _query: str, _retrieved_contents: List[str]
+		) -> List[dict[str, str]]:
+			contents_str = "\n\n".join(_retrieved_contents)
+			result_prompt = copy.deepcopy(_prompt)
+			for lst in result_prompt:
+				if "content" in lst:
+					lst["content"] = lst["content"].format(
+						query=_query, retrieved_contents=contents_str
+					)
+
+			return result_prompt
+
+		return list(
+			map(
+				lambda x: fstring_row(prompt, x[0], x[1]),
+				zip(queries, retrieved_contents),
+			)
+		)

autorag/nodes/promptmaker/fstring.py

@@ -0,0 +1,49 @@
+from typing import List
+
+import pandas as pd
+
+from autorag.nodes.promptmaker.base import BasePromptMaker
+from autorag.utils import result_to_dataframe
+
+
+class Fstring(BasePromptMaker):
+	@result_to_dataframe(["prompts"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		query, retrieved_contents, prompt = self.cast_to_run(
+			previous_result, *args, **kwargs
+		)
+		return self._pure(prompt, query, retrieved_contents)
+
+	def _pure(
+		self, prompt: str, queries: List[str], retrieved_contents: List[List[str]]
+	) -> List[str]:
+		"""
+		Make a prompt using f-string from a query and retrieved_contents.
+		You must type a prompt or prompt list at a config YAML file like this:
+
+		.. Code:: yaml
+		nodes:
+		- node_type: prompt_maker
+		  modules:
+		  - module_type: fstring
+			prompt: [Answer this question: {query} \n\n {retrieved_contents},
+			Read the passages carefully and answer this question: {query} \n\n Passages: {retrieved_contents}]
+
+		:param prompt: A prompt string.
+		:param queries: List of query strings.
+		:param retrieved_contents: List of retrieved contents.
+		:return: Prompts that are made by f-string.
+		"""
+
+		def fstring_row(
+			_prompt: str, _query: str, _retrieved_contents: List[str]
+		) -> str:
+			contents_str = "\n\n".join(_retrieved_contents)
+			return _prompt.format(query=_query, retrieved_contents=contents_str)
+
+		return list(
+			map(
+				lambda x: fstring_row(prompt, x[0], x[1]),
+				zip(queries, retrieved_contents),
+			)
+		)

autorag/nodes/promptmaker/long_context_reorder.py

@@ -0,0 +1,83 @@
+import logging
+from typing import List
+
+import numpy as np
+import pandas as pd
+
+from autorag.nodes.promptmaker.base import BasePromptMaker
+from autorag.utils import result_to_dataframe
+
+logger = logging.getLogger("AutoRAG")
+
+
+class LongContextReorder(BasePromptMaker):
+	@result_to_dataframe(["prompts"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		query, retrieved_contents, prompt = self.cast_to_run(
+			previous_result, *args, **kwargs
+		)
+		assert "retrieve_scores" in previous_result.columns, (
+			"previous_result must have retrieve_scores column."
+		)
+		retrieve_scores = previous_result["retrieve_scores"].tolist()
+		return self._pure(prompt, query, retrieved_contents, retrieve_scores)
+
+	def _pure(
+		self,
+		prompt: str,
+		queries: List[str],
+		retrieved_contents: List[List[str]],
+		retrieve_scores: List[List[float]],
+	) -> List[str]:
+		"""
+		Models struggle to access significant details found
+		in the center of extended contexts. A study
+		(https://arxiv.org/abs/2307.03172) observed that the best
+		performance typically arises when crucial data is positioned
+		at the start or conclusion of the input context. Additionally,
+		as the input context lengthens, performance drops notably, even
+		in models designed for long contexts."
+
+		.. Code:: yaml
+		nodes:
+		- node_type: prompt_maker
+		  modules:
+		  - module_type: long_context_reorder
+		    prompt: [Answer this question: {query} \n\n {retrieved_contents},
+		    Read the passages carefully and answer this question: {query} \n\n Passages: {retrieved_contents}]
+
+		:param prompt: A prompt string.
+		:param queries: List of query strings.
+		:param retrieved_contents: List of retrieved contents.
+		:param retrieve_scores: List of `retrieve scores`.
+		:return: Prompts that are made by long context reorder.
+		"""
+
+		def long_context_reorder_row(
+			_prompt: str,
+			_query: str,
+			_retrieved_contents: List[str],
+			_retrieve_scores: List[float],
+		) -> str:
+			if isinstance(_retrieved_contents, np.ndarray):
+				_retrieved_contents = _retrieved_contents.tolist()
+			if not len(_retrieved_contents) == len(_retrieve_scores):
+				logger.info("If you use a summarizer, the reorder will not proceed.")
+				return _prompt.format(
+					query=_query, retrieved_contents="\n\n".join(_retrieved_contents)
+				)
+			content_scores = list(zip(_retrieved_contents, _retrieve_scores))
+			sorted_content_scores = sorted(
+				content_scores, key=lambda x: x[1], reverse=True
+			)
+			content_result, score_result = zip(*sorted_content_scores)
+			_retrieved_contents.append(content_result[0])
+			contents_str = "\n\n".join(_retrieved_contents)
+			return _prompt.format(query=_query, retrieved_contents=contents_str)
+
+		return list(
+			map(
+				lambda x: long_context_reorder_row(prompt, x[0], x[1], x[2]),
+				zip(queries, retrieved_contents, retrieve_scores),
+			)
+		)

autorag/nodes/promptmaker/run.py

@@ -0,0 +1,283 @@
+import os
+import pathlib
+from copy import deepcopy
+from typing import List, Dict, Optional, Union
+
+import pandas as pd
+import tokenlog
+
+from autorag.evaluation import evaluate_generation
+from autorag.evaluation.util import cast_metrics
+from autorag.schema.metricinput import MetricInput
+from autorag.strategy import measure_speed, filter_by_threshold, select_best
+from autorag.support import get_support_modules
+from autorag.utils import validate_qa_dataset
+from autorag.utils.util import make_combinations, explode, split_dataframe, to_list
+
+
+def run_prompt_maker_node(
+	modules: List,
+	module_params: List[Dict],
+	previous_result: pd.DataFrame,
+	node_line_dir: str,
+	strategies: Dict,
+) -> pd.DataFrame:
+	"""
+	Run prompt maker node.
+	With this function, you can select the best prompt maker module.
+	As default, when you can use only one module, the evaluation will be skipped.
+	If you want to select the best prompt among modules, you can use strategies.
+	When you use them, you must pass 'generator_modules' and its parameters at strategies.
+	Because it uses generator modules and generator metrics for evaluation this module.
+	It is recommended to use one params and modules for evaluation,
+	but you can use multiple params and modules for evaluation.
+	When you don't set generator module at strategies, it will use the default generator module.
+	The default generator module is llama_index_llm with openai gpt-3.5-turbo model.
+
+	:param modules: Prompt maker module classes to run.
+	:param module_params: Prompt maker 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 prompt maker node.
+	:return: The best result dataframe.
+	    It contains previous result columns and prompt maker's result columns which is 'prompts'.
+	"""
+	if not os.path.exists(node_line_dir):
+		os.makedirs(node_line_dir)
+	node_dir = os.path.join(node_line_dir, "prompt_maker")
+	if not os.path.exists(node_dir):
+		os.makedirs(node_dir)
+	project_dir = pathlib.PurePath(node_line_dir).parent.parent
+
+	# 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),
+		)
+	)
+	average_times = list(map(lambda x: x / len(results[0]), execution_times))
+
+	# get average token usage
+	token_usages = []
+	for i, result in enumerate(results):
+		token_logger = tokenlog.getLogger(
+			f"prompt_maker_{i}", strategies.get("tokenizer", "openai/gpt-oss-20b")
+		)
+		token_logger.query_batch(result["prompts"].tolist())
+		token_usages.append(token_logger.get_token_usage() / len(result))
+
+	# save results to folder
+	filepaths = list(
+		map(lambda x: os.path.join(node_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,
+			"average_prompt_token": token_usages,
+		}
+	)
+
+	metric_names, metric_params = cast_metrics(strategies.get("metrics"))
+
+	# Run evaluation when there are more than one module.
+	if len(modules) > 1:
+		# pop general keys from strategies (e.g. metrics, speed_threshold)
+		general_key = ["metrics", "speed_threshold", "token_threshold", "tokenizer"]
+		general_strategy = dict(
+			filter(lambda x: x[0] in general_key, strategies.items())
+		)
+		extra_strategy = dict(
+			filter(lambda x: x[0] not in general_key, strategies.items())
+		)
+
+		# first, filter by threshold if it is enabled.
+		if general_strategy.get("speed_threshold") is not None:
+			results, filenames = filter_by_threshold(
+				results, average_times, general_strategy["speed_threshold"], filenames
+			)
+
+		# Calculate tokens and save to summary
+		if general_strategy.get("token_threshold") is not None:
+			results, filenames = filter_by_threshold(
+				results, token_usages, general_strategy["token_threshold"], filenames
+			)
+
+		# run metrics before filtering
+		if metric_names is None or len(metric_names) <= 0:
+			raise ValueError(
+				"You must at least one metrics for prompt maker evaluation."
+			)
+
+		# get generator modules from strategy
+		generator_callables, generator_params = make_generator_callable_params(
+			extra_strategy
+		)
+
+		# get generation_gt
+		qa_data = pd.read_parquet(
+			os.path.join(project_dir, "data", "qa.parquet"), engine="pyarrow"
+		)
+		validate_qa_dataset(qa_data)
+		generation_gt = qa_data["generation_gt"].tolist()
+		generation_gt = list(map(lambda x: x.tolist(), generation_gt))
+		queries = to_list(qa_data["query"].tolist())
+		metric_inputs = [
+			MetricInput(generation_gt=gen_gt, query=query)
+			for gen_gt, query in zip(generation_gt, queries)
+		]
+
+		all_prompts = []
+		for result in results:
+			all_prompts.extend(result["prompts"].tolist())
+
+		evaluation_result_all = evaluate_one_prompt_maker_node(
+			all_prompts,
+			generator_callables,
+			generator_params,
+			metric_inputs * len(results),
+			general_strategy["metrics"],
+			project_dir,
+			strategy_name=strategies.get("strategy", "mean"),
+		)
+		evaluation_results = split_dataframe(
+			evaluation_result_all, chunk_size=len(results[0])
+		)
+
+		evaluation_df = pd.DataFrame(
+			{
+				"filename": filenames,
+				**{
+					f"prompt_maker_{metric_name}": list(
+						map(lambda x: x[metric_name].mean(), evaluation_results)
+					)
+					for metric_name in metric_names
+				},
+			}
+		)
+		summary_df = pd.merge(
+			on="filename", left=summary_df, right=evaluation_df, how="left"
+		)
+
+		best_result, best_filename = select_best(
+			evaluation_results,
+			metric_names,
+			filenames,
+			strategies.get("strategy", "mean"),
+		)
+		# change metric name columns to prompt_maker_metric_name
+		best_result = best_result.rename(
+			columns={
+				metric_name: f"prompt_maker_{metric_name}"
+				for metric_name in metric_names
+			}
+		)
+		best_result = best_result.drop(columns=["generated_texts"])
+	else:
+		best_result, best_filename = results[0], filenames[0]
+
+	# add 'is_best' column at summary file
+	summary_df["is_best"] = summary_df["filename"] == best_filename
+
+	best_result = pd.concat([previous_result, best_result], axis=1)
+
+	# save files
+	summary_df.to_csv(os.path.join(node_dir, "summary.csv"), index=False)
+	best_result.to_parquet(
+		os.path.join(node_dir, f"best_{os.path.splitext(best_filename)[0]}.parquet"),
+		index=False,
+	)
+
+	return best_result
+
+
+def make_generator_callable_params(strategy_dict: Dict):
+	node_dict = deepcopy(strategy_dict)
+	generator_module_list: Optional[List[Dict]] = node_dict.pop(
+		"generator_modules", None
+	)
+	if generator_module_list is None:
+		generator_module_list = [
+			{
+				"module_type": "llama_index_llm",
+				"llm": "openai",
+				"model": "gpt-3.5-turbo",
+			}
+		]
+	node_params = node_dict
+	modules = list(
+		map(
+			lambda module_dict: get_support_modules(module_dict.pop("module_type")),
+			generator_module_list,
+		)
+	)
+	param_combinations = list(
+		map(
+			lambda module_dict: make_combinations({**module_dict, **node_params}),
+			generator_module_list,
+		)
+	)
+	return explode(modules, param_combinations)
+
+
+def evaluate_one_prompt_maker_node(
+	prompts: List[str],
+	generator_classes: List,
+	generator_params: List[Dict],
+	metric_inputs: List[MetricInput],
+	metrics: Union[List[str], List[Dict]],
+	project_dir,
+	strategy_name: str,
+) -> pd.DataFrame:
+	input_df = pd.DataFrame({"prompts": prompts})
+	generator_results = list(
+		map(
+			lambda x: x[0].run_evaluator(
+				project_dir=project_dir, previous_result=input_df, **x[1]
+			),
+			zip(generator_classes, generator_params),
+		)
+	)
+	evaluation_results = list(
+		map(
+			lambda x: evaluate_generator_result(x[0], metric_inputs, metrics),
+			zip(generator_results, generator_classes),
+		)
+	)
+	metric_names = (
+		list(map(lambda x: x["metric_name"], metrics))
+		if isinstance(metrics[0], dict)
+		else metrics
+	)
+	best_result, _ = select_best(
+		evaluation_results, metric_names, strategy_name=strategy_name
+	)
+	best_result = pd.concat([input_df, best_result], axis=1)
+	return best_result  # it has 'generated_texts' column
+
+
+def evaluate_generator_result(
+	result_df: pd.DataFrame,
+	metric_inputs: List[MetricInput],
+	metrics: Union[List[str], List[Dict]],
+) -> pd.DataFrame:
+	@evaluate_generation(metric_inputs=metric_inputs, metrics=metrics)
+	def evaluate(df):
+		return df["generated_texts"].tolist()
+
+	return evaluate(result_df)