feat: NLP-2 Grammar model integration (Bayan_Arabic_Grammar T5) - New: src/nlp/grammar/grammar_service.py (lazy T5 loading) - Dockerfile: pre-download grammar model (2.14GB safetensors) - Wire /api/analyze Step 2 to grammar service - Update /api/grammar standalone endpoint - Add _grammar_available() health check

Dockerfile

@@ -45,6 +45,16 @@ AutoModelForMaskedLM.from_pretrained('aubmindlab/bert-base-arabertv02'); \
 print('Spelling model + MLM cached!'); \
 "
 
+# 3. Grammar model (T5ForConditionalGeneration)
+RUN python -c "\
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM; \
+print('Downloading Grammar tokenizer...'); \
+AutoTokenizer.from_pretrained('bayan10/Bayan_Arabic_Grammar'); \
+print('Downloading Grammar T5 model...'); \
+AutoModelForSeq2SeqLM.from_pretrained('bayan10/Bayan_Arabic_Grammar'); \
+print('Grammar model cached!'); \
+"
+
 # Copy application code
 COPY src/ ./src/
 COPY .env* ./

src/app.py

@@ -155,7 +155,7 @@ def health_check():
                 'summarization': summarization_model is not None,
                 'spelling': _spelling_available(),
                 'autocomplete': False,
-                'grammar': False,
+                'grammar': _grammar_available(),
                 'punctuation': False
             },
             'note': 'Free tier: summarization local, other models return input unchanged',
@@ -232,6 +232,15 @@ def _spelling_available():
         return False
 
 
+def _grammar_available():
+    """Check if grammar model is loaded (without triggering lazy load)."""
+    try:
+        from nlp.grammar.grammar_service import _grammar_model
+        return _grammar_model is not None
+    except Exception:
+        return False
+
+
 @app.route('/api/spelling', methods=['POST'])
 def spelling_correction():
     """
@@ -443,12 +452,6 @@ def grammar_correction():
         "text": "Arabic text to correct"
     }
     """
-    if not USE_HF_API and grammar_model is None:
-        return jsonify({
-            'error': 'Grammar model not loaded. Please check server logs.',
-            'status': 'error'
-        }), 503
-    
     try:
         if not request.is_json:
             return jsonify({'error': 'Request must be JSON', 'status': 'error'}), 400
@@ -460,17 +463,13 @@ def grammar_correction():
             return jsonify({'error': 'Text is required', 'status': 'error'}), 400
         
         logger.info(f"Correcting grammar for text of length: {len(text)}")
-        if USE_HF_API:
-            # Grammar uses spelling model as proxy (no dedicated grammar model yet)
-            corrected = hf_correct_spelling(text)
-        else:
-            corrected = grammar_model.correct(text)
+        from nlp.grammar.grammar_service import correct_grammar
+        corrected = correct_grammar(text)
         
         return jsonify({
-            'corrected': corrected,
-            'status': 'success',
-            'original_length': len(text),
-            'corrected_length': len(corrected)
+            'original_text': text,
+            'corrected_text': corrected,
+            'status': 'success'
         })
     
     except Exception as e:
@@ -850,15 +849,13 @@ def analyze_text():
                 logger.error(f"[ANALYZE] Spelling failed: {e}")
 
         # 2. Grammar (runs on spelling-corrected text)
-        has_grammar = USE_HF_API or grammar_model
+        has_grammar = True  # Always available via lazy-loaded grammar_service
         if has_grammar:
             try:
                 t0 = time.time()
                 logger.info(f"[ANALYZE] Step 2: Grammar correction starting...")
-                if USE_HF_API:
-                    corrected_grammar = hf_correct_spelling(current_text)
-                else:
-                    corrected_grammar = grammar_model.correct(current_text)
+                from nlp.grammar.grammar_service import correct_grammar
+                corrected_grammar = correct_grammar(current_text)
                 logger.info(f"[ANALYZE] Step 2: Grammar done in {time.time()-t0:.2f}s")
                 if corrected_grammar != current_text:
                     diffs = get_word_diffs(current_text, corrected_grammar)

src/nlp/grammar/__init__.py

@@ -0,0 +1 @@
+# Grammar NLP package

src/nlp/grammar/grammar_service.py

@@ -0,0 +1,118 @@
+"""
+Grammar correction service using Bayan_Arabic_Grammar (T5).
+Lazy-loads the model on first request, keeps it resident in memory.
+"""
+import logging
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+logger = logging.getLogger(__name__)
+
+_grammar_model = None
+_grammar_tokenizer = None
+_device = None
+
+MODEL_ID = "bayan10/Bayan_Arabic_Grammar"
+
+# Generation parameters
+MAX_INPUT_LENGTH = 512
+MAX_OUTPUT_LENGTH = 512
+NUM_BEAMS = 4
+
+
+def get_grammar_model():
+    """
+    Lazy-load the grammar model. Returns (model, tokenizer, device).
+    Loads once, stays in memory forever.
+    """
+    global _grammar_model, _grammar_tokenizer, _device
+
+    if _grammar_model is not None:
+        return _grammar_model, _grammar_tokenizer, _device
+
+    logger.info("Loading Grammar model (lazy init)...")
+
+    _device = "cpu"
+
+    try:
+        logger.info(f"Loading tokenizer: {MODEL_ID}")
+        _grammar_tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_ID, local_files_only=True
+        )
+
+        logger.info(f"Loading T5 model: {MODEL_ID}")
+        _grammar_model = AutoModelForSeq2SeqLM.from_pretrained(
+            MODEL_ID, local_files_only=True
+        )
+        _grammar_model.to(_device)
+        _grammar_model.eval()
+
+        logger.info(f"Grammar model loaded on {_device}")
+
+    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)
+            _grammar_model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
+            _grammar_model.to(_device)
+            _grammar_model.eval()
+            logger.info(f"Grammar model loaded on {_device} (fallback)")
+        except Exception as e2:
+            logger.error(f"Grammar model load failed completely: {e2}")
+            _grammar_model = None
+            _grammar_tokenizer = None
+            raise
+
+    return _grammar_model, _grammar_tokenizer, _device
+
+
+def correct_grammar(text: str) -> str:
+    """
+    Run grammar correction on input text using T5.
+    Returns the corrected text.
+    """
+    if not text or not text.strip():
+        return text
+
+    model, tokenizer, device = get_grammar_model()
+
+    try:
+        # Tokenize input
+        inputs = tokenizer(
+            text,
+            return_tensors="pt",
+            max_length=MAX_INPUT_LENGTH,
+            truncation=True,
+            padding=True
+        )
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+
+        # Generate correction
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=MAX_OUTPUT_LENGTH,
+                num_beams=NUM_BEAMS,
+                early_stopping=True,
+                no_repeat_ngram_size=3,
+            )
+
+        # Decode output
+        corrected = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        # Safety: if model returns empty or much shorter, keep original
+        if not corrected or not corrected.strip():
+            return text
+
+        # If the correction is drastically different (>80% changed), keep original
+        if len(corrected) < len(text) * 0.3:
+            logger.warning(f"Grammar output too short ({len(corrected)} vs {len(text)}), keeping original")
+            return text
+
+        return corrected.strip()
+
+    except Exception as e:
+        logger.error(f"Grammar correction failed: {e}")
+        return text