Complete BERT integration with frontend support

✨ Features:
- Add BERT option in frontend method selection
- Add conditional Mistral model selection (hidden when BERT selected)
- Complete BERT integration in FastAPI app
- Update health check to include BERT service
- Update API info with BERT support

🔧 Technical:
- BERT service initialization in lifespan
- BERT prediction route support
- Frontend JavaScript for method switching
- Updated request/response models

app.py

@@ -44,7 +44,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
+    global mistral_base_service, mistral_finetuned_service, bert_service
     
     # Startup
     logger.info("🚀 Starting PII Masking Demo application...")
@@ -61,7 +61,17 @@ async def lifespan(app: FastAPI):
         logger.info("✅ Fine-tuned Mistral service initialized successfully")
         
     except Exception as e:
-        logger.error(f"Failed to initialize services: {e}")
+        logger.error(f"Failed to initialize Mistral services: {e}")
+        # Don't raise exception - let app start but handle gracefully in endpoints
+    
+    try:
+        # Initialize BERT service
+        logger.info("Initializing BERT service...")
+        bert_service = await create_bert_service(model_path=BERT_MODEL_PATH)
+        logger.info("✅ BERT service initialized successfully")
+        
+    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
         
     yield
@@ -80,8 +90,8 @@ app = FastAPI(
 # Request/Response models
 class PredictionRequest(BaseModel):
     text: str = Field(..., description="Text to analyze for PII", min_length=1, max_length=5000)
-    method: str = Field(default="mistral", description="Method to use (currently only 'mistral')")
-    model: str = Field(default="base", description="Model to use: 'base' for mistral-large-latest or 'finetuned' for fine-tuned model")
+    method: str = Field(default="mistral", description="Method to use: 'mistral' or 'bert'")
+    model: str = Field(default="base", description="Model to use: 'base' for mistral-large-latest or 'finetuned' for fine-tuned model (ignored for BERT)")
 
 class PredictionResponse(BaseModel):
     masked_text: str = Field(description="Text with PII entities masked")
@@ -200,26 +210,41 @@ async def predict(request: PredictionRequest):
     """
     Predict PII entities and return masked text.
     
-    Supports both base and fine-tuned Mistral models.
+    Supports Mistral models (base and fine-tuned) and BERT.
     """
     # Validate method
-    if request.method != "mistral":
+    if request.method not in ["mistral", "bert"]:
         raise HTTPException(
             status_code=400, 
-            detail=f"Method '{request.method}' not supported. Currently only 'mistral' is available."
+            detail=f"Method '{request.method}' not supported. Use 'mistral' or 'bert'."
         )
     
-    # Get the appropriate service
-    service = get_mistral_service(request.model)
-    
     start_time = time.time()
     
     try:
-        model_type = "Fine-tuned" if request.model == "finetuned" else "Base"
-        logger.info(f"🔍 Processing text with {model_type} Mistral model: {request.text[:100]}...")
-        
-        # Call Mistral service
-        prediction = await service.predict(request.text)
+        if request.method == "mistral":
+            # Get the appropriate Mistral service
+            service = get_mistral_service(request.model)
+            model_type = "Fine-tuned" if request.model == "finetuned" else "Base"
+            logger.info(f"🔍 Processing text with {model_type} Mistral model: {request.text[:100]}...")
+            
+            # Call Mistral service
+            prediction = await service.predict(request.text)
+            method_used = f"{request.method}-{request.model}"
+            
+        elif request.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."
+                )
+            
+            logger.info(f"🔍 Processing text with BERT model: {request.text[:100]}...")
+            
+            # Call BERT service
+            prediction = await bert_service.predict(request.text)
+            method_used = "bert"
         
         processing_time = time.time() - start_time
         
@@ -232,7 +257,7 @@ async def predict(request: PredictionRequest):
             masked_text=prediction.masked_text,
             entities=prediction.entities,
             processing_time=processing_time,
-            method_used=f"{request.method}-{request.model}",
+            method_used=method_used,
             num_entities=num_entities
         )
         
@@ -246,7 +271,7 @@ async def predict(request: PredictionRequest):
 @app.get("/health", response_model=HealthResponse)
 async def health_check():
     """Health check endpoint."""
-    global mistral_base_service, mistral_finetuned_service
+    global mistral_base_service, mistral_finetuned_service, bert_service
     
     services_status = {
         "mistral_base": {
@@ -260,16 +285,25 @@ async def health_check():
             "initialized": mistral_finetuned_service.is_initialized if mistral_finetuned_service else False,
             "model": MODELS["finetuned"],
             "info": mistral_finetuned_service.get_service_info() if mistral_finetuned_service else None
+        },
+        "bert": {
+            "available": bert_service is not None,
+            "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
         }
     }
     
     # Overall status
     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
+    
+    healthy_services = sum([base_healthy, finetuned_healthy, bert_healthy])
     
-    if base_healthy and finetuned_healthy:
+    if healthy_services == 3:
         overall_status = "healthy"
-    elif base_healthy or finetuned_healthy:
+    elif healthy_services >= 1:
         overall_status = "partial"
     else:
         overall_status = "degraded"
@@ -287,7 +321,7 @@ async def api_info():
         "name": "PII Masking Demo API",
         "version": "1.0.0",
         "description": "Personal Identifiable Information masking using Mistral AI",
-        "available_methods": ["mistral"],
+        "available_methods": ["mistral", "bert"],
         "available_models": {
             "base": {
                 "name": MODELS["base"],
@@ -296,6 +330,10 @@ async def api_info():
             "finetuned": {
                 "name": MODELS["finetuned"],
                 "description": "Fine-tuned Mistral model specialized for PII detection"
+            },
+            "bert": {
+                "name": BERT_MODEL_PATH,
+                "description": "BERT token classification model for fast PII detection"
             }
         },
         "endpoints": {

static/index.html

@@ -280,16 +280,23 @@ Example: Hi, my name is John Smith and my email is [email protected]. Call
                     <label>Select masking method:</label>
                     <div class="method-selection">
                         <div class="method-option">
-                            <input type="radio" id="mistral" name="method" value="mistral" class="method-radio" checked>
+                            <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-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-desc">Fast local processing</div>
+                            </label>
+                        </div>
                     </div>
                 </div>
 
-                <div class="form-group">
+                <div class="form-group" id="mistralModelSelection">
                     <label>Select Mistral model:</label>
                     <div class="method-selection">
                         <div class="method-option">
@@ -435,6 +442,17 @@ Example: Hi, my name is John Smith and my email is [email protected]. Call
         function hideError() {
             document.getElementById('error').style.display = 'none';
         }
+        
+        function toggleMistralModelSelection() {
+            const mistralSelected = document.getElementById('mistral').checked;
+            const mistralModelSelection = document.getElementById('mistralModelSelection');
+            
+            if (mistralSelected) {
+                mistralModelSelection.style.display = 'block';
+            } else {
+                mistralModelSelection.style.display = 'none';
+            }
+        }
     </script>
 </body>
 </html>