feat: implement context-aware pronoun resolution in translation process

backend/app/models/mt_model.py

@@ -1,7 +1,9 @@
+import re
 import threading
 import time
 from transformers import MarianMTModel, MarianTokenizer
 
+# ── Language pair → HuggingFace model name ───────────────────────────────────
 MODEL_MAP: dict[tuple[str, str], str] = {
     ("en", "fr"): "Helsinki-NLP/opus-mt-en-fr",
     ("fr", "en"): "Helsinki-NLP/opus-mt-fr-en",
@@ -27,6 +29,38 @@ PIVOT_PAIRS: frozenset[tuple[str, str]] = frozenset({
 
 CONTEXT_WINDOW = 3
 
+_PRONOUN_SUBJECT = re.compile(
+    r'\b(he|she|they|it|him|her|them|his|hers|their|its)\b',
+    re.IGNORECASE
+)
+
+_NAME_PATTERN = re.compile(r'\b[A-Z][a-z]{2,}\b')
+
+
+def _resolve_pronouns(text: str, context: list[str]) -> str:
+    if not _PRONOUN_SUBJECT.search(text):
+        return text  
+
+    candidate = None
+    for utt in reversed(context):
+        names = _NAME_PATTERN.findall(utt)
+        names = [n for n in names if n not in (
+            "The", "A", "An", "This", "That", "These", "Those",
+            "I", "You", "We", "They", "He", "She", "It"
+        )]
+        if names:
+            candidate = names[-1]
+            break
+
+    if not candidate:
+        return text 
+
+    resolved = _PRONOUN_SUBJECT.sub(candidate, text)
+    if resolved != text:
+        print(f"[MT] Pronoun resolved: '{text}' → '{resolved}' "
+              f"(candidate: {candidate})")
+    return resolved
+
 
 class HelsinkiTranslator:
     def __init__(self):
@@ -47,16 +81,14 @@ class HelsinkiTranslator:
     def _load_pair(self, src: str, tgt: str) -> None:
         pair = (src, tgt)
         lock = self._get_load_lock(pair)
-
         with lock:
             if pair in self._models:
                 return
-
             model_name = MODEL_MAP.get(pair)
             if model_name is None:
-                raise ValueError(f"[MT] No direct model for {src}→{tgt}. "
-                                 f"Use ensure_pair_loaded for pivot pairs.")
-
+                raise ValueError(
+                    f"[MT] No direct model for {src}→{tgt}."
+                )
             print(f"[MT] Loading {src}→{tgt} ({model_name})...")
             t0 = time.time()
             tokenizer = MarianTokenizer.from_pretrained(model_name)
@@ -71,7 +103,6 @@ class HelsinkiTranslator:
         pair = (src, tgt)
         tokenizer = self._tokenizers[pair]
         model = self._models[pair]
-
         inputs = tokenizer(
             text,
             return_tensors="pt",
@@ -87,15 +118,13 @@ class HelsinkiTranslator:
 
     def ensure_pair_loaded(self, src: str, tgt: str) -> None:
         pair = (src, tgt)
-
         if pair in MODEL_MAP:
             if pair not in self._models:
                 self._load_pair(src, tgt)
         elif pair in PIVOT_PAIRS:
-            # Need src→en and en→tgt
-            if (src, "en") in MODEL_MAP and (src, "en") not in self._models:
+            if (src, "en") not in self._models:
                 self._load_pair(src, "en")
-            if ("en", tgt) in MODEL_MAP and ("en", tgt) not in self._models:
+            if ("en", tgt) not in self._models:
                 self._load_pair("en", tgt)
         else:
             print(f"[MT] ⚠️  No route defined for {src}→{tgt}")
@@ -111,60 +140,33 @@ class HelsinkiTranslator:
         if src == tgt:
             return text
 
-        pair = (src, tgt)
-
-        # ── Context-aware translation ─────────────────────────────────────────
-
-        if use_context and context:
+        # Apply pronoun resolution heuristic (English source only for now)
+        input_text = text
+        if use_context and context and src == "en":
             recent = context[-CONTEXT_WINDOW:]
-            context_str = " ".join(recent) 
-            full_str    = context_str + " " + text 
-
-            if pair in MODEL_MAP:
-                if pair not in self._models:
-                    self._load_pair(src, tgt)
-                context_translated = self._translate_direct(context_str, src, tgt)
-                full_translated    = self._translate_direct(full_str,    src, tgt)
-
-            elif pair in PIVOT_PAIRS:
-                en_pair  = (src, "en")
-                tgt_pair = ("en", tgt)
-                if en_pair  not in self._models: self._load_pair(src, "en")
-                if tgt_pair not in self._models: self._load_pair("en", tgt)
-
-                ctx_en   = self._translate_direct(context_str, src, "en")
-                full_en  = self._translate_direct(full_str,    src, "en")
-                context_translated = self._translate_direct(ctx_en,  "en", tgt)
-                full_translated    = self._translate_direct(full_en, "en", tgt)
-
-            else:
-                print(f"[MT] ⚠️  No route for {src}→{tgt}, returning original")
-                return text
-
-            prefix = context_translated.strip()
-            result = full_translated.strip()
+            input_text = _resolve_pronouns(text, recent)
 
-            if prefix and result.startswith(prefix):
-                result = result[len(prefix):].lstrip(" ,.;")
-
-            return result.strip() if result.strip() else full_translated.strip()
+        pair = (src, tgt)
 
-        # ── Sentence-level translation ────────────────────────────────────────────
-        else:
-            input_text = text
-
-            if pair in MODEL_MAP:
-                if pair not in self._models:
-                    self._load_pair(src, tgt)
-                return self._translate_direct(input_text, src, tgt)
-
-            if pair in PIVOT_PAIRS:
-                en_pair  = (src, "en")
-                tgt_pair = ("en", tgt)
-                if en_pair  not in self._models: self._load_pair(src, "en")
-                if tgt_pair not in self._models: self._load_pair("en", tgt)
-                en_text = self._translate_direct(input_text, src, "en")
-                return self._translate_direct(en_text, "en", tgt)
-
-            print(f"[MT] ⚠️  No route for {src}→{tgt}, returning original")
-            return text
+        # ── Direct translation ────────────────────────────────────────────────
+        if pair in MODEL_MAP:
+            if pair not in self._models:
+                self._load_pair(src, tgt)
+            result = self._translate_direct(input_text, src, tgt)
+            print(f"[MT] context={'on' if use_context and context else 'off'} "
+                  f"input='{input_text[:60]}' output='{result[:60]}'")
+            return result
+
+        # ── Pivot through English ─────────────────────────────────────────────
+        if pair in PIVOT_PAIRS:
+            en_pair  = (src, "en")
+            tgt_pair = ("en", tgt)
+            if en_pair  not in self._models: self._load_pair(src, "en")
+            if tgt_pair not in self._models: self._load_pair("en", tgt)
+            en_text = self._translate_direct(input_text, src, "en")
+            result  = self._translate_direct(en_text, "en", tgt)
+            print(f"[MT] pivot {src}→en→{tgt}: '{result[:60]}'")
+            return result
+
+        print(f"[MT] ⚠️  No route for {src}→{tgt}, returning original")
+        return text