AutoRAG 0.0.0__py3-none-any.whl

autorag/nodes/passageaugmenter/prev_next_augmenter.py

@@ -0,0 +1,155 @@
+from typing import List, Union
+
+import numpy as np
+import pandas as pd
+
+from autorag.embedding.base import EmbeddingModel
+from autorag.evaluation.metric.util import calculate_cosine_similarity
+from autorag.nodes.passageaugmenter.base import BasePassageAugmenter
+from autorag.utils.util import (
+	filter_dict_keys,
+	fetch_contents,
+	embedding_query_content,
+	result_to_dataframe,
+	empty_cuda_cache,
+)
+
+
+class PrevNextPassageAugmenter(BasePassageAugmenter):
+	def __init__(
+		self,
+		project_dir: str,
+		embedding_model: Union[str, dict] = "openai",
+		*args,
+		**kwargs,
+	):
+		"""
+		Initialize the PrevNextPassageAugmenter module.
+
+		:param project_dir:
+		:param embedding_model: The embedding model name to use for calculating cosine similarity
+			Default is openai (text-embedding-ada-002)
+		:param kwargs:
+		"""
+		super().__init__(project_dir, *args, **kwargs)
+		slim_corpus_df = self.corpus_df[["doc_id", "metadata"]]
+		slim_corpus_df.loc[:, "metadata"] = slim_corpus_df["metadata"].apply(
+			filter_dict_keys, keys=["prev_id", "next_id"]
+		)
+		self.slim_corpus_df = slim_corpus_df
+
+		# init embedding model
+		self.embedding_model = EmbeddingModel.load(embedding_model)()
+
+	def __del__(self):
+		del self.embedding_model
+		empty_cuda_cache()
+		super().__del__()
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		"""
+		Run the passage augmenter node - PrevNextPassageAugmenter module.
+
+		:param previous_result: The previous result Dataframe.
+		:param top_k: You must input the top_k value to get the top k results.
+		:param kwargs: Not affected.
+		:return: DataFrame with retrieved_contents, retrieved_ids, and retrieve_scores columns
+		"""
+		top_k = kwargs.pop("top_k")
+
+		ids = self.cast_to_run(previous_result)
+		# find queries columns
+		assert "query" in previous_result.columns, (
+			"previous_result must have query column."
+		)
+		queries = previous_result["query"].tolist()
+
+		mode = kwargs.pop("mode", "both")
+		num_passages = kwargs.pop("num_passages", 1)
+		augmented_ids = self._pure(ids, num_passages, mode)
+
+		# fetch contents from corpus to use augmented ids
+		augmented_contents = fetch_contents(self.corpus_df, augmented_ids)
+
+		query_embeddings, contents_embeddings = embedding_query_content(
+			queries, augmented_contents, self.embedding_model, batch=128
+		)
+
+		# get scores from calculated cosine similarity
+		augmented_scores = [
+			np.array(
+				[
+					calculate_cosine_similarity(query_embedding, x)
+					for x in content_embeddings
+				]
+			).tolist()
+			for query_embedding, content_embeddings in zip(
+				query_embeddings, contents_embeddings
+			)
+		]
+		return self.sort_by_scores(
+			augmented_contents, augmented_ids, augmented_scores, top_k
+		)
+
+	def _pure(
+		self,
+		ids_list: List[List[str]],
+		num_passages: int = 1,
+		mode: str = "both",
+	) -> List[List[str]]:
+		"""
+		Add passages before and/or after the retrieved passage.
+		For more information, visit https://docs.llamaindex.ai/en/stable/examples/node_postprocessor/PrevNextPostprocessorDemo/.
+
+		:param ids_list: The list of lists of ids retrieved
+		:param num_passages: The number of passages to add before and after the retrieved passage
+		    Default is 1.
+		:param mode: The mode of augmentation
+		    'prev': add passages before the retrieved passage
+		    'next': add passages after the retrieved passage
+		    'both': add passages before and after the retrieved passage
+		    Default is 'next'.
+		:return: The list of lists of augmented ids
+		"""
+		if mode not in ["prev", "next", "both"]:
+			raise ValueError(f"mode must be 'prev', 'next', or 'both', but got {mode}")
+
+		augmented_ids = [
+			(
+				lambda ids: prev_next_augmenter_pure(
+					ids, self.slim_corpus_df, mode, num_passages
+				)
+			)(ids)
+			for ids in ids_list
+		]
+
+		return augmented_ids
+
+
+def prev_next_augmenter_pure(
+	ids: List[str], corpus_df: pd.DataFrame, mode: str, num_passages: int
+):
+	def fetch_id_sequence(start_id, key):
+		sequence = []
+		current_id = start_id
+		for _ in range(num_passages):
+			current_id = (
+				corpus_df.loc[corpus_df["doc_id"] == current_id]["metadata"]
+				.values[0]
+				.get(key)
+			)
+			if current_id is None:
+				break
+			sequence.append(current_id)
+		return sequence
+
+	augmented_group = []
+	for id_ in ids:
+		current_ids = [id_]
+		if mode in ["prev", "both"]:
+			current_ids = fetch_id_sequence(id_, "prev_id")[::-1] + current_ids
+		if mode in ["next", "both"]:
+			current_ids += fetch_id_sequence(id_, "next_id")
+		augmented_group.extend(current_ids)
+	return augmented_group

autorag/nodes/passageaugmenter/run.py

@@ -0,0 +1,131 @@
+import logging
+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 apply_recursive, to_list
+
+logger = logging.getLogger("AutoRAG")
+
+
+def run_passage_augmenter_node(
+	modules: List,
+	module_params: List[Dict],
+	previous_result: pd.DataFrame,
+	node_line_dir: str,
+	strategies: Dict,
+) -> pd.DataFrame:
+	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))
+
+	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))
+	metric_inputs = [
+		MetricInput(retrieval_gt=ret_gt, query=query, generation_gt=gen_gt)
+		for ret_gt, query, gen_gt in zip(
+			retrieval_gt,
+			previous_result["query"].tolist(),
+			previous_result["generation_gt"].tolist(),
+		)
+	]
+
+	# run metrics before filtering
+	if strategies.get("metrics") is None:
+		raise ValueError(
+			"You must at least one metrics for passage_augmenter 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_augmenter")  # 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_augmenter_{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"),
+	)
+	# change metric name columns to passage_augmenter_metric_name
+	selected_result = selected_result.rename(
+		columns={
+			metric_name: f"passage_augmenter_{metric_name}"
+			for metric_name in strategies["metrics"]
+		}
+	)
+	# drop retrieval result columns in previous_result
+	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

autorag/nodes/passagecompressor/__init__.py

@@ -0,0 +1,4 @@
+from .longllmlingua import LongLLMLingua
+from .pass_compressor import PassCompressor
+from .refine import Refine
+from .tree_summarize import TreeSummarize

autorag/nodes/passagecompressor/base.py

@@ -0,0 +1,78 @@
+import abc
+import logging
+from typing import Dict
+
+import pandas as pd
+from llama_index.core.llms import LLM
+
+from autorag import generator_models
+from autorag.schema import BaseModule
+from autorag.utils import result_to_dataframe
+from autorag.utils.cast import cast_retrieved_contents
+
+logger = logging.getLogger("AutoRAG")
+
+
+class BasePassageCompressor(BaseModule, metaclass=abc.ABCMeta):
+	def __init__(self, project_dir: str, *args, **kwargs):
+		logger.info(
+			f"Initialize passage compressor node - {self.__class__.__name__} module..."
+		)
+
+	def __del__(self):
+		logger.info(
+			f"Deleting passage compressor node - {self.__class__.__name__} module..."
+		)
+
+	def cast_to_run(self, previous_result: pd.DataFrame, *args, **kwargs):
+		logger.info(
+			f"Running passage compressor node - {self.__class__.__name__} module..."
+		)
+		assert "query" in previous_result.columns, (
+			"previous_result must contain 'query' column."
+		)
+		assert len(previous_result) > 0, "previous_result must have at least one row."
+
+		queries = previous_result["query"].tolist()
+
+		return queries, cast_retrieved_contents(previous_result)
+
+
+class LlamaIndexCompressor(BasePassageCompressor, metaclass=abc.ABCMeta):
+	param_list = ["prompt", "chat_prompt", "batch"]
+
+	def __init__(self, project_dir: str, **kwargs):
+		"""
+		Initialize passage compressor module.
+
+		:param project_dir: The project directory
+		:param llm: The llm name that will be used to summarize.
+			The LlamaIndex LLM model can be used in here.
+		:param kwargs: Extra parameter for init llm
+		"""
+		super().__init__(project_dir)
+		kwargs_dict = dict(
+			filter(lambda x: x[0] not in self.param_list, kwargs.items())
+		)
+		llm_name = kwargs_dict.pop("llm")
+		self.llm: LLM = make_llm(llm_name, kwargs_dict)
+
+	def __del__(self):
+		del self.llm
+		super().__del__()
+
+	@result_to_dataframe(["retrieved_contents"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		queries, retrieved_contents = self.cast_to_run(previous_result)
+		param_dict = dict(filter(lambda x: x[0] in self.param_list, kwargs.items()))
+		result = self._pure(queries, retrieved_contents, **param_dict)
+		return list(map(lambda x: [x], result))
+
+
+def make_llm(llm_name: str, kwargs: Dict) -> LLM:
+	if llm_name not in generator_models:
+		raise KeyError(
+			f"{llm_name} is not supported. "
+			"You can add it manually by calling autorag.generator_models."
+		)
+	return generator_models[llm_name](**kwargs)

autorag/nodes/passagecompressor/longllmlingua.py

@@ -0,0 +1,115 @@
+from typing import List, Optional
+
+import pandas as pd
+
+from autorag.nodes.passagecompressor.base import BasePassageCompressor
+from autorag.utils.util import pop_params, result_to_dataframe, empty_cuda_cache
+
+
+# TODO: Parallel Processing Refactoring at #460
+
+
+class LongLLMLingua(BasePassageCompressor):
+	def __init__(
+		self, project_dir: str, model_name: str = "NousResearch/Llama-2-7b-hf", **kwargs
+	):
+		try:
+			from llmlingua import PromptCompressor
+		except ImportError:
+			raise ImportError(
+				"LongLLMLingua is not installed. Please install it by running `pip install llmlingua`."
+			)
+
+		super().__init__(project_dir)
+		model_init_params = pop_params(PromptCompressor.__init__, kwargs)
+		self.llm_lingua = PromptCompressor(model_name=model_name, **model_init_params)
+
+	def __del__(self):
+		del self.llm_lingua
+		empty_cuda_cache()
+		super().__del__()
+
+	@result_to_dataframe(["retrieved_contents"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		queries, retrieved_contents = self.cast_to_run(previous_result)
+		results = self._pure(queries, retrieved_contents, **kwargs)
+		return list(map(lambda x: [x], results))
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents: List[List[str]],
+		instructions: Optional[str] = None,
+		target_token: int = 300,
+		**kwargs,
+	) -> List[str]:
+		"""
+		Compresses the retrieved texts using LongLLMLingua.
+		For more information, visit https://github.com/microsoft/LLMLingua.
+
+		:param queries: The queries for retrieved passages.
+		:param contents: The contents of retrieved passages.
+		:param model_name: The model name to use for compression.
+		    The default is "NousResearch/Llama-2-7b-hf".
+		:param instructions: The instructions for compression.
+		    Default is None. When it is None, it will use default instructions.
+		:param target_token: The target token for compression.
+		    Default is 300.
+		:param kwargs: Additional keyword arguments.
+		:return: The list of compressed texts.
+		"""
+		if instructions is None:
+			instructions = "Given the context, please answer the final question"
+		results = [
+			llmlingua_pure(
+				query, contents_, self.llm_lingua, instructions, target_token, **kwargs
+			)
+			for query, contents_ in zip(queries, contents)
+		]
+
+		return results
+
+
+def llmlingua_pure(
+	query: str,
+	contents: List[str],
+	llm_lingua,
+	instructions: str,
+	target_token: int = 300,
+	**kwargs,
+) -> str:
+	"""
+	Return the compressed text.
+
+	:param query: The query for retrieved passages.
+	:param contents: The contents of retrieved passages.
+	:param llm_lingua: The llm instance, that will be used to compress.
+	:param instructions: The instructions for compression.
+	:param target_token: The target token for compression.
+	    Default is 300.
+	:param kwargs: Additional keyword arguments.
+	:return: The compressed text.
+	"""
+	try:
+		from llmlingua import PromptCompressor
+	except ImportError:
+		raise ImportError(
+			"LongLLMLingua is not installed. Please install it by running `pip install llmlingua`."
+		)
+	# split by "\n\n" (recommended by LongLLMLingua authors)
+	new_context_texts = [c for context in contents for c in context.split("\n\n")]
+	compress_prompt_params = pop_params(PromptCompressor.compress_prompt, kwargs)
+	compressed_prompt = llm_lingua.compress_prompt(
+		new_context_texts,
+		question=query,
+		instruction=instructions,
+		rank_method="longllmlingua",
+		target_token=target_token,
+		**compress_prompt_params,
+	)
+	compressed_prompt_txt = compressed_prompt["compressed_prompt"]
+
+	# separate out the question and instruction
+	result = "\n\n".join(compressed_prompt_txt.split("\n\n")[1:-1])
+
+	return result

autorag/nodes/passagecompressor/pass_compressor.py

@@ -0,0 +1,16 @@
+from typing import List
+
+import pandas as pd
+
+from autorag.nodes.passagecompressor.base import BasePassageCompressor
+from autorag.utils import result_to_dataframe
+
+
+class PassCompressor(BasePassageCompressor):
+	@result_to_dataframe(["retrieved_contents"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		_, contents = self.cast_to_run(previous_result)
+		return self._pure(contents)
+
+	def _pure(self, contents: List[List[str]]):
+		return contents

autorag/nodes/passagecompressor/refine.py

@@ -0,0 +1,54 @@
+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 Refine as rf
+
+from autorag.nodes.passagecompressor.base import LlamaIndexCompressor
+from autorag.utils.util import get_event_loop, process_batch
+
+
+class Refine(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]:
+		"""
+		Refine a response to a query across text chunks.
+		This function is a wrapper for llama_index.response_synthesizers.Refine.
+		For more information, visit https://docs.llamaindex.ai/en/stable/examples/response_synthesizers/refine/.
+
+		:param queries: The queries for retrieved passages.
+		:param contents: The contents of retrieved passages.
+		:param prompt: The prompt template for refine.
+		    If you want to use chat prompt, you should pass chat_prompt instead.
+		    At prompt, you must specify where to put 'context_msg' 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 refine.
+		    If you want to use normal prompt, you should pass prompt instead.
+		    At prompt, you must specify where to put 'context_msg' 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):
+			refine_template = PromptTemplate(prompt, prompt_type=PromptType.REFINE)
+		elif chat_prompt is not None and is_chat_model(self.llm):
+			refine_template = PromptTemplate(chat_prompt, prompt_type=PromptType.REFINE)
+		else:
+			refine_template = None
+		summarizer = rf(llm=self.llm, refine_template=refine_template, verbose=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