Update app.py

app.py

@@ -12,12 +12,13 @@ import re
 import uuid
 import pymupdf
 
-# 원래 코드에 있던 os.system() 호출들
+# (원래 설치 & 세팅 코드)
 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)
 
@@ -29,8 +30,8 @@ if os.getenv('apikey'):
 with open('/home/user/magic-pdf.json', 'w') as file:
     json.dump(data, file, indent=4)
 
+# (OCR 폴더 복사)
 os.system('cp -r paddleocr /home/user/.paddleocr')
-# from gradio_pdf import PDF  # PDF 미리보기를 위한 컴포넌트이지만, 지금은 숨길 예정
 
 import gradio as gr
 from loguru import logger
@@ -40,22 +41,31 @@ from magic_pdf.libs.hash_utils import compute_sha256
 from magic_pdf.tools.common import do_parse, prepare_env
 
 def create_css():
+    """
+    화면을 넓고 꽉 차게 채우도록 설정.
+    필요하면 배경색/폰트 수정 가능
+    """
     return """
-    /* 전체 스타일 */
+    /* 전체 컨테이너: 화면 가득 채우기 */
     .gradio-container {
+        width: 100vw !important;
+        height: 100vh !important;
+        margin: 0 !important;
+        padding: 0 !important;
+        /* 배경 그라디언트 예시 */
         background: linear-gradient(135deg, #EFF6FF 0%, #F5F3FF 100%);
-        max-width: 1200px !important;
-        margin: 0 auto !important;
-        padding: 2rem !important;
+        display: flex;
+        flex-direction: column;
     }
-    /* 제목 스타일 */
+    /* 타이틀 영역 */
     .title-area {
         text-align: center;
-        margin-bottom: 2rem;
+        margin: 1rem auto;
         padding: 1rem;
         background: white;
         border-radius: 1rem;
         box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
+        max-width: 800px;
     }
     .title-area h1 {
         background: linear-gradient(90deg, #2563EB 0%, #7C3AED 100%);
@@ -69,10 +79,14 @@ def create_css():
         color: #6B7280;
         font-size: 1.1rem;
     }
-    /* 숨길 예정인 컴포넌트 스타일 */
+    /* 숨기고 싶은 컴포넌트 */
     .invisible {
         display: none !important;
     }
+    /* 기본 패딩 줄이기 */
+    .gr-block, .gr-box {
+        padding: 0.5rem !important;
+    }
     """
 
 def read_fn(path):
@@ -136,7 +150,7 @@ def replace_image_with_base64(markdown_text, image_dir_path):
 
 def to_pdf(file_path):
     """
-    PDF가 아닌 경우(예: PNG, JPG 파일)에도 pymupdf를 이용하여 PDF로 변환하기 위한 함수.
+    PDF가 아닌 경우(예: PNG, JPG 파일)에도 pymupdf를 이용하여 PDF로 변환
     """
     with pymupdf.open(file_path) as f:
         if f.is_pdf:
@@ -151,8 +165,8 @@ def to_pdf(file_path):
 
 def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table_enable, language):
     """
-    파일을 받아 최대 end_pages 페이지까지 마크다운 추출 후,
-    base64 이미지가 포함된 md_content를 반환.
+    파일(이미지/PDF)을 받아 최대 end_pages 페이지까지 마크다운 추출 후,
+    base64 이미지가 포함된 최종 md_content 반환
     """
     file_path = to_pdf(file_path)
     if end_pages > 20:
@@ -171,10 +185,10 @@ def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table
         txt_content = f.read()
 
     md_content = replace_image_with_base64(txt_content, local_md_dir)
-    # new_pdf_path = os.path.join(local_md_dir, file_name + "_layout.pdf") # 원래 pdf 미리보기용
 
-    return md_content  # base64 이미지가 포함된 최종 마크다운 텍스트만 반환
+    return md_content  # 최종 마크다운 텍스트
 
+# latex 수식 구분자
 latex_delimiters = [
     {"left": "$$", "right": "$$", "display": True},
     {"left": '$', "right": '$', "display": False}
@@ -182,7 +196,7 @@ latex_delimiters = [
 
 def init_model():
     """
-    magic_pdf의 모델을 미리 초기화.
+    magic_pdf의 모델(레이아웃, OCR 등)을 미리 초기화
     """
     from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
     try:
@@ -223,29 +237,25 @@ all_lang.extend([*other_lang, *latin_lang, *arabic_lang, *cyrillic_lang, *devana
 ##############################
 # 2) Gemini LLM 챗 코드
 ##############################
-
-# (중복 import이지만 "누락 없이" 출력해야 하므로 주석 처리)
-# import os
-# import gradio as gr
+# (중복 import이지만 "누락 없이" 출력해야 하므로 주석 처리 안함)
+import google.generativeai as genai
 from gradio import ChatMessage
 from typing import Iterator
-import google.generativeai as genai
 import time
 
-# get Gemini API Key from the environ variable
+# Gemini API Key 세팅
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 genai.configure(api_key=GEMINI_API_KEY)
 
-# we will be using the Gemini 2.0 Flash model with Thinking capabilities
+# Gemini 2.0 Flash (Thinking)
 model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp-1219")
 
 def format_chat_history(messages: list) -> list:
     """
-    Formats the chat history into a structure Gemini can understand
+    Gemini가 이해할 수 있는 형식으로 변환
     """
     formatted_history = []
     for message in messages:
-        # Skip thinking messages (messages with metadata)
         if not (message.get("role") == "assistant" and "metadata" in message):
             formatted_history.append({
                 "role": "user" if message.get("role") == "user" else "assistant",
@@ -255,11 +265,13 @@ def format_chat_history(messages: list) -> list:
 
 def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
     """
-    Streams thoughts and response with conversation history support for text input only.
+    Gemini 응답을 스트리밍. 
+    ChatMessage 형식의 messages를 받아 최종적으로 Gradio가 
+    인식할 수 있는 (user, assistant) 튜플 리스트로 변환해서 yield
     """
     if not user_message.strip():
-        messages.append(ChatMessage(role="assistant", content="Please provide a non-empty text message. Empty input is not allowed."))
-        yield messages
+        messages.append(ChatMessage(role="assistant", content="Please provide a non-empty text message."))
+        yield convert_chat_messages_to_gradio_format(messages)
         return
 
     try:
@@ -274,6 +286,7 @@ def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
         response_buffer = ""
         thinking_complete = False
 
+        # 우선 "Thinking" 메시지 추가
         messages.append(
             ChatMessage(
                 role="assistant",
@@ -281,225 +294,191 @@ def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
                 metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
             )
         )
+        yield convert_chat_messages_to_gradio_format(messages)
 
         for chunk in response:
             parts = chunk.candidates[0].content.parts
             current_chunk = parts[0].text
 
             if len(parts) == 2 and not thinking_complete:
-                # Complete thought and start response
+                # Complete thought
                 thought_buffer += current_chunk
-                print(f"\n=== Complete Thought ===\n{thought_buffer}")
-
                 messages[-1] = ChatMessage(
                     role="assistant",
                     content=thought_buffer,
                     metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
                 )
-                yield messages
+                yield convert_chat_messages_to_gradio_format(messages)
 
                 # Start response
                 response_buffer = parts[1].text
-                print(f"\n=== Starting Response ===\n{response_buffer}")
-
-                messages.append(
-                    ChatMessage(
-                        role="assistant",
-                        content=response_buffer
-                    )
-                )
+                messages.append(ChatMessage(role="assistant", content=response_buffer))
                 thinking_complete = True
-
             elif thinking_complete:
+                # Response ongoing
                 response_buffer += current_chunk
-                print(f"\n=== Response Chunk ===\n{current_chunk}")
-
-                messages[-1] = ChatMessage(
-                    role="assistant",
-                    content=response_buffer
-                )
-
+                messages[-1] = ChatMessage(role="assistant", content=response_buffer)
             else:
+                # Still in "thinking"
                 thought_buffer += current_chunk
-                print(f"\n=== Thinking Chunk ===\n{current_chunk}")
-
                 messages[-1] = ChatMessage(
                     role="assistant",
                     content=thought_buffer,
                     metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
                 )
 
-            # time.sleep(0.05) #Optional debugging delay
-            yield messages
+            yield convert_chat_messages_to_gradio_format(messages)
 
         print(f"\n=== Final Response ===\n{response_buffer}")
 
     except Exception as e:
         print(f"\n=== Error ===\n{str(e)}")
-        messages.append(
-            ChatMessage(
-                role="assistant",
-                content=f"I apologize, but I encountered an error: {str(e)}"
-            )
-        )
-        yield messages
+        messages.append(ChatMessage(role="assistant", content=f"I encountered an error: {str(e)}"))
+        yield convert_chat_messages_to_gradio_format(messages)
 
-def user_message(msg: str, history: list) -> tuple[str, list]:
-    """Adds user message to chat history"""
-    history.append(ChatMessage(role="user", content=msg))
+def convert_chat_messages_to_gradio_format(messages):
+    """
+    ChatMessage list -> Gradio Chatbot 포맷( (user_str, assistant_str) 튜플 리스트 )
+    """
+    gradio_chat = []
+    user_text, assistant_text = None, None
+
+    for msg in messages:
+        role = msg.get("role")
+        content = msg.get("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
+            assistant_text = None
+        else:
+            # assistant
+            if user_text is None:
+                user_text = ""
+            if assistant_text is None:
+                assistant_text = content
+            else:
+                assistant_text += content  # 스트리밍 시 이어붙임
+    # 마지막 턴
+    if user_text is not None or assistant_text is not None:
+        gradio_chat.append((user_text or "", assistant_text or ""))
+
+    return gradio_chat
+
+def user_message(msg: str, history: list, doc_text: str) -> tuple[str, list]:
+    """
+    사용자가 입력할 때마다 doc_text (마크다운)를 
+    '참고하라'는 프롬프트를 자동으로 추가 (간단 예시)
+    """
+    if not doc_text.strip():
+        # 변환된 문서가 없으면 그냥 입력
+        user_query = msg
+    else:
+        # 문서를 참조해 달라는 지시어 추가
+        user_query = f"다음 문서를 참고해서 답변:\n\n{doc_text}\n\n질문: {msg}"
+
+    history.append(ChatMessage(role="user", content=user_query))
     return "", history
 
 
-######################################################
+##############################
 # 3) 통합 Gradio 앱 구성
-#   - PDF 업로드만 보이게 하고(나머지는 hidden)
-#   - 업로드 후 "변환" 버튼 클릭 시, 마크다운을 만들어
-#     Chatbot과 대화할 수 있도록 전달
-######################################################
-with gr.Blocks(title="통합 OCR & Gemini Chat", css=create_css(), theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate", neutral_hue="neutral")) as demo:
+#    - 화면을 최대한 넓게, 
+#      PDF 업로드와 변환 버튼, 
+#      Gemini 채팅만 보이게
+##############################
+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/이미지 -> 텍스트(마크다운) 변환 후, LLM Gemini와 대화</p>
+        <p>PDF/이미지 -> 텍스트(마크다운) 변환 후, Gemini LLM에 질의응답</p>
     </div>
     """)
-    
-    # 내부 상태(마크다운 텍스트)
-    md_state = gr.State("")
-    chat_history = gr.State([])  # Gemini 챗 기록 상태
 
-    # 1) 파일 업로드 UI
+    md_state = gr.State("")        # 변환된 마크다운 텍스트
+    chat_history = gr.State([])    # ChatMessage 리스트 (Gemini 대화 이력)
+
+    # 1) 파일 업로드 & 변환 버튼
     with gr.Row():
         file = gr.File(
-            label="PDF 또는 이미지 파일 업로드",
+            label="PDF / 이미지 업로드",
             file_types=[".pdf", ".png", ".jpeg", ".jpg"],
             interactive=True
         )
-        convert_btn = gr.Button(
-            "변환",
-            elem_classes="primary-button"
-        )
+        convert_btn = gr.Button("변환하기")
 
-    # 2) 원래 존재하던 슬라이더, 체크박스 등은 전부 hidden
+    # 2) 숨긴 컴포넌트들
     max_pages = gr.Slider(
         1, 20, 10,
         step=1,
         label='최대 변환 페이지 수',
-        elem_classes="invisible",
-        visible=False
+        visible=False,  # 숨김
+        elem_classes="invisible"
     )
     layout_mode = gr.Dropdown(
         ["layoutlmv3", "doclayout_yolo"],
         label="레이아웃 모델",
         value="doclayout_yolo",
-        elem_classes="invisible",
-        visible=False
+        visible=False,
+        elem_classes="invisible"
     )
     language = gr.Dropdown(
         all_lang,
         label="언어",
         value='auto',
-        elem_classes="invisible",
-        visible=False
+        visible=False,
+        elem_classes="invisible"
     )
     formula_enable = gr.Checkbox(
-        label="수식 인식 활성화",
+        label="수식 인식",
         value=True,
-        elem_classes="invisible",
-        visible=False
+        visible=False,
+        elem_classes="invisible"
     )
     is_ocr = gr.Checkbox(
-        label="OCR 강제 활성화",
+        label="OCR 강제",
         value=False,
-        elem_classes="invisible",
-        visible=False
+        visible=False,
+        elem_classes="invisible"
     )
     table_enable = gr.Checkbox(
-        label="표 인식 활성화(테스트)",
+        label="표 인식",
         value=True,
-        elem_classes="invisible",
-        visible=False
+        visible=False,
+        elem_classes="invisible"
     )
 
-    # 3) 출력 결과(파일, 마크다운 등)��� 숨김
-    #    필요하면 주석 해제하여 확인 가능
-    # output_file = gr.File(
-    #     label="변환 결과",
-    #     interactive=False,
-    #     visible=False
-    # )
-    # md = gr.Markdown(
-    #     label="마크다운 렌더링",
-    #     visible=False
-    # )
-    # md_text = gr.TextArea(
-    #     lines=45,
-    #     visible=False
-    # )
-    # pdf_show = PDF(
-    #     label='PDF 미리보기',
-    #     interactive=False,
-    #     visible=False,
-    #     height=800
-    # )
-
-    # 4) 파일 업로드 -> '변환' 버튼 클릭시 동작:
-    #    to_markdown 함수를 통해 md_state에 마크다운 저장
+    # convert_btn 누르면 마크다운 추출 -> md_state
     convert_btn.click(
         fn=to_markdown,
         inputs=[file, max_pages, is_ocr, layout_mode, formula_enable, table_enable, language],
         outputs=md_state
     )
 
-    # ==========================
-    # Gemini Chat 부분
-    # ==========================
-    gr.Markdown("## Gemini 2.0 Flash (With Thinking) Chat")
+    # 3) Gemini Chat 영역
+    gr.Markdown("## Gemini 2.0 Flash (Thinking) Chat")
 
     chatbot = gr.Chatbot(
-        label="Gemini2.0 Chatbot (Streaming Output)",
-        render_markdown=True,
-        height=400
+        label="Gemini2.0 Chatbot (Streaming)",
+        height=600
     )
-
     with gr.Row():
-        chat_input = gr.Textbox(
-            lines=1,
-            label="질문 입력",
-            placeholder="추출된 문서(마크다운 내용)에 대해 궁금한 점을 물어보세요..."
-        )
+        chat_input = gr.Textbox(lines=1, placeholder="질문을 입력하세요...")
         clear_button = gr.Button("대화 초기화")
 
-    # 사용자가 질문 -> user_message -> Gemini 처리 -> stream_gemini_response
-    def user_message_wrapper(msg, history, doc_text):
-        """
-        사용자가 입력할 때마다, doc_text(마크다운)를 참고하도록 
-        질문을 약간 변형해서 history에 추가하는 방식(간단 예시).
-        """
-        if not doc_text:
-            # 아직 변환된 문서가 없다면 그냥 질문
-            user_query = msg
-        else:
-            # 문서 내용(doc_text)을 "참고" 요청하는 간단 프롬프트 예시
-            user_query = f"다음 문서를 참고하여 답변:\n\n{doc_text}\n\n질문: {msg}"
-
-        history.append(ChatMessage(role="user", content=user_query))
-        return "", history
-
+    # 사용자가 질문하면 -> user_message -> Gemini 응답
     chat_input.submit(
-        fn=user_message_wrapper,
+        fn=user_message,
         inputs=[chat_input, chat_history, md_state],
         outputs=[chat_input, chat_history]
     ).then(
         fn=stream_gemini_response,
         inputs=[chat_input, chat_history],
-        outputs=chat_history
-    ).then(
-        fn=lambda h: h,
-        inputs=chat_history,
         outputs=chatbot
     )
 
+    # 초기화 버튼 -> 대화 기록, md_state 모두 비우고 -> chatbot도 초기화
     clear_button.click(
         fn=lambda: ([], ""),
         inputs=[],
@@ -515,14 +494,4 @@ with gr.Blocks(title="통합 OCR & Gemini Chat", css=create_css(), theme=gr.them
 # 4) 실제 실행
 ##############################
 if __name__ == "__main__":
-    # 첫 번째 demo.launch() - 통합 앱 실행
-    demo.launch(ssr_mode=True, debug=True)
-
-
-###############################################
-#  아래는 "Gemini 챗 코드" 원본에 있던 
-#  별도의 demo.launch() 부분 (누락 없이 주석 보존)
-###############################################
-# if __name__ == "__main__":
-#     demo.launch(debug=True)
-   
+    demo.launch(debug=True, server_name="0.0.0.0", server_port=7860)