Update app.py

app.py

@@ -1,3 +1,19 @@
+"""
+Khmer Legal Bridge - Translation API
+=====================================
+
+Flask application with COMETKiwi-based confidence scoring.
+
+Features:
+- Bidirectional EN↔KM translation using fine-tuned NLLB-200
+- Scientific confidence scoring with COMETKiwi
+- PDF text extraction
+- Privacy-first design (zero retention)
+
+Author: Khmer Legal Bridge Project
+License: MIT
+"""
+
 from flask import Flask, render_template, request, jsonify
 from transformers import AutoModelForSeq2SeqLM, NllbTokenizerFast
 import torch
@@ -5,10 +21,30 @@ import fitz
 import re
 import unicodedata
 import time
+import logging
+import os
 from sacremoses import MosesPunctNormalizer
 
+# Import confidence scoring module
+from confidence_scoring_v2 import (
+    TransparencyScorer,
+    DEFAULT_LEGAL_GLOSSARY,
+    ConfidenceResult
+)
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
 app = Flask(__name__)
 
+# ============================================================================
+# Text Preprocessing
+# ============================================================================
+
 mpn = MosesPunctNormalizer(lang="en")
 mpn.substitutions = [(re.compile(r), sub) for r, sub in mpn.substitutions]
 
@@ -23,29 +59,90 @@ def get_non_printing_char_replacer(replace_by: str = " "):
 replace_nonprint = get_non_printing_char_replacer(" ")
 
 def preprocess_text(text: str) -> str:
+    """Clean and normalize text for translation."""
     clean = mpn.normalize(text)
     clean = replace_nonprint(clean)
     clean = unicodedata.normalize("NFKC", clean)
     return clean
 
-print("Loading model...")
+
+# ============================================================================
+# Model Loading
+# ============================================================================
+
+logger.info("Loading translation model...")
 MODEL_ID = "ClaudBarbara/Open_Access_Khmer"
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
 tokenizer = NllbTokenizerFast.from_pretrained(MODEL_ID)
-print("Model loaded!")
+logger.info("Translation model loaded!")
+
+# Configuration
+USE_COMET = os.environ.get("USE_COMET", "true").lower() == "true"
+USE_DETAILED_SCORING = os.environ.get("DETAILED_SCORING", "true").lower() == "true"
+
+# Initialize confidence scorer (lazy loading for COMETKiwi)
+confidence_scorer = None
+
+def get_confidence_scorer():
+    """Lazy initialization of confidence scorer."""
+    global confidence_scorer
+    if confidence_scorer is None:
+        logger.info(f"Initializing confidence scorer (COMETKiwi: {USE_COMET})")
+        confidence_scorer = TransparencyScorer(
+            translator_func=translate_simple,
+            glossary=DEFAULT_LEGAL_GLOSSARY,
+            use_comet=USE_COMET,
+            use_back_translation=True,
+            use_terminology=True
+        )
+    return confidence_scorer
+
 
-CONFIDENCE_THRESHOLD = 70  # Below this = human review needed
+# ============================================================================
+# Translation Functions
+# ============================================================================
 
-def segment_text(text, src_lang):
+def segment_text(text: str, src_lang: str) -> list:
+    """Segment text into sentences for batch processing."""
     if src_lang == "khm_Khmr":
+        # Khmer sentence boundaries
         sentences = re.split(r'(?<=[។៖])\s*', text)
     else:
+        # English sentence boundaries
         sentences = re.split(r'(?<=[.!?])\s+', text)
     return [s.strip() for s in sentences if s.strip()]
 
-def translate_batch(texts, src_lang, tgt_lang):
+
+def translate_simple(text: str, src_lang: str, tgt_lang: str) -> str:
+    """
+    Simple translation without confidence scoring.
+    Used for back-translation verification.
+    """
+    tokenizer.src_lang = src_lang
+    inputs = tokenizer(
+        text,
+        return_tensors='pt',
+        padding=True,
+        truncation=True,
+        max_length=512
+    )
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_lang),
+            max_new_tokens=int(32 + 3 * inputs.input_ids.shape[1]),
+            num_beams=4,
+            early_stopping=True
+        )
+    
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+def translate_batch(texts: list, src_lang: str, tgt_lang: str) -> list:
+    """Translate a batch of texts efficiently."""
     if not texts:
-        return [], []
+        return []
     
     tokenizer.src_lang = src_lang
     inputs = tokenizer(
@@ -62,56 +159,151 @@ def translate_batch(texts, src_lang, tgt_lang):
             forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_lang),
             max_new_tokens=int(32 + 3 * inputs.input_ids.shape[1]),
             num_beams=4,
-            early_stopping=True,
-            return_dict_in_generate=True,
-            output_scores=True
+            early_stopping=True
         )
     
-    translations = tokenizer.batch_decode(outputs.sequences, skip_special_tokens=True)
+    return tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+
+def translate_long(
+    text: str, 
+    src_lang: str, 
+    tgt_lang: str, 
+    batch_size: int = 8,
+    compute_confidence: bool = True
+) -> tuple:
+    """
+    Translate long text with sentence segmentation and confidence scoring.
     
-    # Extract confidence scores
-    if hasattr(outputs, 'sequences_scores') and outputs.sequences_scores is not None:
-        scores = torch.sigmoid(outputs.sequences_scores).tolist()
-    else:
-        scores = [0.85] * len(texts)
+    Args:
+        text: Input text
+        src_lang: Source language code
+        tgt_lang: Target language code
+        batch_size: Batch size for processing
+        compute_confidence: Whether to compute detailed confidence
     
-    return translations, scores
-
-def translate_long(text, src_lang, tgt_lang, batch_size=8):
+    Returns:
+        Tuple of (translation, metrics_dict)
+    """
     start_time = time.time()
     
+    # Preprocess
     clean_text = preprocess_text(text)
     sentences = segment_text(clean_text, src_lang)
     
     if not sentences:
-        return "", {}
+        return "", {"error": "No text to translate"}
     
+    # Translate in batches
     translated_parts = []
-    all_scores = []
-    
     for i in range(0, len(sentences), batch_size):
         batch = sentences[i:i + batch_size]
-        translations, scores = translate_batch(batch, src_lang, tgt_lang)
+        translations = translate_batch(batch, src_lang, tgt_lang)
         translated_parts.extend(translations)
-        all_scores.extend(scores)
     
     result = " ".join(translated_parts)
     elapsed = time.time() - start_time
     
-    avg_confidence = (sum(all_scores) / len(all_scores) * 100) if all_scores else 0
-    min_confidence = (min(all_scores) * 100) if all_scores else 0
+    # Compute confidence score
+    direction = "en2km" if src_lang == "eng_Latn" else "km2en"
     
-    metrics = {
-        "confidence": round(avg_confidence, 1),
-        "min_confidence": round(min_confidence, 1),
-        "needs_review": avg_confidence < CONFIDENCE_THRESHOLD,
-        "time_seconds": round(elapsed, 2),
-        "sentences": len(sentences)
-    }
+    if compute_confidence and USE_COMET:
+        try:
+            scorer = get_confidence_scorer()
+            
+            # For long texts, sample representative sentences for scoring
+            if len(sentences) > 5:
+                # Score first, middle, and last sentences
+                sample_indices = [0, len(sentences)//2, -1]
+                sample_scores = []
+                
+                for idx in sample_indices:
+                    src_sent = sentences[idx]
+                    tgt_sent = translated_parts[idx]
+                    
+                    conf_result = scorer.score(
+                        src_sent, tgt_sent, direction,
+                        detailed=USE_DETAILED_SCORING
+                    )
+                    sample_scores.append(conf_result.overall_score)
+                
+                avg_score = sum(sample_scores) / len(sample_scores)
+                min_score = min(sample_scores)
+                
+                # Use most conservative estimate
+                confidence_score = min(avg_score, min_score + 0.1)
+                
+            else:
+                # Score entire translation
+                conf_result = scorer.score(
+                    clean_text, result, direction,
+                    detailed=USE_DETAILED_SCORING
+                )
+                confidence_score = conf_result.overall_score
+            
+            # Determine review recommendation
+            needs_review = confidence_score < 0.75
+            quality_level = (
+                "excellent" if confidence_score >= 0.85 else
+                "good" if confidence_score >= 0.70 else
+                "acceptable" if confidence_score >= 0.55 else
+                "low" if confidence_score >= 0.40 else
+                "very_low"
+            )
+            
+            metrics = {
+                "confidence": round(confidence_score * 100, 1),
+                "quality_level": quality_level,
+                "needs_review": needs_review,
+                "time_seconds": round(elapsed, 2),
+                "sentences": len(sentences),
+                "method": "comet_kiwi"
+            }
+            
+        except Exception as e:
+            logger.error(f"Confidence scoring failed: {e}")
+            # Fallback to lightweight scoring
+            metrics = compute_lightweight_metrics(
+                clean_text, result, direction, elapsed, len(sentences)
+            )
+    else:
+        # Use lightweight scoring
+        metrics = compute_lightweight_metrics(
+            clean_text, result, direction, elapsed, len(sentences)
+        )
     
     return result, metrics
 
-def extract_pdf_text(pdf_file):
+
+def compute_lightweight_metrics(
+    source: str, 
+    translation: str, 
+    direction: str,
+    elapsed: float,
+    num_sentences: int
+) -> dict:
+    """
+    Compute lightweight confidence metrics without COMETKiwi.
+    """
+    scorer = get_confidence_scorer()
+    conf_result = scorer.score_fast(source, translation, direction)
+    
+    return {
+        "confidence": round(conf_result.overall_score * 100, 1),
+        "quality_level": conf_result.quality_level,
+        "needs_review": conf_result.human_review_recommended,
+        "time_seconds": round(elapsed, 2),
+        "sentences": num_sentences,
+        "method": "lightweight"
+    }
+
+
+# ============================================================================
+# PDF Extraction
+# ============================================================================
+
+def extract_pdf_text(pdf_file) -> str:
+    """Extract text from uploaded PDF file."""
     try:
         pdf_bytes = pdf_file.read()
         doc = fitz.open(stream=pdf_bytes, filetype="pdf")
@@ -120,15 +312,35 @@ def extract_pdf_text(pdf_file):
             text += page.get_text()
         doc.close()
         return text.strip()
-    except:
+    except Exception as e:
+        logger.error(f"PDF extraction failed: {e}")
         return None
 
+
+# ============================================================================
+# API Routes
+# ============================================================================
+
 @app.route("/")
 def index():
+    """Serve the main translation interface."""
     return render_template("index.html")
 
+
 @app.route("/translate", methods=["POST"])
 def translate_endpoint():
+    """
+    Translation API endpoint.
+    
+    Request JSON:
+        - text: str - Text to translate
+        - direction: str - "en-km" or "km-en"
+    
+    Response JSON:
+        - success: bool
+        - translation: str
+        - metrics: dict with confidence scores
+    """
     data = request.json
     text = data.get("text", "")
     direction = data.get("direction", "en-km")
@@ -140,12 +352,27 @@ def translate_endpoint():
     
     try:
         result, metrics = translate_long(text, src_lang, tgt_lang)
-        return jsonify({"success": True, "translation": result, "metrics": metrics})
+        return jsonify({
+            "success": True, 
+            "translation": result, 
+            "metrics": metrics
+        })
     except Exception as e:
-        return jsonify({"success": False, "error": str(e)})
+        logger.error(f"Translation failed: {e}")
+        return jsonify({
+            "success": False, 
+            "error": str(e)
+        })
+
 
 @app.route("/upload-pdf", methods=["POST"])
 def upload_pdf():
+    """
+    PDF upload endpoint.
+    
+    Accepts multipart form with 'file' field.
+    Returns extracted text.
+    """
     if 'file' not in request.files:
         return jsonify({"success": False, "error": "No file uploaded"})
     
@@ -162,5 +389,21 @@ def upload_pdf():
     else:
         return jsonify({"success": False, "error": "Could not extract text"})
 
+
+@app.route("/health", methods=["GET"])
+def health_check():
+    """Health check endpoint for monitoring."""
+    return jsonify({
+        "status": "healthy",
+        "model": MODEL_ID,
+        "comet_enabled": USE_COMET
+    })
+
+
+# ============================================================================
+# Main Entry Point
+# ============================================================================
+
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
+    port = int(os.environ.get("PORT", 7860))
+    app.run(host="0.0.0.0", port=port)