AutoRAG 0.0.0__py3-none-any.whl

autorag/data/parse/run.py

@@ -0,0 +1,141 @@
+import os
+from typing import List, Callable, Dict
+import pandas as pd
+from glob import glob
+
+from autorag.strategy import measure_speed
+from autorag.data.utils.util import get_param_combinations
+
+default_map = {
+	"pdf": {
+		"file_type": "pdf",
+		"module_type": "langchain_parse",
+		"parse_method": "pdfminer",
+	},
+	"csv": {
+		"file_type": "csv",
+		"module_type": "langchain_parse",
+		"parse_method": "csv",
+	},
+	"md": {
+		"file_type": "md",
+		"module_type": "langchain_parse",
+		"parse_method": "unstructuredmarkdown",
+	},
+	"html": {
+		"file_type": "html",
+		"module_type": "langchain_parse",
+		"parse_method": "bshtml",
+	},
+	"xml": {
+		"file_type": "xml",
+		"module_type": "langchain_parse",
+		"parse_method": "unstructuredxml",
+	},
+}
+
+
+def run_parser(
+	modules: List[Callable],
+	module_params: List[Dict],
+	data_path_glob: str,
+	project_dir: str,
+	all_files: bool,
+):
+	if not all_files:
+		# Set the parsing module to default if it is a file type in paths but not set in YAML.
+		data_path_list = glob(data_path_glob)
+		if not data_path_list:
+			raise FileNotFoundError(f"data does not exits in {data_path_glob}")
+
+		file_types = set(
+			[os.path.basename(data_path).split(".")[-1] for data_path in data_path_list]
+		)
+		set_file_types = set([module["file_type"] for module in module_params])
+
+		# Calculate the set difference once
+		file_types_to_remove = set_file_types - file_types
+
+		# Use list comprehension to filter out unwanted elements
+		module_params = [
+			param
+			for param in module_params
+			if param["file_type"] not in file_types_to_remove
+		]
+		modules = [
+			module
+			for module, param in zip(modules, module_params)
+			if param["file_type"] not in file_types_to_remove
+		]
+
+		# create a list of only those file_types that are in file_types but not in set_file_types
+		missing_file_types = list(file_types - set_file_types)
+
+		if missing_file_types:
+			add_modules_list = []
+			for missing_file_type in missing_file_types:
+				if missing_file_type == "json":
+					raise ValueError(
+						"JSON file type must have a jq_schema so you must set it in the YAML file."
+					)
+
+				add_modules_list.append(default_map[missing_file_type])
+
+			add_modules, add_params = get_param_combinations(add_modules_list)
+			modules.extend(add_modules)
+			module_params.extend(add_params)
+
+	results, execution_times = zip(
+		*map(
+			lambda x: measure_speed(x[0], data_path_glob=data_path_glob, **x[1]),
+			zip(modules, module_params),
+		)
+	)
+	average_times = list(map(lambda x: x / len(results[0]), execution_times))
+
+	# save results to parquet files
+	if all_files:
+		if len(module_params) > 1:
+			raise ValueError(
+				"All files is set to True, You can only use one parsing module."
+			)
+		filepaths = [os.path.join(project_dir, "parsed_result.parquet")]
+	else:
+		filepaths = list(
+			map(
+				lambda x: os.path.join(project_dir, f"{x['file_type']}.parquet"),
+				module_params,
+			)
+		)
+
+	_files = {}
+	for result, filepath in zip(results, filepaths):
+		_files[filepath].append(result) if filepath in _files.keys() else _files.update(
+			{filepath: [result]}
+		)
+	# Save files with a specific file type as Parquet files.
+	for filepath, value in _files.items():
+		pd.concat(value).to_parquet(filepath, index=False)
+
+	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,
+		}
+	)
+	summary_df.to_csv(os.path.join(project_dir, "summary.csv"), index=False)
+
+	# concat all parquet files here if not all_files.
+	_filepaths = list(_files.keys())
+	if not all_files:
+		dataframes = [pd.read_parquet(file) for file in _filepaths]
+		combined_df = pd.concat(dataframes, ignore_index=True)
+		combined_df.to_parquet(
+			os.path.join(project_dir, "parsed_result.parquet"), index=False
+		)
+
+	return summary_df

autorag/data/parse/table_hybrid_parse.py

@@ -0,0 +1,134 @@
+import os
+import tempfile
+from glob import glob
+from typing import List, Tuple, Dict
+
+from PyPDF2 import PdfFileReader, PdfFileWriter
+import pdfplumber
+
+from autorag.support import get_support_modules
+from autorag.data.parse.base import parser_node
+
+
+@parser_node
+def table_hybrid_parse(
+	data_path_list: List[str],
+	text_parse_module: str,
+	text_params: Dict,
+	table_parse_module: str,
+	table_params: Dict,
+) -> Tuple[List[str], List[str], List[int]]:
+	"""
+	Parse documents to use table_hybrid_parse method.
+	The table_hybrid_parse method is a hybrid method that combines the parsing results of PDFs with and without tables.
+	It splits the PDF file into pages, separates pages with and without tables, and then parses and merges the results.
+
+	:param data_path_list: The list of data paths to parse.
+	:param text_parse_module: The text parsing module to use. The type should be a string.
+	:param text_params: The extra parameters for the text parsing module. The type should be a dictionary.
+	:param table_parse_module: The table parsing module to use. The type should be a string.
+	:param table_params: The extra parameters for the table parsing module. The type should be a dictionary.
+	:return: tuple of lists containing the parsed texts, path and pages.
+	"""
+	# make save folder directory
+	with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as save_dir:
+		text_dir = os.path.join(save_dir, "text")
+		table_dir = os.path.join(save_dir, "table")
+
+		os.makedirs(text_dir, exist_ok=True)
+		os.makedirs(table_dir, exist_ok=True)
+
+		# Split PDF file into pages and Save PDFs with and without tables
+		path_map_dict_lst = [
+			save_page_by_table(data_path, text_dir, table_dir)
+			for data_path in data_path_list
+		]
+		path_map_dict = {k: v for d in path_map_dict_lst for k, v in d.items()}
+
+		# Extract text pages
+		table_results, table_file_path = get_each_module_result(
+			table_parse_module, table_params, os.path.join(table_dir, "*")
+		)
+
+		# Extract table pages
+		text_results, text_file_path = get_each_module_result(
+			text_parse_module, text_params, os.path.join(text_dir, "*")
+		)
+
+		# Merge parsing results of PDFs with and without tables
+		texts = table_results + text_results
+		temp_path_lst = table_file_path + text_file_path
+
+		# Sort by file names
+		temp_path_lst, texts = zip(*sorted(zip(temp_path_lst, texts)))
+
+		# get original file path
+		path = list(map(lambda temp_path: path_map_dict[temp_path], temp_path_lst))
+
+		# get pages
+		pages = list(map(lambda x: get_page_from_path(x), temp_path_lst))
+
+		return list(texts), path, pages
+
+
+# Save PDFs with and without tables
+def save_page_by_table(data_path: str, text_dir: str, table_dir: str) -> Dict[str, str]:
+	file_name = os.path.basename(data_path).split(".pdf")[0]
+
+	with open(data_path, "rb") as input_data:
+		pdf_reader = PdfFileReader(input_data)
+		num_pages = pdf_reader.getNumPages()
+
+		path_map_dict = {}
+		for page_num in range(num_pages):
+			output_pdf_path = _get_output_path(
+				data_path, page_num, file_name, text_dir, table_dir
+			)
+			_save_single_page(pdf_reader, page_num, output_pdf_path)
+			path_map_dict.update({output_pdf_path: data_path})
+
+	return path_map_dict
+
+
+def _get_output_path(
+	data_path: str, page_num: int, file_name: str, text_dir: str, table_dir: str
+) -> str:
+	with pdfplumber.open(data_path) as pdf:
+		page = pdf.pages[page_num]
+		tables = page.extract_tables()
+		directory = table_dir if tables else text_dir
+		return os.path.join(directory, f"{file_name}_page_{page_num + 1}.pdf")
+
+
+def _save_single_page(pdf_reader: PdfFileReader, page_num: int, output_pdf_path: str):
+	pdf_writer = PdfFileWriter()
+	pdf_writer.addPage(pdf_reader.getPage(page_num))
+
+	with open(output_pdf_path, "wb") as output_file:
+		pdf_writer.write(output_file)
+
+
+def get_each_module_result(
+	module: str, module_params: Dict, data_path_glob: str
+) -> Tuple[List[str], List[str]]:
+	module_params["module_type"] = module
+
+	data_path_list = glob(data_path_glob)
+	if not data_path_list:
+		return [], []
+
+	module_name = module_params.pop("module_type")
+	module_callable = get_support_modules(module_name)
+	module_original = module_callable.__wrapped__
+	texts, path, _ = module_original(data_path_list, **module_params)
+
+	return texts, path
+
+
+def get_page_from_path(file_path: str) -> int:
+	file_name = os.path.basename(file_path)
+	split_result = file_name.rsplit("_page_", -1)
+	page_number_with_extension = split_result[1]
+	page_number, _ = page_number_with_extension.split(".")
+
+	return int(page_number)

autorag/data/qa/__init__.py

@@ -0,0 +1,3 @@
+# This is v2 version, the next version of data creation
+# The legacy (v1) version will be deprecated on AutoRAG version 0.3
+# The legacy (v1) version and new v2 data creation is not compatible with each other

autorag/data/qa/evolve/llama_index_query_evolve.py

@@ -0,0 +1,64 @@
+import itertools
+from typing import Dict, List
+
+from llama_index.core.base.llms.base import BaseLLM
+from llama_index.core.base.llms.types import ChatResponse, ChatMessage, MessageRole
+
+from autorag.data.qa.evolve.prompt import QUERY_EVOLVE_PROMPT
+
+
+async def llama_index_generate_base(
+	row: Dict,
+	llm: BaseLLM,
+	messages: List[ChatMessage],
+) -> Dict:
+	original_query = row["query"]
+	context = list(itertools.chain.from_iterable(row["retrieval_gt_contents"]))
+	context_str = "Text:\n" + "\n".join(
+		[f"{i + 1}. {c}" for i, c in enumerate(context)]
+	)
+	user_prompt = f"Question: {original_query}\nContext: {context_str}\nOutput: "
+	messages.append(ChatMessage(role=MessageRole.USER, content=user_prompt))
+
+	chat_response: ChatResponse = await llm.achat(messages=messages)
+	row["query"] = chat_response.message.content
+	return row
+
+
+async def conditional_evolve_ragas(
+	row: Dict,
+	llm: BaseLLM,
+	lang: str = "en",
+) -> Dict:
+	return await llama_index_generate_base(
+		row,
+		llm,
+		QUERY_EVOLVE_PROMPT["conditional_evolve_ragas"][lang],
+	)
+
+
+async def reasoning_evolve_ragas(
+	row: Dict,
+	llm: BaseLLM,
+	lang: str = "en",
+) -> Dict:
+	return await llama_index_generate_base(
+		row,
+		llm,
+		QUERY_EVOLVE_PROMPT["reasoning_evolve_ragas"][lang],
+	)
+
+
+async def compress_ragas(
+	row: Dict,
+	llm: BaseLLM,
+	lang: str = "en",
+) -> Dict:
+	original_query = row["query"]
+	user_prompt = f"Question: {original_query}\nOutput: "
+	messages = QUERY_EVOLVE_PROMPT["compress_ragas"][lang]
+	messages.append(ChatMessage(role=MessageRole.USER, content=user_prompt))
+
+	chat_response: ChatResponse = await llm.achat(messages=messages)
+	row["query"] = chat_response.message.content
+	return row

autorag/data/qa/evolve/openai_query_evolve.py

@@ -0,0 +1,81 @@
+import itertools
+from typing import Dict, List
+
+from llama_index.core.base.llms.types import ChatMessage, MessageRole
+from llama_index.llms.openai.utils import to_openai_message_dicts
+from openai import AsyncClient
+from pydantic import BaseModel
+
+from autorag.data.qa.evolve.prompt import QUERY_EVOLVE_PROMPT
+
+
+class Response(BaseModel):
+	evolved_query: str
+
+
+async def query_evolve_openai_base(
+	row: Dict,
+	client: AsyncClient,
+	messages: List[ChatMessage],
+	model_name: str = "gpt-4o-2024-08-06",
+):
+	"""
+	Evolve the original query to a new evolved query using OpenAI structured outputs.
+	"""
+	original_query = row["query"]
+	context = list(itertools.chain.from_iterable(row["retrieval_gt_contents"]))
+	context_str = "Text:\n" + "\n".join(
+		[f"{i + 1}. {c}" for i, c in enumerate(context)]
+	)
+	user_prompt = f"Question: {original_query}\nContext: {context_str}\nOutput: "
+	messages.append(ChatMessage(role=MessageRole.USER, content=user_prompt))
+
+	completion = await client.beta.chat.completions.parse(
+		model=model_name,
+		messages=to_openai_message_dicts(messages),
+		response_format=Response,
+	)
+	row["query"] = completion.choices[0].message.parsed.evolved_query
+	return row
+
+
+async def conditional_evolve_ragas(
+	row: Dict,
+	client: AsyncClient,
+	model_name: str = "gpt-4o-2024-08-06",
+	lang: str = "en",
+) -> Dict:
+	return await query_evolve_openai_base(
+		row, client, QUERY_EVOLVE_PROMPT["conditional_evolve_ragas"][lang], model_name
+	)
+
+
+async def reasoning_evolve_ragas(
+	row: Dict,
+	client: AsyncClient,
+	model_name: str = "gpt-4o-2024-08-06",
+	lang: str = "en",
+) -> Dict:
+	return await query_evolve_openai_base(
+		row, client, QUERY_EVOLVE_PROMPT["reasoning_evolve_ragas"][lang], model_name
+	)
+
+
+async def compress_ragas(
+	row: Dict,
+	client: AsyncClient,
+	model_name: str = "gpt-4o-2024-08-06",
+	lang: str = "en",
+) -> Dict:
+	original_query = row["query"]
+	messages = QUERY_EVOLVE_PROMPT["compress_ragas"][lang]
+	user_prompt = f"Question: {original_query}\nOutput: "
+	messages.append(ChatMessage(role=MessageRole.USER, content=user_prompt))
+
+	completion = await client.beta.chat.completions.parse(
+		model=model_name,
+		messages=to_openai_message_dicts(messages),
+		response_format=Response,
+	)
+	row["query"] = completion.choices[0].message.parsed.evolved_query
+	return row