app.py

import os
import shutil
from typing import List

import gradio as gr
import pandas as pd
from autorag.data.parse import langchain_parse
from autorag.data.parse.llamaparse import llama_parse
from autorag.data.qa.schema import Raw
from llama_index.llms.openai import OpenAI

from src.create import default_create, fast_create, advanced_create
from src.util import on_submit_openai_key, on_submit_llama_cloud_key, on_submit_upstage_key

root_dir = os.path.dirname(os.path.realpath(__file__))
FILE_DIR = os.path.join(root_dir, "file_cache")
if not os.path.exists(FILE_DIR):
	os.makedirs(FILE_DIR)
DATA_DIR = os.path.join(root_dir, "data")
if not os.path.exists(DATA_DIR):
	os.makedirs(DATA_DIR)


def change_lang_choice(lang: str) -> str:
	lang_dict = {
		"English": "en",
		"한국어": "ko",
		"日本語": "ja"
	}
	return lang_dict[lang]

def change_visible_status_api_key(parse_method: str):
	if parse_method == "llama-parse":
		return gr.update(visible=True), gr.update(visible=False)
	elif parse_method == "upstage🇰🇷":
		return gr.update(visible=False), gr.update(visible=True)
	else:
		return gr.update(visible=False), gr.update(visible=False)



def run_parse(file_lists: List[str], parse_method: str, progress=gr.Progress()):
	# save an input file to a directory
	for file_path in file_lists:
		shutil.copy(file_path, FILE_DIR)
	progress(0.05)

	if parse_method in ["pdfminer", "pdfplumber", "pypdfium2", "pypdf", "pymupdf"]:
		raw_df: pd.DataFrame = langchain_parse(data_path_glob=os.path.join(FILE_DIR, "*.pdf"), parse_method=parse_method)
	elif parse_method == "llama-parse":
		llama_cloud_api_key = os.getenv("LLAMA_CLOUD_API_KEY")
		if llama_cloud_api_key is None:
			return "Please submit your Llama Cloud API key first."
		raw_df: pd.DataFrame = llama_parse(data_path_glob=os.path.join(FILE_DIR, "*.pdf"))
	elif parse_method == "upstage🇰🇷":
		upstage_api_key = os.getenv("UPSTAGE_API_KEY")
		if upstage_api_key is None:
			return "Please submit your Upstage API key first."
		raw_df: pd.DataFrame = langchain_parse(data_path_glob=os.path.join(FILE_DIR, "*.pdf"), parse_method="upstagedocumentparse")
	else:
		return "Unsupported parse method."
	progress(0.8)

	raw_df.to_parquet(os.path.join(DATA_DIR, "raw.parquet"), index=False)
	return "Parsing Complete. Download at the bottom button."


def run_chunk(use_existed_raw: bool, raw_file: str, chunk_method: str, chunk_size: int, chunk_overlap: int,
			  lang: str = "English", progress=gr.Progress()):
	lang = change_lang_choice(lang)
	if use_existed_raw:
		raw_df_path = os.path.join(DATA_DIR, "raw.parquet")
	else:
		raw_df_path = raw_file

	if not os.path.exists(raw_df_path):
		return "Please upload raw.parquet file first. Or run the parsing stage first."
	raw_df = pd.read_parquet(raw_df_path, engine="pyarrow")
	raw_instance = Raw(raw_df)

	if chunk_method in ["Token", "Sentence"]:
		corpus = raw_instance.chunk("llama_index_chunk", chunk_method=chunk_method, chunk_size=chunk_size,
									chunk_overlap=chunk_overlap, add_file_name=lang)
	elif chunk_method in ["Semantic"]:
		corpus = raw_instance.chunk("llama_index_chunk", chunk_method="Semantic_llama_index",
									embed_model="openai", breakpoint_percnetile_threshold=0.95,
									add_file_name=lang)
	elif chunk_method == "Recursive":
		corpus = raw_instance.chunk("langchain_chunk", chunk_method="recursivecharacter",
									add_file_name=lang, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
	else:
		gr.Error("Unsupported chunk method.")
		return "Unsupported chunk method."
	progress(0.8)
	corpus.to_parquet(os.path.join(DATA_DIR, "corpus.parquet"))
	return "Chunking Complete. Download at the bottom button."


def run_qa(use_existed_corpus: bool, corpus_file: str, qa_method: str,
		   model_name: str, qa_cnt: int, batch_size: int, lang: str = "English", progress=gr.Progress()):
	lang = change_lang_choice(lang)
	if use_existed_corpus:
		corpus_df_path = os.path.join(DATA_DIR, "corpus.parquet")
	else:
		corpus_df_path = corpus_file

	if not os.path.exists(corpus_df_path):
		gr.Error("Please upload corpus.parquet file first. Or run the chunking stage first.")
		return "Please upload corpus.parquet file first. Or run the chunking stage first."
	corpus_df = pd.read_parquet(corpus_df_path, engine="pyarrow")

	if os.getenv("OPENAI_API_KEY") is None:
		gr.Error("Please submit your OpenAI API key first.")
		return "Please submit your OpenAI API key first."
	llm = OpenAI(model=model_name)

	if qa_method == "default":
		qa = default_create(corpus_df, llm=llm, n=qa_cnt, lang=lang, progress=progress, batch_size=batch_size)
	elif qa_method == "fast":
		qa = fast_create(corpus_df, llm=llm, n=qa_cnt, lang=lang, progress=progress, batch_size=batch_size)
	elif qa_method == "advanced":
		qa = advanced_create(corpus_df, llm=llm, n=qa_cnt, lang=lang, progress=progress, batch_size=batch_size)
	else:
		gr.Error("Unsupported QA method.")
		return "Unsupported QA method."

	qa.to_parquet(os.path.join(DATA_DIR, "qa.parquet"), os.path.join(DATA_DIR, "corpus.parquet"))
	return "QA Creation Complete. Download at the bottom button."


def file_reset() -> str:
	shutil.rmtree(FILE_DIR)
	os.makedirs(FILE_DIR)
	return "Files reset complete."

with gr.Blocks(theme="earneleh/paris") as demo:
	gr.HTML("<h1>AutoRAG Data Creation 🛠️</h1>")
	with gr.Row():
		openai_key_textbox = gr.Textbox(label="Please input your OpenAI API key and press Enter.", type="password",
										info="You can get your API key from https://platform.openai.com/account/api-keys\n\n"
											 "AutoRAG do not store your API key.",
										autofocus=True)
		api_key_status_box = gr.Textbox(label="OpenAI API status", value="Not Set", interactive=False)
		lang_choice = gr.Radio(["English", "한국어", "日本語"], label="Language",
									   value="English", info="Choose Langauge. En, Ko, Ja are supported.",
									   interactive=True)

	with gr.Row(visible=False) as llama_cloud_api_key_row:
		llama_key_textbox = gr.Textbox(label="Please input your Llama Cloud API key and press Enter.", type="password",
									   		info="You can get your API key from https://docs.cloud.llamaindex.ai/llamacloud/getting_started/api_key\n\n"
											 "AutoRAG do not store your API key.",)
		llama_key_status_box = gr.Textbox(label="Llama Cloud API status", value="Not Set", interactive=False)

	with gr.Row(visible=False) as upstage_api_key_row:
		upstage_key_textbox = gr.Textbox(label="Please input your Upstage API key and press Enter.", type="password",
									   		info="You can get your API key from https://upstage.ai/\n\n"
											 "AutoRAG do not store your API key.",)
		upstage_key_status_box = gr.Textbox(label="Upstage API status", value="Not Set", interactive=False)

	with gr.Row():
		with gr.Column(scale=1):
			gr.Markdown("## 1. Parse your PDF files\n\nUpload your pdf files and make it to raw.parquet.")
			document_file_input = gr.File(label="Upload Files", type="filepath", file_count="multiple")
			parse_choice = gr.Dropdown(
				["pdfminer", "pdfplumber", "pypdfium2", "pypdf", "pymupdf", "llama-parse", "upstage🇰🇷"],
				label="Parsing Method", info="Choose parsing method that you want")
			parse_button = gr.Button(value="Run Parsing")
			parse_status = gr.Textbox(value="Not Started", interactive=False)
			raw_download_button = gr.Button(value="Download raw.parquet",
											link=f"/file={os.path.join(DATA_DIR, 'raw.parquet')}")
			file_reset_button = gr.Button(value="Reset uploaded files")

		with gr.Column(scale=1):
			gr.Markdown(
				"## 2. Chunk your raw.parquet\n\nUse parsed raw.parquet or upload your own. It will make a corpus.parquet."
			)
			raw_file_input = gr.File(label="Upload raw.parquet", type="filepath", file_count="single", visible=False)
			use_previous_raw_file = gr.Checkbox(label="Use previous raw.parquet", value=True)

			chunk_choice = gr.Dropdown(
				["Token", "Sentence", "Semantic", "Recursive"],
				label="Chunking Method", info="Choose chunking method that you want")
			chunk_size = gr.Slider(minimum=128, maximum=1024, step=128, label="Chunk Size", value=256)
			chunk_overlap = gr.Slider(minimum=16, maximum=256, step=16, label="Chunk Overlap", value=32)
			chunk_button = gr.Button(value="Run Chunking")
			chunk_status = gr.Textbox(value="Not Started", interactive=False)
			corpus_download_button = gr.Button(value="Download corpus.parquet",
											   link=f"/file={os.path.join(DATA_DIR, 'corpus.parquet')}")

		with gr.Column(scale=1):
			gr.Markdown(
				"## 3. Create QA dataset from your corpus.parquet\n\nQA dataset is essential to run AutoRAG. Upload corpus.parquet & select QA method and run.")
			gr.HTML("<b style='color: red; background-color: black; font-weight: bold;'>Warning: QA Creation uses an OpenAI model, which can be costly. Start with a small batch to gauge expenses.</b>")
			corpus_file_input = gr.File(label="Upload corpus.parquet", type="filepath", file_count="single",
										visible=False)
			use_previous_corpus_file = gr.Checkbox(label="Use previous corpus.parquet", value=True)

			qa_choice = gr.Radio(["default", "fast", "advanced"], label="QA Method",
									info="Choose QA method that you want")
			model_choice = gr.Radio(["gpt-4o-mini", "gpt-4o"], label="Select model for data creation",
									)
			qa_cnt = gr.Slider(minimum=20, maximum=150, step=5, label="Number of QA pairs", value=80)
			batch_size = gr.Slider(minimum=1, maximum=16, step=1,
								   label="Batch Size to OpenAI model. If there is an error, decrease this.", value=16)
			run_qa_button = gr.Button(value="Run QA Creation")
			qa_status = gr.Textbox(value="Not Started", interactive=False)
			qa_download_button = gr.Button(value="Download qa.parquet",
										   link=f"/file={os.path.join(DATA_DIR, 'qa.parquet')}")

	#================================================================================================#
	# Logics

	use_previous_raw_file.change(lambda x: gr.update(visible=not x), inputs=[use_previous_raw_file],
								 outputs=[raw_file_input])
	use_previous_corpus_file.change(lambda x: gr.update(visible=not x), inputs=[use_previous_corpus_file],
									outputs=[corpus_file_input])
	openai_key_textbox.submit(on_submit_openai_key, inputs=[openai_key_textbox], outputs=api_key_status_box)

	# Parsing
	parse_button.click(run_parse, inputs=[document_file_input, parse_choice], outputs=parse_status)
	file_reset_button.click(file_reset, outputs=parse_status)

	# Chunking
	chunk_button.click(run_chunk, inputs=[use_previous_raw_file, raw_file_input, chunk_choice, chunk_size, chunk_overlap,
										  lang_choice],
					   outputs=chunk_status)

	# QA Creation
	run_qa_button.click(run_qa, inputs=[use_previous_corpus_file, corpus_file_input, qa_choice, model_choice, qa_cnt,
										batch_size, lang_choice], outputs=qa_status)

	# API Key visibility
	parse_choice.change(change_visible_status_api_key, inputs=[parse_choice],
						outputs=[llama_cloud_api_key_row, upstage_api_key_row])
	llama_key_textbox.submit(on_submit_llama_cloud_key, inputs=[llama_key_textbox], outputs=llama_key_status_box)
	upstage_key_textbox.submit(on_submit_upstage_key, inputs=[upstage_key_textbox], outputs=upstage_key_status_box)


demo.launch(share=False, debug=True, allowed_paths=[FILE_DIR, DATA_DIR])