Update app.py

app.py

@@ -9,7 +9,6 @@ from sacremoses import MosesPunctNormalizer
 
 app = Flask(__name__)
 
-# Preprocessing NLLB
 mpn = MosesPunctNormalizer(lang="en")
 mpn.substitutions = [(re.compile(r), sub) for r, sub in mpn.substitutions]
 
@@ -29,15 +28,15 @@ def preprocess_text(text: str) -> str:
     clean = unicodedata.normalize("NFKC", clean)
     return clean
 
-# Load model
 print("Loading model...")
 MODEL_ID = "ClaudBarbara/Open_Access_Khmer"
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
 tokenizer = NllbTokenizerFast.from_pretrained(MODEL_ID)
 print("Model loaded!")
 
+CONFIDENCE_THRESHOLD = 70  # Below this = human review needed
+
 def segment_text(text, src_lang):
-    """Segment text into sentences"""
     if src_lang == "khm_Khmr":
         sentences = re.split(r'(?<=[។៖])\s*', text)
     else:
@@ -45,9 +44,8 @@ def segment_text(text, src_lang):
     return [s.strip() for s in sentences if s.strip()]
 
 def translate_batch(texts, src_lang, tgt_lang):
-    """Translate a batch of sentences"""
     if not texts:
-        return []
+        return [], []
     
     tokenizer.src_lang = src_lang
     inputs = tokenizer(
@@ -64,13 +62,22 @@ 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
+            early_stopping=True,
+            return_dict_in_generate=True,
+            output_scores=True
         )
     
-    return tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    translations = tokenizer.batch_decode(outputs.sequences, skip_special_tokens=True)
+    
+    # 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)
+    
+    return translations, scores
 
 def translate_long(text, src_lang, tgt_lang, batch_size=8):
-    """Translate long document in batches"""
     start_time = time.time()
     
     clean_text = preprocess_text(text)
@@ -80,20 +87,26 @@ def translate_long(text, src_lang, tgt_lang, batch_size=8):
         return "", {}
     
     translated_parts = []
+    all_scores = []
     
     for i in range(0, len(sentences), batch_size):
         batch = sentences[i:i + batch_size]
-        translations = translate_batch(batch, src_lang, tgt_lang)
+        translations, scores = 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
+    
     metrics = {
-        "sentences": len(sentences),
-        "source_chars": len(text),
-        "target_chars": len(result),
-        "time_seconds": round(elapsed, 2)
+        "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)
     }
     
     return result, metrics