app.py

import os
import gradio as gr
import tempfile
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from unstructured.partition.pdf import partition_pdf
from unstructured.partition.utils.constants import PartitionStrategy
from huggingface_hub import InferenceClient
import base64
from PIL import Image
import io
import requests
from getpass import getpass
import PyPDF2
import fitz  # PyMuPDF
import pytesseract

# # Step 2: Set up Hugging Face Token
# print("🔑 Setting up Hugging Face Token...")
# print("Please enter your Hugging Face token (get it from: https://huggingface.co/settings/tokens)")
# HF_TOKEN = getpass("Hugging Face Token: ")

# # Set environment variable
# os.environ["HUGGINGFACE_HUB_TOKEN"] = HF_TOKEN

# Step 3: Initialize Hugging Face components
print("🚀 Initializing models...")

# Initialize embeddings model (runs locally for better performance)
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={'device': 'cpu'}
)

# Initialize vector store
vector_store = InMemoryVectorStore(embeddings)

# Initialize Hugging Face Inference clients with proper multimodal support
def initialize_multimodal_clients():
    """Initialize clients with proper multimodal capabilities"""
    
    # Vision-Language Models (can understand images AND text together)
    multimodal_models = [
        "microsoft/git-large-coco",           # Best for image+text understanding
        "Salesforce/blip2-opt-2.7b",         # Strong multimodal model
        "microsoft/git-base-coco",            # Lighter alternative
        "Salesforce/blip-image-captioning-large"  # Good image understanding
    ]
    
    # Text-only models for when no images are involved
    text_models = [
        "google/flan-t5-base",                # Excellent for Q&A
        "microsoft/DialoGPT-medium",          # Conversational
        "facebook/blenderbot-400M-distill",   # Another option
    ]
    
    vision_client = None
    text_client = None
    
    # Try to initialize multimodal/vision client
    for model_name in multimodal_models:
        try:
            vision_client = InferenceClient(model=model_name, token=HF_TOKEN)
            print(f"✅ Multimodal client initialized: {model_name}")
            break
        except Exception as e:
            print(f"⚠️  Failed to initialize {model_name}: {e}")
            continue
    
    # Try to initialize text client
    for model_name in text_models:
        try:
            text_client = InferenceClient(model=model_name, token=HF_TOKEN)
            print(f"✅ Text client initialized: {model_name}")
            break
        except Exception as e:
            print(f"⚠️  Failed to initialize {model_name}: {e}")
            continue
    
    return vision_client, text_client

vision_client, text_client = initialize_multimodal_clients()

template = """
You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.

Question: {question} 
Context: {context} 
Answer:
"""

def extract_text_with_multiple_methods(pdf_path):
    """Try multiple methods to extract text from PDF"""
    extracted_text = ""
    methods_tried = []
    
    # Method 1: PyPDF2
    try:
        print("🔍 Trying PyPDF2...")
        with open(pdf_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            text_parts = []
            for page_num, page in enumerate(pdf_reader.pages):
                page_text = page.extract_text()
                if page_text.strip():
                    text_parts.append(f"Page {page_num + 1}:\n{page_text}")
            
            if text_parts:
                extracted_text = "\n\n".join(text_parts)
                methods_tried.append("PyPDF2")
                print(f"✅ PyPDF2 extracted {len(extracted_text)} characters")
    except Exception as e:
        print(f"⚠️ PyPDF2 failed: {e}")
    
    # Method 2: PyMuPDF (fitz) - often better for complex PDFs
    if not extracted_text.strip():
        try:
            print("🔍 Trying PyMuPDF...")
            doc = fitz.open(pdf_path)
            text_parts = []
            for page_num in range(len(doc)):
                page = doc.load_page(page_num)
                page_text = page.get_text()
                if page_text.strip():
                    text_parts.append(f"Page {page_num + 1}:\n{page_text}")
            
            if text_parts:
                extracted_text = "\n\n".join(text_parts)
                methods_tried.append("PyMuPDF")
                print(f"✅ PyMuPDF extracted {len(extracted_text)} characters")
            doc.close()
        except Exception as e:
            print(f"⚠️ PyMuPDF failed: {e}")
    
    # Method 3: OCR with PyMuPDF for image-based PDFs
    if not extracted_text.strip():
        try:
            print("🔍 Trying OCR with PyMuPDF...")
            doc = fitz.open(pdf_path)
            text_parts = []
            for page_num in range(min(len(doc), 5)):  # Limit to first 5 pages for OCR
                page = doc.load_page(page_num)
                # Convert page to image
                pix = page.get_pixmap()
                img_data = pix.tobytes("png")
                img = Image.open(io.BytesIO(img_data))
                
                # Apply OCR
                ocr_text = pytesseract.image_to_string(img)
                if ocr_text.strip():
                    text_parts.append(f"Page {page_num + 1} (OCR):\n{ocr_text}")
            
            if text_parts:
                extracted_text = "\n\n".join(text_parts)
                methods_tried.append("OCR")
                print(f"✅ OCR extracted {len(extracted_text)} characters")
            doc.close()
        except Exception as e:
            print(f"⚠️ OCR failed: {e}")
    
    return extracted_text, methods_tried

def upload_and_process_pdf(pdf_file):
    """Process uploaded PDF file with enhanced error handling"""
    if pdf_file is None:
        return "Please upload a PDF file first."
    
    try:
        # Create temporary directories
        with tempfile.TemporaryDirectory() as temp_dir:
            figures_dir = os.path.join(temp_dir, "figures")
            os.makedirs(figures_dir, exist_ok=True)
            
            # Save uploaded file temporarily
            temp_pdf_path = os.path.join(temp_dir, "uploaded.pdf")
            with open(temp_pdf_path, "wb") as f:
                f.write(pdf_file)
            
            # Check file size and validity
            file_size = os.path.getsize(temp_pdf_path)
            print(f"📄 Processing PDF: {file_size} bytes")
            
            if file_size == 0:
                return "❌ The uploaded file is empty. Please check your PDF file."
            
            if file_size > 50 * 1024 * 1024:  # 50MB limit
                return "❌ File too large (>50MB). Please upload a smaller PDF."
            
            # Try multiple extraction methods
            text, methods = extract_text_with_multiple_methods(temp_pdf_path)
            
            # Process with unstructured as backup/additional method
            unstructured_text = ""
            try:
                print("🔍 Trying unstructured...")
                elements = partition_pdf(
                    temp_pdf_path,
                    strategy=PartitionStrategy.FAST,
                    extract_image_block_types=["Image", "Table"],
                    extract_image_block_output_dir=figures_dir,
                    infer_table_structure=True
                )
                
                # Extract text elements
                text_elements = []
                for element in elements:
                    if hasattr(element, 'text') and element.text and element.category not in ["Image", "Table"]:
                        text_elements.append(element.text)
                
                if text_elements:
                    unstructured_text = "\n\n".join(text_elements)
                    print(f"✅ Unstructured extracted {len(unstructured_text)} characters")
                    
                    # Combine with existing text if available
                    if text.strip():
                        text = f"{text}\n\n--- Additional Content ---\n\n{unstructured_text}"
                    else:
                        text = unstructured_text
                        methods.append("unstructured")
                
            except Exception as unstructured_error:
                print(f"⚠️ Unstructured processing failed: {unstructured_error}")
            
            # Process images
            image_text = ""
            image_count = 0
            if os.path.exists(figures_dir):
                for file in os.listdir(figures_dir):
                    if file.lower().endswith(('.png', '.jpg', '.jpeg')):
                        try:
                            extracted_image_text = extract_text_from_image(os.path.join(figures_dir, file))
                            image_text += f"\n\n{extracted_image_text}"
                            image_count += 1
                        except Exception as e:
                            print(f"⚠️ Error processing image {file}: {e}")
            
            # Also try to extract images directly from PDF using PyMuPDF
            try:
                doc = fitz.open(temp_pdf_path)
                for page_num in range(min(len(doc), 10)):  # Process first 10 pages
                    page = doc.load_page(page_num)
                    image_list = page.get_images(full=True)
                    
                    for img_index, img in enumerate(image_list[:3]):  # Max 3 images per page
                        try:
                            xref = img[0]
                            pix = fitz.Pixmap(doc, xref)
                            if pix.n - pix.alpha < 4:  # GRAY or RGB
                                img_data = pix.tobytes("png")
                                img_path = os.path.join(figures_dir, f"page_{page_num}_img_{img_index}.png")
                                with open(img_path, "wb") as img_file:
                                    img_file.write(img_data)
                                
                                extracted_image_text = extract_text_from_image(img_path)
                                image_text += f"\n\n{extracted_image_text}"
                                image_count += 1
                            pix = None
                        except Exception as img_error:
                            print(f"⚠️ Error extracting image: {img_error}")
                            continue
                doc.close()
            except Exception as e:
                print(f"⚠️ Error extracting images from PDF: {e}")
            
            # Combine all text
            full_text = text
            if image_text.strip():
                full_text += f"\n\n--- Image Content ---\n{image_text}"
            
            if not full_text.strip():
                return (f"⚠️ No text could be extracted from the PDF using any method. "
                       f"This might be a scanned PDF without OCR text, or the file might be corrupted. "
                       f"Methods tried: {', '.join(['PyPDF2', 'PyMuPDF', 'OCR', 'unstructured']) if not methods else ', '.join(methods)}")
            
            # Split and index the text
            chunked_texts = split_text(full_text)
            
            if not chunked_texts:
                return "⚠️ Text was extracted but could not be split into chunks."
            
            # Clear existing vector store and add new documents
            global vector_store
            vector_store = InMemoryVectorStore(embeddings)
            index_docs(chunked_texts)
            
            success_msg = (f"✅ PDF processed successfully!\n"
                          f"📊 Statistics:\n"
                          f"- Text chunks: {len(chunked_texts)}\n"
                          f"- Images processed: {image_count}\n"
                          f"- Methods used: {', '.join(methods)}\n"
                          f"- Total characters: {len(full_text)}")
            
            return success_msg
    
    except Exception as e:
        return f"❌ Error processing PDF: {str(e)}\n\nTroubleshooting tips:\n- Ensure the PDF is not password protected\n- Try a different PDF file\n- Check if the file is corrupted"

def load_pdf(file_path, figures_directory):
    """Legacy function - now handled by upload_and_process_pdf"""
    return extract_text_with_multiple_methods(file_path)[0]

def extract_text_from_image(image_path):
    """Extract text description from image using Hugging Face Vision model"""
    try:
        # First try OCR for any text in the image
        ocr_text = ""
        try:
            img = Image.open(image_path)
            ocr_text = pytesseract.image_to_string(img)
            if ocr_text.strip():
                ocr_text = f"Text in image: {ocr_text.strip()}"
        except Exception as ocr_error:
            print(f"⚠️ OCR failed for image: {ocr_error}")
        
        # Then use vision model for description
        vision_description = ""
        if vision_client:
            try:
                with open(image_path, "rb") as img_file:
                    image_data = img_file.read()
                
                response = vision_client.image_to_text(image_data)
                
                if isinstance(response, list) and len(response) > 0:
                    vision_description = response[0].get('generated_text', '')
                elif isinstance(response, dict):
                    vision_description = response.get('generated_text', '')
                else:
                    vision_description = str(response)
                    
            except Exception as vision_error:
                print(f"⚠️ Vision model failed: {vision_error}")
        
        # Combine OCR and vision results
        combined_result = []
        if ocr_text:
            combined_result.append(ocr_text)
        if vision_description:
            combined_result.append(f"Image description: {vision_description}")
        
        if combined_result:
            return "\n".join(combined_result)
        else:
            return "Image content: Visual element present but could not be processed"
            
    except Exception as e:
        print(f"⚠️ Error extracting text from image: {e}")
        return "Image content: Visual element present but could not be processed"

def split_text(text):
    """Split text into chunks"""
    if not text or not text.strip():
        return []
    
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        add_start_index=True
    )
    return text_splitter.split_text(text)

def index_docs(texts):
    """Index documents in vector store"""
    if texts:
        vector_store.add_texts(texts)
        print(f"📚 Indexed {len(texts)} text chunks")

def retrieve_docs(query, k=4):
    """Retrieve relevant documents"""
    try:
        return vector_store.similarity_search(query, k=k)
    except Exception as e:
        print(f"⚠️ Error retrieving documents: {e}")
        return []

def answer_question_hf(question):
    """Answer question using Hugging Face multimodal models"""
    try:
        # Retrieve relevant documents
        related_documents = retrieve_docs(question)
        
        if not related_documents:
            return "❓ No relevant documents found. Please upload and process a PDF first."
        
        # Prepare context
        context = "\n\n".join([doc.page_content for doc in related_documents])
        
        # Limit context length for better performance
        if len(context) > 1500:
            context = context[:1500] + "..."
        
        # Check if we have image content in the context
        has_image_content = "Image content:" in context or "Image description:" in context
        
        if has_image_content and vision_client:
            # Use multimodal approach for questions involving images
            try:
                # For multimodal models, we can send both text and image context
                multimodal_prompt = f"""
                Based on the document content below (including text and image descriptions), answer this question: {question}
                
                Document content:
                {context}
                
                Please provide a clear, concise answer in 2-3 sentences.
                """
                
                response = vision_client.text_generation(
                    multimodal_prompt,
                    max_new_tokens=150,
                    temperature=0.7,
                    do_sample=True,
                    return_full_text=False,
                    stop=["Question:", "Document content:", "\n\n\n"]
                )
                
                if isinstance(response, dict):
                    answer = response.get('generated_text', '')
                elif isinstance(response, str):
                    answer = response
                else:
                    answer = str(response)
                
                if answer.strip():
                    return f"🖼️ {answer.strip()}"
                    
            except Exception as multimodal_error:
                print(f"⚠️ Multimodal model failed: {multimodal_error}")
        
        # Fall back to text-only approach
        if text_client:
            try:
                text_prompt = f"""
                Question: {question}
                
                Based on the following information from the document, provide a clear and concise answer:
                
                {context}
                
                Answer:"""
                
                response = text_client.text_generation(
                    text_prompt,
                    max_new_tokens=150,
                    temperature=0.7,
                    do_sample=True,
                    return_full_text=False,
                    stop=["Question:", "Answer:", "\n\n\n"]
                )
                
                if isinstance(response, dict):
                    answer = response.get('generated_text', '')
                elif isinstance(response, str):
                    answer = response
                else:
                    answer = str(response)
                
                # Clean up the answer
                answer = answer.strip()
                if answer:
                    return f"📄 {answer}"
                    
            except Exception as text_error:
                print(f"⚠️ Text model failed: {text_error}")
        
        # Last resort: Return extracted context
        if context:
            return f"📋 Based on the document, here's the relevant information:\n\n{context[:500]}{'...' if len(context) > 500 else ''}"
        else:
            return "❌ Unable to find relevant information in the document."
    
    except Exception as e:
        return f"❌ Error generating answer: {str(e)}"

def create_colab_interface():
    """Create Gradio interface optimized for Colab"""
    
    with gr.Blocks(
        title="Enhanced Multimodal RAG with Hugging Face", 
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px !important;
        }
        """
    ) as iface:
        
        gr.HTML("""
        <div style="text-align: center; padding: 20px;">
            <h1>📚 Enhanced Multimodal RAG with Hugging Face</h1>
            <p>Upload a PDF document and ask questions about its content, including images and tables!</p>
            <p><em>Now with improved PDF processing and multiple extraction methods</em></p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                # PDF Upload Section
                gr.Markdown("### 📤 Upload Document")
                pdf_input = gr.File(
                    label="Upload PDF Document",
                    file_types=[".pdf"],
                    type="binary",
                    height=100
                )
                
                upload_btn = gr.Button("🔄 Process PDF", variant="primary", size="lg")
                upload_status = gr.Textbox(
                    label="Processing Status",
                    interactive=False,
                    lines=6,
                    placeholder="Upload a PDF and click 'Process PDF' to begin..."
                )
            
            with gr.Column(scale=2):
                # Chat Interface
                gr.Markdown("### 💬 Chat Interface")
                chatbot = gr.Chatbot(
                    label="Chat with your document",
                    height=400,
                    show_label=False
                )
                
                with gr.Row():
                    question_input = gr.Textbox(
                        label="Ask a question",
                        placeholder="What is this document about?",
                        lines=1,
                        scale=4
                    )
                    ask_btn = gr.Button("Ask", variant="secondary", scale=1)
        
        # Example questions
        gr.Markdown("### 💡 Example Questions")
        example_questions = [
            "What is the main topic of this document?",
            "Can you summarize the key points?",
            "What information is shown in the images or tables?",
            "What are the conclusions or recommendations?"
        ]
        
        with gr.Row():
            for i, eq in enumerate(example_questions):
                example_btn = gr.Button(eq, size="sm")
                example_btn.click(
                    lambda x=eq: x,
                    outputs=[question_input]
                )
        
        # Event handlers
        def process_pdf_and_update(pdf_file):
            if pdf_file is None:
                return "Please select a PDF file first."
            return upload_and_process_pdf(pdf_file)
        
        def ask_and_update_chat(question, chat_history):
            if not question.strip():
                return chat_history, ""
            
            # Get answer
            answer = answer_question_hf(question)
            
            # Update chat history
            if chat_history is None:
                chat_history = []
            
            chat_history.append([question, answer])
            
            return chat_history, ""
        
        def clear_chat():
            return []
        
        # Connect events
        upload_btn.click(
            fn=process_pdf_and_update,
            inputs=[pdf_input],
            outputs=[upload_status]
        )
        
        ask_btn.click(
            fn=ask_and_update_chat,
            inputs=[question_input, chatbot],
            outputs=[chatbot, question_input]
        )
        
        question_input.submit(
            fn=ask_and_update_chat,
            inputs=[question_input, chatbot],
            outputs=[chatbot, question_input]
        )
        
        # Clear chat button
        clear_btn = gr.Button("🗑️ Clear Chat", variant="stop", size="sm")
        clear_btn.click(
            fn=clear_chat,
            outputs=[chatbot]
        )
        
        # Enhanced Instructions
        gr.Markdown("""
        ---
        ### 📋 Instructions:
        1. **Get HF Token**: Visit [Hugging Face Settings](https://huggingface.co/settings/tokens) to get your token
        2. **Upload PDF**: Click "Choose File" and select your PDF document
        3. **Process Document**: Click "Process PDF" and wait for confirmation
        4. **Ask Questions**: Type questions or use example prompts
        
        ### ✨ Enhanced Features:
        - 📄 **Multiple Text Extraction Methods**: PyPDF2, PyMuPDF, OCR, and Unstructured
        - 🖼️ **Advanced Image Processing**: Direct PDF image extraction + vision models
        - 🔍 **Robust PDF Handling**: Works with scanned PDFs, complex layouts, and image-heavy documents
        - 💬 **Interactive Chat**: Conversation history with multimodal understanding
        - ⚡ **Error Recovery**: Graceful fallbacks when one extraction method fails
        - 📊 **Processing Statistics**: Detailed feedback on what was extracted
        
        ### 🔧 Models Used:
        - **🎭 Multimodal**: Microsoft GIT-Large (understands images + text together)
        - **📝 Text Generation**: Google FLAN-T5-Base (optimized for Q&A)
        - **👁️ Vision**: Salesforce BLIP (image captioning and understanding)
        - **🔍 Embeddings**: Sentence Transformers all-MiniLM-L6-v2
        - **📖 OCR**: Tesseract for text recognition in images
        
        ### 🎯 Multimodal Capabilities:
        - **Text + Images**: Can answer questions about both text content and visual elements
        - **Image Understanding**: Describes charts, diagrams, photos in your PDFs
        - **OCR Integration**: Extracts text from images within PDFs
        - **Context Awareness**: Combines text and visual information for comprehensive answers
        - **Fallback Strategy**: Uses multiple methods to ensure successful text extraction
        
        ### 🛠️ Troubleshooting:
        - **No text extracted**: Try different PDF files, ensure not password-protected
        - **Large files**: Keep PDFs under 50MB for optimal performance
        - **Scanned PDFs**: OCR will automatically process image-based text
        - **Complex layouts**: Multiple extraction methods handle various PDF formats
        """)
    
    return iface

# Step 4: Launch the application
print("✅ Setup complete! Launching Enhanced Gradio interface...")

# Create and launch interface
iface = create_colab_interface()

# Launch with public link for Colab
iface.launch(
    debug=True,
    share=True,  # Creates public link
    server_name="0.0.0.0",
    server_port=7860,
    show_error=True
)

print("🎉 Enhanced Application launched successfully!")
print("📱 Use the public link above to access your app from anywhere!")