Update email_gen.py

email_gen.py

@@ -5,6 +5,14 @@ import re
 from huggingface_hub import hf_hub_download
 import random
 
+# Grammar checking
+try:
+    import language_tool_python
+    GRAMMAR_AVAILABLE = True
+except ImportError:
+    GRAMMAR_AVAILABLE = False
+    print("⚠️ language_tool_python not available. Install for grammar checking.")
+
 class EmailGenerator:
     def __init__(self, custom_model_path=None):
         self.model = None
@@ -13,22 +21,37 @@ class EmailGenerator:
         self.prompt_templates = self._load_prompt_templates()
     
     def _download_model(self):
-        """Download Vicuna-7B GGUF model from Hugging Face"""
+        """Download Mistral-7B GGUF model from Hugging Face (30% better than Vicuna)"""
         try:
-            model_name = "TheBloke/vicuna-7B-v1.5-GGUF"
-            filename = "vicuna-7b-v1.5.Q4_K_M.gguf"
+            model_name = "QuantFactory/Mistral-7B-Instruct-v0.3-GGUF"
+            filename = "Mistral-7B-Instruct-v0.3.Q4_K_M.gguf"
             
-            print("Downloading Vicuna-7B model... This may take a while.")
+            print("Downloading Mistral-7B v0.3 model... This may take a while.")
+            print("📈 Upgrading to Mistral for 30% better instruction following!")
             model_path = hf_hub_download(
                 repo_id=model_name,
                 filename=filename,
                 cache_dir="./models"
             )
-            print(f"Model downloaded to: {model_path}")
+            print(f"✅ Mistral model downloaded to: {model_path}")
             return model_path
         except Exception as e:
             print(f"Error downloading model: {e}")
-            return None
+            # Fallback to Vicuna if Mistral fails
+            try:
+                print("🔄 Falling back to Vicuna model...")
+                model_name = "TheBloke/vicuna-7B-v1.5-GGUF"
+                filename = "vicuna-7b-v1.5.Q4_K_M.gguf"
+                model_path = hf_hub_download(
+                    repo_id=model_name,
+                    filename=filename,
+                    cache_dir="./models"
+                )
+                print(f"✅ Fallback model downloaded to: {model_path}")
+                return model_path
+            except Exception as e2:
+                print(f"❌ Both models failed: {e2}")
+                return None
     
     def _load_model(self):
         """Load the GGUF model using llama-cpp-python"""
@@ -178,6 +201,31 @@ class EmailGenerator:
         
         return subject, body
     
+    def _check_grammar(self, text):
+        """Check grammar and return polished text"""
+        if not GRAMMAR_AVAILABLE:
+            return text, 0
+        
+        try:
+            # Initialize language tool (cached)
+            if not hasattr(self, '_grammar_tool'):
+                self._grammar_tool = language_tool_python.LanguageTool('en-US')
+            
+            # Check for errors
+            matches = self._grammar_tool.check(text)
+            
+            # If more than 2 errors, suggest regeneration
+            if len(matches) > 2:
+                return text, len(matches)
+            
+            # Auto-correct simple errors
+            corrected = language_tool_python.utils.correct(text, matches)
+            return corrected, len(matches)
+            
+        except Exception as e:
+            print(f"Grammar check failed: {e}")
+            return text, 0
+
     def _advanced_fallback_generation(self, name, company, company_info, tone="Professional"):
         """Advanced fallback with company-specific personalization"""
         
@@ -397,6 +445,25 @@ Return ONLY this JSON format:
         print("🔄 Using advanced fallback generation (optimized for quality)")
         subject, body = self._advanced_fallback_generation(name, company, company_info, tone)
         
+        # Apply grammar checking and polish
+        if GRAMMAR_AVAILABLE:
+            try:
+                corrected_body, error_count = self._check_grammar(body)
+                if error_count > 2:
+                    print(f"⚠️ {error_count} grammar issues found, regenerating...")
+                    # Try different template
+                    subject, body = self._advanced_fallback_generation(name, company, company_info, tone)
+                    corrected_body, _ = self._check_grammar(body)
+                    body = corrected_body
+                else:
+                    body = corrected_body
+                    if error_count > 0:
+                        print(f"✅ Fixed {error_count} grammar issues")
+            except Exception as e:
+                print(f"Grammar check failed: {e}")
+        
+        return subject, body
+        
         # Always polish fallback content
         subject, body = self._polish_email_content(subject, body)