Update app.py

app.py

@@ -8,9 +8,7 @@ import re
 import uuid
 import pymupdf
 
-# =======================================================
-# magic-pdf & MinerU 설치 관련 (기존 코드 그대로 예시)
-# =======================================================
+# (기존 magic-pdf 설치 및 설정 로직)
 os.system('pip uninstall -y magic-pdf')
 os.system('pip install git+https://github.com/opendatalab/MinerU.git@dev')
 os.system('wget https://github.com/opendatalab/MinerU/raw/dev/scripts/download_models_hf.py -O download_models_hf.py')
@@ -29,6 +27,10 @@ with open('/home/user/magic-pdf.json', 'w') as file:
 
 os.system('cp -r paddleocr /home/user/.paddleocr')
 
+
+###############################
+# Gradio 및 기타 라이브러리
+###############################
 import gradio as gr
 from loguru import logger
 
@@ -36,14 +38,11 @@ from magic_pdf.data.data_reader_writer import FileBasedDataReader
 from magic_pdf.libs.hash_utils import compute_sha256
 from magic_pdf.tools.common import do_parse, prepare_env
 
-
-###########################################
-# 1) UI 스타일(CSS) + PDF처리 관련 함수들
-###########################################
+#######################
+# 1) UI CSS + PDF funcs
+#######################
 def create_css():
-    """
-    화면을 가득 채우고 스크롤 가능하도록 설정
-    """
+    """화면 가득 사용 + 스크롤 허용"""
     return """
     .gradio-container {
         width: 100vw !important;
@@ -152,10 +151,7 @@ def to_pdf(file_path):
             return tmp_file_path
 
 def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table_enable, language, progress=gr.Progress(track_tqdm=False)):
-    """
-    - PDF 변환 과정에서 'progress(...)' 로 진행률 업데이트
-    - Gradio 버전이 낮아도 'with progress:' 를 쓰지 않으면 __enter__ 오류가 안뜸
-    """
+    """PDF 변환 함수 (프로그레스바 표시)"""
     progress(0, "PDF로 변환 중...")
     file_path = to_pdf(file_path)
     time.sleep(0.5)
@@ -177,13 +173,13 @@ def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table
         logger.error("압축 실패")
     time.sleep(0.5)
 
-    progress(70, "마크다운 로드 중...")
+    progress(70, "마크다운 읽는 중...")
     md_path = os.path.join(local_md_dir, file_name + ".md")
     with open(md_path, 'r', encoding='utf-8') as f:
         txt_content = f.read()
     time.sleep(0.5)
 
-    progress(90, "이미지(base64) 변환 중...")
+    progress(90, "이미지 base64 변환 중...")
     md_content = replace_image_with_base64(txt_content, local_md_dir)
     time.sleep(0.5)
 
@@ -191,9 +187,9 @@ def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table
     return md_content
 
 
-###############################
+################################
 # magic_pdf 모델 초기화
-###############################
+################################
 def init_model():
     from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
     try:
@@ -210,6 +206,9 @@ def init_model():
 model_init = init_model()
 logger.info(f"model_init: {model_init}")
 
+##################
+# 언어 목록
+##################
 latin_lang = [
     'af','az','bs','cs','cy','da','de','es','et','fr','ga','hr','hu','id','is','it','ku',
     'la','lt','lv','mi','ms','mt','nl','no','oc','pi','pl','pt','ro','rs_latin','sk','sl',
@@ -224,9 +223,9 @@ all_lang = ['', 'auto']
 all_lang.extend([*other_lang, *latin_lang, *arabic_lang, *cyrillic_lang, *devanagari_lang])
 
 
-#################################
+##################################
 # 2) Gemini (google.generativeai)
-#################################
+##################################
 import google.generativeai as genai
 from gradio import ChatMessage
 from typing import Iterator
@@ -239,7 +238,7 @@ model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp-1219")
 
 def format_chat_history(messages: list) -> list:
     """
-    Gemini가 이해할 수 있는 형식 (role, content)
+    Gemini가 이해할 수 있는 (role, parts[]) 형식으로 변환
     """
     formatted_history = []
     for message in messages:
@@ -252,15 +251,13 @@ def format_chat_history(messages: list) -> list:
 
 def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
     """
-    Gemini 응답 스트리밍: user_message가 비어있으면 기본 문구로 대체
+    Gemini 응답 스트리밍 (user_message가 공백이면 기본 문구로 교체)
     """
-    # 빈 문자열이면 기본 문구로 교체(오류 방지)
     if not user_message.strip():
-        user_message = "…(No content from user)…"
+        user_message = "...(No content from user)..."
 
     try:
         print(f"\n=== [Gemini] New Request ===\nUser message: '{user_message}'")
-
         chat_history = format_chat_history(messages)
         chat = model.start_chat(history=chat_history)
         response = chat.send_message(user_message, stream=True)
@@ -302,7 +299,7 @@ def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
                 response_buffer += current_chunk
                 messages[-1] = ChatMessage(role="assistant", content=response_buffer)
             else:
-                # Still thinking
+                # Still in "thinking"
                 thought_buffer += current_chunk
                 messages[-1] = ChatMessage(
                     role="assistant",
@@ -328,7 +325,6 @@ def convert_chat_messages_to_gradio_format(messages):
 
     for msg in messages:
         if msg.role == "user":
-            # 이전 턴 저장
             if user_text is not None or assistant_text is not None:
                 gradio_chat.append((user_text or "", assistant_text or ""))
             user_text = msg.content
@@ -340,9 +336,8 @@ def convert_chat_messages_to_gradio_format(messages):
             if assistant_text is None:
                 assistant_text = msg.content
             else:
-                assistant_text += msg.content  # 스트리밍 시 누적
+                assistant_text += msg.content
 
-    # 마지막 턴
     if user_text is not None or assistant_text is not None:
         gradio_chat.append((user_text or "", assistant_text or ""))
 
@@ -350,7 +345,7 @@ def convert_chat_messages_to_gradio_format(messages):
 
 def user_message(msg: str, history: list, doc_text: str) -> tuple[str, list]:
     """
-    doc_text(마크다운) 참고 문구를 자동 삽입
+    doc_text(마크다운) 사용해 질문 자동 변형
     """
     if doc_text.strip():
         user_query = f"다음 문서를 참고하여 답변:\n\n{doc_text}\n\n질문: {msg}"
@@ -360,34 +355,60 @@ def user_message(msg: str, history: list, doc_text: str) -> tuple[str, list]:
     history.append(ChatMessage(role="user", content=user_query))
     return "", history
 
+########################
+# 새 파일 업로드시 초기화
+########################
+def reset_states(_):
+    """
+    파일이 변경되면(새 파일 업로드되면) 
+    - 대화 이력(chat_history)
+    - 마크다운 상태(md_state)
+    - 챗봇(기존 대화)
+    모두 초기화
+    """
+    return [], "", []
 
 ################################
-# 3) 통합 Gradio 앱 구성 & 실행
+# 3) 통합 Gradio 앱 구성
 ################################
 with gr.Blocks(title="OCR FLEX + Gemini Chat", css=create_css()) as demo:
     gr.HTML("""
     <div class="title-area">
         <h1>OCR FLEX + Gemini Chat</h1>
-        <p>PDF/이미지 -> 텍스트(마크다운) 변환 후, Gemini LLM 대화</p>
+        <p>PDF/이미지 -> 텍스트(마크다운) 변환 후, Gemini LLM과 대화</p>
     </div>
     """)
 
+    # 상태
     md_state = gr.State("")
     chat_history = gr.State([])
 
+    # 1) 업로드 & 변환
     with gr.Row():
-        file = gr.File(label="PDF/이미지 업로드", file_types=[".pdf", ".png", ".jpeg", ".jpg"], interactive=True)
+        file = gr.File(
+            label="PDF/이미지 업로드",
+            file_types=[".pdf", ".png", ".jpeg", ".jpg"],
+            interactive=True
+        )
         convert_btn = gr.Button("변환하기")
 
-    # 숨긴 컴포넌트들
-    max_pages = gr.Slider(1, 20, 10, visible=False, elem_classes="invisible")
-    layout_mode = gr.Dropdown(["layoutlmv3","doclayout_yolo"], value="doclayout_yolo", visible=False, elem_classes="invisible")
+    # 새 파일 업로드 시: 이전 대화/마크다운/챗봇 초기화
+    # => 만약 "파일 변경 이벤트"로 초기화하길 원치 않으면 주석 처리
+    file.change(
+        fn=reset_states,
+        inputs=file,
+        outputs=[chat_history, md_state, "chatbot"]
+    )
+
+    # 숨긴 컴포넌트
+    max_pages = gr.Slider(1,20,10, visible=False, elem_classes="invisible")
+    layout_mode = gr.Dropdown(["layoutlmv3","doclayout_yolo"],value="doclayout_yolo",visible=False,elem_classes="invisible")
     language = gr.Dropdown(all_lang, value='auto', visible=False, elem_classes="invisible")
     formula_enable = gr.Checkbox(value=True, visible=False, elem_classes="invisible")
     is_ocr = gr.Checkbox(value=False, visible=False, elem_classes="invisible")
     table_enable = gr.Checkbox(value=True, visible=False, elem_classes="invisible")
 
-    # 변환 클릭 -> to_markdown (progress)
+    # 변환 버튼 → to_markdown
     convert_btn.click(
         fn=to_markdown,
         inputs=[file, max_pages, is_ocr, layout_mode, formula_enable, table_enable, language],
@@ -395,14 +416,15 @@ with gr.Blocks(title="OCR FLEX + Gemini Chat", css=create_css()) as demo:
         show_progress=True  # 프로그레스바+로딩 표시
     )
 
-    # Gemini Chat
+    # 2) Gemini Chat
     gr.Markdown("## Gemini 2.0 Flash (Thinking) Chat")
-    chatbot = gr.Chatbot(height=600)
+    chatbot = gr.Chatbot(height=600, elem_id="chatbot")  # elem_id 지정 (reset에서 outputs로 사용)
+
     with gr.Row():
         chat_input = gr.Textbox(lines=1, placeholder="질문을 입력하세요...")
         clear_btn = gr.Button("대화 초기화")
 
-    # 프롬프트 전송 -> user_message -> stream_gemini_response
+    # 채팅 흐름
     chat_input.submit(
         fn=user_message,
         inputs=[chat_input, chat_history, md_state],
@@ -413,19 +435,15 @@ with gr.Blocks(title="OCR FLEX + Gemini Chat", css=create_css()) as demo:
         outputs=chatbot
     )
 
+    # 초기화 버튼: 대화/상태/챗봇 모두 초기화
     def clear_states():
-        return [], ""
+        return [], "", []
 
     clear_btn.click(
         fn=clear_states,
         inputs=[],
-        outputs=[chat_history, md_state]
-    ).then(
-        fn=lambda: [],
-        inputs=[],
-        outputs=chatbot
+        outputs=[chat_history, md_state, chatbot]
     )
 
-
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)