fix: Switch grammar to greedy decoding + stronger anti-hallucination - num_beams=1 (greedy, most stable for seq2seq correction) - repetition_penalty=2.0 (strong anti-repeat) - Log raw model output before hallucination check - Tightened bigram duplicate detection

src/nlp/grammar/grammar_service.py

@@ -17,7 +17,6 @@ MODEL_ID = "bayan10/Bayan_Arabic_Grammar"
 
 # Generation parameters
 MAX_INPUT_LENGTH = 512
-NUM_BEAMS = 2  # Lower beams = less hallucination on CPU
 
 
 def get_grammar_model():
@@ -51,7 +50,6 @@ def get_grammar_model():
 
     except Exception as e:
         logger.error(f"Failed to load grammar model: {e}")
-        # Try without local_files_only as fallback
         try:
             logger.info("Retrying without local_files_only...")
             _grammar_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -68,52 +66,44 @@ def get_grammar_model():
     return _grammar_model, _grammar_tokenizer, _device
 
 
-def _word_set(text: str) -> set:
-    """Extract set of Arabic words from text."""
-    return set(re.findall(r'[\u0600-\u06FF]+', text))
-
-
 def _is_hallucinated(original: str, corrected: str) -> bool:
     """
     Check if the grammar output is hallucinated / garbled.
     Returns True if output should be discarded.
     """
-    # 1. Output too long (grammar should NOT add lots of text)
+    orig_words = re.findall(r'[\u0600-\u06FF]+', original)
+    corr_words = re.findall(r'[\u0600-\u06FF]+', corrected)
+
+    # 1. Output too long (grammar corrections are minor edits)
     if len(corrected) > len(original) * 1.5:
-        logger.warning(f"Grammar hallucination: output too long ({len(corrected)} vs {len(original)})")
+        logger.warning(f"Hallucination: output too long ({len(corrected)} vs {len(original)} chars)")
         return True
 
     # 2. Output too short
     if len(corrected) < len(original) * 0.5:
-        logger.warning(f"Grammar hallucination: output too short ({len(corrected)} vs {len(original)})")
+        logger.warning(f"Hallucination: output too short ({len(corrected)} vs {len(original)} chars)")
         return True
 
-    # 3. Check word order preservation — grammar should NOT rearrange sentences
-    orig_words = re.findall(r'[\u0600-\u06FF]+', original)
-    corr_words = re.findall(r'[\u0600-\u06FF]+', corrected)
-
-    if len(corr_words) > len(orig_words) * 1.5:
-        logger.warning(f"Grammar hallucination: too many words ({len(corr_words)} vs {len(orig_words)})")
+    # 3. Too many words added
+    if len(corr_words) > len(orig_words) * 1.4:
+        logger.warning(f"Hallucination: too many words ({len(corr_words)} vs {len(orig_words)})")
         return True
 
-    # 4. Check for duplicated segments (repeated phrases)
-    if len(corr_words) > 3:
-        # Check if any 3-word sequence appears more than once
-        trigrams = [' '.join(corr_words[i:i+3]) for i in range(len(corr_words)-2)]
-        if len(trigrams) != len(set(trigrams)):
-            logger.warning("Grammar hallucination: repeated phrases detected")
+    # 4. Check for duplicated 2-word phrases (bigrams)
+    if len(corr_words) > 4:
+        bigrams = [f"{corr_words[i]} {corr_words[i+1]}" for i in range(len(corr_words)-1)]
+        if len(bigrams) != len(set(bigrams)):
+            logger.warning("Hallucination: repeated bigrams detected")
             return True
 
-    # 5. First word of output should be close to first word of input
-    #    (grammar shouldn't rearrange sentence start)
+    # 5. Sentence start rearranged (grammar shouldn't move first word to end)
     if orig_words and corr_words:
-        if orig_words[0] != corr_words[0]:
-            # Check if the first few words are completely rearranged
-            orig_start = set(orig_words[:3])
-            corr_start = set(corr_words[:3])
-            if not orig_start.intersection(corr_start):
-                logger.warning("Grammar hallucination: sentence completely rearranged")
-                return True
+        # Check if first 2 original words appear in first 4 corrected words
+        orig_head = set(orig_words[:2])
+        corr_head = set(corr_words[:4])
+        if len(orig_head) > 0 and not orig_head.intersection(corr_head):
+            logger.warning("Hallucination: sentence start completely rearranged")
+            return True
 
     return False
 
@@ -139,30 +129,31 @@ def correct_grammar(text: str) -> str:
         )
         inputs = {k: v.to(device) for k, v in inputs.items()}
 
-        # Calculate dynamic max_length: slightly more than input to allow small expansions
+        # Dynamic max_length: close to input to prevent hallucination
         input_len = inputs['input_ids'].shape[1]
-        dynamic_max_len = min(int(input_len * 1.3) + 10, MAX_INPUT_LENGTH)
+        dynamic_max_len = min(int(input_len * 1.3) + 5, MAX_INPUT_LENGTH)
 
-        # Generate correction with conservative params to avoid hallucination
+        # Strategy: greedy decoding (most stable for grammar correction)
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
                 max_length=dynamic_max_len,
-                num_beams=NUM_BEAMS,
-                early_stopping=True,
-                no_repeat_ngram_size=3,
-                length_penalty=1.0,
-                repetition_penalty=1.5,
+                num_beams=1,          # Greedy decoding — most stable
+                do_sample=False,
+                repetition_penalty=2.0,  # Strong penalty against repeating
             )
 
         # Decode output
-        corrected = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        raw_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        logger.info(f"[GRAMMAR-RAW] input({len(text)}): '{text}'")
+        logger.info(f"[GRAMMAR-RAW] output({len(raw_output)}): '{raw_output}'")
 
         # Safety: if model returns empty, keep original
-        if not corrected or not corrected.strip():
+        if not raw_output or not raw_output.strip():
+            logger.warning("Grammar returned empty output")
             return text
 
-        corrected = corrected.strip()
+        corrected = raw_output.strip()
 
         # Check for hallucination
         if _is_hallucinated(text, corrected):