Add PDF drag-and-drop functionality with Mistral OCR integration

app.py

@@ -11,7 +11,7 @@ import logging
 from contextlib import asynccontextmanager
 from typing import Dict, Any, List
 
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request, File, UploadFile
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
 from pydantic import BaseModel, Field
@@ -19,6 +19,7 @@ from pydantic import BaseModel, Field
 # Import our inference services
 from inference.mistral_prompting import create_mistral_service, MistralPromptingService
 from inference.bert_classif import create_bert_service, BERTInferenceService
+from inference.ocr_service import create_ocr_service, OCRService
 
 # Setup logging
 logging.basicConfig(
@@ -31,6 +32,7 @@ logger = logging.getLogger(__name__)
 mistral_base_service: MistralPromptingService = None
 mistral_finetuned_service: MistralPromptingService = None
 bert_service: BERTInferenceService = None
+ocr_service: OCRService = None
 
 # Model configurations
 MODELS = {
@@ -44,7 +46,7 @@ BERT_MODEL_PATH = "SoelMgd/bert-pii-detection"
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Manage application lifespan - startup and shutdown."""
-    global mistral_base_service, mistral_finetuned_service, bert_service
+    global mistral_base_service, mistral_finetuned_service, bert_service, ocr_service
     
     # Startup
     logger.info("🚀 Starting PII Masking Demo application...")
@@ -73,6 +75,16 @@ async def lifespan(app: FastAPI):
     except Exception as e:
         logger.error(f"Failed to initialize BERT service: {e}")
         # Don't raise exception - let app start but handle gracefully in endpoints
+    
+    try:
+        # Initialize OCR service
+        logger.info("Initializing OCR service...")
+        ocr_service = await create_ocr_service()
+        logger.info("✅ OCR service initialized successfully")
+        
+    except Exception as e:
+        logger.error(f"Failed to initialize OCR service: {e}")
+        # Don't raise exception - let app start but handle gracefully in endpoints
         
     yield
     
@@ -271,10 +283,109 @@ async def predict(request: PredictionRequest):
             detail=f"Prediction failed: {str(e)}"
         )
 
+@app.post("/predict-pdf", response_model=PredictionResponse)
+async def predict_pdf(
+    file: UploadFile = File(...),
+    method: str = "mistral",
+    model: str = "base",
+    pii_entities: str = "[]"
+):
+    """
+    Extract text from PDF using OCR, then predict PII entities and return masked text.
+    
+    Supports the same methods as /predict: Mistral (base/fine-tuned) and BERT.
+    """
+    # Validate file type
+    if not file.filename.lower().endswith('.pdf'):
+        raise HTTPException(
+            status_code=400,
+            detail="Only PDF files are supported"
+        )
+    
+    # Check OCR service availability
+    if ocr_service is None:
+        raise HTTPException(
+            status_code=503,
+            detail="OCR service not available. Please check API key configuration."
+        )
+    
+    # Validate method
+    if method not in ["mistral", "bert"]:
+        raise HTTPException(
+            status_code=400, 
+            detail=f"Method '{method}' not supported. Use 'mistral' or 'bert'."
+        )
+    
+    try:
+        # Parse PII entities list
+        import json
+        pii_entities_list = json.loads(pii_entities) if pii_entities else []
+        
+        start_time = time.time()
+        
+        # Read PDF content
+        pdf_content = await file.read()
+        logger.info(f"📄 Received PDF file: {file.filename} ({len(pdf_content)} bytes)")
+        
+        # Extract text using OCR
+        logger.info("🔍 Extracting text from PDF using Mistral OCR...")
+        extracted_text = await ocr_service.extract_text_from_pdf(pdf_content)
+        
+        if not extracted_text or len(extracted_text.strip()) < 10:
+            raise HTTPException(
+                status_code=400,
+                detail="Could not extract sufficient text from PDF"
+            )
+        
+        logger.info(f"📝 Extracted {len(extracted_text)} characters from PDF")
+        
+        # Now process the extracted text with the selected method
+        if method == "mistral":
+            # Get the appropriate Mistral service
+            service = get_mistral_service(model)
+            prediction = await service.predict(extracted_text, pii_entities_list)
+            method_used = f"{method}-{model}"
+            
+        elif method == "bert":
+            # Check BERT service availability
+            if bert_service is None:
+                raise HTTPException(
+                    status_code=503, 
+                    detail="BERT service not available. Please check model configuration."
+                )
+            
+            prediction = await bert_service.predict(extracted_text, pii_entities_list)
+            method_used = "bert"
+        
+        processing_time = time.time() - start_time
+        
+        # Count total entities
+        num_entities = sum(len(entities) for entities in prediction.entities.values())
+        
+        logger.info(f"✅ PDF processing completed in {processing_time:.3f}s - found {num_entities} entities")
+        
+        return PredictionResponse(
+            masked_text=prediction.masked_text,
+            entities=prediction.entities,
+            processing_time=processing_time,
+            method_used=f"pdf-{method_used}",
+            num_entities=num_entities,
+            selected_entities=pii_entities_list
+        )
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"❌ PDF processing failed: {e}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"PDF processing failed: {str(e)}"
+        )
+
 @app.get("/health", response_model=HealthResponse)
 async def health_check():
     """Health check endpoint."""
-    global mistral_base_service, mistral_finetuned_service, bert_service
+    global mistral_base_service, mistral_finetuned_service, bert_service, ocr_service
     
     services_status = {
         "mistral_base": {
@@ -294,6 +405,12 @@ async def health_check():
             "initialized": bert_service.is_initialized if bert_service else False,
             "model": BERT_MODEL_PATH,
             "info": bert_service.get_service_info() if bert_service else None
+        },
+        "ocr": {
+            "available": ocr_service is not None,
+            "initialized": ocr_service.is_initialized if ocr_service else False,
+            "model": "mistral-ocr-latest",
+            "info": ocr_service.get_service_info() if ocr_service else None
         }
     }
     
@@ -301,12 +418,13 @@ async def health_check():
     base_healthy = mistral_base_service and mistral_base_service.is_initialized
     finetuned_healthy = mistral_finetuned_service and mistral_finetuned_service.is_initialized
     bert_healthy = bert_service and bert_service.is_initialized
+    ocr_healthy = ocr_service and ocr_service.is_initialized
     
-    healthy_services = sum([base_healthy, finetuned_healthy, bert_healthy])
+    healthy_services = sum([base_healthy, finetuned_healthy, bert_healthy, ocr_healthy])
     
-    if healthy_services == 3:
+    if healthy_services == 4:
         overall_status = "healthy"
-    elif healthy_services >= 1:
+    elif healthy_services >= 2:
         overall_status = "partial"
     else:
         overall_status = "degraded"

inference/ocr_service.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+"""
+OCR Service for PDF processing using Mistral OCR API.
+"""
+
+import os
+import sys
+import logging
+import base64
+from pathlib import Path
+from typing import Optional, Dict, Any
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+from mistralai import Mistral
+
+# Setup logging
+logger = logging.getLogger(__name__)
+
+class OCRService:
+    """
+    OCR service for extracting text from PDF documents using Mistral OCR API.
+    """
+    
+    def __init__(self, api_key: Optional[str] = None):
+        """
+        Initialize the OCR service.
+        
+        Args:
+            api_key: Mistral API key (if None, will read from environment)
+        """
+        self.api_key = api_key or os.environ.get("MISTRAL_API_KEY")
+        if not self.api_key:
+            raise ValueError("MISTRAL_API_KEY not found in environment variables")
+        
+        self.client = Mistral(api_key=self.api_key)
+        self.is_initialized = True
+        
+        logger.info("🔧 OCR service initialized with Mistral API")
+    
+    async def extract_text_from_pdf(self, pdf_content: bytes) -> str:
+        """
+        Extract text from PDF content using Mistral OCR.
+        
+        Args:
+            pdf_content: Raw PDF file content as bytes
+            
+        Returns:
+            Extracted text content
+        """
+        try:
+            # Encode PDF content to base64
+            base64_pdf = base64.b64encode(pdf_content).decode('utf-8')
+            
+            logger.info(f"📄 Processing PDF ({len(pdf_content)} bytes) with Mistral OCR...")
+            
+            # Process the PDF with OCR
+            ocr_response = self.client.ocr.process(
+                model="mistral-ocr-latest",
+                document={
+                    "type": "document_url",
+                    "document_url": f"data:application/pdf;base64,{base64_pdf}"
+                },
+                include_image_base64=False  # Don't include images to save bandwidth
+            )
+            
+            logger.info("✅ OCR processing completed")
+            
+            # Extract text from all pages
+            extracted_text = ""
+            
+            if hasattr(ocr_response, 'pages') and ocr_response.pages:
+                logger.info(f"📄 Found {len(ocr_response.pages)} pages")
+                
+                for i, page in enumerate(ocr_response.pages):
+                    if hasattr(page, 'markdown') and page.markdown:
+                        page_text = page.markdown
+                        extracted_text += page_text + "\n\n"
+                        logger.debug(f"📝 Page {i+1}: {len(page_text)} characters")
+                
+                logger.info(f"📄 Total extracted text: {len(extracted_text)} characters")
+                
+                if not extracted_text.strip():
+                    logger.warning("⚠️  No text extracted from PDF")
+                    return "No text could be extracted from this PDF."
+                
+                return extracted_text.strip()
+                
+            else:
+                logger.warning("⚠️  No pages found in OCR response")
+                return "No text could be extracted from this PDF."
+                
+        except Exception as e:
+            logger.error(f"❌ OCR processing failed: {e}")
+            raise RuntimeError(f"Failed to extract text from PDF: {str(e)}")
+    
+    def get_service_info(self) -> Dict[str, Any]:
+        """Get service information for monitoring."""
+        return {
+            "service_name": "OCRService",
+            "is_initialized": self.is_initialized,
+            "api_provider": "Mistral",
+            "model": "mistral-ocr-latest",
+            "description": "PDF text extraction using Mistral OCR API"
+        }
+
+# Factory function for easy initialization
+async def create_ocr_service(api_key: Optional[str] = None) -> OCRService:
+    """
+    Factory function to create and initialize OCR service.
+    
+    Args:
+        api_key: Mistral API key (if None, will read from environment)
+        
+    Returns:
+        Initialized OCRService
+    """
+    try:
+        service = OCRService(api_key)
+        logger.info("✅ OCR service created successfully")
+        return service
+    except Exception as e:
+        logger.error(f"❌ Failed to create OCR service: {e}")
+        raise 

static/index.html

@@ -298,6 +298,131 @@
             background: #4338ca;
         }
 
+        .input-method-selection {
+            margin-bottom: 25px;
+        }
+
+        .input-tabs {
+            display: flex;
+            border-bottom: 2px solid #e5e7eb;
+            margin-bottom: 20px;
+        }
+
+        .input-tab {
+            padding: 12px 24px;
+            background: none;
+            border: none;
+            cursor: pointer;
+            font-size: 16px;
+            font-weight: 500;
+            color: #6b7280;
+            border-bottom: 2px solid transparent;
+            transition: all 0.3s ease;
+        }
+
+        .input-tab.active {
+            color: #4f46e5;
+            border-bottom-color: #4f46e5;
+        }
+
+        .input-tab:hover {
+            color: #4f46e5;
+        }
+
+        .input-content {
+            display: none;
+        }
+
+        .input-content.active {
+            display: block;
+        }
+
+        .dropzone {
+            border: 2px dashed #d1d5db;
+            border-radius: 12px;
+            padding: 40px 20px;
+            text-align: center;
+            background: #f9fafb;
+            transition: all 0.3s ease;
+            cursor: pointer;
+        }
+
+        .dropzone.dragover {
+            border-color: #4f46e5;
+            background: #f0f9ff;
+        }
+
+        .dropzone-icon {
+            font-size: 48px;
+            color: #9ca3af;
+            margin-bottom: 16px;
+        }
+
+        .dropzone.dragover .dropzone-icon {
+            color: #4f46e5;
+        }
+
+        .dropzone-text {
+            color: #374151;
+            font-size: 16px;
+            margin-bottom: 8px;
+        }
+
+        .dropzone-subtext {
+            color: #6b7280;
+            font-size: 14px;
+        }
+
+        .file-info {
+            display: none;
+            background: #f0f9ff;
+            border: 1px solid #bfdbfe;
+            border-radius: 8px;
+            padding: 12px 16px;
+            margin-top: 12px;
+        }
+
+        .file-info.show {
+            display: flex;
+            align-items: center;
+            gap: 12px;
+        }
+
+        .file-icon {
+            color: #3b82f6;
+            font-size: 20px;
+        }
+
+        .file-details {
+            flex: 1;
+        }
+
+        .file-name {
+            font-weight: 500;
+            color: #1f2937;
+        }
+
+        .file-size {
+            font-size: 14px;
+            color: #6b7280;
+        }
+
+        .file-remove {
+            background: none;
+            border: none;
+            color: #6b7280;
+            cursor: pointer;
+            font-size: 18px;
+            padding: 4px;
+            border-radius: 4px;
+            transition: all 0.2s ease;
+        }
+
+        .file-remove:hover {
+            background: #fee2e2;
+            color: #dc2626;
+        }
+
         .error {
             background: #fef2f2;
             color: #dc2626;
@@ -331,16 +456,43 @@
 
         <div class="content">
             <form id="piiForm">
-                <div class="form-group">
-                    <label for="inputText">Enter your text:</label>
-                    <textarea 
-                        id="inputText" 
-                        class="input-textarea"
-                        placeholder="Enter text containing PII information...
+                <div class="form-group input-method-selection">
+                    <label>Choose input method:</label>
+                    <div class="input-tabs">
+                        <button type="button" class="input-tab active" onclick="switchInputMethod('text')">
+                            📝 Text Input
+                        </button>
+                        <button type="button" class="input-tab" onclick="switchInputMethod('pdf')">
+                            📄 PDF Upload
+                        </button>
+                    </div>
+                    
+                    <div id="textInput" class="input-content active">
+                        <textarea 
+                            id="inputText" 
+                            class="input-textarea"
+                            placeholder="Enter text containing PII information...
 
 Example: Hi, my name is John Smith and my email is [email protected]. Call me at 555-1234."
-                        required
-                    ></textarea>
+                        ></textarea>
+                    </div>
+                    
+                    <div id="pdfInput" class="input-content">
+                        <div class="dropzone" id="dropzone" onclick="document.getElementById('fileInput').click()">
+                            <div class="dropzone-icon">📄</div>
+                            <div class="dropzone-text">Drop your PDF here or click to browse</div>
+                            <div class="dropzone-subtext">Supports PDF files up to 10MB</div>
+                        </div>
+                        <input type="file" id="fileInput" accept=".pdf" style="display: none;">
+                        <div id="fileInfo" class="file-info">
+                            <div class="file-icon">📄</div>
+                            <div class="file-details">
+                                <div class="file-name" id="fileName"></div>
+                                <div class="file-size" id="fileSize"></div>
+                            </div>
+                            <button type="button" class="file-remove" onclick="removeFile()">✕</button>
+                        </div>
+                    </div>
                 </div>
 
                 <div class="form-group">
@@ -349,14 +501,14 @@ Example: Hi, my name is John Smith and my email is [email protected]. Call
                         <div class="method-option">
                             <input type="radio" id="mistral" name="method" value="mistral" class="method-radio" checked onchange="toggleMistralModelSelection()">
                             <label for="mistral" class="method-label">
-                                <div class="method-title">🧠 Mistral AI</div>
+                                <div class="method-title">Mistral AI</div>
                                 <div class="method-desc">High accuracy via API</div>
                             </label>
                         </div>
                         <div class="method-option">
                             <input type="radio" id="bert" name="method" value="bert" class="method-radio" onchange="toggleMistralModelSelection()">
                             <label for="bert" class="method-label">
-                                <div class="method-title">🤖 BERT</div>
+                                <div class="method-title">BERT</div>
                                 <div class="method-desc">Fast local processing</div>
                             </label>
                         </div>
@@ -369,14 +521,14 @@ Example: Hi, my name is John Smith and my email is [email protected]. Call
                         <div class="method-option">
                             <input type="radio" id="base" name="model" value="base" class="method-radio" checked>
                             <label for="base" class="method-label">
-                                <div class="method-title">🎯 Base Model</div>
+                                <div class="method-title">Base Model</div>
                                 <div class="method-desc">mistral-large-latest with detailed prompting</div>
                             </label>
                         </div>
                         <div class="method-option">
                             <input type="radio" id="finetuned" name="model" value="finetuned" class="method-radio">
                             <label for="finetuned" class="method-label">
-                                <div class="method-title">🎯 Fine-tuned Model</div>
+                                <div class="method-title">Fine-tuned Model</div>
                                 <div class="method-desc">Specialized PII detection model</div>
                             </label>
                         </div>
@@ -384,7 +536,7 @@ Example: Hi, my name is John Smith and my email is [email protected]. Call
                 </div>
 
                 <div class="form-group pii-selection">
-                    <label>🎯 Select PII entities to mask:</label>
+                    <label>Select PII entities to mask:</label>
                     <div class="pii-controls">
                         <button type="button" class="pii-btn primary" onclick="selectAllPII()">Select All</button>
                         <button type="button" class="pii-btn" onclick="selectNonePII()">Select None</button>
@@ -396,7 +548,7 @@ Example: Hi, my name is John Smith and my email is [email protected]. Call
                 </div>
 
                 <button type="submit" class="process-btn" id="processBtn">
-                    🚀 Process Text
+                    Process Text
                 </button>
             </form>
 
@@ -407,12 +559,12 @@ Example: Hi, my name is John Smith and my email is [email protected]. Call
 
             <div class="result" id="result">
                 <div class="result-section">
-                    <div class="result-title">🎭 Masked Text</div>
+                    <div class="result-title">Masked Text</div>
                     <div class="result-content" id="maskedText"></div>
                 </div>
 
                 <div class="result-section">
-                    <div class="result-title">📋 Detected Entities</div>
+                    <div class="result-title">Detected Entities</div>
                     <div class="result-content" id="entities"></div>
                 </div>
 
@@ -462,9 +614,14 @@ Example: Hi, my name is John Smith and my email is [email protected]. Call
             'ADDRESS', 'CITY', 'STATE', 'ZIPCODE', 'DOB', 'AGE', 'IP'
         ];
 
+        // Global variables
+        let selectedFile = null;
+        let currentInputMethod = 'text';
+
         // Initialize PII selection on page load
         document.addEventListener('DOMContentLoaded', function() {
             initializePIISelection();
+            initializeFileUpload();
         });
 
         function initializePIISelection() {
@@ -519,19 +676,114 @@ Example: Hi, my name is John Smith and my email is [email protected]. Call
             return Array.from(checkboxes).map(cb => cb.value);
         }
 
+        function switchInputMethod(method) {
+            currentInputMethod = method;
+            
+            // Update tab appearance
+            document.querySelectorAll('.input-tab').forEach(tab => tab.classList.remove('active'));
+            event.target.classList.add('active');
+            
+            // Update content visibility
+            document.querySelectorAll('.input-content').forEach(content => content.classList.remove('active'));
+            document.getElementById(method + 'Input').classList.add('active');
+        }
+
+        function initializeFileUpload() {
+            const dropzone = document.getElementById('dropzone');
+            const fileInput = document.getElementById('fileInput');
+            
+            // Prevent default drag behaviors
+            ['dragenter', 'dragover', 'dragleave', 'drop'].forEach(eventName => {
+                dropzone.addEventListener(eventName, preventDefaults, false);
+                document.body.addEventListener(eventName, preventDefaults, false);
+            });
+            
+            // Highlight drop area when item is dragged over it
+            ['dragenter', 'dragover'].forEach(eventName => {
+                dropzone.addEventListener(eventName, highlight, false);
+            });
+            
+            ['dragleave', 'drop'].forEach(eventName => {
+                dropzone.addEventListener(eventName, unhighlight, false);
+            });
+            
+            // Handle dropped files
+            dropzone.addEventListener('drop', handleDrop, false);
+            
+            // Handle file input change
+            fileInput.addEventListener('change', function(e) {
+                handleFiles(e.target.files);
+            });
+            
+            function preventDefaults(e) {
+                e.preventDefault();
+                e.stopPropagation();
+            }
+            
+            function highlight(e) {
+                dropzone.classList.add('dragover');
+            }
+            
+            function unhighlight(e) {
+                dropzone.classList.remove('dragover');
+            }
+            
+            function handleDrop(e) {
+                const dt = e.dataTransfer;
+                const files = dt.files;
+                handleFiles(files);
+            }
+            
+            function handleFiles(files) {
+                if (files.length > 0) {
+                    const file = files[0];
+                    
+                    // Validate file type
+                    if (!file.name.toLowerCase().endsWith('.pdf')) {
+                        showError('Please select a PDF file.');
+                        return;
+                    }
+                    
+                    // Validate file size (10MB limit)
+                    const maxSize = 10 * 1024 * 1024; // 10MB
+                    if (file.size > maxSize) {
+                        showError('File size must be less than 10MB.');
+                        return;
+                    }
+                    
+                    selectedFile = file;
+                    showFileInfo(file);
+                }
+            }
+        }
+
+        function showFileInfo(file) {
+            document.getElementById('fileName').textContent = file.name;
+            document.getElementById('fileSize').textContent = formatFileSize(file.size);
+            document.getElementById('fileInfo').classList.add('show');
+        }
+
+        function removeFile() {
+            selectedFile = null;
+            document.getElementById('fileInfo').classList.remove('show');
+            document.getElementById('fileInput').value = '';
+        }
+
+        function formatFileSize(bytes) {
+            if (bytes === 0) return '0 Bytes';
+            const k = 1024;
+            const sizes = ['Bytes', 'KB', 'MB', 'GB'];
+            const i = Math.floor(Math.log(bytes) / Math.log(k));
+            return parseFloat((bytes / Math.pow(k, i)).toFixed(2)) + ' ' + sizes[i];
+        }
+
         document.getElementById('piiForm').addEventListener('submit', async function(e) {
             e.preventDefault();
             
-            const text = document.getElementById('inputText').value.trim();
             const method = document.querySelector('input[name="method"]:checked').value;
             const model = document.querySelector('input[name="model"]:checked').value;
             const selectedPIIEntities = getSelectedPIIEntities();
             
-            if (!text) {
-                showError('Please enter some text to analyze.');
-                return;
-            }
-            
             if (selectedPIIEntities.length === 0) {
                 showError('Please select at least one PII entity to mask.');
                 return;
@@ -543,18 +795,50 @@ Example: Hi, my name is John Smith and my email is [email protected]. Call
             hideResult();
             
             try {
-                const response = await fetch('/predict', {
-                    method: 'POST',
-                    headers: {
-                        'Content-Type': 'application/json',
-                    },
-                    body: JSON.stringify({
-                        text: text,
-                        method: method,
-                        model: model,
-                        pii_entities: selectedPIIEntities
-                    })
-                });
+                let response;
+                
+                if (currentInputMethod === 'text') {
+                    // Text input processing
+                    const text = document.getElementById('inputText').value.trim();
+                    
+                    if (!text) {
+                        showError('Please enter some text to analyze.');
+                        setLoading(false);
+                        return;
+                    }
+                    
+                    response = await fetch('/predict', {
+                        method: 'POST',
+                        headers: {
+                            'Content-Type': 'application/json',
+                        },
+                        body: JSON.stringify({
+                            text: text,
+                            method: method,
+                            model: model,
+                            pii_entities: selectedPIIEntities
+                        })
+                    });
+                    
+                } else if (currentInputMethod === 'pdf') {
+                    // PDF processing
+                    if (!selectedFile) {
+                        showError('Please select a PDF file to analyze.');
+                        setLoading(false);
+                        return;
+                    }
+                    
+                    const formData = new FormData();
+                    formData.append('file', selectedFile);
+                    formData.append('method', method);
+                    formData.append('model', model);
+                    formData.append('pii_entities', JSON.stringify(selectedPIIEntities));
+                    
+                    response = await fetch('/predict-pdf', {
+                        method: 'POST',
+                        body: formData
+                    });
+                }
                 
                 const result = await response.json();
                 
@@ -578,11 +862,11 @@ Example: Hi, my name is John Smith and my email is [email protected]. Call
             if (loading) {
                 loadingEl.style.display = 'block';
                 processBtn.disabled = true;
-                processBtn.textContent = '⏳ Processing...';
+                processBtn.textContent = 'Processing...';
             } else {
                 loadingEl.style.display = 'none';
                 processBtn.disabled = false;
-                processBtn.textContent = '🚀 Process Text';
+                processBtn.textContent = 'Process Text';
             }
         }
         

test_ocr.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""
+Test script for Mistral OCR functionality.
+"""
+
+import os
+import sys
+import asyncio
+import logging
+from pathlib import Path
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Add the inference directory to Python path
+sys.path.insert(0, str(Path(__file__).parent / "inference"))
+
+from mistralai import Mistral
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+async def test_mistral_ocr():
+    """Test Mistral OCR with a sample PDF."""
+    
+    # Check if API key is available
+    api_key = os.environ.get("MISTRAL_API_KEY")
+    if not api_key:
+        logger.error("❌ MISTRAL_API_KEY not found in environment variables")
+        return
+    
+    try:
+        # Initialize Mistral client
+        client = Mistral(api_key=api_key)
+        logger.info("✅ Mistral client initialized")
+        
+        # Test with a sample PDF from arXiv (Mistral paper)
+        test_pdf_url = "https://arxiv.org/pdf/2201.04234"
+        
+        logger.info(f"🔍 Testing OCR with PDF: {test_pdf_url}")
+        
+        # Process the PDF with OCR
+        ocr_response = client.ocr.process(
+            model="mistral-ocr-latest",
+            document={
+                "type": "document_url",
+                "document_url": test_pdf_url
+            },
+            include_image_base64=False  # Don't include images for testing
+        )
+        
+        logger.info("✅ OCR processing completed")
+        
+        # Extract the text content
+        logger.info(f"📊 OCR Response structure:")
+        logger.info(f"   - Type: {type(ocr_response)}")
+        logger.info(f"   - Has pages: {hasattr(ocr_response, 'pages')}")
+        logger.info(f"   - Has content: {hasattr(ocr_response, 'content')}")
+        
+        extracted_text = ""
+        
+        if hasattr(ocr_response, 'pages') and ocr_response.pages:
+            logger.info(f"📄 Found {len(ocr_response.pages)} pages")
+            
+            # Extract text from all pages
+            for i, page in enumerate(ocr_response.pages):
+                logger.info(f"📃 Page {i+1} structure: {dir(page)}")
+                
+                if hasattr(page, 'markdown') and page.markdown:
+                    page_text = page.markdown
+                    extracted_text += page_text + "\n\n"
+                    logger.info(f"📝 Page {i+1} text length: {len(page_text)} characters")
+                    logger.info(f"📝 Page {i+1} preview: {page_text[:200]}...")
+                elif hasattr(page, 'content'):
+                    page_text = page.content
+                    extracted_text += page_text + "\n\n"
+                    logger.info(f"📝 Page {i+1} text length: {len(page_text)} characters")
+                    logger.info(f"📝 Page {i+1} preview: {page_text[:200]}...")
+                elif hasattr(page, 'text'):
+                    page_text = page.text
+                    extracted_text += page_text + "\n\n"
+                    logger.info(f"📝 Page {i+1} text length: {len(page_text)} characters")
+                    logger.info(f"📝 Page {i+1} preview: {page_text[:200]}...")
+                else:
+                    logger.info(f"⚠️  Page {i+1} attributes: {[attr for attr in dir(page) if not attr.startswith('_')]}")
+            
+            if extracted_text:
+                logger.info(f"📄 Total extracted text length: {len(extracted_text)} characters")
+                
+                # Test if we can use this text for PII detection
+                if len(extracted_text) > 50:
+                    logger.info("✅ OCR extraction successful - text is suitable for PII detection")
+                    
+                    # Try to import and test PII detection
+                    try:
+                        from mistral_prompting import create_mistral_service
+                        
+                        logger.info("🔍 Testing PII detection on OCR text...")
+                        service = await create_mistral_service()
+                        
+                        # Use a small sample of the text for testing
+                        sample_text = extracted_text[:500]  # First 500 characters
+                        prediction = await service.predict(sample_text)
+                        
+                        logger.info(f"📊 PII detection results:")
+                        logger.info(f"   - Entities found: {len(prediction.entities)}")
+                        logger.info(f"   - Spans detected: {len(prediction.spans)}")
+                        logger.info(f"   - Masked text preview: {prediction.masked_text[:100]}...")
+                        
+                    except Exception as e:
+                        logger.warning(f"⚠️  PII detection test failed: {e}")
+                        logger.info("💡 OCR works, but PII detection needs API key setup")
+                else:
+                    logger.warning("⚠️  Extracted text too short")
+            else:
+                logger.warning("⚠️  No text extracted from pages")
+                
+        elif hasattr(ocr_response, 'content'):
+            extracted_text = ocr_response.content
+            logger.info(f"📄 Extracted text length: {len(extracted_text)} characters")
+            logger.info(f"📝 First 200 characters: {extracted_text[:200]}...")
+        else:
+            logger.error("❌ No content found in OCR response")
+            logger.info(f"Available attributes: {[attr for attr in dir(ocr_response) if not attr.startswith('_')]}")
+        
+    except Exception as e:
+        logger.error(f"❌ OCR test failed: {e}")
+        logger.info("💡 Make sure MISTRAL_API_KEY is set correctly")
+
+if __name__ == "__main__":
+    asyncio.run(test_mistral_ocr())