Delete data persistent on the huggingface space

app.py

@@ -1,25 +1,24 @@
 import os
-import shutil
-from typing import List
+import tempfile
+from typing import List, Callable
 
 import gradio as gr
 import pandas as pd
 from autorag.data.parse import langchain_parse
+from autorag.data.parse.base import _add_last_modified_datetime
 from autorag.data.parse.llamaparse import llama_parse
 from autorag.data.qa.schema import Raw
+from autorag.utils import result_to_dataframe
 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)
-
+@result_to_dataframe(["texts", "path", "page", "last_modified_datetime"])
+def original_parse(fn: Callable, **kwargs):
+	result = fn(**kwargs)
+	result = _add_last_modified_datetime(result)
+	return result
 
 def change_lang_choice(lang: str) -> str:
 	lang_dict = {
@@ -39,43 +38,38 @@ def change_visible_status_api_key(parse_method: str):
 
 
 
-def run_parse(file_lists: List[str], parse_method: str, progress=gr.Progress()):
+def run_parse(file_lists: List[str], parse_method: str, original_raw_df, 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)
+	langchain_parse_original = langchain_parse.__wrapped__
 
 	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)
+		raw_df: pd.DataFrame = original_parse(langchain_parse_original,
+											  data_path_list=file_lists, 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"))
+			return "Please submit your Llama Cloud API key first.", original_raw_df
+		raw_df: pd.DataFrame = original_parse(llama_parse.__wrapped__, data_path_list=file_lists)
 	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")
+			return "Please submit your Upstage API key first.", original_raw_df
+		raw_df: pd.DataFrame = original_parse(langchain_parse_original,
+											  data_path_list=file_lists, parse_method="upstagedocumentparse")
 	else:
-		return "Unsupported parse method."
+		return "Unsupported parse method.", original_raw_df
 	progress(0.8)
 
-	raw_df.to_parquet(os.path.join(DATA_DIR, "raw.parquet"), index=False)
-	return "Parsing Complete. Download at the bottom button."
+	return "Parsing Complete. Download at the bottom button.", raw_df
 
 
-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()):
+def run_chunk(use_existed_raw: bool, raw_df: pd.DataFrame, raw_file: str, chunk_method: str, chunk_size: int, chunk_overlap: int,
+			  lang: str = "English", original_corpus_df = None, 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")
+	if not use_existed_raw:
+		raw_df = pd.read_parquet(raw_file, engine="pyarrow")
 	raw_instance = Raw(raw_df)
 
 	if chunk_method in ["Token", "Sentence"]:
@@ -90,28 +84,21 @@ def run_chunk(use_existed_raw: bool, raw_file: str, chunk_method: str, chunk_siz
 									add_file_name=lang, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
 	else:
 		gr.Error("Unsupported chunk method.")
-		return "Unsupported chunk method."
+		return "Unsupported chunk method.", original_corpus_df
 	progress(0.8)
-	corpus.to_parquet(os.path.join(DATA_DIR, "corpus.parquet"))
-	return "Chunking Complete. Download at the bottom button."
+	return "Chunking Complete. Download at the bottom button.", corpus.data
 
 
-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()):
+def run_qa(use_existed_corpus: bool, corpus_df: pd.DataFrame, corpus_file: str, qa_method: str,
+		   model_name: str, qa_cnt: int, batch_size: int, lang: str = "English", original_qa_df = None,
+		   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 not use_existed_corpus:
+		corpus_df = pd.read_parquet(corpus_file, 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."
+		return "Please submit your OpenAI API key first.", original_qa_df
 	llm = OpenAI(model=model_name)
 
 	if qa_method == "default":
@@ -122,18 +109,25 @@ def run_qa(use_existed_corpus: bool, corpus_file: str, qa_method: str,
 		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."
+		return "Unsupported QA method.", original_qa_df
+
+	return "QA Creation Complete. Download at the bottom button.", qa.data
 
-	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 download_state(state: pd.DataFrame, change_name: str):
+	if state is None:
+		gr.Error("No data to download.")
+		return ""
+	with tempfile.TemporaryDirectory() as temp_dir:
+		filename = os.path.join(temp_dir, f"{change_name}.parquet")
+		state.to_parquet(filename, engine="pyarrow")
+		yield filename
 
-def file_reset() -> str:
-	shutil.rmtree(FILE_DIR)
-	os.makedirs(FILE_DIR)
-	return "Files reset complete."
 
 with gr.Blocks(theme="earneleh/paris") as demo:
+	raw_df_state = gr.State()
+	corpus_df_state = gr.State()
+	qa_df_state = gr.State()
 	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",
@@ -166,9 +160,8 @@ with gr.Blocks(theme="earneleh/paris") as demo:
 				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.DownloadButton(value=os.path.join(DATA_DIR, 'raw.parquet'),
+			raw_download_button = gr.DownloadButton(value=download_state, inputs=[raw_df_state, gr.State("raw")],
 				label="Download raw.parquet")
-			file_reset_button = gr.Button(value="Reset uploaded files")
 
 		with gr.Column(scale=1):
 			gr.Markdown(
@@ -185,7 +178,7 @@ with gr.Blocks(theme="earneleh/paris") as demo:
 			chunk_button = gr.Button(value="Run Chunking")
 			chunk_status = gr.Textbox(value="Not Started", interactive=False)
 			corpus_download_button = gr.DownloadButton(label="Download corpus.parquet",
-											   value=os.path.join(DATA_DIR, 'corpus.parquet'))
+											   value=download_state, inputs=[corpus_df_state, gr.State("corpus")])
 
 		with gr.Column(scale=1):
 			gr.Markdown(
@@ -205,7 +198,7 @@ with gr.Blocks(theme="earneleh/paris") as demo:
 			run_qa_button = gr.Button(value="Run QA Creation")
 			qa_status = gr.Textbox(value="Not Started", interactive=False)
 			qa_download_button = gr.DownloadButton(label="Download qa.parquet",
-										   value=os.path.join(DATA_DIR, 'qa.parquet'))
+										   value=download_state, inputs=[qa_df_state, gr.State("qa")])
 
 	#================================================================================================#
 	# Logics
@@ -217,17 +210,19 @@ with gr.Blocks(theme="earneleh/paris") as demo:
 	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)
+	parse_button.click(run_parse, inputs=[document_file_input, parse_choice, raw_df_state],
+					   outputs=[parse_status, raw_df_state])
 
 	# 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)
+	chunk_button.click(run_chunk, inputs=[use_previous_raw_file, raw_df_state, raw_file_input, chunk_choice, chunk_size, chunk_overlap,
+										  lang_choice, corpus_df_state],
+					   outputs=[chunk_status, corpus_df_state])
 
 	# 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)
+	run_qa_button.click(run_qa, inputs=[use_previous_corpus_file, corpus_df_state, corpus_file_input, qa_choice,
+										model_choice, qa_cnt, batch_size, lang_choice,
+										qa_df_state],
+						outputs=[qa_status, qa_df_state])
 
 	# API Key visibility
 	parse_choice.change(change_visible_status_api_key, inputs=[parse_choice],
@@ -236,4 +231,6 @@ with gr.Blocks(theme="earneleh/paris") as demo:
 	upstage_key_textbox.submit(on_submit_upstage_key, inputs=[upstage_key_textbox], outputs=upstage_key_status_box)
 
 
-demo.launch(share=False, debug=False, allowed_paths=[FILE_DIR, DATA_DIR])
+# if __name__ == "__main__":
+# 	demo.launch(share=False, debug=True)
+demo.launch(share=False, debug=False)