AutoRAG 0.0.0__py3-none-any.whl

autorag/nodes/semanticretrieval/vectordb.py

@@ -0,0 +1,339 @@
+import itertools
+import logging
+import os
+from typing import List, Tuple, Optional
+
+import numpy as np
+import pandas as pd
+from llama_index.core.embeddings import BaseEmbedding
+from llama_index.embeddings.openai import OpenAIEmbedding
+
+from autorag.evaluation.metric.util import (
+	calculate_l2_distance,
+	calculate_inner_product,
+	calculate_cosine_similarity,
+)
+from autorag.nodes.retrieval.base import evenly_distribute_passages, BaseRetrieval
+from autorag.utils import (
+	validate_corpus_dataset,
+	cast_corpus_dataset,
+	cast_qa_dataset,
+	validate_qa_dataset,
+)
+from autorag.utils.util import (
+	get_event_loop,
+	process_batch,
+	openai_truncate_by_token,
+	flatten_apply,
+	result_to_dataframe,
+	pop_params,
+	fetch_contents,
+	empty_cuda_cache,
+	convert_inputs_to_list,
+	make_batch,
+)
+from autorag.vectordb import load_vectordb_from_yaml
+from autorag.vectordb.base import BaseVectorStore
+
+logger = logging.getLogger("AutoRAG")
+
+
+class VectorDB(BaseRetrieval):
+	def __init__(self, project_dir: str, vectordb: str = "default", **kwargs):
+		"""
+		Initialize VectorDB retrieval node.
+
+		:param project_dir: The project directory path.
+		:param vectordb: The vectordb name.
+			You must configure the vectordb name in the config.yaml file.
+			If you don't configure, it uses the default vectordb.
+		:param kwargs: The optional arguments.
+			Not affected in the init method.
+		"""
+		super().__init__(project_dir)
+
+		vectordb_config_path = os.path.join(self.resources_dir, "vectordb.yaml")
+		self.vector_store = load_vectordb_from_yaml(
+			vectordb_config_path, vectordb, project_dir
+		)
+
+		self.embedding_model = self.vector_store.embedding
+
+	def __del__(self):
+		del self.vector_store
+		del self.embedding_model
+		empty_cuda_cache()
+		super().__del__()
+
+	@result_to_dataframe(
+		[
+			"retrieved_contents_semantic",
+			"retrieved_ids_semantic",
+			"retrieve_scores_semantic",
+		]
+	)
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		queries = self.cast_to_run(previous_result)
+		pure_params = pop_params(self._pure, kwargs)
+		ids, scores = self._pure(queries, **pure_params)
+		contents = fetch_contents(self.corpus_df, ids)
+		return contents, ids, scores
+
+	def _pure(
+		self,
+		queries: List[List[str]],
+		top_k: int,
+		embedding_batch: int = 128,
+		ids: Optional[List[List[str]]] = None,
+	) -> Tuple[List[List[str]], List[List[float]]]:
+		"""
+		VectorDB retrieval function.
+		You have to get a chroma collection that is already ingested.
+		You have to get an embedding model that is already used in ingesting.
+
+		:param queries: 2-d list of query strings.
+		    Each element of the list is a query strings of each row.
+		:param top_k: The number of passages to be retrieved.
+		:param embedding_batch: The number of queries to be processed in parallel.
+		    This is used to prevent API error at the query embedding.
+		    Default is 128.
+		:param ids: The optional list of ids that you want to retrieve.
+		    You don't need to specify this in the general use cases.
+		    Default is None.
+
+		:return: The 2-d list contains a list of passage ids that retrieved from vectordb and 2-d list of its scores.
+		    It will be a length of queries. And each element has a length of top_k.
+		"""
+		# if ids are specified, fetch the ids score from Chroma
+		if ids is not None:
+			return self.__get_ids_scores(queries, ids, embedding_batch)
+
+		# run async vector_db_pure function
+		tasks = [
+			vectordb_pure(query_list, top_k, self.vector_store)
+			for query_list in queries
+		]
+		loop = get_event_loop()
+		results = loop.run_until_complete(
+			process_batch(tasks, batch_size=embedding_batch)
+		)
+		id_result = list(map(lambda x: x[0], results))
+		score_result = list(map(lambda x: x[1], results))
+		return id_result, score_result
+
+	def __get_ids_scores(self, queries, ids, embedding_batch: int):
+		# truncate queries and embedding execution here.
+		openai_embedding_limit = 8000
+		if isinstance(self.embedding_model, OpenAIEmbedding):
+			queries = list(
+				map(
+					lambda query_list: openai_truncate_by_token(
+						query_list,
+						openai_embedding_limit,
+						self.embedding_model.model_name,
+					),
+					queries,
+				)
+			)
+
+		query_embeddings = flatten_apply(
+			run_query_embedding_batch,
+			queries,
+			embedding_model=self.embedding_model,
+			batch_size=embedding_batch,
+		)
+
+		loop = get_event_loop()
+
+		async def run_fetch(ids):
+			final_result = []
+			for id_list in ids:
+				if len(id_list) == 0:
+					final_result.append([])
+				else:
+					result = await self.vector_store.fetch(id_list)
+					final_result.append(result)
+			return final_result
+
+		content_embeddings = loop.run_until_complete(run_fetch(ids))
+
+		score_result = list(
+			map(
+				lambda query_embedding_list, content_embedding_list: get_id_scores(
+					query_embedding_list,
+					content_embedding_list,
+					similarity_metric=self.vector_store.similarity_metric,
+				),
+				query_embeddings,
+				content_embeddings,
+			)
+		)
+		return ids, score_result
+
+
+async def vectordb_pure(
+	queries: List[str], top_k: int, vectordb: BaseVectorStore
+) -> Tuple[List[str], List[float]]:
+	"""
+	Async VectorDB retrieval function.
+	Its usage is for async retrieval of vector_db row by row.
+
+	:param query_embeddings: A list of query embeddings.
+	:param top_k: The number of passages to be retrieved.
+	:param vectordb: The vector store instance.
+	:return: The tuple contains a list of passage ids that are retrieved from vectordb and a list of its scores.
+	"""
+	id_result, score_result = await vectordb.query(queries=queries, top_k=top_k)
+
+	# Distribute passages evenly
+	id_result, score_result = evenly_distribute_passages(id_result, score_result, top_k)
+	# sort id_result and score_result by score
+	result = [
+		(_id, score)
+		for score, _id in sorted(
+			zip(score_result, id_result), key=lambda pair: pair[0], reverse=True
+		)
+	]
+	id_result, score_result = zip(*result)
+	return list(id_result), list(score_result)
+
+
+async def filter_exist_ids(
+	vectordb: BaseVectorStore,
+	corpus_data: pd.DataFrame,
+) -> pd.DataFrame:
+	corpus_data = cast_corpus_dataset(corpus_data)
+	validate_corpus_dataset(corpus_data)
+	ids = corpus_data["doc_id"].tolist()
+
+	# Query the collection to check if IDs already exist
+	existed_bool_list = await vectordb.is_exist(ids=ids)
+	# Assuming 'ids' is the key in the response
+	new_passage = corpus_data[~pd.Series(existed_bool_list)]
+	return new_passage
+
+
+async def filter_exist_ids_from_retrieval_gt(
+	vectordb: BaseVectorStore,
+	qa_data: pd.DataFrame,
+	corpus_data: pd.DataFrame,
+) -> pd.DataFrame:
+	qa_data = cast_qa_dataset(qa_data)
+	validate_qa_dataset(qa_data)
+	corpus_data = cast_corpus_dataset(corpus_data)
+	validate_corpus_dataset(corpus_data)
+	retrieval_gt = (
+		qa_data["retrieval_gt"]
+		.apply(lambda x: list(itertools.chain.from_iterable(x)))
+		.tolist()
+	)
+	retrieval_gt = list(itertools.chain.from_iterable(retrieval_gt))
+	retrieval_gt = list(set(retrieval_gt))
+
+	existed_bool_list = await vectordb.is_exist(ids=retrieval_gt)
+	add_ids = []
+	for ret_gt, is_exist in zip(retrieval_gt, existed_bool_list):
+		if not is_exist:
+			add_ids.append(ret_gt)
+	new_passage = corpus_data[corpus_data["doc_id"].isin(add_ids)]
+	return new_passage
+
+
+async def vectordb_ingest_api(
+	vectordb: BaseVectorStore,
+	corpus_data: pd.DataFrame,
+):
+	"""
+	Ingest given corpus data to the vectordb.
+	It truncates corpus content when the embedding model is OpenAIEmbedding to the 8000 tokens.
+	Plus, when the corpus content is empty (whitespace), it will be ignored.
+	And if there is a document id that already exists in the collection, it will be ignored.
+
+	:param vectordb: A vector stores instance that you want to ingest.
+	:param corpus_data: The corpus data that contains doc_id and contents columns.
+	"""
+	embedding_batch = vectordb.embedding_batch
+	if not corpus_data.empty:
+		new_contents = corpus_data["contents"].tolist()
+		new_ids = corpus_data["doc_id"].tolist()
+		content_batches = make_batch(new_contents, embedding_batch)
+		id_batches = make_batch(new_ids, embedding_batch)
+		for content_batch, id_batch in zip(content_batches, id_batches):
+			await vectordb.add(ids=id_batch, texts=content_batch)
+
+
+def vectordb_ingest_huggingface(
+	vectordb: BaseVectorStore,
+	corpus_data: pd.DataFrame,
+):
+	"""
+	Ingest given corpus data to the vectordb using local model.
+	When the corpus content is empty (whitespace), it will be ignored.
+	And if there is a document id that already exists in the collection, it will be ignored.
+
+	:param vectordb: A vector stores instance that you want to ingest.
+	:param corpus_data: The corpus data that contains doc_id and contents columns.
+	"""
+	embedding_batch_size = vectordb.embedding_batch
+	embedding_model = vectordb.embedding._model
+	if corpus_data.empty:
+		logger.warning("The corpus data is empty. Nothing to ingest.")
+		return
+	new_contents = corpus_data["contents"].tolist()
+	new_ids = corpus_data["doc_id"].tolist()
+	logger.info("Start embedding corpus data with huggingface model.")
+	embeddings = embedding_model.encode(
+		new_contents,
+		batch_size=embedding_batch_size,
+		normalize_embeddings=vectordb.embedding.normalize,
+		show_progress_bar=True,
+	)
+	vectordb.add_embedding(new_ids, embeddings)
+	logger.info("Finish embedding & ingesting corpus data with huggingface model.")
+
+
+def run_query_embedding_batch(
+	queries: List[str], embedding_model: BaseEmbedding, batch_size: int
+) -> List[List[float]]:
+	result = []
+	for i in range(0, len(queries), batch_size):
+		batch = queries[i : i + batch_size]
+		embeddings = embedding_model.get_text_embedding_batch(batch)
+		result.extend(embeddings)
+	return result
+
+
+@convert_inputs_to_list
+def get_id_scores(  # To find the uncalculated score when fuse the scores for the hybrid retrieval
+	query_embeddings: List[
+		List[float]
+	],  # `queries` is input. This is one user input query.
+	content_embeddings: List[List[float]],
+	similarity_metric: str,
+) -> List[
+	float
+]:  # The most high scores among each query. The length of a result is the same as the contents length.
+	"""
+	Calculate the highest similarity scores between query embeddings and content embeddings.
+
+	:param query_embeddings: A list of lists containing query embeddings.
+	:param content_embeddings: A list of lists containing content embeddings.
+	:param similarity_metric: The similarity metric to use ('l2', 'ip', or 'cosine').
+	:return: A list of the highest similarity scores for each content embedding.
+	"""
+	metric_func_dict = {
+		"l2": lambda x, y: 1 - calculate_l2_distance(x, y),
+		"ip": calculate_inner_product,
+		"cosine": calculate_cosine_similarity,
+	}
+	metric_func = metric_func_dict[similarity_metric]
+
+	result = []
+	for content_embedding in content_embeddings:
+		scores = []
+		for query_embedding in query_embeddings:
+			scores.append(
+				metric_func(np.array(query_embedding), np.array(content_embedding))
+			)
+		result.append(max(scores))
+	return result

autorag/nodes/util.py

@@ -0,0 +1,16 @@
+from typing import Optional, Dict
+
+from autorag.support import get_support_modules
+
+
+def make_generator_callable_param(generator_dict: Optional[Dict]):
+	if "generator_module_type" not in generator_dict.keys():
+		generator_dict = {
+			"generator_module_type": "llama_index_llm",
+			"llm": "openai",
+			"model": "gpt-4o-mini",
+		}
+	module_str = generator_dict.pop("generator_module_type")
+	module_class = get_support_modules(module_str)
+	module_param = generator_dict
+	return module_class, module_param

autorag/parser.py

@@ -0,0 +1,37 @@
+import logging
+import os
+import shutil
+from typing import Optional
+
+from autorag.data.parse.run import run_parser
+from autorag.data.utils.util import load_yaml, get_param_combinations
+
+logger = logging.getLogger("AutoRAG")
+
+
+class Parser:
+	def __init__(self, data_path_glob: str, project_dir: Optional[str] = None):
+		self.data_path_glob = data_path_glob
+		self.project_dir = project_dir if project_dir is not None else os.getcwd()
+
+	def start_parsing(self, yaml_path: str, all_files: bool = False):
+		if not os.path.exists(self.project_dir):
+			os.makedirs(self.project_dir)
+
+		# copy yaml file to project directory
+		shutil.copy(yaml_path, os.path.join(self.project_dir, "parse_config.yaml"))
+
+		# load yaml file
+		modules = load_yaml(yaml_path)
+
+		input_modules, input_params = get_param_combinations(modules)
+
+		logger.info("Parsing Start...")
+		run_parser(
+			modules=input_modules,
+			module_params=input_params,
+			data_path_glob=self.data_path_glob,
+			project_dir=self.project_dir,
+			all_files=all_files,
+		)
+		logger.info("Parsing Done!")

autorag/schema/__init__.py

@@ -0,0 +1,3 @@
+from .module import Module
+from .node import Node
+from .base import BaseModule

autorag/schema/base.py

@@ -0,0 +1,35 @@
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Union
+
+import pandas as pd
+
+
+class BaseModule(metaclass=ABCMeta):
+	@abstractmethod
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		pass
+
+	@abstractmethod
+	def _pure(self, *args, **kwargs):
+		pass
+
+	@classmethod
+	def run_evaluator(
+		cls,
+		project_dir: Union[str, Path],
+		previous_result: pd.DataFrame,
+		*args,
+		**kwargs,
+	):
+		instance = cls(project_dir, *args, **kwargs)
+		result = instance.pure(previous_result, *args, **kwargs)
+		del instance
+		return result
+
+	@abstractmethod
+	def cast_to_run(self, previous_result: pd.DataFrame, *args, **kwargs):
+		"""
+		This function is for cast function (a.k.a decorator) only for pure function in the whole node.
+		"""
+		pass

autorag/schema/metricinput.py

@@ -0,0 +1,99 @@
+from dataclasses import dataclass
+from typing import Optional, List, Dict, Callable, Any, Union
+
+import numpy as np
+import pandas as pd
+
+
+@dataclass
+class MetricInput:
+	query: Optional[str] = None
+	queries: Optional[List[str]] = None
+	retrieval_gt_contents: Optional[List[List[str]]] = None
+	retrieved_contents: Optional[List[str]] = None
+	retrieval_gt: Optional[List[List[str]]] = None
+	retrieved_ids: Optional[List[str]] = None
+	prompt: Optional[str] = None
+	generated_texts: Optional[str] = None
+	generation_gt: Optional[List[str]] = None
+	generated_log_probs: Optional[List[float]] = None
+
+	def is_fields_notnone(self, fields_to_check: List[str]) -> bool:
+		for field in fields_to_check:
+			actual_value = getattr(self, field)
+
+			if actual_value is None:
+				return False
+
+			try:
+				if not type_checks.get(type(actual_value), lambda _: False)(
+					actual_value
+				):
+					return False
+			except Exception:
+				return False
+
+		return True
+
+	@classmethod
+	def from_dataframe(cls, qa_data: pd.DataFrame) -> List["MetricInput"]:
+		"""
+		Convert a pandas DataFrame into a list of MetricInput instances.
+		qa_data: pd.DataFrame: qa_data DataFrame containing metric data.
+
+		:returns: List[MetricInput]: List of MetricInput objects created from DataFrame rows.
+		"""
+		instances = []
+
+		for _, row in qa_data.iterrows():
+			instance = cls()
+
+			for attr_name in cls.__annotations__:
+				if attr_name in row:
+					value = row[attr_name]
+
+					if isinstance(value, str):
+						setattr(
+							instance,
+							attr_name,
+							value.strip() if value.strip() != "" else None,
+						)
+					elif isinstance(value, list):
+						setattr(instance, attr_name, value if len(value) > 0 else None)
+					else:
+						setattr(instance, attr_name, value)
+
+			instances.append(instance)
+
+		return instances
+
+	@staticmethod
+	def _check_list(lst_or_arr: Union[List[Any], np.ndarray]) -> bool:
+		if isinstance(lst_or_arr, np.ndarray):
+			lst_or_arr = lst_or_arr.flatten().tolist()
+
+		if len(lst_or_arr) == 0:
+			return False
+
+		for item in lst_or_arr:
+			if item is None:
+				return False
+
+			item_type = type(item)
+
+			if item_type in type_checks:
+				if not type_checks[item_type](item):
+					return False
+			else:
+				return False
+
+		return True
+
+
+type_checks: Dict[type, Callable[[Any], bool]] = {
+	str: lambda x: len(x.strip()) > 0,
+	list: MetricInput._check_list,
+	np.ndarray: MetricInput._check_list,
+	int: lambda _: True,
+	float: lambda _: True,
+}

autorag/schema/module.py

@@ -0,0 +1,24 @@
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import Callable, Dict
+
+from autorag.support import get_support_modules
+
+
+@dataclass
+class Module:
+	module_type: str
+	module_param: Dict
+	module: Callable = field(init=False)
+
+	def __post_init__(self):
+		self.module = get_support_modules(self.module_type)
+		if self.module is None:
+			raise ValueError(f"Module type {self.module_type} is not supported.")
+
+	@classmethod
+	def from_dict(cls, module_dict: Dict) -> "Module":
+		_module_dict = deepcopy(module_dict)
+		module_type = _module_dict.pop("module_type")
+		module_params = _module_dict
+		return cls(module_type, module_params)

autorag/schema/node.py

@@ -0,0 +1,144 @@
+import itertools
+import logging
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import Dict, List, Callable, Tuple, Any
+
+import pandas as pd
+
+from autorag.schema.module import Module
+from autorag.support import get_support_nodes
+from autorag.utils.util import make_combinations, explode, find_key_values
+
+logger = logging.getLogger("AutoRAG")
+
+
+@dataclass
+class Node:
+	node_type: str
+	strategy: Dict
+	node_params: Dict
+	modules: List[Module]
+	run_node: Callable = field(init=False)
+
+	def __post_init__(self):
+		self.run_node = get_support_nodes(self.node_type)
+		if self.run_node is None:
+			raise ValueError(f"Node type {self.node_type} is not supported.")
+
+	def get_param_combinations(self) -> Tuple[List[Callable], List[Dict]]:
+		"""
+		This method returns a combination of module and node parameters, also corresponding modules.
+
+		:return: Each module and its module parameters.
+		:rtype: Tuple[List[Callable], List[Dict]]
+		"""
+
+		def make_single_combination(module: Module) -> List[Dict]:
+			input_dict = {**self.node_params, **module.module_param}
+			return make_combinations(input_dict)
+
+		combinations = list(map(make_single_combination, self.modules))
+		module_list, combination_list = explode(self.modules, combinations)
+		return list(map(lambda x: x.module, module_list)), combination_list
+
+	@classmethod
+	def from_dict(cls, node_dict: Dict) -> "Node":
+		_node_dict = deepcopy(node_dict)
+		node_type = _node_dict.pop("node_type")
+		strategy = _node_dict.pop("strategy")
+		modules = list(map(lambda x: Module.from_dict(x), _node_dict.pop("modules")))
+		node_params = _node_dict
+		return cls(node_type, strategy, node_params, modules)
+
+	def run(self, previous_result: pd.DataFrame, node_line_dir: str) -> pd.DataFrame:
+		logger.info(f"Running node {self.node_type}...")
+		input_modules, input_params = self.get_param_combinations()
+		return self.run_node(
+			modules=input_modules,
+			module_params=input_params,
+			previous_result=previous_result,
+			node_line_dir=node_line_dir,
+			strategies=self.strategy,
+		)
+
+
+def extract_values(node: Node, key: str) -> List[str]:
+	"""
+	This function extract values from node's modules' module_param.
+
+	:param node: The node you want to extract values from.
+	:param key: The key of module_param that you want to extract.
+	:return: The list of extracted values.
+	    It removes duplicated elements automatically.
+	"""
+
+	def extract_module_values(module: Module):
+		if key not in module.module_param:
+			return []
+		value = module.module_param[key]
+		if isinstance(value, str) or isinstance(value, int):
+			return [value]
+		elif isinstance(value, list):
+			return value
+		else:
+			raise ValueError(f"{key} must be str,list or int, but got {type(value)}")
+
+	values = list(map(extract_module_values, node.modules))
+	return list(set(list(itertools.chain.from_iterable(values))))
+
+
+def extract_values_from_nodes(nodes: List[Node], key: str) -> List[str]:
+	"""
+	This function extract values from nodes' modules' module_param.
+
+	:param nodes: The nodes you want to extract values from.
+	:param key: The key of module_param that you want to extract.
+	:return: The list of extracted values.
+	    It removes duplicated elements automatically.
+	"""
+	values = list(map(lambda node: extract_values(node, key), nodes))
+	return list(set(list(itertools.chain.from_iterable(values))))
+
+
+def extract_values_from_nodes_strategy(nodes: List[Node], key: str) -> List[Any]:
+	"""
+	This function extract values from nodes' strategy.
+
+	:param nodes: The nodes you want to extract values from.
+	:param key: The key string that you want to extract.
+	:return: The list of extracted values.
+	    It removes duplicated elements automatically.
+	"""
+	values = []
+	for node in nodes:
+		value_list = find_key_values(node.strategy, key)
+		if value_list:
+			values.extend(value_list)
+	return values
+
+
+def module_type_exists(nodes: List[Node], module_type: str) -> bool:
+	"""
+	This function check if the module type exists in the nodes.
+
+	:param nodes: The nodes you want to check.
+	:param module_type: The module type you want to check.
+	:return: True if the module type exists in the nodes.
+	"""
+	return any(
+		list(
+			map(
+				lambda node: any(
+					list(
+						map(
+							lambda module: module.module_type.lower()
+							== module_type.lower(),
+							node.modules,
+						)
+					)
+				),
+				nodes,
+			)
+		)
+	)