AutoRAG 0.0.0__py3-none-any.whl

autorag/nodes/lexicalretrieval/bm25.py

@@ -0,0 +1,381 @@
+import asyncio
+import os
+import pickle
+import re
+from typing import List, Dict, Tuple, Callable, Union, Iterable, Optional
+
+import numpy as np
+import pandas as pd
+from llama_index.core.indices.keyword_table.utils import simple_extract_keywords
+from nltk import PorterStemmer
+from rank_bm25 import BM25Okapi
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
+from autorag.nodes.retrieval.base import (
+	evenly_distribute_passages,
+	BaseRetrieval,
+	get_bm25_pkl_name,
+)
+from autorag.utils import validate_corpus_dataset, fetch_contents
+from autorag.utils.util import (
+	get_event_loop,
+	normalize_string,
+	result_to_dataframe,
+	pop_params,
+)
+
+
+def tokenize_ko_kiwi(texts: List[str]) -> List[List[str]]:
+	try:
+		from kiwipiepy import Kiwi, Token
+	except ImportError:
+		raise ImportError(
+			"You need to install kiwipiepy to use 'ko_kiwi' tokenizer. "
+			"Please install kiwipiepy by running 'pip install kiwipiepy'. "
+			"Or install Korean version of AutoRAG by running 'pip install AutoRAG[ko]'."
+		)
+	texts = list(map(lambda x: x.strip().lower(), texts))
+	kiwi = Kiwi()
+	tokenized_list: Iterable[List[Token]] = kiwi.tokenize(texts)
+
+	def extract_form_safe(x):
+		try:
+			return x.form
+		except UnicodeDecodeError:
+			return " "
+
+	return [
+		list(map(lambda x: extract_form_safe(x), token_list))
+		for token_list in tokenized_list
+	]
+
+
+def tokenize_ko_kkma(texts: List[str]) -> List[List[str]]:
+	try:
+		from konlpy.tag import Kkma
+	except ImportError:
+		raise ImportError(
+			"You need to install konlpy to use 'ko_kkma' tokenizer. "
+			"Please install konlpy by running 'pip install konlpy'. "
+			"Or install Korean version of AutoRAG by running 'pip install AutoRAG[ko]'."
+		)
+	tokenizer = Kkma()
+	tokenized_list: List[List[str]] = list(map(lambda x: tokenizer.morphs(x), texts))
+	return tokenized_list
+
+
+def tokenize_ko_okt(texts: List[str]) -> List[List[str]]:
+	try:
+		from konlpy.tag import Okt
+	except ImportError:
+		raise ImportError(
+			"You need to install konlpy to use 'ko_kkma' tokenizer. "
+			"Please install konlpy by running 'pip install konlpy'. "
+			"Or install Korean version of AutoRAG by running 'pip install AutoRAG[ko]'."
+		)
+	tokenizer = Okt()
+	tokenized_list: List[List[str]] = list(map(lambda x: tokenizer.morphs(x), texts))
+	return tokenized_list
+
+
+def tokenize_porter_stemmer(texts: List[str]) -> List[List[str]]:
+	def tokenize_remove_stopword(text: str, stemmer) -> List[str]:
+		text = text.lower()
+		words = list(simple_extract_keywords(text))
+		return [stemmer.stem(word) for word in words]
+
+	stemmer = PorterStemmer()
+	tokenized_list: List[List[str]] = list(
+		map(lambda x: tokenize_remove_stopword(x, stemmer), texts)
+	)
+	return tokenized_list
+
+
+def tokenize_space(texts: List[str]) -> List[List[str]]:
+	def tokenize_space_text(text: str) -> List[str]:
+		text = normalize_string(text)
+		return re.split(r"\s+", text.strip())
+
+	return list(map(tokenize_space_text, texts))
+
+
+def load_bm25_corpus(bm25_path: str) -> Dict:
+	if bm25_path is None:
+		return {}
+	with open(bm25_path, "rb") as f:
+		bm25_corpus = pickle.load(f)
+	return bm25_corpus
+
+
+def tokenize_ja_sudachipy(texts: List[str]) -> List[List[str]]:
+	try:
+		from sudachipy import dictionary, tokenizer
+	except ImportError:
+		raise ImportError(
+			"You need to install SudachiPy to use 'sudachipy' tokenizer. "
+			"Please install SudachiPy by running 'pip install sudachipy'."
+		)
+
+	# Initialize SudachiPy with the default tokenizer
+	tokenizer_obj = dictionary.Dictionary(dict="core").create()
+
+	# Choose the tokenizer mode: NORMAL, SEARCH, A
+	mode = tokenizer.Tokenizer.SplitMode.A
+
+	# Tokenize the input texts
+	tokenized_list = []
+	for text in texts:
+		tokens = tokenizer_obj.tokenize(text, mode)
+		tokenized_list.append([token.surface() for token in tokens])
+
+	return tokenized_list
+
+
+BM25_TOKENIZER = {
+	"porter_stemmer": tokenize_porter_stemmer,
+	"ko_kiwi": tokenize_ko_kiwi,
+	"space": tokenize_space,
+	"ko_kkma": tokenize_ko_kkma,
+	"ko_okt": tokenize_ko_okt,
+	"sudachipy": tokenize_ja_sudachipy,
+}
+
+
+class BM25(BaseRetrieval):
+	def __init__(self, project_dir: str, *args, **kwargs):
+		"""
+		Initialize BM25 module.
+		(Retrieval)
+
+		:param project_dir: The project directory path.
+		:param bm25_tokenizer: The tokenizer name that is used to the BM25.
+		    It supports 'porter_stemmer', 'ko_kiwi', and huggingface `AutoTokenizer`.
+		    You can pass huggingface tokenizer name.
+		    Default is porter_stemmer.
+		:param kwargs: The optional arguments.
+		"""
+
+		super().__init__(project_dir)
+		# check if bm25_path and file exist
+		bm25_tokenizer = kwargs.get("bm25_tokenizer", None)
+		if bm25_tokenizer is None:
+			bm25_tokenizer = "porter_stemmer"
+		bm25_path = os.path.join(self.resources_dir, get_bm25_pkl_name(bm25_tokenizer))
+
+		assert bm25_path is not None, (
+			"bm25_path must be specified for using bm25 retrieval."
+		)
+		assert os.path.exists(bm25_path), (
+			f"bm25_path {bm25_path} does not exist. Please ingest first."
+		)
+
+		self.bm25_corpus = load_bm25_corpus(bm25_path)
+		assert "tokens" and "passage_id" in list(self.bm25_corpus.keys()), (
+			"bm25_corpus must contain tokens and passage_id. Please check you ingested bm25 corpus correctly."
+		)
+		self.tokenizer = select_bm25_tokenizer(bm25_tokenizer)
+		assert self.bm25_corpus["tokenizer_name"] == bm25_tokenizer, (
+			f"The bm25 corpus tokenizer is {self.bm25_corpus['tokenizer_name']}, but your input is {bm25_tokenizer}. "
+			f"You need to ingest again. Delete bm25 pkl file and re-ingest it."
+		)
+		self.bm25_instance = BM25Okapi(self.bm25_corpus["tokens"])
+
+	@result_to_dataframe(
+		[
+			"retrieved_contents_lexical",
+			"retrieved_ids_lexical",
+			"retrieve_scores_lexical",
+		]
+	)
+	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, *args, **pure_params)
+		contents = fetch_contents(self.corpus_df, ids)
+		return contents, ids, scores
+
+	def _pure(
+		self,
+		queries: List[List[str]],
+		top_k: int,
+		ids: Optional[List[List[str]]] = None,
+	) -> Tuple[List[List[str]], List[List[float]]]:
+		"""
+		BM25 retrieval function.
+		You have to load a pickle file that is already ingested.
+
+		: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 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 bm25 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 is not None:
+			score_result = list(
+				map(
+					lambda query_list, id_list: get_bm25_scores(
+						query_list,
+						id_list,
+						self.tokenizer,
+						self.bm25_instance,
+						self.bm25_corpus,
+					),
+					queries,
+					ids,
+				)
+			)
+			return ids, score_result
+
+		# run async bm25_pure function
+		tasks = [
+			bm25_pure(
+				input_queries,
+				top_k,
+				self.tokenizer,
+				self.bm25_instance,
+				self.bm25_corpus,
+			)
+			for input_queries in queries
+		]
+		loop = get_event_loop()
+		results = loop.run_until_complete(asyncio.gather(*tasks))
+		id_result = list(map(lambda x: x[0], results))
+		score_result = list(map(lambda x: x[1], results))
+		return id_result, score_result
+
+
+async def bm25_pure(
+	queries: List[str], top_k: int, tokenizer, bm25_api: BM25Okapi, bm25_corpus: Dict
+) -> Tuple[List[str], List[float]]:
+	"""
+	Async BM25 retrieval function.
+	Its usage is for async retrieval of bm25 row by row.
+
+	:param queries: A list of query strings.
+	:param top_k: The number of passages to be retrieved.
+	:param tokenizer: A tokenizer that will be used to tokenize queries.
+	:param bm25_api: A bm25 api instance that will be used to retrieve passages.
+	:param bm25_corpus: A dictionary containing the bm25 corpus, which is doc_id from corpus and tokenized corpus.
+	    Its data structure looks like this:
+
+	    .. Code:: python
+
+	        {
+	            "tokens": [], # 2d list of tokens
+	            "passage_id": [], # 2d list of passage_id. Type must be str.
+	        }
+	:return: The tuple contains a list of passage ids that retrieved from bm25 and its scores.
+	"""
+	# I don't make queries operation to async, because queries length might be small, so it will occur overhead.
+	tokenized_queries = tokenize(queries, tokenizer)
+	id_result = []
+	score_result = []
+	for query in tokenized_queries:
+		scores = bm25_api.get_scores(query)
+		sorted_scores = sorted(scores, reverse=True)
+		top_n_index = np.argsort(scores)[::-1][:top_k]
+		ids = [bm25_corpus["passage_id"][i] for i in top_n_index]
+		id_result.append(ids)
+		score_result.append(sorted_scores[:top_k])
+
+	# make a total result to top_k
+	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)
+
+
+def get_bm25_scores(
+	queries: List[str],
+	ids: List[str],
+	tokenizer,
+	bm25_api: BM25Okapi,
+	bm25_corpus: Dict,
+) -> List[float]:
+	if len(ids) == 0 or not bool(ids):
+		return []
+	tokenized_queries = tokenize(queries, tokenizer)
+	result_dict = {id_: [] for id_ in ids}
+	for query in tokenized_queries:
+		scores = bm25_api.get_scores(query)
+		for i, id_ in enumerate(ids):
+			result_dict[id_].append(scores[bm25_corpus["passage_id"].index(id_)])
+	result_df = pd.DataFrame(result_dict)
+	return result_df.max(axis=0).tolist()
+
+
+def tokenize(queries: List[str], tokenizer) -> List[List[int]]:
+	if isinstance(tokenizer, PreTrainedTokenizerBase):
+		tokenized_queries = tokenizer(queries).input_ids
+	else:
+		tokenized_queries = tokenizer(queries)
+	return tokenized_queries
+
+
+def bm25_ingest(
+	corpus_path: str, corpus_data: pd.DataFrame, bm25_tokenizer: str = "porter_stemmer"
+):
+	if not corpus_path.endswith(".pkl"):
+		raise ValueError(f"Corpus path {corpus_path} is not a pickle file.")
+	validate_corpus_dataset(corpus_data)
+	ids = corpus_data["doc_id"].tolist()
+
+	# Initialize bm25_corpus
+	bm25_corpus = pd.DataFrame()
+
+	# Load the BM25 corpus if it exists and get the passage ids
+	if os.path.exists(corpus_path) and os.path.getsize(corpus_path) > 0:
+		with open(corpus_path, "rb") as r:
+			corpus = pickle.load(r)
+			bm25_corpus = pd.DataFrame.from_dict(corpus)
+		duplicated_passage_rows = bm25_corpus[bm25_corpus["passage_id"].isin(ids)]
+		new_passage = corpus_data[
+			~corpus_data["doc_id"].isin(duplicated_passage_rows["passage_id"])
+		]
+	else:
+		new_passage = corpus_data
+
+	if not new_passage.empty:
+		tokenizer = select_bm25_tokenizer(bm25_tokenizer)
+		if isinstance(tokenizer, PreTrainedTokenizerBase):
+			tokenized_corpus = tokenizer(new_passage["contents"].tolist()).input_ids
+		else:
+			tokenized_corpus = tokenizer(new_passage["contents"].tolist())
+		new_bm25_corpus = pd.DataFrame(
+			{
+				"tokens": tokenized_corpus,
+				"passage_id": new_passage["doc_id"].tolist(),
+			}
+		)
+
+		if not bm25_corpus.empty:
+			bm25_corpus_updated = pd.concat(
+				[bm25_corpus, new_bm25_corpus], ignore_index=True
+			)
+			bm25_dict = bm25_corpus_updated.to_dict("list")
+		else:
+			bm25_dict = new_bm25_corpus.to_dict("list")
+
+		# add tokenizer name to bm25_dict
+		bm25_dict["tokenizer_name"] = bm25_tokenizer
+
+		with open(corpus_path, "wb") as w:
+			pickle.dump(bm25_dict, w)
+
+
+def select_bm25_tokenizer(
+	bm25_tokenizer: str,
+) -> Callable[[str], List[Union[int, str]]]:
+	if bm25_tokenizer in list(BM25_TOKENIZER.keys()):
+		return BM25_TOKENIZER[bm25_tokenizer]
+
+	return AutoTokenizer.from_pretrained(bm25_tokenizer, use_fast=False)

autorag/nodes/lexicalretrieval/run.py

@@ -0,0 +1,148 @@
+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_lexical_retrieval_node(
+	modules: List,
+	module_params: List[Dict],
+	previous_result: pd.DataFrame,
+	node_line_dir: str,
+	strategies: Dict,
+) -> pd.DataFrame:
+	"""
+	Run the semantic retrieval node.
+
+	:param modules: Retrieval modules to run.
+	:param module_params: Retrieval 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 retrieval node.
+	:return: The best result dataframe.
+	    It contains previous result columns and retrieval node's 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()
+		)
+	]
+
+	save_dir = os.path.join(node_line_dir, "lexical_retrieval")
+	if not os.path.exists(save_dir):
+		os.makedirs(save_dir)
+
+	# Run the modules
+	lexical_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),
+		)
+	)
+	lexical_times = list(map(lambda x: x / len(lexical_results[0]), execution_times))
+
+	# run metrics
+	if strategies.get("metrics") is None:
+		raise ValueError("You must at least one metrics for retrieval evaluation.")
+	lexical_results = list(
+		map(
+			lambda x: evaluate_lexical_retrieval_node(
+				x,
+				metric_inputs,
+				strategies.get("metrics"),
+			),
+			lexical_results,
+		)
+	)
+
+	lexical_summary_df = save_and_summary(
+		modules,
+		module_params,
+		lexical_results,
+		lexical_times,
+		0,
+		save_dir,
+		strategies,
+	)
+	lexical_selected_result, lexical_selected_filename = find_best(
+		lexical_results,
+		lexical_times,
+		lexical_summary_df["filename"].tolist(),
+		strategies,
+	)
+	lexical_summary_df["is_best"] = (
+		lexical_summary_df["filename"] == lexical_selected_filename
+	)
+	previous_result.drop(
+		columns=list(RETRIEVAL_METRIC_FUNC_DICT.keys()), inplace=True, errors="ignore"
+	)
+	lexical_selected_result.rename(
+		columns={
+			"retrieved_contents": "retrieved_contents_lexical",
+			"retrieved_ids": "retrieved_ids_lexical",
+			"retrieve_scores": "retrieve_scores_lexical",
+		},
+		inplace=True,
+	)
+	best_result = pd.concat([previous_result, lexical_selected_result], axis=1)
+	best_result.to_parquet(
+		os.path.join(
+			save_dir, f"best_{os.path.splitext(lexical_selected_filename)[0]}.parquet"
+		),
+		index=False,
+	)
+	lexical_summary_df.to_csv(os.path.join(save_dir, "summary.csv"), index=False)
+	return best_result
+
+
+def evaluate_lexical_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_lexical', 'retrieved_ids_lexical', 'retrieve_scores_lexical', and metric names.
+	"""
+
+	@evaluate_retrieval(
+		metric_inputs=metric_inputs,
+		metrics=metrics,
+	)
+	def evaluate_this_module(df: pd.DataFrame):
+		return (
+			df["retrieved_contents_lexical"].tolist(),
+			df["retrieved_ids_lexical"].tolist(),
+			df["retrieve_scores_lexical"].tolist(),
+		)
+
+	return evaluate_this_module(result_df)

autorag/nodes/passageaugmenter/__init__.py

@@ -0,0 +1,2 @@
+from .pass_passage_augmenter import PassPassageAugmenter
+from .prev_next_augmenter import PrevNextPassageAugmenter

autorag/nodes/passageaugmenter/base.py

@@ -0,0 +1,76 @@
+import abc
+import logging
+import os
+
+import pandas as pd
+
+from autorag.schema import BaseModule
+from autorag.utils import (
+	validate_qa_dataset,
+	sort_by_scores,
+	validate_corpus_dataset,
+	cast_corpus_dataset,
+)
+from autorag.utils.cast import cast_retrieved_ids
+from autorag.utils.util import select_top_k
+
+logger = logging.getLogger("AutoRAG")
+
+
+class BasePassageAugmenter(BaseModule, metaclass=abc.ABCMeta):
+	def __init__(self, project_dir: str, *args, **kwargs):
+		logger.info(
+			f"Initialize passage augmenter node - {self.__class__.__name__} module..."
+		)
+		data_dir = os.path.join(project_dir, "data")
+		corpus_df = pd.read_parquet(
+			os.path.join(data_dir, "corpus.parquet"), engine="pyarrow"
+		)
+		validate_corpus_dataset(corpus_df)
+		corpus_df = cast_corpus_dataset(corpus_df)
+		self.corpus_df = corpus_df
+
+	def __del__(self):
+		logger.info(
+			f"Initialize passage augmenter node - {self.__class__.__name__} module..."
+		)
+
+	def cast_to_run(self, previous_result: pd.DataFrame, *args, **kwargs):
+		logger.info(
+			f"Running passage augmenter node - {self.__class__.__name__} module..."
+		)
+		validate_qa_dataset(previous_result)
+
+		# find ids columns
+		return cast_retrieved_ids(previous_result)
+
+	@staticmethod
+	def sort_by_scores(
+		augmented_contents,
+		augmented_ids,
+		augmented_scores,
+		top_k: int,
+		reverse: bool = True,
+	):
+		# sort by scores
+		df = pd.DataFrame(
+			{
+				"contents": augmented_contents,
+				"ids": augmented_ids,
+				"scores": augmented_scores,
+			}
+		)
+		df[["contents", "ids", "scores"]] = df.apply(
+			lambda row: sort_by_scores(row, reverse=reverse),
+			axis=1,
+			result_type="expand",
+		)
+
+		# select by top_k
+		results = select_top_k(df, ["contents", "ids", "scores"], top_k)
+
+		return (
+			results["contents"].tolist(),
+			results["ids"].tolist(),
+			results["scores"].tolist(),
+		)

autorag/nodes/passageaugmenter/pass_passage_augmenter.py

@@ -0,0 +1,43 @@
+from typing import List
+
+import pandas as pd
+
+from autorag.nodes.passageaugmenter.base import BasePassageAugmenter
+from autorag.utils import result_to_dataframe
+
+
+class PassPassageAugmenter(BasePassageAugmenter):
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		"""
+		Run the passage augmenter node - PassPassageAugmenter 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)
+		contents = previous_result["retrieved_contents"].tolist()
+		scores = previous_result["retrieve_scores"].tolist()
+
+		augmented_ids, augmented_contents, augmented_scores = self._pure(
+			ids, contents, scores
+		)
+		return self.sort_by_scores(
+			augmented_contents, augmented_ids, augmented_scores, top_k
+		)
+
+	def _pure(
+		self,
+		ids_list: List[List[str]],
+		contents_list: List[List[str]],
+		scores_list: List[List[float]],
+	):
+		"""
+		Do not perform augmentation.
+		Return given passages, scores, and ids as is.
+		"""
+		return ids_list, contents_list, scores_list