Create app.py

app.py

@@ -0,0 +1,680 @@
+import spaces
+import json
+import subprocess
+import os
+from llama_cpp import Llama
+from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
+from llama_cpp_agent.providers import LlamaCppPythonProvider
+from llama_cpp_agent.chat_history import BasicChatHistory
+from llama_cpp_agent.chat_history.messages import Roles
+import gradio as gr
+from huggingface_hub import hf_hub_download
+import tempfile
+from typing import List, Tuple, Optional
+
+# PDF 처리 라이브러리 조건부 import
+try:
+    from docling.document_converter import DocumentConverter
+    DOCLING_AVAILABLE = True
+except ImportError:
+    DOCLING_AVAILABLE = False
+    print("Docling not available, using alternative PDF processing")
+    try:
+        import PyPDF2
+        import pdfplumber
+    except ImportError:
+        print("Warning: PDF processing libraries not fully installed")
+
+# 환경 변수에서 HF_TOKEN 가져오기
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# 전역 변수 초기화 (중요!)
+llm = None
+llm_model = None
+document_context = ""  # PDF에서 추출한 문서 컨텍스트 저장
+document_filename = ""  # 현재 로드된 문서의 파일명
+
+print("Global variables initialized")
+print(f"document_context initial value: '{document_context}'")
+print(f"document_filename initial value: '{document_filename}'")
+
+# 모델 이름과 경로를 정의
+MISTRAL_MODEL_NAME = "Private-BitSix-Mistral-Small-3.1-24B-Instruct-2503.gguf"
+
+# 모델 다운로드 (HF_TOKEN 사용)
+model_path = hf_hub_download(
+    repo_id="ginigen/Private-BitSix-Mistral-Small-3.1-24B-Instruct-2503",
+    filename=MISTRAL_MODEL_NAME,
+    local_dir="./models",
+    token=HF_TOKEN
+)
+
+print(f"Downloaded model path: {model_path}")
+
+css = """
+.bubble-wrap {
+    padding-top: calc(var(--spacing-xl) * 3) !important;
+}
+.message-row {
+    justify-content: space-evenly !important;
+    width: 100% !important;
+    max-width: 100% !important;
+    margin: calc(var(--spacing-xl)) 0 !important;
+    padding: 0 calc(var(--spacing-xl) * 3) !important;
+}
+.flex-wrap.user {
+    border-bottom-right-radius: var(--radius-lg) !important;
+}
+.flex-wrap.bot {
+    border-bottom-left-radius: var(--radius-lg) !important;
+}
+.message.user{
+    padding: 10px;
+}
+.message.bot{
+    text-align: right;
+    width: 100%;
+    padding: 10px;
+    border-radius: 10px;
+}
+.message-bubble-border {
+    border-radius: 6px !important;
+}
+.message-buttons {
+    justify-content: flex-end !important;
+}
+.message-buttons-left {
+    align-self: end !important;
+}
+.message-buttons-bot, .message-buttons-user {
+    right: 10px !important;
+    left: auto !important;
+    bottom: 2px !important;
+}
+.dark.message-bubble-border {
+    border-color: #343140 !important;
+}
+.dark.user {
+    background: #1e1c26 !important;
+}
+.dark.assistant.dark, .dark.pending.dark {
+    background: #16141c !important;
+}
+.upload-container {
+    margin-bottom: 20px;
+    padding: 15px;
+    border: 2px dashed #666;
+    border-radius: 10px;
+    background-color: #f0f0f0;
+}
+.dark .upload-container {
+    background-color: #292733;
+    border-color: #444;
+}
+"""
+
+def get_messages_formatter_type(model_name):
+    if "Mistral" in model_name or "BitSix" in model_name:
+        return MessagesFormatterType.MISTRAL  # CHATML 대신 MISTRAL 형식 사용
+    else:
+        raise ValueError(f"Unsupported model: {model_name}")
+
+@spaces.GPU
+def convert_pdf_to_markdown(file):
+    """Convert PDF file to Markdown"""
+    global document_context, document_filename
+    
+    if file is None:
+        return "No file uploaded.", {}
+    
+    try:
+        print(f"\n=== PDF Conversion Started ===")
+        print(f"File path: {file.name}")
+        
+        # DocumentConverter 인스턴스 생성
+        converter = DocumentConverter()
+        
+        # 파일 변환
+        result = converter.convert(file.name)
+        
+        # Markdown으로 내보내기
+        markdown_content = result.document.export_to_markdown()
+        
+        # 문서 컨텍스트 업데이트 (중요!)
+        document_context = markdown_content
+        document_filename = os.path.basename(file.name)
+        
+        # 메타데이터 추출
+        metadata = {
+            "filename": document_filename,
+            "conversion_status": "success",
+            "content_length": len(markdown_content),
+            "preview": markdown_content[:500] + "..." if len(markdown_content) > 500 else markdown_content
+        }
+        
+        print(f"✅ PDF conversion successful!")
+        print(f"📄 Filename: {document_filename}")
+        print(f"📏 Document length: {len(markdown_content)} characters")
+        print(f"📝 Document preview (first 300 chars):\n{markdown_content[:300]}...")
+        print(f"=== PDF Conversion Complete ===\n")
+        
+        # 전역 변수 확인 및 강제 설정
+        print(f"\n=== Before setting global variables ===")
+        print(f"global document_context length: {len(document_context)}")
+        print(f"global document_filename: {document_filename}")
+        
+        # globals() 함수를 사용하여 강제로 전역 변수 설정
+        globals()['document_context'] = markdown_content
+        globals()['document_filename'] = document_filename
+        
+        print(f"\n=== After setting global variables ===")
+        print(f"global document_context length: {len(globals()['document_context'])}")
+        print(f"global document_filename: {globals()['document_filename']}")
+        
+        return markdown_content, metadata
+        
+    except Exception as e:
+        error_msg = f"Error during PDF conversion: {str(e)}"
+        print(f"❌ {error_msg}")
+        document_context = ""
+        document_filename = ""
+        return error_msg, {"error": str(e)}
+
+def find_relevant_chunks(document, query, chunk_size=1500, overlap=300):
+    """Find relevant chunks from document based on query"""
+    if not document:
+        return ""
+    
+    print(f"Finding relevant chunks for query: {query}")
+    
+    # 간단한 키워드 기반 검색
+    query_words = query.lower().split()
+    chunks = []
+    
+    # 문서를 청크로 나누기
+    for i in range(0, len(document), chunk_size - overlap):
+        chunk = document[i:i + chunk_size]
+        chunks.append((i, chunk))
+    
+    print(f"Document split into {len(chunks)} chunks")
+    
+    # 각 청크의 관련성 점수 계산
+    scored_chunks = []
+    for idx, chunk in chunks:
+        chunk_lower = chunk.lower()
+        score = sum(1 for word in query_words if word in chunk_lower)
+        if score > 0:
+            scored_chunks.append((score, idx, chunk))
+    
+    # 상위 2개 청크 선택 (메모리 절약)
+    scored_chunks.sort(reverse=True, key=lambda x: x[0])
+    relevant_chunks = scored_chunks[:2]
+    
+    if relevant_chunks:
+        result = ""
+        for score, idx, chunk in relevant_chunks:
+            result += f"\n[Extracted from position {idx} - relevance score: {score}]\n{chunk}\n"
+        print(f"Found {len(relevant_chunks)} relevant chunks")
+        return result
+    else:
+        # 관련 청크를 찾지 못한 경우 문서 시작 부분 반환
+        print("No relevant chunks found, returning document beginning")
+        return document[:2000]
+
+@spaces.GPU(duration=120)
+def respond(
+    message,
+    history: list[dict],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+    top_k,
+    repeat_penalty,
+):
+    global llm, llm_model
+    
+    # globals()를 사용하여 전역 변수에 접근
+    document_context = globals().get('document_context', '')
+    document_filename = globals().get('document_filename', '')
+    
+    # 디버깅을 위한 상세 로그
+    print(f"\n=== RESPOND Function Started ===")
+    print(f"User message: {message}")
+    print(f"Document context exists: {bool(document_context)}")
+    if document_context:
+        print(f"Document length: {len(document_context)}")
+        print(f"Document filename: {document_filename}")
+        print(f"Document preview (first 100 chars): {document_context[:100]}...")
+    else:
+        print("⚠️ document_context is empty!")
+        print(f"globals() keys (first 20): {list(globals().keys())[:20]}...")
+    
+    chat_template = get_messages_formatter_type(MISTRAL_MODEL_NAME)
+    
+    # 모델 파일 경로 확인
+    model_path_local = os.path.join("./models", MISTRAL_MODEL_NAME)
+    
+    if llm is None or llm_model != MISTRAL_MODEL_NAME:
+        print("Loading LLM model...")
+        llm = Llama(
+            model_path=model_path_local,
+            flash_attn=True,
+            n_gpu_layers=81,
+            n_batch=1024,
+            n_ctx=16384,  # 컨텍스트 크기
+            verbose=True  # 디버깅을 위한 상세 로그
+        )
+        llm_model = MISTRAL_MODEL_NAME
+        print("LLM model loaded successfully!")
+    
+    provider = LlamaCppPythonProvider(llm)
+    
+    # 기본 시스템 메시지 사용
+    system_prompt = system_message
+    
+    # 문서 컨텍스트가 있으면 시스템 메시지와 사용자 메시지 모두에 포함
+    if document_context and len(document_context) > 0:
+        doc_length = len(document_context)
+        print(f"📄 Including document context in message: {doc_length} characters")
+        
+        # 시스템 메시지에도 문서 정보 추가
+        system_prompt += f"\n\nCurrently loaded document: '{document_filename}'. You must reference this document content when answering all user questions."
+        
+        # 문서 내용을 적절한 크기로 제한
+        max_doc_length = 4000  # 최대 4000자로 제한
+        if doc_length > max_doc_length:
+            # 문서가 너무 긴 경우 처음과 끝 부분만 포함
+            doc_snippet = document_context[:2000] + "\n\n[... middle content omitted ...]\n\n" + document_context[-1500:]
+            enhanced_message = f"""Uploaded PDF document information:
+- Filename: {document_filename}
+- Document length: {doc_length} characters
+
+Document content (excerpt):
+{doc_snippet}
+
+User question: {message}
+
+Please answer the question based on the document above."""
+        else:
+            # 짧은 문서는 전체 포함
+            enhanced_message = f"""Uploaded PDF document information:
+- Filename: {document_filename}
+- Document length: {doc_length} characters
+
+Document content:
+{document_context}
+
+User question: {message}
+
+Please answer the question based on the document above."""
+        
+        print(f"Enhanced message length: {len(enhanced_message)}")
+        print(f"Message preview (first 300 chars):\n{enhanced_message[:300]}...")
+        
+        # 디버그: 최종 메시지 파일로 저장 (확인용)
+        with open("debug_last_message.txt", "w", encoding="utf-8") as f:
+            f.write(f"=== Debug Information ===\n")
+            f.write(f"Document length: {len(document_context)}\n")
+            f.write(f"Filename: {document_filename}\n")
+            f.write(f"User question: {message}\n")
+            f.write(f"\n=== Message to be sent ===\n")
+            f.write(enhanced_message)
+    else:
+        # 문서가 없는 경우
+        enhanced_message = message
+        if any(keyword in message.lower() for keyword in ["document", "pdf", "upload", "file", "content", "summary", "문서", "요약"]):
+            enhanced_message = f"{message}\n\n[System message: No PDF document is currently uploaded. Please upload a PDF file first.]"
+            print("Document-related question but no document loaded")
+        
+        # 디버그 메시지
+        print("⚠️ Warning: document_context is empty!")
+        print(f"document_context type: {type(document_context)}")
+        print(f"document_context value: {repr(document_context)}")
+        print(f"document_filename: {document_filename}")
+
+    settings = provider.get_provider_default_settings()
+    settings.temperature = temperature
+    settings.top_k = top_k
+    settings.top_p = top_p
+    settings.max_tokens = max_tokens
+    settings.repeat_penalty = repeat_penalty
+    settings.stream = True
+
+    # 시스템 프롬프트에 문서 내용 직접 포함 (문서가 있는 경우)
+    if document_context and len(document_context) > 0:
+        doc_snippet = document_context[:3000]  # 처음 3000자만 사용
+        enhanced_system_prompt = f"""{system_prompt}
+
+Currently loaded PDF document:
+Filename: {document_filename}
+Document content:
+{doc_snippet}
+{'' if len(document_context) <= 3000 else '... (remainder omitted)'}
+
+Answer user questions based on the document content above."""
+        
+        # 사용자 메시지는 단순하게
+        final_message = message
+    else:
+        enhanced_system_prompt = system_prompt
+        final_message = enhanced_message
+
+    agent = LlamaCppAgent(
+        provider,
+        system_prompt=enhanced_system_prompt,
+        predefined_messages_formatter_type=chat_template,
+        debug_output=True
+    )
+
+    messages = BasicChatHistory()
+
+    # 이전 대화 기록 추가 (수정됨)
+    for i in range(0, len(history)):
+        # 현재 메시지는 제외
+        if i < len(history) - 1 and history[i][1] is not None:
+            # 사용자 메시지
+            messages.add_message({
+                'role': Roles.user,
+                'content': history[i][0]
+            })
+            # 어시스턴트 메시지
+            messages.add_message({
+                'role': Roles.assistant,
+                'content': history[i][1]
+            })
+    
+    print(f"Sending final message: {final_message}")
+    
+    # 스트림 응답 생성
+    try:
+        stream = agent.get_chat_response(
+            final_message,  # 단순한 메시지 사용
+            llm_sampling_settings=settings,
+            chat_history=messages,
+            returns_streaming_generator=True,
+            print_output=False
+        )
+        
+        outputs = ""
+        for output in stream:
+            outputs += output
+            yield outputs
+    except Exception as e:
+        print(f"Error during stream generation: {e}")
+        yield "Sorry, an error occurred while generating the response. Please try again."
+
+def clear_document_context():
+    """Clear document context"""
+    global document_context, document_filename
+    document_context = ""
+    document_filename = ""
+    return "📭 Document context has been cleared. Please upload a new PDF."
+
+def check_document_status():
+    """Check current document status"""
+    global document_context, document_filename
+    print(f"\n=== Document Status Check ===")
+    print(f"document_context type: {type(document_context)}")
+    print(f"document_context length: {len(document_context) if document_context else 0}")
+    print(f"document_filename: '{document_filename}'")
+    
+    if document_context and len(document_context) > 0:
+        status = f"✅ Document loaded successfully.\n📄 Filename: {document_filename}\n📏 Document length: {len(document_context):,} characters"
+        print(f"Document first 100 chars: {document_context[:100]}")
+        return status
+    else:
+        return "📭 No document loaded. Please upload a PDF file."
+
+# Gradio 인터페이스 구성
+with gr.Blocks(theme=gr.themes.Soft(
+        primary_hue="blue",
+        secondary_hue="cyan",
+        neutral_hue="gray",
+        font=[gr.themes.GoogleFont("Exo"), "ui-sans-serif", "system-ui", "sans-serif"]
+    ).set(
+        body_background_fill="#f8f9fa",
+        block_background_fill="#ffffff",
+        block_border_width="1px",
+        block_title_background_fill="#e9ecef",
+        input_background_fill="#ffffff",
+        button_secondary_background_fill="#e9ecef",
+        border_color_accent="#dee2e6",
+        border_color_primary="#ced4da",
+        background_fill_secondary="#f8f9fa",
+        color_accent_soft="transparent",
+        code_background_fill="#f1f3f5",
+    ), css=css) as demo:
+    
+    gr.Markdown("# On-Premise Optimized 'LLM+RAG Model' Service by VIDraft")
+    gr.Markdown("🔍 **Advanced document analysis with state-of-the-art language model for accurate Q&A based on your PDF content**")
+    gr.Markdown("🌐 **Supports both English and Korean languages seamlessly with context-aware responses**")
+    gr.Markdown("💡 **How to use**: 1) Upload PDF below → 2) Ask questions about the document → 3) Get AI-powered answers")
+    
+    # 채팅 인터페이스를 위쪽에 배치
+    with gr.Row():
+        with gr.Column():
+            # 채팅 인터페이스
+            chatbot = gr.Chatbot(elem_id="chatbot", height=500)
+            msg = gr.Textbox(
+                label="Message Input",
+                placeholder="Enter your question... (Upload a PDF to ask questions about its content)",
+                lines=2
+            )
+            with gr.Row():
+                submit = gr.Button("Send", variant="primary")
+                clear_chat = gr.Button("Clear Chat")
+    
+    # 예제를 중간에 배치 - 영어와 한국어 예제 모두 포함
+    gr.Examples(
+        examples=[
+            ["What is this document about?"],
+            ["Please summarize the main contents of the uploaded PDF document."],
+            ["What are the key dates or deadlines mentioned in the document?"],
+            ["What are the 3 most important key points in this document?"],
+            ["이 문서의 주요 내용을 요약해주세요."],
+            ["문서에 나온 중요한 일정이나 날짜를 알려주세요."],
+            ["이 문서에서 가장 중요한 3가지 핵심 포인트는 무엇인가요?"]
+        ],
+        inputs=msg
+    )
+    
+    # PDF 업로드 섹션을 아래쪽에 배치
+    with gr.Accordion("📄 PDF Document Upload", open=True):
+        with gr.Row():
+            with gr.Column(scale=1):
+                file_input = gr.File(
+                    label="Select PDF Document",
+                    file_types=[".pdf"],
+                    type="filepath"
+                )
+                with gr.Row():
+                    convert_button = gr.Button("Convert Document", variant="primary")
+                    clear_button = gr.Button("Clear Document", variant="secondary")
+                    test_button = gr.Button("Test Document", variant="secondary")
+                
+                status_text = gr.Textbox(
+                    label="Document Status", 
+                    interactive=False,
+                    value=check_document_status(),
+                    lines=3
+                )
+            
+            with gr.Column(scale=1):
+                with gr.Accordion("Converted Document Preview", open=False):
+                    converted_text = gr.Textbox(
+                        label="Markdown Conversion Result",
+                        lines=10,
+                        max_lines=20,
+                        interactive=False
+                    )
+                    metadata_output = gr.JSON(label="Metadata")
+    
+    # 고급 설정을 가장 아래에 배치
+    with gr.Accordion("⚙️ Advanced Settings", open=False):
+        system_message = gr.Textbox(
+            value="You are an AI assistant that can answer in both English and Korean. When a PDF document is provided, analyze its content accurately and provide detailed answers. Respond in the same language as the user's question.",
+            label="System Message",
+            lines=3
+        )
+        max_tokens = gr.Slider(minimum=1, maximum=4096, value=2048, step=1, label="Max Tokens")
+        temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.3, step=0.1, label="Temperature (lower = more consistent)")
+        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.90, step=0.05, label="Top-p")
+        top_k = gr.Slider(minimum=0, maximum=100, value=40, step=1, label="Top-k")
+        repeat_penalty = gr.Slider(minimum=0.0, maximum=2.0, value=1.1, step=0.1, label="Repetition penalty")
+    
+    # 이벤트 핸들러
+    def user_submit(message, history):
+        return "", history + [[message, None]]
+    
+    def bot_response(history, system_msg, max_tok, temp, top_p_val, top_k_val, rep_pen):
+        if history and history[-1][1] is None:
+            user_message = history[-1][0]
+            
+            # 디버깅: 문서 컨텍스트 상태 확인
+            global document_context, document_filename
+            print(f"\n=== BOT RESPONSE Started ===")
+            print(f"User message: {user_message}")
+            if document_context:
+                print(f"📄 Document context active: {document_filename} ({len(document_context)} chars)")
+                print(f"Document preview (first 200 chars): {document_context[:200]}...")
+            else:
+                print("📭 No document context")
+            
+            # 단순한 형식 사용 - [user_message, assistant_message]
+            previous_history = []
+            for i in range(len(history) - 1):
+                if history[i][1] is not None:
+                    previous_history.append({
+                        "user": history[i][0],
+                        "assistant": history[i][1]
+                    })
+            
+            print(f"Previous conversations: {len(previous_history)}")
+            
+            # 문서가 있는 경우 특별 처리
+            if document_context and len(document_context) > 0:
+                print(f"📄 Generating document-based response... (Document length: {len(document_context)})")
+            
+            bot_message = ""
+            try:
+                for token in respond(
+                    user_message,
+                    previous_history,
+                    system_msg,
+                    max_tok,
+                    temp,
+                    top_p_val,
+                    top_k_val,
+                    rep_pen
+                ):
+                    bot_message = token
+                    history[-1][1] = bot_message
+                    yield history
+            except Exception as e:
+                print(f"❌ Error during response generation: {e}")
+                import traceback
+                traceback.print_exc()
+                history[-1][1] = "Sorry, an error occurred while generating the response. Please try again."
+                yield history
+    
+    # PDF 변환 이벤트
+    def on_pdf_convert(file):
+        """PDF conversion and status update"""
+        global document_context, document_filename
+        
+        if file is None:
+            return "", {}, "❌ No file selected."
+        
+        markdown_content, metadata = convert_pdf_to_markdown(file)
+        
+        if "error" in metadata:
+            status = f"❌ Conversion failed: {metadata['error']}"
+        else:
+            # 전역 변수 다시 한번 확인 및 설정 (globals() 사용)
+            globals()['document_context'] = markdown_content
+            globals()['document_filename'] = metadata['filename']
+            
+            status = f"✅ PDF document converted successfully!\n📄 Filename: {metadata['filename']}\n📏 Document length: {metadata['content_length']:,} characters\n\nYou can now ask questions about the document content in English or Korean.\n\nExample questions:\n- What is the main topic of this document?\n- Summarize the key points\n- 이 문서의 핵심 내용을 설명해주세요"
+            
+            print(f"\n✅ Document loading confirmed:")
+            print(f"- globals()['document_context'] length: {len(globals()['document_context'])}")
+            print(f"- globals()['document_filename']: {globals()['document_filename']}")
+            
+            # 최종 확인
+            if len(globals()['document_context']) > 0:
+                print("✅ Document successfully saved to global variables!")
+            else:
+                print("❌ Warning: Document not saved to global variables!")
+        
+        return markdown_content, metadata, status
+    
+    # 파일 업로드 시 자동 변환
+    file_input.change(
+        fn=on_pdf_convert,
+        inputs=[file_input],
+        outputs=[converted_text, metadata_output, status_text]
+    )
+    
+    # 수동 변환 버튼
+    convert_button.click(
+        fn=on_pdf_convert,
+        inputs=[file_input],
+        outputs=[converted_text, metadata_output, status_text]
+    )
+    
+    # 문서 테스트 함수
+    def test_document():
+        """Test currently loaded document"""
+        global document_context, document_filename
+        if document_context:
+            test_msg = f"✅ Document test results:\n"
+            test_msg += f"📄 Filename: {document_filename}\n"
+            test_msg += f"📏 Total length: {len(document_context):,} characters\n"
+            test_msg += f"📝 First 500 characters:\n{document_context[:500]}..."
+            return test_msg
+        else:
+            return "❌ No document currently loaded."
+    
+    test_button.click(
+        fn=test_document,
+        outputs=[status_text]
+    )
+    
+    clear_button.click(
+        fn=clear_document_context,
+        outputs=[status_text]
+    ).then(
+        fn=lambda: ("", {}, check_document_status()),
+        outputs=[converted_text, metadata_output, status_text]
+    )
+    
+    # 채팅 이벤트
+    msg.submit(user_submit, [msg, chatbot], [msg, chatbot]).then(
+        bot_response, 
+        [chatbot, system_message, max_tokens, temperature, top_p, top_k, repeat_penalty], 
+        chatbot
+    )
+    
+    submit.click(user_submit, [msg, chatbot], [msg, chatbot]).then(
+        bot_response, 
+        [chatbot, system_message, max_tokens, temperature, top_p, top_k, repeat_penalty], 
+        chatbot
+    )
+    
+    clear_chat.click(lambda: [], None, chatbot)
+
+if __name__ == "__main__":
+    # 필요한 디렉토리 생성
+    os.makedirs("./models", exist_ok=True)
+    
+    # 환경 변수 확인
+    if not HF_TOKEN:
+        print("⚠️  Warning: HF_TOKEN not set. Model download may be restricted.")
+        print("To set environment variable: export HF_TOKEN='your_huggingface_token'")
+    
+    demo.launch(
+        server_name="0.0.0.0",  # Accessible from local network
+        server_port=7860,
+        share=False  # Disabled for on-premise environment
+    )