Update app.py

app.py

@@ -12,13 +12,11 @@ import re
 import uuid
 import pymupdf
 
-# (원래 설치 & 세팅)
 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')
 os.system('python download_models_hf.py')
 
-# magic-pdf.json 수정
 with open('/home/user/magic-pdf.json', 'r') as file:
     data = json.load(file)
 
@@ -41,22 +39,19 @@ from magic_pdf.tools.common import do_parse, prepare_env
 
 def create_css():
     """
-    스크롤 가능하도록 .gradio-container에서 overflow-y를 auto로,
-    높이는 최소 100vh로 설정
+    화면을 꽉 채우고 스크롤 가능하도록 설정
     """
     return """
-    /* 화면 전체를 차지하면서, 스크롤이 가능하도록 설정 */
     .gradio-container {
         width: 100vw !important;
-        min-height: 100vh !important; 
+        min-height: 100vh !important;
         margin: 0 !important;
         padding: 0 !important;
         background: linear-gradient(135deg, #EFF6FF 0%, #F5F3FF 100%);
         display: flex;
         flex-direction: column;
-        overflow-y: auto !important; /* 스크롤 활성화 */
+        overflow-y: auto !important; 
     }
-    /* 타이틀 영역 */
     .title-area {
         text-align: center;
         margin: 1rem auto;
@@ -78,11 +73,9 @@ def create_css():
         color: #6B7280;
         font-size: 1.1rem;
     }
-    /* 숨기고 싶은 컴포넌트 */
     .invisible {
         display: none !important;
     }
-    /* 기타 padding 조정 */
     .gr-block, .gr-box {
         padding: 0.5rem !important;
     }
@@ -98,10 +91,7 @@ def parse_pdf(doc_path, output_dir, end_page_id, is_ocr, layout_mode, formula_en
     try:
         file_name = f"{str(Path(doc_path).stem)}_{time.time()}"
         pdf_data = read_fn(doc_path)
-        if is_ocr:
-            parse_method = "ocr"
-        else:
-            parse_method = "auto"
+        parse_method = "ocr" if is_ocr else "auto"
         local_image_dir, local_md_dir = prepare_env(output_dir, file_name, parse_method)
         do_parse(
             output_dir,
@@ -148,9 +138,6 @@ def replace_image_with_base64(markdown_text, image_dir_path):
     return re.sub(pattern, replace, markdown_text)
 
 def to_pdf(file_path):
-    """
-    PDF가 아닌 파일(이미지 등)을 pymupdf로 PDF로 변환
-    """
     with pymupdf.open(file_path) as f:
         if f.is_pdf:
             return file_path
@@ -162,37 +149,53 @@ def to_pdf(file_path):
                 tmp_pdf_file.write(pdf_bytes)
             return tmp_file_path
 
-def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table_enable, language):
-    file_path = to_pdf(file_path)
-    if end_pages > 20:
-        end_pages = 20
-    local_md_dir, file_name = parse_pdf(file_path, './output', end_pages - 1, is_ocr,
-                                        layout_mode, formula_enable, table_enable, language)
-    archive_zip_path = os.path.join("./output", compute_sha256(local_md_dir) + ".zip")
-    zip_archive_success = compress_directory_to_zip(local_md_dir, archive_zip_path)
-    if zip_archive_success == 0:
-        logger.info("압축 성공")
-    else:
-        logger.error("압축 실패")
+def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table_enable, language, progress=gr.Progress(track_tqdm=False)):
+    """
+    - PDF 변환 과정 단계별로 progress 표시
+    - 반환값: 최종 마크다운 텍스트
+    """
+    with progress:
+        progress(0, "파일을 PDF로 변환 중...")
+        file_path = to_pdf(file_path)
+        time.sleep(0.5)  # 예시로 딜레이
+
+        if end_pages > 20:
+            end_pages = 20
+        progress(30, "PDF 파싱 중...")
+        local_md_dir, file_name = parse_pdf(file_path, './output', end_pages - 1, is_ocr,
+                                            layout_mode, formula_enable, table_enable, language)
+        time.sleep(0.5)
+
+        progress(50, "압축(Zip) 생성 중...")
+        archive_zip_path = os.path.join("./output", compute_sha256(local_md_dir) + ".zip")
+        zip_archive_success = compress_directory_to_zip(local_md_dir, archive_zip_path)
+        if zip_archive_success == 0:
+            logger.info("압축 성공")
+        else:
+            logger.error("압축 실패")
+        time.sleep(0.5)
+
+        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)
 
-    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()
+        progress(90, "이미지 base64 변환 중...")
+        md_content = replace_image_with_base64(txt_content, local_md_dir)
+        time.sleep(0.5)
 
-    md_content = replace_image_with_base64(txt_content, local_md_dir)
-    return md_content
+        progress(100, "변환 완료!")
+        return md_content
 
 def init_model():
-    """
-    magic_pdf 모델(레이아웃, OCR 등) 미리 초기화
-    """
     from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
     try:
         model_manager = ModelSingleton()
         txt_model = model_manager.get_model(False, False)
-        logger.info(f"txt_model init final")
+        logger.info("txt_model init final")
         ocr_model = model_manager.get_model(True, False)
-        logger.info(f"ocr_model init final")
+        logger.info("ocr_model init final")
         return 0
     except Exception as e:
         logger.exception(e)
@@ -202,21 +205,14 @@ 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', 'sq', 'sv',
-    'sw', 'tl', 'tr', 'uz', 'vi', 'french', 'german'
+    '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',
+    'sq','sv','sw','tl','tr','uz','vi','french','german'
 ]
-arabic_lang = ['ar', 'fa', 'ug', 'ur']
-cyrillic_lang = [
-    'ru', 'rs_cyrillic', 'be', 'bg', 'uk', 'mn', 'abq', 'ady', 'kbd', 'ava',
-    'dar', 'inh', 'che', 'lbe', 'lez', 'tab'
-]
-devanagari_lang = [
-    'hi', 'mr', 'ne', 'bh', 'mai', 'ang', 'bho', 'mah', 'sck', 'new', 'gom',
-    'sa', 'bgc'
-]
-other_lang = ['ch', 'en', 'korean', 'japan', 'chinese_cht', 'ta', 'te', 'ka']
+arabic_lang = ['ar','fa','ug','ur']
+cyrillic_lang = ['ru','rs_cyrillic','be','bg','uk','mn','abq','ady','kbd','ava','dar','inh','che','lbe','lez','tab']
+devanagari_lang = ['hi','mr','ne','bh','mai','ang','bho','mah','sck','new','gom','sa','bgc']
+other_lang = ['ch','en','korean','japan','chinese_cht','ta','te','ka']
 
 all_lang = ['', 'auto']
 all_lang.extend([*other_lang, *latin_lang, *arabic_lang, *cyrillic_lang, *devanagari_lang])
@@ -232,6 +228,7 @@ import time
 
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 genai.configure(api_key=GEMINI_API_KEY)
+
 model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp-1219")
 
 def format_chat_history(messages: list) -> list:
@@ -240,7 +237,9 @@ def format_chat_history(messages: list) -> list:
     """
     formatted_history = []
     for message in messages:
-        # ChatMessage.role, ChatMessage.content 활용
+        # ChatMessage => role, content
+        # metadata는 "thinking" 표시 용도
+        # Gemini input에서는 제외
         if not (message.role == "assistant" and hasattr(message, "metadata")):
             formatted_history.append({
                 "role": "user" if message.role == "user" else "assistant",
@@ -250,16 +249,17 @@ def format_chat_history(messages: list) -> list:
 
 def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
     """
-    Gemini 응답을 스트리밍
-    => 최종적으로 (유저, 봇) 튜플 리스트로 변환하여 yield
+    Gemini 응답 스트리밍
+    - user_message가 공백이어도 오류 X
     """
-    if not user_message.strip():
-        messages.append(ChatMessage(role="assistant", content="Please provide a non-empty text message."))
-        yield convert_chat_messages_to_gradio_format(messages)
-        return
+    # 1) 공백 입력도 그냥 통과(오류 안내 메시지 제거)
+    # if not user_message.strip():
+    #     messages.append(ChatMessage(role="assistant", content="Please provide a non-empty text message."))
+    #     yield convert_chat_messages_to_gradio_format(messages)
+    #     return
 
     try:
-        print(f"\n=== New Request (Text) ===\nUser message: {user_message}")
+        print(f"\n=== [Gemini] New Request ===\nUser message: '{user_message}'")
 
         chat_history = format_chat_history(messages)
         chat = model.start_chat(history=chat_history)
@@ -302,7 +302,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 in "thinking"
+                # Still in "Thinking"
                 thought_buffer += current_chunk
                 messages[-1] = ChatMessage(
                     role="assistant",
@@ -312,16 +312,16 @@ def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
 
             yield convert_chat_messages_to_gradio_format(messages)
 
-        print(f"\n=== Final Response ===\n{response_buffer}")
+        print(f"\n=== [Gemini] Final Response ===\n{response_buffer}")
 
     except Exception as e:
-        print(f"\n=== Error ===\n{str(e)}")
+        print(f"\n=== [Gemini] Error ===\n{str(e)}")
         messages.append(ChatMessage(role="assistant", content=f"I encountered an error: {str(e)}"))
         yield convert_chat_messages_to_gradio_format(messages)
 
 def convert_chat_messages_to_gradio_format(messages):
     """
-    ChatMessage list -> Gradio Chatbot용 [(유저발화, 봇응답), (유저발화2, 봇응답2), ...]
+    ChatMessage 리스트 => Gradio (유저, 봇) 튜플 리스트
     """
     gradio_chat = []
     user_text, assistant_text = None, None
@@ -330,7 +330,6 @@ def convert_chat_messages_to_gradio_format(messages):
         role = msg.role
         content = msg.content
         if 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 = content
@@ -351,12 +350,13 @@ def convert_chat_messages_to_gradio_format(messages):
 
 def user_message(msg: str, history: list, doc_text: str) -> tuple[str, list]:
     """
-    사용자 입력 -> doc_text(마크다운) 참고 문구를 첨가해서 history에 ChatMessage 추가
+    - doc_text(마크다운) 참고 문구를 자동으로 추가
+    - 공백 입력도 에러 없이 진행
     """
-    if not doc_text.strip():
-        user_query = msg
-    else:
+    if doc_text.strip():
         user_query = f"다음 문서를 참고하여 답변:\n\n{doc_text}\n\n질문: {msg}"
+    else:
+        user_query = msg
 
     history.append(ChatMessage(role="user", content=user_query))
     return "", history
@@ -369,23 +369,24 @@ 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("")        # PDF -> 마크다운 변환 결과
-    chat_history = gr.State([])    # ChatMessage 리스트
+    # 변환된 마크다운, 채팅 이력
+    md_state = gr.State("")
+    chat_history = gr.State([])
 
     # 업로드 & 변환
     with gr.Row():
         file = gr.File(
-            label="PDF / 이미지 업로드",
+            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")
     language = gr.Dropdown(all_lang, value='auto', visible=False, elem_classes="invisible")
@@ -393,21 +394,22 @@ with gr.Blocks(title="OCR FLEX + Gemini Chat", css=create_css()) as demo:
     is_ocr = gr.Checkbox(value=False, visible=False, elem_classes="invisible")
     table_enable = gr.Checkbox(value=True, visible=False, elem_classes="invisible")
 
-    # 변환 버튼 → md_state 저장
+    # 변환 버튼 → to_markdown + progress
     convert_btn.click(
         fn=to_markdown,
         inputs=[file, max_pages, is_ocr, layout_mode, formula_enable, table_enable, language],
-        outputs=md_state
+        outputs=md_state,
+        show_progress=True  # Gradio 자체 로딩 스피너도 표시
     )
 
-    # Gemini 채팅
+    # Gemini Chat
     gr.Markdown("## Gemini 2.0 Flash (Thinking) Chat")
     chatbot = gr.Chatbot(height=600)
     with gr.Row():
         chat_input = gr.Textbox(lines=1, placeholder="질문을 입력하세요...")
         clear_button = gr.Button("대화 초기화")
 
-    # 채팅 흐름
+    # 프롬프트 입력 -> user_message -> stream_gemini_response
     chat_input.submit(
         fn=user_message,
         inputs=[chat_input, chat_history, md_state],
@@ -418,7 +420,7 @@ with gr.Blocks(title="OCR FLEX + Gemini Chat", css=create_css()) as demo:
         outputs=chatbot
     )
 
-    # 초기화
+    # 초기화: 상태 리셋 + 챗봇 초기화
     def clear_states():
         return [], ""
     clear_button.click(
@@ -432,8 +434,5 @@ with gr.Blocks(title="OCR FLEX + Gemini Chat", css=create_css()) as demo:
     )
 
 
-##############################
-# 4) 앱 실행
-##############################
 if __name__ == "__main__":
     demo.launch(debug=True, server_name="0.0.0.0", server_port=7860)