AutoRAG 0.0.0__py3-none-any.whl

autorag/__init__.py

@@ -0,0 +1,82 @@
+import logging
+import sys
+from typing import Any
+
+from llama_index.core.base.llms.types import CompletionResponse
+from llama_index.core.llms.mock import MockLLM
+from llama_index.llms.bedrock import Bedrock
+
+from llama_index.llms.openai import OpenAI
+from llama_index.llms.openai_like import OpenAILike
+from rich.logging import RichHandler
+
+
+class LazyInit:
+	def __init__(self, factory, *args, **kwargs):
+		self._factory = factory
+		self._args = args
+		self._kwargs = kwargs
+		self._instance = None
+
+	def __call__(self):
+		if self._instance is None:
+			self._instance = self._factory(*self._args, **self._kwargs)
+		return self._instance
+
+	def __getattr__(self, name):
+		if self._instance is None:
+			self._instance = self._factory(*self._args, **self._kwargs)
+		return getattr(self._instance, name)
+
+
+rich_format = "[%(filename)s:%(lineno)s] >> %(message)s"
+logging.basicConfig(
+	level="INFO", format=rich_format, handlers=[RichHandler(rich_tracebacks=True)]
+)
+logger = logging.getLogger("AutoRAG")
+
+
+def handle_exception(exc_type, exc_value, exc_traceback):
+	logger = logging.getLogger("AutoRAG")
+	logger.error("Unexpected exception", exc_info=(exc_type, exc_value, exc_traceback))
+
+
+sys.excepthook = handle_exception
+
+
+class AutoRAGBedrock(Bedrock):
+	async def acomplete(
+		self, prompt: str, formatted: bool = False, **kwargs: Any
+	) -> CompletionResponse:
+		return self.complete(prompt, formatted=formatted, **kwargs)
+
+
+generator_models = {
+	"openai": OpenAI,
+	"openailike": OpenAILike,
+	"mock": MockLLM,
+	"bedrock": AutoRAGBedrock,
+}
+
+try:
+	from llama_index.llms.huggingface import HuggingFaceLLM
+	from llama_index.llms.ollama import Ollama
+
+	generator_models["huggingfacellm"] = HuggingFaceLLM
+	generator_models["ollama"] = Ollama
+
+except ImportError:
+	logger.info(
+		"You are using API version of AutoRAG."
+		"To use local version, run pip install 'AutoRAG[gpu]'"
+	)
+
+try:
+	import transformers
+
+	transformers.logging.set_verbosity_error()
+except ImportError:
+	logger.info(
+		"You are using API version of AutoRAG."
+		"To use local version, run pip install 'AutoRAG[gpu]'"
+	)

autorag/chunker.py

@@ -0,0 +1,51 @@
+import logging
+import os
+import shutil
+from typing import Optional
+
+import pandas as pd
+
+from autorag.data.chunk.run import run_chunker
+from autorag.data.utils.util import load_yaml, get_param_combinations
+
+logger = logging.getLogger("AutoRAG")
+
+
+class Chunker:
+	def __init__(self, raw_df: pd.DataFrame, project_dir: Optional[str] = None):
+		self.parsed_raw = raw_df
+		self.project_dir = project_dir if project_dir is not None else os.getcwd()
+
+	@classmethod
+	def from_parquet(
+		cls, parsed_data_path: str, project_dir: Optional[str] = None
+	) -> "Chunker":
+		if not os.path.exists(parsed_data_path):
+			raise ValueError(f"parsed_data_path {parsed_data_path} does not exist.")
+		if not parsed_data_path.endswith("parquet"):
+			raise ValueError(
+				f"parsed_data_path {parsed_data_path} is not a parquet file."
+			)
+		parsed_result = pd.read_parquet(parsed_data_path, engine="pyarrow")
+		return cls(parsed_result, project_dir)
+
+	def start_chunking(self, yaml_path: str):
+		if not os.path.exists(self.project_dir):
+			os.makedirs(self.project_dir)
+
+		# Copy YAML file to the trial directory
+		shutil.copy(yaml_path, os.path.join(self.project_dir, "chunk_config.yaml"))
+
+		# load yaml file
+		modules = load_yaml(yaml_path)
+
+		input_modules, input_params = get_param_combinations(modules)
+
+		logger.info("Chunking Start...")
+		run_chunker(
+			modules=input_modules,
+			module_params=input_params,
+			parsed_result=self.parsed_raw,
+			project_dir=self.project_dir,
+		)
+		logger.info("Chunking Done!")

autorag/cli.py

@@ -0,0 +1,209 @@
+import importlib.resources
+import logging
+import os
+import pathlib
+import subprocess
+from pathlib import Path
+from typing import Optional
+
+import click
+import nest_asyncio
+
+from autorag import dashboard
+from autorag.deploy import extract_best_config as original_extract_best_config
+from autorag.deploy.api import ApiRunner
+from autorag.evaluator import Evaluator
+from autorag.validator import Validator
+
+logger = logging.getLogger("AutoRAG")
+
+autorag_dir = os.path.dirname(os.path.realpath(__file__))
+version_file = os.path.join(autorag_dir, "VERSION")
+with open(version_file, "r") as f:
+	__version__ = f.read().strip()
+
+
+@click.group()
+@click.version_option(__version__)
+def cli():
+	pass
+
+
+@click.command()
+@click.option(
+	"--config",
+	"-c",
+	help="Path to config yaml file. Must be yaml or yml file.",
+	type=str,
+)
+@click.option(
+	"--qa_data_path", help="Path to QA dataset. Must be parquet file.", type=str
+)
+@click.option(
+	"--corpus_data_path", help="Path to corpus dataset. Must be parquet file.", type=str
+)
+@click.option(
+	"--project_dir", help="Path to project directory.", type=str, default=None
+)
+@click.option(
+	"--skip_validation",
+	help="Skip validation or not. Default is False.",
+	type=bool,
+	default=False,
+)
+def evaluate(config, qa_data_path, corpus_data_path, project_dir, skip_validation):
+	if not config.endswith(".yaml") and not config.endswith(".yml"):
+		raise ValueError(f"Config file {config} is not a yaml or yml file.")
+	if not os.path.exists(config):
+		raise ValueError(f"Config file {config} does not exist.")
+	evaluator = Evaluator(qa_data_path, corpus_data_path, project_dir=project_dir)
+	evaluator.start_trial(config, skip_validation=skip_validation)
+
+
+@click.command()
+@click.option(
+	"--config_path", type=str, help="Path to extracted config yaml file.", default=None
+)
+@click.option("--host", type=str, default="0.0.0.0", help="Host address")
+@click.option("--port", type=int, default=8000, help="Port number")
+@click.option(
+	"--trial_dir",
+	type=click.Path(file_okay=False, dir_okay=True, exists=True),
+	default=None,
+	help="Path to trial directory.",
+)
+@click.option(
+	"--project_dir", help="Path to project directory.", type=str, default=None
+)
+@click.option(
+	"--remote", help="Run the API server in remote mode.", type=bool, default=False
+)
+def run_api(config_path, host, port, trial_dir, project_dir, remote: bool):
+	if trial_dir is None:
+		runner = ApiRunner.from_yaml(config_path, project_dir=project_dir)
+	else:
+		runner = ApiRunner.from_trial_folder(trial_dir)
+	logger.info(f"Running API server at {host}:{port}...")
+	nest_asyncio.apply()
+	runner.run_api_server(host, port, remote=remote)
+
+
+@click.command()
+@click.option(
+	"--yaml_path", type=click.Path(path_type=Path), help="Path to the YAML file."
+)
+@click.option(
+	"--project_dir",
+	type=click.Path(path_type=Path),
+	help="Path to the project directory.",
+)
+@click.option(
+	"--trial_path", type=click.Path(path_type=Path), help="Path to the trial directory."
+)
+def run_web(
+	yaml_path: Optional[str], project_dir: Optional[str], trial_path: Optional[str]
+):
+	try:
+		with importlib.resources.path("autorag", "web.py") as web_path:
+			web_py_path = str(web_path)
+	except ImportError:
+		raise ImportError(
+			"Could not locate the web.py file within the autorag package."
+			" Please ensure that autorag is correctly installed."
+		)
+
+	if not yaml_path and not trial_path:
+		raise ValueError("yaml_path or trial_path must be given.")
+	elif yaml_path and trial_path:
+		raise ValueError("yaml_path and trial_path cannot be given at the same time.")
+	elif yaml_path and not project_dir:
+		subprocess.run(
+			["streamlit", "run", web_py_path, "--", "--yaml_path", yaml_path]
+		)
+	elif yaml_path and project_dir:
+		subprocess.run(
+			[
+				"streamlit",
+				"run",
+				web_py_path,
+				"--",
+				"--yaml_path",
+				yaml_path,
+				"--project_dir",
+				project_dir,
+			]
+		)
+	elif trial_path:
+		subprocess.run(
+			["streamlit", "run", web_py_path, "--", "--trial_path", trial_path]
+		)
+
+
+@click.command()
+@click.option(
+	"--trial_dir",
+	type=click.Path(dir_okay=True, file_okay=False, exists=True),
+	required=True,
+)
+@click.option(
+	"--port", type=int, default=7690, help="Port number. The default is 7690."
+)
+def run_dashboard(trial_dir: str, port: int):
+	dashboard.run(trial_dir, port=port)
+
+
+@click.command()
+@click.option("--trial_path", type=click.Path(), help="Path to the trial directory.")
+@click.option(
+	"--output_path",
+	type=click.Path(),
+	help="Path to the output directory. Must be .yaml or .yml file.",
+)
+def extract_best_config(trial_path: str, output_path: str):
+	original_extract_best_config(trial_path, output_path)
+
+
+@click.command()
+@click.option("--trial_path", help="Path to trial directory.", type=str)
+def restart_evaluate(trial_path):
+	if not os.path.exists(trial_path):
+		raise ValueError(f"trial_path {trial_path} does not exist.")
+	project_dir = str(pathlib.PurePath(trial_path).parent)
+	qa_data_path = os.path.join(project_dir, "data", "qa.parquet")
+	corpus_data_path = os.path.join(project_dir, "data", "corpus.parquet")
+	evaluator = Evaluator(qa_data_path, corpus_data_path, project_dir)
+	evaluator.restart_trial(trial_path)
+
+
+@click.command()
+@click.option(
+	"--config",
+	"-c",
+	help="Path to config yaml file. Must be yaml or yml file.",
+	type=str,
+)
+@click.option(
+	"--qa_data_path", help="Path to QA dataset. Must be parquet file.", type=str
+)
+@click.option(
+	"--corpus_data_path", help="Path to corpus dataset. Must be parquet file.", type=str
+)
+def validate(config, qa_data_path, corpus_data_path):
+	if not config.endswith(".yaml") and not config.endswith(".yml"):
+		raise ValueError(f"Config file {config} is not a parquet file.")
+	if not os.path.exists(config):
+		raise ValueError(f"Config file {config} does not exist.")
+	validator = Validator(qa_data_path=qa_data_path, corpus_data_path=corpus_data_path)
+	validator.validate(config)
+
+
+cli.add_command(evaluate, "evaluate")
+cli.add_command(run_api, "run_api")
+cli.add_command(run_web, "run_web")
+cli.add_command(run_dashboard, "dashboard")
+cli.add_command(extract_best_config, "extract_best_config")
+cli.add_command(restart_evaluate, "restart_evaluate")
+cli.add_command(validate, "validate")
+
+if __name__ == "__main__":
+	cli()

autorag/dashboard.py

@@ -0,0 +1,199 @@
+import ast
+import logging
+import os
+from typing import Dict, List
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import panel as pn
+import seaborn as sns
+import yaml
+from bokeh.models import NumberFormatter, BooleanFormatter
+
+from autorag.utils.util import dict_to_markdown, dict_to_markdown_table
+
+pn.extension(
+	"terminal",
+	"tabulator",
+	"mathjax",
+	"ipywidgets",
+	console_output="disable",
+	sizing_mode="stretch_width",
+	css_files=[
+		"https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/all.min.css"
+	],
+)
+logger = logging.getLogger("AutoRAG")
+
+
+def find_node_dir(trial_dir: str) -> List[str]:
+	trial_summary_df = pd.read_csv(os.path.join(trial_dir, "summary.csv"))
+	result_paths = []
+	for idx, row in trial_summary_df.iterrows():
+		node_line_name = row["node_line_name"]
+		node_type = row["node_type"]
+		result_paths.append(os.path.join(trial_dir, node_line_name, node_type))
+	return result_paths
+
+
+def get_metric_values(node_summary_df: pd.DataFrame) -> Dict:
+	non_metric_column_names = [
+		"filename",
+		"module_name",
+		"module_params",
+		"execution_time",
+		"average_output_token",
+		"is_best",
+	]
+	best_row = node_summary_df.loc[node_summary_df["is_best"]].drop(
+		columns=non_metric_column_names, errors="ignore"
+	)
+	assert len(best_row) == 1, "The best module must be only one."
+	return best_row.iloc[0].to_dict()
+
+
+def make_trial_summary_md(trial_dir):
+	markdown_text = f"""# Trial Result Summary
+- Trial Directory : {trial_dir}
+
+"""
+	node_dirs = find_node_dir(trial_dir)
+	for node_dir in node_dirs:
+		node_summary_filepath = os.path.join(node_dir, "summary.csv")
+		node_type = os.path.basename(node_dir)
+		node_summary_df = pd.read_csv(node_summary_filepath)
+		best_row = node_summary_df.loc[node_summary_df["is_best"]].iloc[0]
+		metric_dict = get_metric_values(node_summary_df)
+		markdown_text += f"""---
+
+## {node_type} best module
+
+### Module Name
+
+{best_row["module_name"]}
+
+### Module Params
+
+{dict_to_markdown(ast.literal_eval(best_row["module_params"]), level=3)}
+
+### Metric Values
+
+{dict_to_markdown_table(metric_dict, key_column_name="metric_name", value_column_name="metric_value")}
+
+"""
+
+	return markdown_text
+
+
+def node_view(node_dir: str):
+	non_metric_column_names = [
+		"filename",
+		"module_name",
+		"module_params",
+		"execution_time",
+		"average_output_token",
+		"is_best",
+	]
+	summary_df = pd.read_csv(os.path.join(node_dir, "summary.csv"))
+	bokeh_formatters = {
+		"float": NumberFormatter(format="0.000"),
+		"bool": BooleanFormatter(),
+	}
+	first_df = pd.read_parquet(os.path.join(node_dir, "0.parquet"), engine="pyarrow")
+
+	each_module_df_widget = pn.widgets.Tabulator(
+		pd.DataFrame(columns=first_df.columns),
+		name="Module DataFrame",
+		formatters=bokeh_formatters,
+		pagination="local",
+		page_size=20,
+		widths=150,
+	)
+
+	def change_module_widget(event):
+		if event.column == "detail":
+			filename = summary_df["filename"].iloc[event.row]
+			filepath = os.path.join(node_dir, filename)
+			each_module_df = pd.read_parquet(filepath, engine="pyarrow")
+			each_module_df_widget.value = each_module_df
+
+	df_widget = pn.widgets.Tabulator(
+		summary_df,
+		name="Summary DataFrame",
+		formatters=bokeh_formatters,
+		buttons={"detail": '<i class="fa fa-eye"></i>'},
+		widths=150,
+	)
+	df_widget.on_click(change_module_widget)
+
+	try:
+		fig, ax = plt.subplots(figsize=(10, 5))
+		metric_df = summary_df.drop(columns=non_metric_column_names, errors="ignore")
+		sns.stripplot(data=metric_df, ax=ax)
+		strip_plot_pane = pn.pane.Matplotlib(fig, tight=True)
+
+		fig2, ax2 = plt.subplots(figsize=(10, 5))
+		sns.boxplot(data=metric_df, ax=ax2)
+		box_plot_pane = pn.pane.Matplotlib(fig2, tight=True)
+		plot_pane = pn.Row(strip_plot_pane, box_plot_pane)
+
+		layout = pn.Column(
+			"## Summary distribution plot",
+			plot_pane,
+			"## Summary DataFrame",
+			df_widget,
+			"## Module Result DataFrame",
+			each_module_df_widget,
+		)
+	except Exception as e:
+		logger.error(f"Skipping make boxplot and stripplot with error {e}")
+		layout = pn.Column("## Summary DataFrame", df_widget)
+	layout.servable()
+	return layout
+
+
+CSS = """
+div.card-margin:nth-child(1) {
+    max-height: 300px;
+}
+div.card-margin:nth-child(2) {
+    max-height: 400px;
+}
+"""
+
+
+def yaml_to_markdown(yaml_filepath):
+	markdown_content = ""
+	with open(yaml_filepath, "r", encoding="utf-8") as file:
+		try:
+			content = yaml.safe_load(file)
+			markdown_content += f"## {os.path.basename(yaml_filepath)}\n```yaml\n{yaml.safe_dump(content, allow_unicode=True)}\n```\n\n"
+		except yaml.YAMLError as exc:
+			print(f"Error in {yaml_filepath}: {exc}")
+	return markdown_content
+
+
+def run(trial_dir: str, port: int = 7690):
+	trial_summary_md = make_trial_summary_md(trial_dir=trial_dir)
+	trial_summary_tab = pn.pane.Markdown(trial_summary_md, sizing_mode="stretch_width")
+
+	node_views = [
+		(str(os.path.basename(node_dir)), node_view(node_dir))
+		for node_dir in find_node_dir(trial_dir)
+	]
+
+	yaml_file_markdown = yaml_to_markdown(os.path.join(trial_dir, "config.yaml"))
+
+	yaml_file = pn.pane.Markdown(yaml_file_markdown, sizing_mode="stretch_width")
+
+	tabs = pn.Tabs(
+		("Summary", trial_summary_tab),
+		*node_views,
+		("Used YAML file", yaml_file),
+		dynamic=True,
+	)
+
+	template = pn.template.FastListTemplate(
+		site="AutoRAG", title="Dashboard", main=[tabs], raw_css=[CSS]
+	).servable()
+	template.show(port=port)

autorag/data/__init__.py

@@ -0,0 +1,109 @@
+import logging
+from typing import List, Callable
+
+from langchain_community.document_loaders import (
+	PDFMinerLoader,
+	PDFPlumberLoader,
+	PyPDFium2Loader,
+	PyPDFLoader,
+	PyMuPDFLoader,
+	UnstructuredPDFLoader,
+	CSVLoader,
+	JSONLoader,
+	UnstructuredMarkdownLoader,
+	BSHTMLLoader,
+	UnstructuredXMLLoader,
+	DirectoryLoader,
+)
+from langchain_unstructured import UnstructuredLoader
+from langchain_upstage import UpstageLayoutAnalysisLoader
+
+from llama_index.core.node_parser import (
+	TokenTextSplitter,
+	SentenceSplitter,
+	SentenceWindowNodeParser,
+	SemanticSplitterNodeParser,
+	SemanticDoubleMergingSplitterNodeParser,
+	SimpleFileNodeParser,
+)
+from langchain.text_splitter import (
+	RecursiveCharacterTextSplitter,
+	CharacterTextSplitter,
+	KonlpyTextSplitter,
+	SentenceTransformersTokenTextSplitter,
+)
+
+from autorag import LazyInit
+
+logger = logging.getLogger("AutoRAG")
+
+parse_modules = {
+	# PDF
+	"pdfminer": PDFMinerLoader,
+	"pdfplumber": PDFPlumberLoader,
+	"pypdfium2": PyPDFium2Loader,
+	"pypdf": PyPDFLoader,
+	"pymupdf": PyMuPDFLoader,
+	"unstructuredpdf": UnstructuredPDFLoader,
+	# Common File Types
+	# 1. CSV
+	"csv": CSVLoader,
+	# 2. JSON
+	"json": JSONLoader,
+	# 3. Markdown
+	"unstructuredmarkdown": UnstructuredMarkdownLoader,
+	# 4. HTML
+	"bshtml": BSHTMLLoader,
+	# 5. XML
+	"unstructuredxml": UnstructuredXMLLoader,
+	# 6. All files
+	"directory": DirectoryLoader,
+	"unstructured": UnstructuredLoader,
+	"upstagedocumentparse": UpstageLayoutAnalysisLoader,
+}
+
+chunk_modules = {
+	# Llama Index
+	# Token
+	"token": TokenTextSplitter,
+	# Sentence
+	"sentence": SentenceSplitter,
+	# window
+	"sentencewindow": SentenceWindowNodeParser,
+	# Semantic
+	"semantic_llama_index": SemanticSplitterNodeParser,
+	"semanticdoublemerging": SemanticDoubleMergingSplitterNodeParser,
+	# Simple
+	"simplefile": SimpleFileNodeParser,
+	# LangChain
+	# Token
+	"sentencetransformerstoken": SentenceTransformersTokenTextSplitter,
+	# Character
+	"recursivecharacter": RecursiveCharacterTextSplitter,
+	"character": CharacterTextSplitter,
+	# Sentence
+	"konlpy": KonlpyTextSplitter,
+}
+
+
+def split_by_sentence_kiwi() -> Callable[[str], List[str]]:
+	try:
+		from kiwipiepy import Kiwi
+	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]'."
+		)
+	kiwi = Kiwi()
+
+	def split(text: str) -> List[str]:
+		kiwi_result = kiwi.split_into_sents(text)
+		sentences = list(map(lambda x: x.text, kiwi_result))
+
+		return sentences
+
+	return split
+
+
+sentence_splitter_modules = {"kiwi": LazyInit(split_by_sentence_kiwi)}

autorag/data/chunk/__init__.py

@@ -0,0 +1,2 @@
+from .llama_index_chunk import llama_index_chunk
+from .langchain_chunk import langchain_chunk