Initial deployment of Unsubscriber app with AI model

README.md

@@ -1,12 +1,28 @@
 ---
-title: Unsubscriber
-emoji: 📊
-colorFrom: red
-colorTo: purple
+title: Email Unsubscribe Classifier
+emoji: 📧
+colorFrom: blue
+colorTo: green
 sdk: gradio
-sdk_version: 5.31.0
+sdk_version: 4.19.0
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Email Unsubscribe Classifier
+
+This AI model classifies emails to determine if they are unsubscribe confirmations or important emails that should not be automatically processed.
+
+## Features
+- Classifies emails as 'unsubscribe' or 'important'
+- Provides confidence scores for predictions
+- Based on fine-tuned DeBERTa-v3-small model
+- Trained on 20,000 email samples
+
+## Usage
+Simply paste your email content (including subject line) into the text box and click "Submit" to get a classification.
+
+## Model Performance
+- Accuracy: 100% on test set
+- F1 Score: 1.0 for both classes
+- Model size: 552MB

app.py

@@ -0,0 +1,118 @@
+import os
+import sys
+import gradio as gr
+import json
+import base64
+from pathlib import Path
+
+# Add parent directory to path to import the main app modules
+parent_dir = str(Path(__file__).parent.parent)
+sys.path.insert(0, parent_dir)
+
+# Import the ML components from the main app
+from ml_suite.predictor import initialize_predictor, get_ai_prediction_for_email, is_predictor_ready, get_model_status
+
+# Initialize the predictor once when the app starts
+print("Initializing AI model...")
+initialize_predictor()
+
+def parse_email_content(email_text):
+    """Parse email content to extract subject and body"""
+    lines = email_text.strip().split('\n')
+    subject = ""
+    body = ""
+    
+    # Simple parsing - look for Subject: line
+    for i, line in enumerate(lines):
+        if line.lower().startswith('subject:'):
+            subject = line[8:].strip()
+            body = '\n'.join(lines[i+1:]).strip()
+            break
+    
+    if not subject and not body:
+        # If no subject line found, treat entire text as body
+        body = email_text
+    
+    return subject, body
+
+def classify_email(email_content):
+    """Classify email using the AI model"""
+    if not email_content.strip():
+        return "Please enter email content to analyze."
+    
+    # Check if model is ready
+    if not is_predictor_ready():
+        status = get_model_status()
+        return f"Model is not ready. Status: {status}"
+    
+    # Parse email content
+    subject, body = parse_email_content(email_content)
+    
+    # Create email data structure similar to the main app
+    email_data = {
+        'snippet': body[:200],  # Gmail API typically provides snippets
+        'subject': subject,
+        'body': body,
+        'sender': 'demo@example.com',  # Placeholder
+        'id': 'demo_id'
+    }
+    
+    try:
+        # Get prediction
+        result = get_ai_prediction_for_email(email_data)
+        
+        # Format the response
+        prediction = result.get('prediction', 'Unknown')
+        confidence = result.get('confidence', 0)
+        
+        # Create formatted output
+        output = f"""
+## Classification Result
+
+**Category:** {prediction}
+**Confidence:** {confidence:.2%}
+
+### Analysis:
+"""
+        
+        if prediction == 'unsubscribe':
+            output += "✅ This email appears to be an unsubscribe confirmation or related to subscription management."
+        elif prediction == 'important':
+            output += "⚠️ This email appears to be important and should not be automatically processed."
+        else:
+            output += "❓ Unable to classify this email with high confidence."
+        
+        # Add confidence interpretation
+        if confidence > 0.9:
+            output += f"\n\n*High confidence prediction ({confidence:.2%})*"
+        elif confidence > 0.7:
+            output += f"\n\n*Moderate confidence prediction ({confidence:.2%})*"
+        else:
+            output += f"\n\n*Low confidence prediction ({confidence:.2%})*"
+        
+        return output
+        
+    except Exception as e:
+        return f"Error during classification: {str(e)}"
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=classify_email,
+    inputs=gr.Textbox(
+        lines=10,
+        placeholder="Paste email content here...\n\nFormat:\nSubject: Your subscription has been cancelled\nBody text goes here...",
+        label="Email Content"
+    ),
+    outputs=gr.Markdown(label="Classification Result"),
+    title="Email Unsubscribe Classifier",
+    description="This AI model classifies emails as either 'unsubscribe' confirmations or 'important' emails that should not be auto-processed.",
+    examples=[
+        ["Subject: Your subscription has been cancelled\n\nHi there,\n\nWe're sorry to see you go! Your subscription to our newsletter has been successfully cancelled. You will no longer receive emails from us.\n\nBest regards,\nThe Team"],
+        ["Subject: Important: Your account security update\n\nDear Customer,\n\nWe've detected unusual activity on your account. Please review your recent transactions and update your password immediately.\n\nThank you,\nSecurity Team"],
+        ["Subject: You've been unsubscribed\n\nYou have been removed from our mailing list and will not receive any further emails from us."],
+    ],
+    theme=gr.themes.Soft()
+)
+
+if __name__ == "__main__":
+    demo.launch()

final_optimized_model/added_tokens.json

@@ -0,0 +1,3 @@
+{
+  "[MASK]": 128000
+}

final_optimized_model/config.json

@@ -0,0 +1,35 @@
+{
+  "_name_or_path": "microsoft/deberta-v3-small",
+  "architectures": [
+    "DebertaV2ForSequenceClassification"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-07,
+  "max_position_embeddings": 512,
+  "max_relative_positions": -1,
+  "model_type": "deberta-v2",
+  "norm_rel_ebd": "layer_norm",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "pooler_dropout": 0,
+  "pooler_hidden_act": "gelu",
+  "pooler_hidden_size": 768,
+  "pos_att_type": [
+    "p2c",
+    "c2p"
+  ],
+  "position_biased_input": false,
+  "position_buckets": 256,
+  "relative_attention": true,
+  "share_att_key": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.36.2",
+  "type_vocab_size": 0,
+  "vocab_size": 128100
+}

final_optimized_model/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c39041f172b403117f415e2947236e097c2c5d4f5607cd332b81a602bdeea05
+size 567598552

final_optimized_model/special_tokens_map.json

@@ -0,0 +1,15 @@
+{
+  "bos_token": "[CLS]",
+  "cls_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

final_optimized_model/spm.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c679fbf93643d19aab7ee10c0b99e460bdbc02fedf34b92b05af343b4af586fd
+size 2464616

final_optimized_model/tokenizer_config.json

@@ -0,0 +1,58 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128000": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "[CLS]",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_lower_case": false,
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "sp_model_kwargs": {},
+  "split_by_punct": false,
+  "tokenizer_class": "DebertaV2Tokenizer",
+  "unk_token": "[UNK]",
+  "vocab_type": "spm"
+}

final_optimized_model/training_info.json

@@ -0,0 +1,21 @@
+{
+  "model_name": "microsoft/deberta-v3-small",
+  "training_date": "2025-05-27T02:45:24.492640",
+  "training_hours": 7.459452472222223,
+  "samples": 20000,
+  "final_metrics": {
+    "eval_loss": 7.772326353006065e-05,
+    "eval_accuracy": 1.0,
+    "eval_precision_important": 1.0,
+    "eval_recall_important": 1.0,
+    "eval_f1_important": 1.0,
+    "eval_precision_unsub": 1.0,
+    "eval_recall_unsub": 1.0,
+    "eval_f1_unsub": 1.0,
+    "eval_false_positives": 0,
+    "eval_runtime": 433.8438,
+    "eval_samples_per_second": 6.915,
+    "eval_steps_per_second": 0.217,
+    "epoch": 3.01
+  }
+}

ml_suite/__init__.py

@@ -0,0 +1,19 @@
+"""
+Machine Learning module for Gmail Unsubscriber application.
+
+This module provides AI-powered classification of emails as 'unsubscribable' or 'important'.
+It includes components for data preparation, model training, and prediction, with user
+control over the AI lifecycle through an in-app AI panel.
+
+The design emphasizes:
+- User control over AI data preparation and training
+- Seamless integration with the existing application
+- Transparency in AI operations and decisions
+- Graceful degradation when AI components are unavailable
+"""
+
+# Import configuration to ensure directories are created
+from . import config
+
+# Version information
+__version__ = "0.1.0"

ml_suite/advanced_predictor.py

@@ -0,0 +1,283 @@
+"""
+Advanced predictor with support for multiple models and ensemble predictions
+"""
+
+import os
+import torch
+import numpy as np
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    TextClassificationPipeline
+)
+import logging
+from typing import Dict, List, Tuple, Optional
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class AdvancedPredictor:
+    """Advanced predictor with ensemble support and confidence calibration"""
+    
+    def __init__(self, model_paths: List[str], weights: Optional[List[float]] = None):
+        """
+        Initialize with multiple models for ensemble prediction
+        
+        Args:
+            model_paths: List of paths to model directories
+            weights: Optional weights for each model (must sum to 1.0)
+        """
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.models = []
+        self.tokenizers = []
+        self.pipelines = []
+        
+        # Load all models
+        for path in model_paths:
+            if os.path.exists(path):
+                logger.info(f"Loading model from {path}")
+                tokenizer = AutoTokenizer.from_pretrained(path)
+                model = AutoModelForSequenceClassification.from_pretrained(path)
+                model = model.to(self.device)
+                model.eval()
+                
+                # Create pipeline
+                pipeline = TextClassificationPipeline(
+                    model=model,
+                    tokenizer=tokenizer,
+                    device=0 if torch.cuda.is_available() else -1,
+                    top_k=None,
+                    function_to_apply="sigmoid"
+                )
+                
+                self.models.append(model)
+                self.tokenizers.append(tokenizer)
+                self.pipelines.append(pipeline)
+            else:
+                logger.warning(f"Model path not found: {path}")
+        
+        if not self.models:
+            raise ValueError("No models loaded successfully")
+        
+        # Set weights
+        if weights:
+            assert len(weights) == len(self.models), "Number of weights must match number of models"
+            assert abs(sum(weights) - 1.0) < 1e-6, "Weights must sum to 1.0"
+            self.weights = weights
+        else:
+            # Equal weights by default
+            self.weights = [1.0 / len(self.models)] * len(self.models)
+        
+        logger.info(f"Initialized with {len(self.models)} models")
+    
+    def predict(self, text: str, return_all_scores: bool = False) -> Dict:
+        """
+        Make ensemble prediction
+        
+        Args:
+            text: Email text to classify
+            return_all_scores: Whether to return individual model scores
+            
+        Returns:
+            Dictionary with prediction results
+        """
+        # Preprocess text
+        text = self._preprocess_email(text)
+        
+        # Get predictions from all models
+        all_predictions = []
+        for pipeline in self.pipelines:
+            try:
+                result = pipeline(text)
+                all_predictions.append(result)
+            except Exception as e:
+                logger.error(f"Error in prediction: {e}")
+                continue
+        
+        if not all_predictions:
+            return {
+                "label": "IMPORTANT",
+                "score": 0.5,
+                "confidence": "low",
+                "error": "Prediction failed"
+            }
+        
+        # Aggregate predictions
+        ensemble_scores = self._aggregate_predictions(all_predictions)
+        
+        # Determine final prediction
+        unsub_score = ensemble_scores.get("UNSUBSCRIBABLE", 0.5)
+        important_score = ensemble_scores.get("IMPORTANT", 0.5)
+        
+        # Apply confidence calibration
+        calibrated_unsub = self._calibrate_confidence(unsub_score)
+        
+        # Determine label
+        if calibrated_unsub > 0.75:  # High confidence threshold
+            label = "UNSUBSCRIBABLE"
+            score = calibrated_unsub
+        else:
+            label = "IMPORTANT"
+            score = important_score
+        
+        # Confidence level
+        if score > 0.9:
+            confidence = "high"
+        elif score > 0.7:
+            confidence = "medium"
+        else:
+            confidence = "low"
+        
+        result = {
+            "label": label,
+            "score": float(score),
+            "confidence": confidence,
+            "raw_scores": {
+                "UNSUBSCRIBABLE": float(unsub_score),
+                "IMPORTANT": float(important_score)
+            }
+        }
+        
+        if return_all_scores:
+            result["model_predictions"] = all_predictions
+        
+        return result
+    
+    def _preprocess_email(self, text: str) -> str:
+        """Advanced email preprocessing"""
+        # Handle subject extraction
+        if "Subject:" in text:
+            parts = text.split("Subject:", 1)
+            if len(parts) > 1:
+                subject = parts[1].split("\n")[0].strip()
+                body = parts[1][len(subject):].strip()
+                # Emphasize subject
+                text = f"Email Subject: {subject}. Email Body: {body}"
+        
+        # Clean text
+        text = text.replace("\\n", " ").replace("\\t", " ")
+        text = " ".join(text.split())
+        
+        # Truncate if too long
+        if len(text) > 2000:
+            text = text[:2000] + "..."
+        
+        return text
+    
+    def _aggregate_predictions(self, predictions: List) -> Dict[str, float]:
+        """Aggregate predictions from multiple models using weighted voting"""
+        aggregated = {"UNSUBSCRIBABLE": 0.0, "IMPORTANT": 0.0}
+        
+        for i, pred_list in enumerate(predictions):
+            weight = self.weights[i]
+            
+            # Handle different prediction formats
+            if isinstance(pred_list, list) and pred_list:
+                for pred in pred_list:
+                    label = pred.get("label", "").upper()
+                    score = pred.get("score", 0.5)
+                    
+                    if label in aggregated:
+                        aggregated[label] += score * weight
+        
+        # Normalize
+        total = sum(aggregated.values())
+        if total > 0:
+            for key in aggregated:
+                aggregated[key] /= total
+        
+        return aggregated
+    
+    def _calibrate_confidence(self, score: float, temperature: float = 1.2) -> float:
+        """Apply temperature scaling for confidence calibration"""
+        # Convert to logit
+        epsilon = 1e-7
+        score = np.clip(score, epsilon, 1 - epsilon)
+        logit = np.log(score / (1 - score))
+        
+        # Apply temperature scaling
+        calibrated_logit = logit / temperature
+        
+        # Convert back to probability
+        calibrated_score = 1 / (1 + np.exp(-calibrated_logit))
+        
+        return float(calibrated_score)
+    
+    def predict_batch(self, texts: List[str]) -> List[Dict]:
+        """Predict multiple emails efficiently"""
+        results = []
+        
+        # Process in batches for efficiency
+        batch_size = 8
+        for i in range(0, len(texts), batch_size):
+            batch = texts[i:i + batch_size]
+            batch_results = [self.predict(text) for text in batch]
+            results.extend(batch_results)
+        
+        return results
+    
+    def get_feature_importance(self, text: str) -> Dict:
+        """Get feature importance for explainability"""
+        # This is a simplified version - in production, use SHAP or LIME
+        important_keywords = [
+            "unsubscribe", "opt out", "preferences", "newsletter",
+            "promotional", "marketing", "deal", "offer", "sale"
+        ]
+        
+        text_lower = text.lower()
+        found_keywords = [kw for kw in important_keywords if kw in text_lower]
+        
+        return {
+            "important_features": found_keywords,
+            "feature_count": len(found_keywords)
+        }
+
+
+def create_advanced_predictor():
+    """Factory function to create predictor with best available models"""
+    model_paths = []
+    
+    # Check for advanced model first
+    if os.path.exists("./advanced_unsubscriber_model"):
+        model_paths.append("./advanced_unsubscriber_model")
+    
+    # Check for optimized model
+    if os.path.exists("./optimized_model"):
+        model_paths.append("./optimized_model")
+    
+    # Fallback to original model
+    if os.path.exists("./ml_suite/models/fine_tuned_unsubscriber"):
+        model_paths.append("./ml_suite/models/fine_tuned_unsubscriber")
+    
+    if not model_paths:
+        raise ValueError("No trained models found")
+    
+    # Use ensemble if multiple models available
+    if len(model_paths) > 1:
+        logger.info(f"Creating ensemble predictor with {len(model_paths)} models")
+        # Give higher weight to advanced model
+        weights = [0.6, 0.4] if len(model_paths) == 2 else None
+        return AdvancedPredictor(model_paths, weights)
+    else:
+        logger.info(f"Creating single model predictor")
+        return AdvancedPredictor(model_paths)
+
+
+# Example usage
+if __name__ == "__main__":
+    # Test the predictor
+    predictor = create_advanced_predictor()
+    
+    test_emails = [
+        "Subject: 50% OFF Everything! Limited time offer. Click here to shop now. Unsubscribe from promotional emails.",
+        "Subject: Security Alert: New login detected. We noticed a login from a new device. If this wasn't you, secure your account.",
+        "Subject: Your monthly newsletter is here! Check out our latest articles and tips. Manage your email preferences.",
+    ]
+    
+    for email in test_emails:
+        result = predictor.predict(email, return_all_scores=True)
+        print(f"\nEmail: {email[:100]}...")
+        print(f"Prediction: {result['label']}")
+        print(f"Confidence: {result['confidence']} ({result['score']:.2%})")
+        print(f"Raw scores: {result['raw_scores']}")

ml_suite/config.py

@@ -0,0 +1,149 @@
+"""
+Centralized configuration for Gmail Unsubscriber AI Suite.
+
+This module defines all configuration parameters for the ML components, including:
+- Directory paths for models, datasets, and task status
+- Hugging Face cache configuration
+- Model specifications
+- Data preparation parameters
+- Training hyperparameters
+- User data collection and personalization parameters
+
+All directories are automatically created when this module is imported.
+"""
+
+import os
+
+# --- Base Path Configuration ---
+ML_SUITE_DIR = os.path.dirname(os.path.abspath(__file__))
+PROJECT_ROOT = os.path.dirname(ML_SUITE_DIR)
+
+# --- Cache and Model Storage ---
+MODELS_DIR = os.path.join(ML_SUITE_DIR, "models")
+BASE_TRANSFORMER_CACHE_DIR = os.path.join(MODELS_DIR, "base_transformer_cache")
+# FINE_TUNED_MODEL_DIR = os.path.join(MODELS_DIR, "fine_tuned_unsubscriber")  # Old model
+FINE_TUNED_MODEL_DIR = os.path.join(PROJECT_ROOT, "final_optimized_model")  # New trained model
+
+# Set Hugging Face environment variables to use project-local cache
+os.environ['HF_HOME'] = BASE_TRANSFORMER_CACHE_DIR
+os.environ['TRANSFORMERS_CACHE'] = BASE_TRANSFORMER_CACHE_DIR
+os.environ['HF_DATASETS_CACHE'] = os.path.join(BASE_TRANSFORMER_CACHE_DIR, 'datasets')
+os.environ['HF_METRICS_CACHE'] = os.path.join(BASE_TRANSFORMER_CACHE_DIR, 'metrics')
+
+# --- Dataset Storage ---
+DATASETS_DIR = os.path.join(ML_SUITE_DIR, "datasets")
+RAW_DATASETS_DIR = os.path.join(DATASETS_DIR, "raw")
+EXTRACTED_DATASETS_DIR = os.path.join(DATASETS_DIR, "extracted")
+PROCESSED_DATASETS_DIR = os.path.join(DATASETS_DIR, "processed")
+PREPARED_DATA_FILE = os.path.join(PROCESSED_DATASETS_DIR, "unsubscriber_training_data.csv")
+DATA_COLUMNS_SCHEMA = ['text', 'label']  # Schema for the training CSV
+
+# --- Task Status Storage ---
+TASK_STATUS_DIR = os.path.join(ML_SUITE_DIR, "task_status")
+DATA_PREP_STATUS_FILE = os.path.join(TASK_STATUS_DIR, "data_preparation_status.json")
+MODEL_TRAIN_STATUS_FILE = os.path.join(TASK_STATUS_DIR, "model_training_status.json")
+PERSONALIZED_TRAIN_STATUS_FILE = os.path.join(TASK_STATUS_DIR, "personalized_training_status.json")
+
+# --- User Data Collection and Personalization ---
+USER_DATA_DIR = os.path.join(ML_SUITE_DIR, "user_data")
+USER_FEEDBACK_DIR = os.path.join(USER_DATA_DIR, "feedback")
+USER_MODELS_DIR = os.path.join(USER_DATA_DIR, "models")
+USER_DATASETS_DIR = os.path.join(USER_DATA_DIR, "datasets")
+
+# User feedback collection configuration
+USER_FEEDBACK_FILE = os.path.join(USER_FEEDBACK_DIR, "user_feedback.csv")
+FEEDBACK_COLUMNS_SCHEMA = ['email_id', 'text', 'predicted_label', 'predicted_confidence', 'user_feedback', 'timestamp', 'session_id']
+
+# Personalized model configuration
+PERSONALIZED_MODEL_DIR_TEMPLATE = os.path.join(USER_MODELS_DIR, "{user_id}")
+PERSONALIZED_MODEL_FILE_TEMPLATE = os.path.join(PERSONALIZED_MODEL_DIR_TEMPLATE, "model.pt")
+PERSONALIZED_MODEL_INFO_TEMPLATE = os.path.join(PERSONALIZED_MODEL_DIR_TEMPLATE, "model_info.json")
+PERSONALIZED_DATASET_FILE_TEMPLATE = os.path.join(USER_DATASETS_DIR, "{user_id}_training_data.csv")
+
+# Personalization hyperparameters
+MIN_FEEDBACK_ENTRIES_FOR_PERSONALIZATION = 10  # Minimum number of user feedback entries required for personalization
+PERSONALIZATION_WEIGHT = 0.7  # Weight given to user feedback vs. base model (higher = more personalized)
+PERSONALIZATION_EPOCHS = 2  # Number of epochs for fine-tuning a personalized model
+
+# --- Directory Creation (Updated with User Data directories) ---
+for dir_path in [MODELS_DIR, BASE_TRANSFORMER_CACHE_DIR, FINE_TUNED_MODEL_DIR,
+                 RAW_DATASETS_DIR, EXTRACTED_DATASETS_DIR, PROCESSED_DATASETS_DIR, TASK_STATUS_DIR,
+                 USER_DATA_DIR, USER_FEEDBACK_DIR, USER_MODELS_DIR, USER_DATASETS_DIR]:
+    os.makedirs(dir_path, exist_ok=True)
+
+# --- Transformer Model Configuration ---
+# Choice: DistilBERT offers a good balance of performance and resource efficiency.
+# Other candidates: 'bert-base-uncased', 'roberta-base', 'google/electra-small-discriminator'.
+# The choice impacts download size, training time, and inference speed.
+PRE_TRAINED_MODEL_NAME = "distilbert-base-uncased"
+
+# --- Data Preparation Parameters ---
+# Define sources for public email data. URLs and types guide the preparator.
+PUBLIC_DATASETS_INFO = {
+    "spamassassin_easy_ham_2003": {
+        "url": "https://spamassassin.apache.org/publiccorpus/20030228_easy_ham.tar.bz2",
+        "type": "important_leaning",  # Expected dominant class after heuristic application
+        "extract_folder_name": "spamassassin_easy_ham_2003"
+    },
+    "spamassassin_spam_2003": {
+        "url": "https://spamassassin.apache.org/publiccorpus/20030228_spam.tar.bz2",
+        "type": "unsubscribable_leaning",
+        "extract_folder_name": "spamassassin_spam_2003"
+    },
+    # Consider adding more diverse datasets like:
+    # - Enron (requires significant parsing and ethical review for a suitable subset)
+    # - Public mailing list archives (e.g., from Apache Software Foundation, carefully selected for relevance)
+}
+MIN_TEXT_LENGTH_FOR_TRAINING = 60  # Emails shorter than this (after cleaning) are likely not useful.
+MAX_SAMPLES_PER_RAW_DATASET = 7500  # Limits processing time for initial data prep. Can be increased.
+EMAIL_SNIPPET_LENGTH_FOR_MODEL = 1024  # Max characters from email body to combine with subject for model input.
+
+# --- Training Hyperparameters & Configuration ---
+NUM_LABELS = 2  # Binary classification: Unsubscribable vs. Important
+LABEL_IMPORTANT_ID = 0
+LABEL_UNSUBSCRIBABLE_ID = 1
+ID_TO_LABEL_MAP = {LABEL_IMPORTANT_ID: "IMPORTANT", LABEL_UNSUBSCRIBABLE_ID: "UNSUBSCRIBABLE"}
+LABEL_TO_ID_MAP = {"IMPORTANT": LABEL_IMPORTANT_ID, "UNSUBSCRIBABLE": LABEL_UNSUBSCRIBABLE_ID}
+
+MAX_SEQ_LENGTH = 512        # Max token sequence length for Transformer. Impacts memory and context window.
+TRAIN_BATCH_SIZE = 16        # Batch size for training. Reduced for GTX 1650 (4GB VRAM)
+EVAL_BATCH_SIZE = 32        # Batch size for evaluation. Reduced for GTX 1650
+NUM_TRAIN_EPOCHS = 8        # Number of full passes through the training data (increased for better learning).
+LEARNING_RATE = 1e-5        # AdamW optimizer learning rate, slightly reduced for more stable training.
+WEIGHT_DECAY = 0.02         # Regularization parameter.
+WARMUP_STEPS_RATIO = 0.15    # Ratio of total training steps for learning rate warmup.
+TEST_SPLIT_SIZE = 0.2       # Proportion of data for the evaluation set (increased for better validation).
+
+# Hugging Face Trainer Arguments
+EVALUATION_STRATEGY = "epoch"  # Evaluate at the end of each epoch.
+SAVE_STRATEGY = "epoch"       # Save model checkpoint at the end of each epoch.
+LOAD_BEST_MODEL_AT_END = True  # Reload the best model (based on metric_for_best_model) at the end of training.
+METRIC_FOR_BEST_MODEL = "f1_unsub"  # Focus on F1 for the "unsubscribable" class.
+FP16_TRAINING = True          # Enable mixed-precision training if a CUDA GPU is available and supports it.
+EARLY_STOPPING_PATIENCE = 3   # Stop training if metric_for_best_model doesn't improve for this many epochs.
+EARLY_STOPPING_THRESHOLD = 0.001  # Minimum change to be considered an improvement.
+
+# --- AI User Preferences (Defaults stored in JS, but can be defined here for reference) ---
+DEFAULT_AI_ENABLED_ON_SCAN = True
+DEFAULT_AI_CONFIDENCE_THRESHOLD = 0.5  # (50%) - Balanced threshold for optimal precision/recall
+
+# --- API Endpoint Configuration for Backend Integration ---
+API_ENDPOINTS = {
+    "submit_feedback": "/api/ai/feedback",
+    "get_feedback_stats": "/api/ai/feedback/stats",
+    "train_personalized": "/api/ai/train_personalized",
+    "reset_user_data": "/api/ai/user_data/reset",
+    "export_user_data": "/api/ai/user_data/export",
+    "import_user_data": "/api/ai/user_data/import"
+}
+# --- Advanced Transformer Configuration (2024 Research) ---
+# Based on 2024 research showing RoBERTa and DistilBERT achieve 99%+ accuracy
+TRANSFORMER_MODEL_NAME = "distilbert-base-uncased"  # Optimal balance of speed and accuracy
+USE_MIXED_PRECISION = True  # FP16 training for efficiency
+GRADIENT_ACCUMULATION_STEPS = 4  # Increased for GTX 1650 to simulate larger batch size
+MAX_GRAD_NORM = 1.0  # Gradient clipping for stability
+LABEL_SMOOTHING_FACTOR = 0.1  # Reduce overconfidence
+SAVE_TOTAL_LIMIT = 3  # Keep only best 3 checkpoints
+LOGGING_STEPS = 50  # Frequent logging for monitoring
+EVAL_STEPS = 100  # Regular evaluation during training
+DATALOADER_NUM_WORKERS = 2  # Reduced for GTX 1650 to avoid memory issues

ml_suite/data_preparator.py

@@ -0,0 +1,458 @@
+"""
+Enhanced data preparation module for the Gmail Unsubscriber AI suite.
+
+This module provides robust data preparation with fallback capabilities:
+- Primary: Download and process public email datasets
+- Fallback: Create comprehensive synthetic dataset when downloads fail
+- Robust error handling and recovery
+- Quality validation and balancing
+"""
+
+import os
+import re
+import csv
+import time
+import json
+import shutil
+import tarfile
+import requests
+import tempfile
+import random
+import email
+import email.parser
+import email.policy
+from typing import Dict, List, Tuple, Optional, Any, Union, Set
+from collections import Counter, defaultdict
+from bs4 import BeautifulSoup
+
+# Local imports
+from . import config
+from . import utils
+from .task_utils import AiTaskLogger
+
+
+def create_comprehensive_fallback_dataset(task_logger: AiTaskLogger) -> List[Tuple[str, int]]:
+    """
+    Create a comprehensive fallback training dataset when external downloads fail.
+    
+    Returns balanced, high-quality training examples for email classification.
+    """
+    task_logger.info("Creating comprehensive fallback training dataset...")
+    
+    # Unsubscribable emails (marketing, promotional, newsletters)
+    unsubscribable_emails = [
+        # Marketing newsletters
+        ("Subject: Weekly Newsletter - Best Deals & Updates Inside!\n\nDiscover this week's top deals and special offers exclusively for our subscribers. Check out our latest product launches and limited-time promotions. If you no longer wish to receive these weekly newsletters, click here to unsubscribe from all marketing communications.", 1),
+        ("Subject: Flash Sale Alert - 48 Hours Only!\n\nDon't miss our biggest flash sale of the season! Save up to 60% on thousands of items with free shipping on orders over $50. This exclusive offer expires in 48 hours. To opt out of flash sale notifications, update your email preferences here.", 1),
+        ("Subject: New Arrivals Handpicked Just For You\n\nBased on your browsing history and previous purchases, our style experts have curated these new arrivals specifically for your taste. Discover the latest trends and exclusive styles before anyone else. Unsubscribe from personalized product recommendations here.", 1),
+        ("Subject: Member Exclusive - VIP Early Access Sale\n\nAs a valued VIP member, enjoy exclusive early access to our annual clearance sale before it opens to the public. Get first pick of the best deals with an additional 20% off already reduced prices. Manage your VIP membership communications here.", 1),
+        ("Subject: Limited Weekend Offer - Free Premium Shipping\n\nThis weekend only! Enjoy free premium shipping on all orders, no minimum purchase required. Perfect time to stock up on your favorites or try something new with expedited delivery. Stop weekend promotion emails here.", 1),
+        
+        # Product promotions and recommendations
+        ("Subject: MASSIVE CLEARANCE EVENT - Up to 80% OFF Everything\n\nOur biggest clearance event ever! Save up to 80% on thousands of items across all categories. Everything must go to make room for new seasonal inventory. Don't miss these incredible once-a-year deals. Unsubscribe from clearance sale alerts here.", 1),
+        ("Subject: Back in Stock Alert - Your Wishlist Items Available\n\nGreat news! Multiple items from your wishlist are now back in stock and ready for immediate shipping. These popular products tend to sell out quickly, so we recommend ordering soon. Turn off wishlist availability notifications here.", 1),
+        ("Subject: Similar Customers Also Purchased These Items\n\nBased on your recent order, here are products that customers with similar purchasing patterns have been loving. Discover new favorites and trending items curated just for you. Opt out of 'customers also bought' recommendation emails here.", 1),
+        ("Subject: Complete Your Collection - Matching Set Available\n\nWe noticed you recently purchased part of our bestselling collection. Complete your set with these perfectly coordinating pieces now available at a special bundle price with free gift wrapping. Unsubscribe from collection completion suggestions here.", 1),
+        ("Subject: Reward Points Expiring Soon - Redeem Now!\n\nDon't let your 750 reward points expire! Use them before the end of this month for instant discounts, free products, or exclusive member perks. Your points are worth $37.50 in savings. Manage reward points notifications here.", 1),
+        
+        # Subscription services and content
+        ("Subject: Your Weekly Meal Plan - Delicious Recipes Inside\n\nThis week's curated meal plan features 7 delicious, nutritionist-approved recipes designed to save you time and money. Includes shopping lists and prep instructions for busy weeknight dinners. Unsubscribe from weekly meal plan emails here.", 1),
+        ("Subject: Daily Market Brief - Today's Financial Highlights\n\nStay informed with today's most important market movements and financial news. Our expert analysts provide insights on stocks, crypto, and economic indicators that matter to your portfolio. Stop daily market briefing emails here.", 1),
+        ("Subject: Fitness Challenge Week 4 - You're Crushing It!\n\nAmazing progress! Week 4 of your personalized fitness challenge includes new strength training exercises and advanced nutrition tips to accelerate your results. Keep up the fantastic work! Opt out of fitness challenge updates here.", 1),
+        ("Subject: Learning Path Update - Your Next Course is Ready\n\nBased on your completed coursework and learning goals, we've prepared your next recommended course in advanced digital marketing. Continue building your professional skills with expert-led content. Manage learning recommendation preferences here.", 1),
+        ("Subject: Podcast Weekly - New Episodes You'll Love\n\nThis week's podcast episodes feature industry leaders discussing the latest trends in technology, business, and personal development. Plus exclusive interviews and behind-the-scenes content. Unsubscribe from podcast update emails here.", 1),
+        
+        # Events, webinars, and community
+        ("Subject: Exclusive Masterclass Invitation - Limited Seats Available\n\nJoin renowned industry expert Dr. Sarah Johnson for an exclusive masterclass on 'Advanced Digital Strategy for 2025.' Only 50 seats available for this interactive session with live Q&A. Opt out of masterclass invitations here.", 1),
+        ("Subject: Local Community Meetup - Connect in Your City\n\nWe're hosting an exclusive meetup in your city next month! Network with like-minded professionals, enjoy complimentary refreshments, and get insider access to upcoming product launches. Stop local community event notifications here.", 1),
+        ("Subject: Annual Conference Early Bird - Save 45% on Registration\n\nSecure your spot at this year's premier industry conference with super early bird pricing. Three days of cutting-edge workshops, keynote speakers, and unparalleled networking opportunities. Unsubscribe from conference promotion emails here.", 1),
+        
+        # Surveys, feedback, and reviews
+        ("Subject: Your Opinion Matters - Quick 3-Minute Survey\n\nHelp us improve your experience with a quick 3-minute survey about our recent website updates. Your feedback directly influences our product development and customer service improvements. Opt out of customer survey requests here.", 1),
+        ("Subject: Share Your Experience - Product Review Request\n\nYou recently purchased our bestselling wireless headphones. Would you mind sharing a quick review to help other customers make informed decisions? Your honest feedback is invaluable to our community. Turn off product review requests here.", 1),
+        ("Subject: Customer Feedback Spotlight - You're Featured!\n\nWe loved your recent review so much that we'd like to feature it in our customer spotlight newsletter! See how your feedback is helping other customers discover their perfect products. Manage customer spotlight communications here.", 1),
+        
+        # Seasonal and holiday promotions
+        ("Subject: Holiday Gift Guide 2025 - Perfect Presents for Everyone\n\nTake the guesswork out of holiday shopping with our expertly curated gift guide. Find perfect presents for everyone on your list, organized by interests, age, and budget ranges. Stop holiday marketing communications here.", 1),
+        ("Subject: Summer Collection Preview - Beach Ready Essentials\n\nGet beach ready with our new summer collection featuring the latest swimwear, accessories, and vacation essentials. Plus early access to summer sale prices for newsletter subscribers only. Opt out of seasonal collection previews here.", 1),
+        ("Subject: Spring Cleaning Sale - Organize Your Space & Save\n\nSpring cleaning season is here! Save big on organization solutions, storage systems, and home improvement essentials. Plus get expert tips and tricks from our home organization specialists. Unsubscribe from seasonal promotion emails here.", 1),
+    ]
+    
+    # Important emails (security, billing, personal, urgent)
+    important_emails = [
+        # Security and account management
+        ("Subject: Critical Security Alert - Immediate Action Required\n\nWe detected unauthorized login attempts to your account from an unrecognized device in Moscow, Russia. Your account has been temporarily secured. Please verify your identity and change your password immediately to restore full access.", 0),
+        ("Subject: Password Successfully Updated - Security Confirmation\n\nYour account password was successfully changed on May 22, 2025 at 4:30 PM EST from IP address 192.168.1.100. If you made this change, no further action is needed. If you didn't authorize this change, contact our security team immediately.", 0),
+        ("Subject: Two-Factor Authentication Setup Required\n\nFor enhanced account security, two-factor authentication setup is now required for all accounts. Please complete the setup process within 7 days to maintain uninterrupted access to your account. This security measure protects against unauthorized access.", 0),
+        ("Subject: Account Verification Required - Expires in 24 Hours\n\nTo comply with security regulations and keep your account active, please verify your email address by clicking the secure link below. This verification link expires in 24 hours for your protection. Unverified accounts will be temporarily suspended.", 0),
+        ("Subject: Suspicious Activity Detected - Account Review in Progress\n\nOur fraud detection system has flagged unusual activity on your account. As a precautionary measure, certain features have been temporarily limited while we conduct a security review. Please contact support if you have questions.", 0),
+        
+        # Financial, billing, and payments
+        ("Subject: Monthly Account Statement Now Available\n\nYour detailed account statement for May 2025 is now available for review in your secure account portal. This statement includes all transactions, fees, and important account activity from the past month. Access your statement securely online.", 0),
+        ("Subject: Payment Method Declined - Update Required Within 48 Hours\n\nYour automatic payment of $129.99 scheduled for today was declined by your bank. To avoid service interruption, please update your payment method or contact your bank within 48 hours. Update payment details securely in your account.", 0),
+        ("Subject: Refund Approved and Processed - $287.46\n\nGood news! Your refund request for order #ORD-2025-5678 has been approved and processed. The amount of $287.46 has been credited to your original payment method and should appear within 3-5 business days depending on your bank.", 0),
+        ("Subject: Urgent: Invoice #INV-2025-4321 Overdue - Payment Required\n\nYour invoice #INV-2025-4321 for $445.00 is now 15 days overdue. To avoid late fees and potential service suspension, please submit payment immediately. Pay securely through your account portal or contact our billing department.", 0),
+        ("Subject: Annual Tax Documents Available for Download\n\nYour 2024 tax documents (Form 1099-MISC) are now available for download in your account's tax documents section. These forms are required for your tax filing and should be retained for your records. Download deadline: December 31, 2025.", 0),
+        
+        # Orders, shipping, and logistics
+        ("Subject: Order Confirmation #ORD-2025-9876 - Processing Started\n\nThank you for your order! Order #ORD-2025-9876 totaling $156.78 has been received and payment successfully processed. Your items are now being prepared for shipment. You'll receive tracking information once dispatched within 1-2 business days.", 0),
+        ("Subject: Shipment Notification - Your Order is on the Way\n\nExcellent news! Your order #ORD-2025-9876 has been shipped via UPS Ground and is currently in transit. Track your package using number 1Z999AA1012345675. Estimated delivery: May 24-26, 2025 between 9 AM and 7 PM.", 0),
+        ("Subject: Delivery Failed - Recipient Not Available\n\nUPS attempted delivery of your package today at 3:15 PM but no one was available to receive it. Your package is now at the local UPS Customer Center. Please arrange redelivery or pickup within 5 business days to avoid return to sender.", 0),
+        ("Subject: Package Delivered Successfully - Confirmation Required\n\nYour package was delivered today at 11:30 AM and left at your front door as per delivery instructions. If you haven't received your package or notice any damage, please contact customer service within 48 hours for immediate assistance.", 0),
+        ("Subject: Return Authorization Approved - Instructions Included\n\nYour return request for order #ORD-2025-9876 has been approved. Use the prepaid return label attached to this email. Package items securely and drop off at any UPS location. Refund will be processed within 5-7 business days upon receipt.", 0),
+        
+        # Service notifications and updates
+        ("Subject: Scheduled System Maintenance - Service Interruption Notice\n\nImportant: We will be performing critical system maintenance on May 25, 2025 from 2:00 AM to 6:00 AM EST. During this time, our website and mobile app will be temporarily unavailable. We apologize for any inconvenience this may cause.", 0),
+        ("Subject: Service Fully Restored - All Systems Operational\n\nOur system maintenance has been completed successfully and all services are now fully operational. Thank you for your patience during the temporary service interruption. If you experience any issues, please contact our technical support team.", 0),
+        ("Subject: Critical: Terms of Service Update - Action Required\n\nImportant changes to our Terms of Service and Privacy Policy will take effect on June 15, 2025. Please review these updates carefully as continued use of our services after this date constitutes acceptance of the new terms.", 0),
+        ("Subject: Subscription Expiration Warning - Renew Within 5 Days\n\nYour premium subscription expires on May 30, 2025. To continue enjoying uninterrupted access to premium features and priority support, please renew your subscription within the next 5 days. Renew now to avoid service interruption.", 0),
+        ("Subject: Account Upgrade Successful - Welcome to Premium\n\nWelcome to Premium! Your account upgrade is now active and you have full access to all premium features including priority customer support, advanced analytics, and exclusive content. Explore your new benefits in your account dashboard.", 0),
+        
+        # Personal, professional, and health
+        ("Subject: Urgent: Medical Appointment Reminder - Tomorrow 9:30 AM\n\nReminder: You have a scheduled appointment with Dr. Jennifer Martinez tomorrow, May 23, 2025 at 9:30 AM. Location: Downtown Medical Center, Suite 402. Please arrive 15 minutes early and bring your insurance card and photo ID.", 0),
+        ("Subject: Lab Results Available - Please Review\n\nYour recent laboratory test results are now available in your patient portal. Please log in to review your results and schedule a follow-up appointment if recommended by your healthcare provider. Results were processed on May 22, 2025.", 0),
+        ("Subject: Emergency Contact Update Required - HR Notice\n\nOur HR records show that your emergency contact information was last updated over two years ago. Please update your emergency contacts in the employee portal within 30 days to ensure we can reach someone in case of workplace emergencies.", 0),
+        ("Subject: Project Deadline Reminder - Report Due Tomorrow\n\nReminder: The quarterly project report for the Johnson account is due tomorrow, May 23, 2025 by 5:00 PM. Please submit your completed report to the project management system and copy all stakeholders. Contact me if you need an extension.", 0),
+        ("Subject: Contract Renewal Notice - Legal Review Required\n\nYour service contract #CON-2025-789 expires on June 30, 2025. Please review the attached renewal terms and return the signed agreement by June 15, 2025. Contact our legal department if you have questions about the updated terms and conditions.", 0),
+        
+        # Family, personal, and community
+        ("Subject: Family Emergency - Please Call Home Immediately\n\nFamily emergency - please call home as soon as you receive this message. Mom is at St. Mary's Hospital, room 304. She's stable but asking for you. Call Dad's cell phone 555-0123 for more details. Drive safely.", 0),
+        ("Subject: School Notification - Parent Conference Scheduled\n\nThis is to inform you that a parent-teacher conference has been scheduled for your child, Emma Johnson, on May 28, 2025 at 3:00 PM. Please contact the school office at 555-0187 if you cannot attend at the scheduled time.", 0),
+        ("Subject: Prescription Ready for Pickup - Pharmacy Notice\n\nYour prescription for Lisinopril 10mg is ready for pickup at CVS Pharmacy, 123 Main Street. Pharmacy hours: Monday-Friday 8 AM-10 PM, Saturday-Sunday 9 AM-8 PM. Please bring photo ID when picking up your medication.", 0),
+    ]
+    
+    # Combine datasets and ensure balance
+    all_examples = unsubscribable_emails + important_emails
+    random.shuffle(all_examples)
+    
+    task_logger.info(f"Created fallback dataset with {len(all_examples)} examples")
+    task_logger.info(f"Unsubscribable examples: {len(unsubscribable_emails)}")
+    task_logger.info(f"Important examples: {len(important_emails)}")
+    
+    return all_examples
+
+
+def download_and_extract_dataset(
+    dataset_key: str, 
+    info: Dict[str, Any], 
+    task_logger: AiTaskLogger
+) -> Tuple[bool, str]:
+    """
+    Download and extract a dataset archive with robust error handling.
+    """
+    utils.ensure_directory_exists(config.RAW_DATASETS_DIR)
+    utils.ensure_directory_exists(config.EXTRACTED_DATASETS_DIR)
+    
+    url = info["url"]
+    extract_folder_name = info["extract_folder_name"]
+    extracted_dir = os.path.join(config.EXTRACTED_DATASETS_DIR, extract_folder_name)
+    
+    # Skip if already exists
+    if os.path.exists(extracted_dir) and os.listdir(extracted_dir):
+        task_logger.info(f"Dataset {dataset_key} already exists, skipping download.")
+        return True, extracted_dir
+    
+    archive_path = os.path.join(config.RAW_DATASETS_DIR, f"{dataset_key}.tar.bz2")
+    
+    try:
+        task_logger.info(f"Downloading dataset {dataset_key} from {url}")
+        
+        # Download with timeout and error handling
+        response = requests.get(url, stream=True, timeout=30)
+        response.raise_for_status()
+        
+        total_size = int(response.headers.get('content-length', 0))
+        if total_size == 0:
+            task_logger.warning(f"Unknown file size for {dataset_key}")
+        
+        downloaded = 0
+        with open(archive_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+                    downloaded += len(chunk)
+                    if total_size > 0:
+                        progress = downloaded / total_size
+                        task_logger.update_progress(
+                            progress, 
+                            f"Downloading {dataset_key}: {downloaded/1024/1024:.1f} MB"
+                        )
+        
+        # Verify download
+        if downloaded == 0:
+            raise ValueError("Downloaded file is empty")
+        
+        task_logger.info(f"Extracting dataset {dataset_key}")
+        utils.ensure_directory_exists(extracted_dir)
+        
+        # Extract with error handling
+        with tarfile.open(archive_path, 'r:bz2') as tar:
+            tar.extractall(path=extracted_dir)
+        
+        # Verify extraction
+        if not os.listdir(extracted_dir):
+            raise ValueError("Extracted directory is empty")
+        
+        task_logger.info(f"Successfully downloaded and extracted {dataset_key}")
+        return True, extracted_dir
+        
+    except Exception as e:
+        task_logger.error(f"Error downloading or extracting dataset {dataset_key}: {str(e)}")
+        
+        # Cleanup failed downloads
+        for path in [archive_path, extracted_dir]:
+            if os.path.exists(path):
+                try:
+                    if os.path.isdir(path):
+                        shutil.rmtree(path)
+                    else:
+                        os.remove(path)
+                except:
+                    pass
+        
+        return False, ""
+
+
+def process_email_content(email_text: str, expected_label_type: str) -> Optional[Tuple[str, int]]:
+    """
+    Process a single email and return cleaned text with label.
+    """
+    try:
+        # Parse email
+        msg = email.message_from_string(email_text, policy=email.policy.default)
+        
+        # Extract subject and body
+        subject = msg.get('Subject', '').strip()
+        body = ""
+        
+        if msg.is_multipart():
+            for part in msg.walk():
+                if part.get_content_type() == "text/plain":
+                    body = part.get_content()
+                    break
+        else:
+            if msg.get_content_type() == "text/plain":
+                body = msg.get_content()
+        
+        if isinstance(body, bytes):
+            body = body.decode('utf-8', errors='ignore')
+        
+        # Clean and combine
+        cleaned_subject = utils.clean_text_for_model(subject, max_length=200)
+        cleaned_body = utils.clean_text_for_model(body, max_length=800)
+        
+        combined_text = f"Subject: {cleaned_subject}\n\n{cleaned_body}"
+        
+        # Skip if too short
+        if len(combined_text.strip()) < config.MIN_TEXT_LENGTH_FOR_TRAINING:
+            return None
+        
+        # Assign label based on heuristics
+        label = 1 if expected_label_type == "unsubscribable_leaning" else 0
+        
+        # Apply some heuristics to improve labeling
+        text_lower = combined_text.lower()
+        
+        # Marketing indicators (more likely unsubscribable)
+        marketing_indicators = [
+            'unsubscribe', 'opt out', 'newsletter', 'promotional', 'sale', 'offer',
+            'deal', 'discount', 'marketing', 'advertisement', 'subscribe'
+        ]
+        
+        # Important indicators (more likely important)
+        important_indicators = [
+            'urgent', 'security', 'password', 'account', 'payment', 'bill',
+            'order', 'shipping', 'delivered', 'confirmation', 'receipt'
+        ]
+        
+        marketing_score = sum(1 for indicator in marketing_indicators if indicator in text_lower)
+        important_score = sum(1 for indicator in important_indicators if indicator in text_lower)
+        
+        # Adjust label based on content analysis
+        if marketing_score > important_score + 1:
+            label = 1  # Unsubscribable
+        elif important_score > marketing_score + 1:
+            label = 0  # Important
+        
+        return (combined_text, label)
+        
+    except Exception as e:
+        return None
+
+
+def download_accessible_datasets(task_logger: AiTaskLogger) -> List[Tuple[str, int]]:
+    """
+    Download and process accessible public datasets from reliable sources.
+    
+    Based on 2024 research, we'll use the most accessible and reliable datasets:
+    1. UCI Spambase patterns (synthetic but research-based)
+    2. Comprehensive promotional email patterns
+    3. Modern phishing and security email patterns
+    """
+    task_logger.info("Downloading accessible public datasets...")
+    
+    accessible_examples = []
+    
+    # UCI Spambase-inspired examples with high-frequency spam words
+    task_logger.info("Creating UCI Spambase-inspired examples...")
+    uci_examples = [
+        ("Subject: FREE MONEY!!! Make $$$ from HOME!!!\n\nCREDIT problems? NO PROBLEM! Our BUSINESS opportunity will REMOVE all your financial worries! RECEIVE MONEY via INTERNET! MAIL us for FREE REPORT! PEOPLE OVER the world are making MONEY with OUR TECHNOLOGY! ORDER now!", 1),
+        ("Subject: URGENT BUSINESS PROPOSAL - FREE MONEY\n\nDear friend, I am writing to seek your assistance in a business proposal involving the transfer of $15 MILLION. This is 100% RISK FREE and will bring you substantial financial reward. Please provide your email and personal details to receive full details.", 1),
+        ("Subject: You've WON $50,000!!! CLAIM NOW!!!\n\nCongratulations! Your email address has been selected in our random drawing! You've won FIFTY THOUSAND DOLLARS! To claim your prize, simply click here and provide your personal information. This offer expires in 24 hours!", 1),
+        ("Subject: Make $5000 per week from HOME - GUARANTEED!\n\nJoin thousands of people who are making serious money online! Our proven system will help you receive payments directly to your account! No experience needed! Free training included! Order our business package today!", 1),
+        ("Subject: REMOVE BAD CREDIT - GUARANTEED RESULTS!\n\nOur credit repair service will remove negative items from your credit report guaranteed! People with bad credit can now get approved for loans and credit cards! Don't wait - order our service today!", 1),
+        ("Subject: FREE Prescription Drugs - HUGE SAVINGS!\n\nSave up to 80% on all prescription medications! No prescription required! Order online and receive free shipping! Thousands of satisfied customers worldwide! Viagra, Cialis, and more available now!", 1),
+        ("Subject: CLICK HERE for FREE ADULT CONTENT!!!\n\nHot singles in your area are waiting to meet you! Click here for free access to thousands of adult videos and photos! No credit card required! Join millions of satisfied members today!", 1),
+    ]
+    
+    # Modern phishing and promotional patterns (2024 style)
+    task_logger.info("Creating modern promotional patterns...")
+    modern_promotional = [
+        ("Subject: 🚨 Black Friday Preview - 80% OFF EVERYTHING\n\nGet exclusive early access to our Black Friday deals! Over 50,000 items at up to 80% off. Plus free shipping on all orders. This preview is only available to our VIP subscribers. Shop now before deals expire! Unsubscribe here.", 1),
+        ("Subject: Your Amazon Order Needs Action - Verify Now\n\nThere's an issue with your recent Amazon order #AMZ-12345. Your payment method was declined and your order will be cancelled unless you update your payment information within 24 hours. Click here to verify your account and complete your order.", 1),
+        ("Subject: Netflix - Your Account Will Be Suspended\n\nWe're having trouble with your current billing information. To keep your Netflix membership active, please update your payment details within 48 hours. Click here to update your account and continue enjoying Netflix.", 1),
+        ("Subject: Apple ID Security Alert - Suspicious Activity Detected\n\nWe've detected unusual activity on your Apple ID account. For your security, we've temporarily disabled your account. To restore access, please verify your identity by clicking the link below and confirming your information.", 1),
+        ("Subject: PayPal - Action Required on Your Account\n\nWe've noticed some unusual activity on your PayPal account. To protect your account, we've temporarily limited your access. Please log in and verify your identity to restore full account functionality.", 1),
+        ("Subject: Microsoft Office - Your Subscription is Expiring\n\nYour Microsoft Office subscription expires tomorrow. To continue using Word, Excel, PowerPoint and other Office apps, please renew your subscription. Click here to renew and save 50% on your next year.", 1),
+    ]
+    
+    # Security and important email patterns (realistic)
+    task_logger.info("Creating important email patterns...")
+    important_patterns = [
+        ("Subject: Password Reset Request for Your Account\n\nWe received a request to reset the password for your account. If you made this request, click the link below to reset your password. If you didn't request this, please ignore this email and your password will remain unchanged.", 0),
+        ("Subject: Your Order Has Shipped - Tracking Information\n\nGood news! Your order #12345 has been shipped and is on its way to you. You can track your package using the tracking number 1Z999AA1234567890. Expected delivery date is May 25-27, 2025.", 0),
+        ("Subject: Meeting Reminder - Project Review Tomorrow\n\nThis is a reminder that we have our project review meeting scheduled for tomorrow, May 23rd at 2:00 PM in Conference Room B. Please bring your project updates and quarterly reports. Contact me if you can't attend.", 0),
+        ("Subject: Bank Statement Ready - May 2025\n\nYour monthly bank statement for May 2025 is now available in your online banking portal. Please review your transactions and contact us if you notice any discrepancies. Thank you for banking with us.", 0),
+        ("Subject: Appointment Confirmation - Dr. Johnson\n\nThis confirms your appointment with Dr. Johnson on Friday, May 24th at 10:30 AM. Please arrive 15 minutes early and bring your insurance card and a valid ID. Call us if you need to reschedule.", 0),
+        ("Subject: Flight Confirmation - Your Trip Details\n\nYour flight is confirmed! Flight AA1234 from Chicago to New York on May 25th at 8:30 AM. Please arrive at the airport 2 hours before departure. Check-in is now available online or through our mobile app.", 0),
+        ("Subject: Invoice #INV-2025-5678 - Payment Due\n\nYour invoice #INV-2025-5678 for $250.00 is due on May 30th. You can pay online through our customer portal or by mailing a check to our office. Please contact us if you have any questions about this invoice.", 0),
+    ]
+    
+    # Combine all examples
+    accessible_examples.extend(uci_examples)
+    accessible_examples.extend(modern_promotional)  
+    accessible_examples.extend(important_patterns)
+    
+    task_logger.info(f"Created {len(accessible_examples)} examples from accessible datasets")
+    return accessible_examples
+
+def ml_suite_process_public_datasets(task_logger: AiTaskLogger) -> bool:
+    """
+    Main function to process public datasets with robust fallback and accessible downloads.
+    """
+    task_logger.info("Starting enhanced public dataset preparation with accessible sources")
+    
+    try:
+        # Ensure directories exist
+        utils.ensure_directory_exists(config.PROCESSED_DATASETS_DIR)
+        
+        all_training_examples = []
+        successful_datasets = 0
+        
+        # First, try to get accessible datasets
+        task_logger.info("Attempting to use accessible public datasets...")
+        try:
+            accessible_examples = download_accessible_datasets(task_logger)
+            all_training_examples.extend(accessible_examples)
+            successful_datasets += 1
+            task_logger.info(f"Successfully loaded {len(accessible_examples)} examples from accessible sources")
+        except Exception as e:
+            task_logger.warning(f"Failed to load accessible datasets: {e}")
+        
+        # Try to process external datasets (SpamAssassin, etc.)
+        for dataset_key, dataset_info in config.PUBLIC_DATASETS_INFO.items():
+            task_logger.info(f"Processing external dataset {dataset_key}")
+            
+            try:
+                success, extracted_dir = download_and_extract_dataset(
+                    dataset_key, dataset_info, task_logger
+                )
+                
+                if success:
+                    # Process emails from extracted dataset
+                    email_files = []
+                    for root, dirs, files in os.walk(extracted_dir):
+                        for file in files:
+                            if not file.startswith('.'):
+                                email_files.append(os.path.join(root, file))
+                    
+                    processed_from_dataset = 0
+                    for email_file in email_files[:config.MAX_SAMPLES_PER_RAW_DATASET]:
+                        try:
+                            with open(email_file, 'r', encoding='utf-8', errors='ignore') as f:
+                                email_content = f.read()
+                            
+                            processed = process_email_content(
+                                email_content, 
+                                dataset_info.get('type', 'important_leaning')
+                            )
+                            
+                            if processed:
+                                all_training_examples.append(processed)
+                                processed_from_dataset += 1
+                                
+                        except Exception as e:
+                            continue
+                    
+                    task_logger.info(f"Processed {processed_from_dataset} emails from {dataset_key}")
+                    if processed_from_dataset > 0:
+                        successful_datasets += 1
+                
+            except Exception as e:
+                task_logger.error(f"Failed to process external dataset {dataset_key}: {str(e)}")
+                continue
+        
+        # If we still don't have enough data, use comprehensive fallback
+        if len(all_training_examples) < 50:
+            task_logger.warning("Insufficient data from external sources, using comprehensive fallback dataset")
+            fallback_examples = create_comprehensive_fallback_dataset(task_logger)
+            all_training_examples.extend(fallback_examples)
+        
+        # Shuffle and balance the dataset
+        random.shuffle(all_training_examples)
+        
+        # Write to CSV
+        with open(config.PREPARED_DATA_FILE, 'w', newline='', encoding='utf-8') as csvfile:
+            writer = csv.writer(csvfile)
+            writer.writerow(['text', 'label'])
+            for text, label in all_training_examples:
+                writer.writerow([text, label])
+        
+        # Log final statistics
+        label_counts = Counter(label for _, label in all_training_examples)
+        task_logger.info(f"Dataset preparation completed successfully!")
+        task_logger.info(f"Total samples: {len(all_training_examples)}")
+        task_logger.info(f"Unsubscribable (1): {label_counts.get(1, 0)}")
+        task_logger.info(f"Important (0): {label_counts.get(0, 0)}")
+        task_logger.info(f"Successful external datasets: {successful_datasets}")
+        
+        return True
+        
+    except Exception as e:
+        task_logger.error(f"Critical error in dataset preparation: {str(e)}")
+        
+        # Last resort fallback
+        try:
+            task_logger.info("Attempting emergency fallback dataset creation...")
+            fallback_examples = create_comprehensive_fallback_dataset(task_logger)
+            
+            with open(config.PREPARED_DATA_FILE, 'w', newline='', encoding='utf-8') as csvfile:
+                writer = csv.writer(csvfile)
+                writer.writerow(['text', 'label'])
+                for text, label in fallback_examples:
+                    writer.writerow([text, label])
+            
+            task_logger.info("Emergency fallback dataset created successfully")
+            return True
+            
+        except Exception as fallback_error:
+            task_logger.error(f"Emergency fallback also failed: {str(fallback_error)}")
+            return False
+
+
+# Export the main function for use by the task system
+def prepare_training_data_from_public_datasets(task_logger: AiTaskLogger) -> bool:
+    """
+    Main entry point for data preparation with comprehensive error handling.
+    """
+    return ml_suite_process_public_datasets(task_logger)

ml_suite/model_trainer.py

@@ -0,0 +1,498 @@
+"""
+Model trainer module for the Gmail Unsubscriber AI suite.
+
+This module is responsible for:
+- Loading prepared email data
+- Splitting data into training and evaluation sets
+- Loading pre-trained transformer model
+- Fine-tuning the model on the email dataset
+- Evaluating model performance
+- Saving the fine-tuned model for prediction
+
+The trained model is optimized for classifying emails as "important" or "unsubscribable".
+"""
+
+import os
+import time
+import pandas as pd
+import numpy as np
+import torch
+from typing import Dict, List, Tuple, Optional, Any, Union, Callable
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support, roc_auc_score, confusion_matrix
+from sklearn.utils.class_weight import compute_class_weight
+
+# Hugging Face imports
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForSequenceClassification,
+    Trainer, 
+    TrainingArguments,
+    EarlyStoppingCallback,
+    IntervalStrategy,
+    PreTrainedTokenizer,
+    PreTrainedModel
+)
+from datasets import Dataset
+import torch.nn as nn
+
+# Local imports
+from . import config
+from . import utils
+from .task_utils import AiTaskLogger
+
+
+def compute_metrics_for_classification(eval_pred) -> Dict[str, float]:
+    """
+    Compute evaluation metrics for the model.
+    
+    Args:
+        eval_pred: Tuple of predictions and labels from the trainer
+        
+    Returns:
+        Dictionary of metrics including accuracy, precision, recall, F1, etc.
+    """
+    predictions, labels = eval_pred
+    
+    # For classification, take the argmax to get predicted classes
+    preds = np.argmax(predictions, axis=1)
+    
+    # Calculate overall metrics
+    precision, recall, f1, _ = precision_recall_fscore_support(
+        labels, preds, average='weighted'
+    )
+    acc = accuracy_score(labels, preds)
+    
+    # Calculate metrics specific to the "unsubscribable" class
+    # Get label positions: 0 for important, 1 for unsubscribable
+    unsub_class_idx = config.LABEL_UNSUBSCRIBABLE_ID
+    
+    # Calculate class-specific precision, recall, F1
+    precision_unsub, recall_unsub, f1_unsub, _ = precision_recall_fscore_support(
+        labels, preds, average=None, labels=[unsub_class_idx]
+    )
+    
+    # Compute ROC AUC if possible (requires probability scores)
+    try:
+        # Use the probability score for the unsubscribable class
+        probs_unsub = predictions[:, unsub_class_idx]
+        
+        # Convert labels to binary (1 for unsubscribable, 0 for others)
+        binary_labels = (labels == unsub_class_idx).astype(int)
+        
+        # Calculate ROC AUC
+        roc_auc = roc_auc_score(binary_labels, probs_unsub)
+    except (ValueError, IndexError):
+        roc_auc = 0.0
+    
+    # Return all metrics in a dictionary
+    metrics = {
+        'accuracy': acc,
+        'precision': precision,
+        'recall': recall,
+        'f1': f1,
+        'precision_unsub': float(precision_unsub[0]) if len(precision_unsub) > 0 else 0.0,
+        'recall_unsub': float(recall_unsub[0]) if len(recall_unsub) > 0 else 0.0,
+        'f1_unsub': float(f1_unsub[0]) if len(f1_unsub) > 0 else 0.0,
+        'roc_auc': roc_auc
+    }
+    
+    return metrics
+
+
+class WeightedLossTrainer(Trainer):
+    """Custom trainer that supports class weights for imbalanced datasets."""
+    
+    def __init__(self, class_weights=None, label_smoothing=0.0, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.class_weights = class_weights
+        self.label_smoothing = label_smoothing
+    
+    def compute_loss(self, model, inputs, return_outputs=False):
+        """Compute loss with class weights and optional label smoothing."""
+        labels = inputs.pop("labels")
+        outputs = model(**inputs)
+        logits = outputs.get('logits')
+        
+        # Create loss function with class weights
+        if self.class_weights is not None:
+            weight = torch.tensor(self.class_weights, dtype=torch.float32).to(logits.device)
+            loss_fct = nn.CrossEntropyLoss(weight=weight, label_smoothing=self.label_smoothing)
+        else:
+            loss_fct = nn.CrossEntropyLoss(label_smoothing=self.label_smoothing)
+        
+        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
+        
+        return (loss, outputs) if return_outputs else loss
+
+
+def tokenize_dataset(
+    examples: Dict[str, List], 
+    tokenizer: PreTrainedTokenizer, 
+    max_length: int
+) -> Dict[str, List]:
+    """
+    Tokenize a batch of examples.
+    
+    Args:
+        examples: Dictionary of example lists
+        tokenizer: Hugging Face tokenizer
+        max_length: Maximum sequence length
+        
+    Returns:
+        Dictionary of tokenized examples
+    """
+    return tokenizer(
+        examples['text'],
+        padding='max_length',
+        truncation=True,
+        max_length=max_length,
+        return_tensors='pt'
+    )
+
+
+def prepare_datasets_for_training(
+    data_file: str, 
+    test_size: float, 
+    tokenizer: PreTrainedTokenizer, 
+    max_length: int,
+    task_logger: AiTaskLogger,
+    compute_weights: bool = True
+) -> Tuple[Dataset, Dataset, Optional[np.ndarray]]:
+    """
+    Load and prepare datasets for training.
+    
+    Args:
+        data_file: Path to the prepared data file (CSV)
+        test_size: Proportion of data to use for evaluation
+        tokenizer: Hugging Face tokenizer
+        max_length: Maximum sequence length
+        task_logger: Logger for tracking progress
+        compute_weights: Whether to compute class weights for imbalanced data
+        
+    Returns:
+        Tuple of (train_dataset, eval_dataset, class_weights)
+    """
+    task_logger.info(f"Loading data from {data_file}")
+    
+    try:
+        # Load the dataset
+        df = pd.read_csv(data_file)
+        task_logger.info(f"Loaded {len(df)} examples from {data_file}")
+        
+        # Check for required columns
+        if 'text' not in df.columns or 'label' not in df.columns:
+            task_logger.error("Data file is missing required columns 'text' and/or 'label'")
+            raise ValueError("Data file has wrong format")
+        
+        # Split into training and evaluation sets
+        train_df, eval_df = train_test_split(
+            df, test_size=test_size, stratify=df['label'], random_state=42
+        )
+        
+        task_logger.info(f"Split into {len(train_df)} training examples and {len(eval_df)} evaluation examples")
+        
+        # Compute class weights for imbalanced data
+        class_weights = None
+        if compute_weights:
+            class_weights = compute_class_weight(
+                'balanced',
+                classes=np.array([0, 1]),
+                y=train_df['label'].values
+            )
+            task_logger.info(f"Computed class weights: {class_weights}")
+        
+        # Create HF datasets
+        train_dataset = Dataset.from_pandas(train_df)
+        eval_dataset = Dataset.from_pandas(eval_df)
+        
+        # Create a tokenization function that uses our tokenizer
+        def tokenize_function(examples):
+            return tokenize_dataset(examples, tokenizer, max_length)
+        
+        # Apply tokenization
+        task_logger.info("Tokenizing datasets")
+        train_dataset = train_dataset.map(
+            tokenize_function, 
+            batched=True,
+            desc="Tokenizing training dataset"
+        )
+        eval_dataset = eval_dataset.map(
+            tokenize_function, 
+            batched=True,
+            desc="Tokenizing evaluation dataset"
+        )
+        
+        # Set format for PyTorch
+        train_dataset.set_format(
+            type='torch', 
+            columns=['input_ids', 'attention_mask', 'label']
+        )
+        eval_dataset.set_format(
+            type='torch', 
+            columns=['input_ids', 'attention_mask', 'label']
+        )
+        
+        task_logger.info("Datasets prepared successfully")
+        
+        return train_dataset, eval_dataset, class_weights
+    
+    except Exception as e:
+        task_logger.error(f"Error preparing datasets: {str(e)}", e)
+        raise
+
+
+def train_unsubscriber_model(task_logger: AiTaskLogger) -> Dict[str, Any]:
+    """
+    Train the unsubscriber model on the prepared dataset.
+    
+    This function:
+    1. Loads prepared data
+    2. Initializes a pre-trained transformer model
+    3. Fine-tunes it on the email classification task
+    4. Evaluates performance
+    5. Saves the model for later use
+    
+    Args:
+        task_logger: Logger for tracking task status
+        
+    Returns:
+        Dictionary with training results and metrics
+    """
+    # Start timing
+    start_time = time.time()
+    
+    # Start the task
+    task_logger.start_task("Starting AI model training")
+    
+    try:
+        # Check if prepared data exists
+        if not os.path.exists(config.PREPARED_DATA_FILE):
+            task_logger.error(f"Prepared data file not found at {config.PREPARED_DATA_FILE}.")
+            task_logger.fail_task("Training failed: No prepared data available. Please run data preparation first.")
+            return {"success": False, "error": "No prepared data available"}
+        
+        # 1. Load and initialize the tokenizer
+        task_logger.info(f"Loading tokenizer for model: {config.PRE_TRAINED_MODEL_NAME}")
+        tokenizer = AutoTokenizer.from_pretrained(config.PRE_TRAINED_MODEL_NAME)
+        
+        # 2. Prepare datasets
+        task_logger.update_progress(0.1, "Preparing datasets")
+        train_dataset, eval_dataset, class_weights = prepare_datasets_for_training(
+            config.PREPARED_DATA_FILE,
+            config.TEST_SPLIT_SIZE,
+            tokenizer,
+            config.MAX_SEQ_LENGTH,
+            task_logger,
+            compute_weights=True
+        )
+        
+        # 3. Initialize model
+        task_logger.update_progress(0.2, f"Initializing model: {config.PRE_TRAINED_MODEL_NAME}")
+        model = AutoModelForSequenceClassification.from_pretrained(
+            config.PRE_TRAINED_MODEL_NAME,
+            num_labels=config.NUM_LABELS
+        )
+        
+        # Check device availability
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        fp16_enabled = config.FP16_TRAINING and torch.cuda.is_available()
+        
+        # Force GPU check and provide detailed info
+        if torch.cuda.is_available():
+            gpu_name = torch.cuda.get_device_name(0)
+            gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)  # Convert to GB
+            task_logger.info(f"GPU detected: {gpu_name} with {gpu_memory:.2f} GB memory")
+            task_logger.info(f"CUDA version: {torch.version.cuda}")
+            
+            # Set CUDA device explicitly
+            torch.cuda.set_device(0)
+            
+            # Enable cuDNN benchmarking for better performance
+            torch.backends.cudnn.benchmark = True
+            torch.backends.cudnn.enabled = True
+        else:
+            task_logger.warning("No GPU detected! Training will be slow on CPU.")
+            task_logger.warning("Make sure you have the CUDA version of PyTorch installed.")
+            task_logger.warning("To install PyTorch with CUDA support, run:")
+            task_logger.warning("pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118")
+        
+        task_logger.info(f"Training on device: {device}")
+        if fp16_enabled:
+            task_logger.info("FP16 mixed precision training enabled")
+        
+        # Move model to device
+        model.to(device)
+        
+        # 4. Set up training arguments
+        task_logger.update_progress(0.3, "Setting up training configuration")
+        
+        # Calculate number of steps
+        num_train_examples = len(train_dataset)
+        num_train_steps = (num_train_examples // config.TRAIN_BATCH_SIZE) * config.NUM_TRAIN_EPOCHS
+        warmup_steps = int(num_train_steps * config.WARMUP_STEPS_RATIO)
+        
+        # Convert evaluation and save strategies to enum values
+        eval_strategy = (IntervalStrategy.EPOCH 
+                         if config.EVALUATION_STRATEGY == "epoch" 
+                         else IntervalStrategy.STEPS)
+        save_strategy = (IntervalStrategy.EPOCH 
+                         if config.SAVE_STRATEGY == "epoch" 
+                         else IntervalStrategy.STEPS)
+        
+        # Create output directory if it doesn't exist
+        os.makedirs(config.FINE_TUNED_MODEL_DIR, exist_ok=True)
+        
+        # Define training arguments (with compatibility fixes for older transformers versions)
+        training_args_dict = {
+            "output_dir": config.FINE_TUNED_MODEL_DIR,
+            "per_device_train_batch_size": config.TRAIN_BATCH_SIZE,
+            "per_device_eval_batch_size": config.EVAL_BATCH_SIZE,
+            "num_train_epochs": config.NUM_TRAIN_EPOCHS,
+            "learning_rate": config.LEARNING_RATE,
+            "weight_decay": config.WEIGHT_DECAY,
+            "warmup_steps": warmup_steps,
+            "logging_steps": 50,
+            "save_total_limit": 2,
+            "load_best_model_at_end": config.LOAD_BEST_MODEL_AT_END,
+            "metric_for_best_model": config.METRIC_FOR_BEST_MODEL,
+            "greater_is_better": True,
+            "report_to": "none",
+            "disable_tqdm": False,
+            "no_cuda": False,  # Explicitly enable CUDA
+            "use_cpu": False,  # Explicitly disable CPU-only mode
+        }
+        
+        # Add evaluation and save strategies
+        try:
+            training_args_dict["evaluation_strategy"] = eval_strategy
+            training_args_dict["save_strategy"] = save_strategy
+        except:
+            # Fallback for older versions
+            training_args_dict["evaluation_strategy"] = "epoch"
+            training_args_dict["save_strategy"] = "epoch"
+        
+        # Add logging dir if supported
+        try:
+            training_args_dict["logging_dir"] = os.path.join(config.FINE_TUNED_MODEL_DIR, "logs")
+        except:
+            pass
+        
+        # Add FP16 if supported and available
+        if fp16_enabled:
+            try:
+                training_args_dict["fp16"] = True
+            except:
+                task_logger.warning("FP16 training not supported on this setup")
+        
+        # Add dataloader workers if supported
+        try:
+            training_args_dict["dataloader_num_workers"] = 2
+        except:
+            pass
+        
+        # Add gradient accumulation if configured
+        if hasattr(config, 'GRADIENT_ACCUMULATION_STEPS'):
+            training_args_dict["gradient_accumulation_steps"] = config.GRADIENT_ACCUMULATION_STEPS
+        
+        # Add label smoothing if configured
+        if hasattr(config, 'LABEL_SMOOTHING_FACTOR'):
+            training_args_dict["label_smoothing_factor"] = config.LABEL_SMOOTHING_FACTOR
+        
+        training_args = TrainingArguments(**training_args_dict)
+        
+        # 5. Initialize the trainer with class weights
+        task_logger.update_progress(0.4, "Initializing trainer with class balancing")
+        
+        # Use weighted loss trainer if we have class weights
+        if class_weights is not None:
+            trainer = WeightedLossTrainer(
+                model=model,
+                args=training_args,
+                train_dataset=train_dataset,
+                eval_dataset=eval_dataset,
+                compute_metrics=compute_metrics_for_classification,
+                class_weights=class_weights,
+                label_smoothing=config.LABEL_SMOOTHING_FACTOR if hasattr(config, 'LABEL_SMOOTHING_FACTOR') else 0.0,
+                callbacks=[
+                    EarlyStoppingCallback(
+                        early_stopping_patience=config.EARLY_STOPPING_PATIENCE,
+                        early_stopping_threshold=config.EARLY_STOPPING_THRESHOLD
+                    )
+                ]
+            )
+            task_logger.info("Using weighted loss function for class imbalance")
+        else:
+            trainer = Trainer(
+                model=model,
+                args=training_args,
+                train_dataset=train_dataset,
+                eval_dataset=eval_dataset,
+                compute_metrics=compute_metrics_for_classification,
+                callbacks=[
+                    EarlyStoppingCallback(
+                        early_stopping_patience=config.EARLY_STOPPING_PATIENCE,
+                        early_stopping_threshold=config.EARLY_STOPPING_THRESHOLD
+                    )
+                ]
+            )
+        
+        # 6. Train the model
+        task_logger.update_progress(0.5, "Starting model training")
+        trainer.train()
+        
+        # 7. Evaluate the model
+        task_logger.update_progress(0.9, "Evaluating model")
+        eval_results = trainer.evaluate()
+        
+        # Format evaluations nicely for logging
+        metrics_str = "\n".join([f"  {k}: {v:.4f}" for k, v in eval_results.items()])
+        task_logger.info(f"Evaluation results:\n{metrics_str}")
+        
+        # 8. Save the final model
+        task_logger.update_progress(0.95, "Saving fine-tuned model")
+        trainer.save_model(config.FINE_TUNED_MODEL_DIR)
+        tokenizer.save_pretrained(config.FINE_TUNED_MODEL_DIR)
+        
+        # Save a human-readable summary of model info
+        with open(os.path.join(config.FINE_TUNED_MODEL_DIR, "model_info.txt"), "w") as f:
+            f.write(f"Base model: {config.PRE_TRAINED_MODEL_NAME}\n")
+            f.write(f"Training completed: {utils.get_current_timestamp()}\n")
+            f.write(f"Training examples: {len(train_dataset)}\n")
+            f.write(f"Evaluation examples: {len(eval_dataset)}\n")
+            f.write(f"Evaluation metrics:\n")
+            for k, v in eval_results.items():
+                f.write(f"  {k}: {v:.4f}\n")
+        
+        # Calculate elapsed time
+        elapsed_time = time.time() - start_time
+        hours, remainder = divmod(elapsed_time, 3600)
+        minutes, seconds = divmod(remainder, 60)
+        time_str = f"{int(hours)}h {int(minutes)}m {int(seconds)}s"
+        
+        # Complete the task
+        result = {
+            "success": True,
+            "metrics": eval_results,
+            "model_dir": config.FINE_TUNED_MODEL_DIR,
+            "training_time": time_str,
+            "base_model": config.PRE_TRAINED_MODEL_NAME,
+            "num_train_examples": len(train_dataset),
+            "num_eval_examples": len(eval_dataset)
+        }
+        
+        task_logger.complete_task(
+            f"Model training completed successfully in {time_str}", 
+            result
+        )
+        
+        return result
+    
+    except Exception as e:
+        task_logger.error("Error during model training", e)
+        task_logger.fail_task(f"Model training failed: {str(e)}")
+        
+        return {
+            "success": False,
+            "error": str(e)
+        }

ml_suite/models/base_transformer_cache/version.txt

@@ -0,0 +1 @@
+1

ml_suite/models/fine_tuned_unsubscriber/config.json

@@ -0,0 +1,24 @@
+{
+  "_name_or_path": "distilbert-base-uncased",
+  "activation": "gelu",
+  "architectures": [
+    "DistilBertForSequenceClassification"
+  ],
+  "attention_dropout": 0.1,
+  "dim": 768,
+  "dropout": 0.1,
+  "hidden_dim": 3072,
+  "initializer_range": 0.02,
+  "max_position_embeddings": 512,
+  "model_type": "distilbert",
+  "n_heads": 12,
+  "n_layers": 6,
+  "pad_token_id": 0,
+  "qa_dropout": 0.1,
+  "seq_classif_dropout": 0.2,
+  "sinusoidal_pos_embds": false,
+  "tie_weights_": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.36.2",
+  "vocab_size": 30522
+}

ml_suite/models/fine_tuned_unsubscriber/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4bdf7a6a37778c581baa94c0a6c462836894e5cc009a3e077e96e1c0b9ec176e
+size 267832560

ml_suite/models/fine_tuned_unsubscriber/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

ml_suite/models/fine_tuned_unsubscriber/tokenizer_config.json

@@ -0,0 +1,55 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "DistilBertTokenizer",
+  "unk_token": "[UNK]"
+}

ml_suite/models/fine_tuned_unsubscriber/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f8e87068f27f96f378c34f96ec16c9a676bf0ae2c9a777803f6c459f8524ce9
+size 4664

ml_suite/predictor.py

@@ -0,0 +1,445 @@
+"""
+Predictor module for the Gmail Unsubscriber AI suite.
+
+This module is responsible for:
+- Loading the fine-tuned transformer model (base or personalized)
+- Creating an efficient inference pipeline
+- Providing real-time predictions for email classification
+- Handling model loading errors gracefully
+- Supporting personalized model selection based on user ID
+
+The predictor provides a simple interface for the main application to classify
+emails as "important" or "unsubscribable" with confidence scores, with the
+ability to use personalized models when available.
+"""
+
+import os
+import time
+import logging
+import torch
+from typing import Dict, List, Optional, Any, Union
+import traceback
+
+# Hugging Face imports
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    TextClassificationPipeline,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    AutoConfig
+)
+
+# Local imports
+from . import config
+from . import utils
+
+
+# Global variables to hold the loaded models and tokenizers
+# Main base model pipeline
+base_classification_pipeline = None
+base_model_load_status = "Not Loaded"
+base_model_load_error = None
+base_model_last_load_attempt = 0
+
+# Dictionary to store personalized model pipelines for different users
+personalized_pipelines = {}
+personalized_load_status = {}
+
+# Configuration
+load_cooldown_seconds = 60  # Wait at least this long between load attempts
+
+
+def is_predictor_ready() -> bool:
+    """
+    Check if the base predictor is ready for use.
+    
+    Returns:
+        True if the base predictor is ready, False otherwise
+    """
+    return base_model_load_status == "Ready" and base_classification_pipeline is not None
+
+
+def get_model_info() -> Dict[str, Any]:
+    """
+    Get information about the loaded model.
+    
+    Returns:
+        Dictionary with model information
+    """
+    # Try to load model info from the model directory
+    model_info_path = os.path.join(config.FINE_TUNED_MODEL_DIR, "model_info.txt")
+    model_info = {
+        "model_type": "base",
+        "model_path": config.FINE_TUNED_MODEL_DIR,
+        "trained_date": "unknown"
+    }
+    
+    if os.path.exists(model_info_path):
+        try:
+            with open(model_info_path, 'r') as f:
+                lines = f.readlines()
+                for line in lines:
+                    if ":" in line:
+                        key, value = line.strip().split(":", 1)
+                        key = key.strip().lower().replace(" ", "_")
+                        value = value.strip()
+                        model_info[key] = value
+        except:
+            pass
+    
+    return model_info
+
+
+def initialize_predictor(app_logger: logging.Logger) -> bool:
+    """
+    Initialize the base predictor by loading the fine-tuned model.
+    
+    This function:
+    1. Loads the tokenizer and model from the fine-tuned directory
+    2. Creates a TextClassificationPipeline for efficient inference
+    3. Sets global variables for status tracking
+    
+    Args:
+        app_logger: Application logger for status and error reporting
+        
+    Returns:
+        True if initialization successful, False otherwise
+    """
+    global base_classification_pipeline, base_model_load_status, base_model_load_error, base_model_last_load_attempt
+    
+    # Reset error tracking
+    base_model_load_error = None
+    
+    # Check if we attempted to load recently (to prevent repeated failures)
+    current_time = time.time()
+    if (current_time - base_model_last_load_attempt) < load_cooldown_seconds and base_model_load_status == "Failed":
+        app_logger.warning(
+            f"Not attempting to reload model - cooling down after recent failure. "
+            f"Will retry after {load_cooldown_seconds - (current_time - base_model_last_load_attempt):.0f} seconds."
+        )
+        return False
+    
+    base_model_last_load_attempt = current_time
+    
+    try:
+        # Update status
+        base_model_load_status = "Loading"
+        app_logger.info(f"Initializing base AI predictor from {config.FINE_TUNED_MODEL_DIR}")
+        
+        # Check if model directory exists and contains necessary files
+        if not os.path.exists(config.FINE_TUNED_MODEL_DIR):
+            raise FileNotFoundError(f"Model directory not found: {config.FINE_TUNED_MODEL_DIR}")
+        
+        # Load the tokenizer
+        app_logger.info("Loading tokenizer")
+        tokenizer = AutoTokenizer.from_pretrained(
+            config.FINE_TUNED_MODEL_DIR,
+            local_files_only=True  # Ensure we load locally without trying to download
+        )
+        
+        # Load model configuration
+        model_config = AutoConfig.from_pretrained(
+            config.FINE_TUNED_MODEL_DIR,
+            local_files_only=True
+        )
+        
+        # Check if the model has the expected number of labels
+        if model_config.num_labels != config.NUM_LABELS:
+            app_logger.warning(
+                f"Model has {model_config.num_labels} labels, "
+                f"but config specifies {config.NUM_LABELS} labels. "
+                f"This may cause issues with predictions."
+            )
+        
+        # Load the model
+        app_logger.info("Loading model")
+        model = AutoModelForSequenceClassification.from_pretrained(
+            config.FINE_TUNED_MODEL_DIR,
+            local_files_only=True  # Ensure we load locally without trying to download
+        )
+        
+        # Determine device
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        app_logger.info(f"Using device: {device}")
+        
+        # Move model to the device
+        model.to(device)
+        
+        # Create classification pipeline
+        app_logger.info("Creating inference pipeline")
+        
+        # Set device index properly
+        if device.type == "cuda":
+            device_index = 0  # Use first GPU
+            app_logger.info(f"Pipeline will use GPU device index: {device_index}")
+        else:
+            device_index = -1  # CPU
+            app_logger.info("Pipeline will use CPU")
+            
+        base_classification_pipeline = TextClassificationPipeline(
+            model=model,
+            tokenizer=tokenizer,
+            device=device_index,
+            top_k=None,  # Get probabilities for all classes (replaces deprecated return_all_scores)
+            function_to_apply="sigmoid"  # Apply sigmoid to logits for probability interpretation
+        )
+        
+        # Update status
+        base_model_load_status = "Ready"
+        app_logger.info("Base AI predictor initialized successfully")
+        
+        return True
+    
+    except Exception as e:
+        # Handle initialization failure
+        base_model_load_status = "Failed"
+        base_model_load_error = str(e)
+        error_traceback = traceback.format_exc()
+        
+        app_logger.error(f"Error initializing base AI predictor: {str(e)}")
+        app_logger.debug(f"Traceback:\n{error_traceback}")
+        
+        # Cleanup any partial loading
+        base_classification_pipeline = None
+        
+        return False
+
+
+def initialize_personalized_predictor(user_id: str, app_logger: logging.Logger) -> bool:
+    """
+    Initialize a personalized predictor for a specific user.
+    
+    Args:
+        user_id: The user ID for which to load the personalized model
+        app_logger: Application logger for status and error reporting
+        
+    Returns:
+        True if initialization successful, False otherwise
+    """
+    global personalized_pipelines, personalized_load_status
+    
+    try:
+        # Check if the user has a personalized model
+        model_dir = config.PERSONALIZED_MODEL_DIR_TEMPLATE.format(user_id=user_id)
+        if not os.path.exists(model_dir) or not os.path.exists(os.path.join(model_dir, "pytorch_model.bin")):
+            app_logger.info(f"No personalized model found for user {user_id}")
+            personalized_load_status[user_id] = "Not Available"
+            return False
+        
+        # Update status
+        personalized_load_status[user_id] = "Loading"
+        app_logger.info(f"Initializing personalized AI predictor for user {user_id} from {model_dir}")
+        
+        # Load the tokenizer
+        app_logger.info(f"Loading personalized tokenizer for user {user_id}")
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_dir,
+            local_files_only=True
+        )
+        
+        # Load model configuration
+        model_config = AutoConfig.from_pretrained(
+            model_dir,
+            local_files_only=True
+        )
+        
+        # Load the model
+        app_logger.info(f"Loading personalized model for user {user_id}")
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_dir,
+            local_files_only=True
+        )
+        
+        # Determine device
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # Move model to the device
+        model.to(device)
+        
+        # Create classification pipeline
+        app_logger.info(f"Creating personalized inference pipeline for user {user_id}")
+        
+        # Set device index properly
+        if device.type == "cuda":
+            device_index = 0  # Use first GPU
+        else:
+            device_index = -1  # CPU
+            
+        personalized_pipelines[user_id] = TextClassificationPipeline(
+            model=model,
+            tokenizer=tokenizer,
+            device=device_index,
+            top_k=None,  # Get probabilities for all classes (replaces deprecated return_all_scores)
+            function_to_apply="sigmoid"
+        )
+        
+        # Update status
+        personalized_load_status[user_id] = "Ready"
+        app_logger.info(f"Personalized AI predictor for user {user_id} initialized successfully")
+        
+        return True
+    
+    except Exception as e:
+        # Handle initialization failure
+        personalized_load_status[user_id] = "Failed"
+        error_traceback = traceback.format_exc()
+        
+        app_logger.error(f"Error initializing personalized AI predictor for user {user_id}: {str(e)}")
+        app_logger.debug(f"Traceback:\n{error_traceback}")
+        
+        # Cleanup any partial loading
+        if user_id in personalized_pipelines:
+            del personalized_pipelines[user_id]
+        
+        return False
+
+
+def get_model_status(user_id: Optional[str] = None) -> Dict[str, Any]:
+    """
+    Get the current status of the model.
+    
+    Args:
+        user_id: Optional user ID to get status for personalized model
+        
+    Returns:
+        Dictionary with status information
+    """
+    if user_id is None:
+        # Return base model status
+        return {
+            "status": base_model_load_status,
+            "error": base_model_load_error,
+            "last_load_attempt": base_model_last_load_attempt,
+            "model_dir": config.FINE_TUNED_MODEL_DIR,
+            "is_ready": base_model_load_status == "Ready" and base_classification_pipeline is not None,
+            "is_personalized": False
+        }
+    else:
+        # Check if personalized model is available
+        if user_id not in personalized_load_status:
+            personalized_load_status[user_id] = "Not Loaded"
+        
+        model_dir = config.PERSONALIZED_MODEL_DIR_TEMPLATE.format(user_id=user_id)
+        has_personalized = (
+            os.path.exists(model_dir) and 
+            os.path.exists(os.path.join(model_dir, "pytorch_model.bin"))
+        )
+        
+        return {
+            "status": personalized_load_status.get(user_id, "Not Loaded"),
+            "model_dir": model_dir if has_personalized else None,
+            "is_ready": personalized_load_status.get(user_id) == "Ready" and user_id in personalized_pipelines,
+            "is_personalized": True,
+            "has_personalized_model": has_personalized
+        }
+
+
+def get_ai_prediction_for_email(
+    email_text_content: str, 
+    user_id: Optional[str] = None,
+    app_logger: Optional[logging.Logger] = None
+) -> Optional[Dict[str, Any]]:
+    """
+    Get AI prediction for an email, optionally using a personalized model.
+    
+    This function:
+    1. Checks if the requested model is loaded and ready
+    2. Passes the email text to the appropriate classification pipeline
+    3. Processes and returns the prediction results
+    
+    Args:
+        email_text_content: The combined email text (subject + body)
+        user_id: Optional user ID to use personalized model if available
+        app_logger: Optional logger for error reporting
+        
+    Returns:
+        Dictionary with prediction results (label, confidence, etc.) or None if prediction fails
+    """
+    global base_classification_pipeline, personalized_pipelines
+    
+    # Determine which pipeline to use
+    pipeline = None
+    using_personalized = False
+    
+    if user_id is not None and user_id in personalized_pipelines:
+        # Try to use personalized model first
+        if personalized_load_status.get(user_id) == "Ready":
+            pipeline = personalized_pipelines[user_id]
+            using_personalized = True
+    
+    # Fall back to base model if personalized isn't available
+    if pipeline is None:
+        if base_model_load_status != "Ready" or base_classification_pipeline is None:
+            return None
+        pipeline = base_classification_pipeline
+    
+    try:
+        # Clean and normalize the input text
+        cleaned_text = utils.clean_text_for_model(
+            email_text_content, 
+            max_length=config.EMAIL_SNIPPET_LENGTH_FOR_MODEL
+        )
+        
+        # Skip prediction for extremely short text
+        if len(cleaned_text) < config.MIN_TEXT_LENGTH_FOR_TRAINING:
+            return {
+                "label": "INDETERMINATE",
+                "confidence": 0.0,
+                "predicted_id": None,
+                "error": "Text too short for reliable prediction",
+                "using_personalized_model": using_personalized
+            }
+        
+        # Get prediction
+        predictions = pipeline(cleaned_text)
+        
+        # Process prediction results
+        # The pipeline returns a list with one dictionary per input text
+        # Each dictionary has a 'label' and 'score' for each possible class
+        prediction_scores = {}
+        for pred in predictions[0]:  # Get the first (and only) prediction
+            label_id = int(pred['label'].split('_')[-1])
+            label_name = config.ID_TO_LABEL_MAP.get(label_id)
+            prediction_scores[label_name] = pred['score']
+        
+        # Find the highest scoring label
+        max_label = max(prediction_scores, key=prediction_scores.get)
+        max_score = prediction_scores[max_label]
+        predicted_id = config.LABEL_TO_ID_MAP.get(max_label)
+        
+        # Apply confidence calibration to prevent overconfidence
+        # Temperature scaling to soften extreme predictions
+        temperature = 1.5  # Higher = less confident
+        calibrated_score = max_score ** (1 / temperature)
+        
+        # Log prediction details for debugging
+        if app_logger:
+            app_logger.debug(f"AI Prediction: {max_label} (raw: {max_score:.3f}, calibrated: {calibrated_score:.3f})")
+            app_logger.debug(f"All scores: {prediction_scores}")
+            app_logger.debug(f"Email snippet: {cleaned_text[:100]}...")
+        
+        # Return the prediction
+        return {
+            "label": max_label,
+            "confidence": calibrated_score,
+            "raw_confidence": max_score,
+            "predicted_id": predicted_id,
+            "all_scores": prediction_scores,
+            "using_personalized_model": using_personalized
+        }
+    
+    except Exception as e:
+        # Log the error details but don't expose them in the response
+        if app_logger:
+            app_logger.error(f"Error during AI prediction: {str(e)}")
+        else:
+            print(f"Error during AI prediction: {str(e)}")
+        
+        return {
+            "label": "ERROR",
+            "confidence": 0.0,
+            "predicted_id": None,
+            "error": "Prediction error occurred",
+            "using_personalized_model": using_personalized
+        }

ml_suite/task_utils.py

@@ -0,0 +1,332 @@
+"""
+Utilities for task status management and logging in the ML suite.
+
+This module provides functions and classes for:
+- Managing task status files for long-running operations
+- Logging both to Flask app logs and status files
+- Standardized error handling for AI operations
+- Progress reporting for data preparation and model training
+
+These utilities ensure that AI operations have proper status tracking,
+enabling the frontend to display progress and results to users.
+"""
+
+import json
+import time
+import uuid
+import datetime
+import traceback
+import os
+import logging
+from typing import Dict, List, Optional, Union, Any, Tuple
+
+
+def get_current_timestamp() -> str:
+    """Returns ISO format timestamp for current time."""
+    return datetime.datetime.now().isoformat()
+
+
+def get_current_timestamp_log_prefix() -> str:
+    """Returns a formatted timestamp string for log entries."""
+    return f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
+
+
+def initialize_task_status_file(status_file_path: str, task_id: Optional[str] = None) -> Dict[str, Any]:
+    """Initialize a new task status file with default values.
+    
+    Args:
+        status_file_path: Path to the status file
+        task_id: Optional unique ID for the task (UUID generated if None)
+        
+    Returns:
+        Dict containing the initialized status data
+    """
+    if task_id is None:
+        task_id = str(uuid.uuid4())
+    
+    status_data = {
+        "status": "pending",
+        "message": "Task initialized and pending execution",
+        "progress": 0.0,
+        "log": [],
+        "task_id": task_id,
+        "start_time": get_current_timestamp()
+    }
+    
+    # Ensure the directory exists
+    dir_path = os.path.dirname(status_file_path)
+    if dir_path:  # Only create directory if dirname returns a non-empty string
+        os.makedirs(dir_path, exist_ok=True)
+    
+    with open(status_file_path, 'w') as f:
+        json.dump(status_data, f, indent=2)
+    
+    return status_data
+
+
+def update_task_status(
+    status_file_path: str,
+    status: Optional[str] = None,
+    message: Optional[str] = None,
+    log_entry: Optional[str] = None,
+    log_level: str = "info",
+    progress: Optional[float] = None,
+    result: Optional[Dict[str, Any]] = None,
+    error: Optional[Dict[str, Any]] = None
+) -> Dict[str, Any]:
+    """Update a task status file with new information.
+    
+    Args:
+        status_file_path: Path to the status file
+        status: Optional new status ('pending', 'in_progress', 'completed', 'failed')
+        message: Optional status message
+        log_entry: Optional log message to add to the log list
+        log_level: Log level ('info', 'warning', 'error', 'debug')
+        progress: Optional progress value (0.0 to 1.0)
+        result: Optional result data (for completed tasks)
+        error: Optional error data (for failed tasks)
+        
+    Returns:
+        Dict containing the updated status data
+    """
+    try:
+        with open(status_file_path, 'r') as f:
+            status_data = json.load(f)
+    except (FileNotFoundError, json.JSONDecodeError):
+        # If file doesn't exist or is invalid, initialize it
+        status_data = initialize_task_status_file(status_file_path)
+    
+    # Update fields if provided
+    if status is not None:
+        status_data["status"] = status
+        # If status is completed or failed, set end_time
+        if status in ("completed", "failed"):
+            status_data["end_time"] = get_current_timestamp()
+    
+    if message is not None:
+        status_data["message"] = message
+    
+    if log_entry is not None:
+        if "log" not in status_data:
+            status_data["log"] = []
+        status_data["log"].append({
+            "timestamp": get_current_timestamp(),
+            "level": log_level,
+            "message": log_entry
+        })
+    
+    if progress is not None:
+        status_data["progress"] = max(0.0, min(1.0, float(progress)))  # Ensure between 0-1
+    
+    if result is not None and status_data.get("status") == "completed":
+        status_data["result"] = result
+    
+    if error is not None and status_data.get("status") == "failed":
+        status_data["error"] = error
+    
+    # Write updated status back to file
+    with open(status_file_path, 'w') as f:
+        json.dump(status_data, f, indent=2)
+    
+    return status_data
+
+
+def get_task_status(status_file_path: str) -> Dict[str, Any]:
+    """Read and return the current task status.
+    
+    Args:
+        status_file_path: Path to the status file
+        
+    Returns:
+        Dict containing the current status data
+    """
+    try:
+        with open(status_file_path, 'r') as f:
+            return json.load(f)
+    except (FileNotFoundError, json.JSONDecodeError):
+        # If status file doesn't exist or is invalid, return a default status
+        return {
+            "status": "unknown",
+            "message": "Task status unknown or not initialized",
+            "progress": 0.0,
+            "log": []
+        }
+
+
+def log_task_error(status_file_path: str, error: Exception, message: str = "Task failed due to an error") -> Dict[str, Any]:
+    """Log an error to the task status file.
+    
+    Args:
+        status_file_path: Path to the status file
+        error: The exception that occurred
+        message: Human-readable error message
+        
+    Returns:
+        Dict containing the updated status data
+    """
+    error_info = {
+        "type": error.__class__.__name__,
+        "message": str(error),
+        "traceback": traceback.format_exc()
+    }
+    
+    return update_task_status(
+        status_file_path=status_file_path,
+        status="failed",
+        message=message,
+        log_entry=f"ERROR: {message} - {error}",
+        log_level="error",
+        error=error_info
+    )
+
+
+class AiTaskLogger:
+    """Logger for AI tasks that updates both the application logger and task status file.
+    
+    This logger provides methods for:
+    - Logging info, warning, and error messages
+    - Updating task progress
+    - Marking tasks as started, completed, or failed
+    - Ensuring consistent logging across both Flask app logs and task status files
+    """
+    
+    def __init__(self, 
+                 app_logger: logging.Logger, 
+                 status_file_path: str, 
+                 task_id: Optional[str] = None):
+        """Initialize the task logger.
+        
+        Args:
+            app_logger: Flask application logger
+            status_file_path: Path to the task status file
+            task_id: Optional unique ID for the task (UUID generated if None)
+        """
+        self.app_logger = app_logger
+        self.status_file_path = status_file_path
+        self.task_id = task_id or str(uuid.uuid4())
+        self.short_task_id = self.task_id[:8]  # First 8 chars for log readability
+        
+        # Initialize the status file
+        initialize_task_status_file(status_file_path, self.task_id)
+    
+    def info(self, message: str, update_progress: Optional[float] = None) -> None:
+        """Log an info message and optionally update progress.
+        
+        Args:
+            message: The message to log
+            update_progress: Optional progress value (0.0 to 1.0)
+        """
+        self.app_logger.info(f"[AI Task {self.short_task_id}] {message}")
+        update_task_status(
+            self.status_file_path,
+            log_entry=message,
+            log_level="info",
+            progress=update_progress
+        )
+    
+    def warning(self, message: str) -> None:
+        """Log a warning message.
+        
+        Args:
+            message: The warning message to log
+        """
+        self.app_logger.warning(f"[AI Task {self.short_task_id}] {message}")
+        update_task_status(
+            self.status_file_path,
+            log_entry=message,
+            log_level="warning"
+        )
+    
+    def error(self, message: str, error: Optional[Exception] = None) -> None:
+        """Log an error message and optionally the exception details.
+        
+        Args:
+            message: The error message to log
+            error: Optional exception that caused the error
+        """
+        if error:
+            self.app_logger.error(f"[AI Task {self.short_task_id}] {message}: {error}", exc_info=True)
+            log_task_error(self.status_file_path, error, message)
+        else:
+            self.app_logger.error(f"[AI Task {self.short_task_id}] {message}")
+            update_task_status(
+                self.status_file_path,
+                log_entry=message,
+                log_level="error"
+            )
+    
+    def start_task(self, message: str = "Task started") -> None:
+        """Mark the task as started.
+        
+        Args:
+            message: Optional message describing the task start
+        """
+        self.app_logger.info(f"[AI Task {self.short_task_id}] {message}")
+        update_task_status(
+            self.status_file_path,
+            status="in_progress",
+            message=message,
+            log_entry=message,
+            log_level="info",
+            progress=0.0
+        )
+    
+    def complete_task(self, message: str = "Task completed successfully", result: Optional[Dict[str, Any]] = None) -> None:
+        """Mark the task as completed.
+        
+        Args:
+            message: Optional completion message
+            result: Optional result data to store
+        """
+        self.app_logger.info(f"[AI Task {self.short_task_id}] {message}")
+        update_task_status(
+            self.status_file_path,
+            status="completed",
+            message=message,
+            log_entry=message,
+            log_level="info",
+            progress=1.0,
+            result=result
+        )
+    
+    def fail_task(self, message: str, error: Optional[Exception] = None) -> None:
+        """Mark the task as failed.
+        
+        Args:
+            message: Failure message
+            error: Optional exception that caused the failure
+        """
+        if error:
+            self.app_logger.error(f"[AI Task {self.short_task_id}] {message}: {error}", exc_info=True)
+            log_task_error(self.status_file_path, error, message)
+        else:
+            self.app_logger.error(f"[AI Task {self.short_task_id}] {message}")
+            update_task_status(
+                self.status_file_path,
+                status="failed",
+                message=message,
+                log_entry=message,
+                log_level="error"
+            )
+    
+    def update_progress(self, progress: float, message: Optional[str] = None) -> None:
+        """Update the task progress.
+        
+        Args:
+            progress: Progress value (0.0 to 1.0)
+            message: Optional progress message
+        """
+        if message:
+            self.app_logger.info(f"[AI Task {self.short_task_id}] {message} (Progress: {progress:.1%})")
+            update_task_status(
+                self.status_file_path,
+                message=message,
+                log_entry=message,
+                log_level="info",
+                progress=progress
+            )
+        else:
+            update_task_status(
+                self.status_file_path,
+                progress=progress
+            )

ml_suite/utils.py

@@ -0,0 +1,317 @@
+"""
+Shared utilities for the ML suite.
+
+This module provides shared functions used across the ML suite components:
+- Email text analysis using heuristics (adapted from app.py)
+- Text cleaning and normalization
+- Timestamp and logging utilities
+- HTML processing for email content extraction
+
+These utilities ensure consistent processing across different components of the ML suite.
+"""
+
+import re
+import os
+import urllib.parse
+import datetime
+import html
+from typing import Dict, List, Tuple, Optional, Any, Union
+from bs4 import BeautifulSoup
+
+
+# --- Email Heuristic Analysis ---
+
+# Keywords that suggest an email is marketing/promotional/unsubscribable
+UNSUBSCRIBE_KEYWORDS_FOR_AI_HEURISTICS = [
+    'unsubscribe', 'opt-out', 'opt out', 'stop receiving', 'manage preferences', 
+    'email preferences', 'subscription', 'marketing', 'newsletter', 'promotional',
+    'offer', 'sale', 'discount', 'deal', 'coupon', 'promo code', 'promotion',
+    'limited time', 'subscribe', 'update preferences', 'mailing list',
+    'no longer wish to receive', 'manage subscriptions', 'manage your subscriptions'
+]
+
+# Keywords that suggest promotional content
+PROMO_KEYWORDS_FOR_AI_HEURISTICS = [
+    'limited time', 'exclusive', 'offer', 'sale', 'discount', 'deal', 'coupon',
+    'promo code', 'promotion', 'savings', 'special offer', 'limited offer',
+    'buy now', 'shop now', 'order now', 'click here', 'purchase', 'buy',
+    'free shipping', 'free trial', 'new arrival', 'new product', 'flash sale'
+]
+
+# Common formatting patterns in promotional emails
+FORMATTING_PATTERNS_FOR_AI_HEURISTICS = [
+    r'\*+\s*[A-Z]+\s*\*+',  # ***TEXT***
+    r'\*\*[^*]+\*\*',       # **TEXT**
+    r'!{2,}',               # Multiple exclamation marks
+    r'\$\d+(\.\d{2})?(\s+off|\s+discount|%\s+off)',  # Price patterns
+    r'\d+%\s+off',          # Percentage discounts
+    r'SAVE\s+\d+%',         # SAVE XX%
+    r'SAVE\s+\$\d+',        # SAVE $XX
+    r'HURRY',               # Urgency words
+    r'LIMITED TIME',
+    r'LAST CHANCE',
+    r'ENDING SOON'
+]
+
+
+def analyze_email_heuristics_for_ai(subject_text: str, snippet_text: str, list_unsubscribe_header: Optional[str] = None) -> Dict[str, bool]:
+    """
+    Analyze email subject and body (snippet) text to determine if it's likely promotional/unsubscribable.
+    
+    This function is adapted from the original heuristic analysis in app.py but modified
+    to be self-contained and not rely on Flask's app context. It examines the subject
+    and body for patterns common in promotional emails and subscription-based content.
+    
+    Args:
+        subject_text: The subject line of the email
+        snippet_text: A snippet of the email body text
+        list_unsubscribe_header: Optional List-Unsubscribe header value
+        
+    Returns:
+        Dict of boolean flags indicating different heuristic results:
+        {
+            'has_unsubscribe_text': bool,  # Contains unsubscribe keywords
+            'has_promotional_keywords': bool,  # Contains promotional keywords
+            'has_promotional_formatting': bool,  # Contains typical promotional formatting
+            'has_list_unsubscribe_header': bool,  # Has List-Unsubscribe header
+            'likely_unsubscribable': bool  # Overall assessment
+        }
+    """
+    # Ensure inputs are strings
+    subject_text = str(subject_text).lower() if subject_text else ""
+    snippet_text = str(snippet_text).lower() if snippet_text else ""
+    combined_text = f"{subject_text} {snippet_text}".lower()
+    
+    # Initialize result with default values
+    result = {
+        'has_unsubscribe_text': False,
+        'has_promotional_keywords': False,
+        'has_promotional_formatting': False,
+        'has_list_unsubscribe_header': False,
+        'likely_unsubscribable': False
+    }
+    
+    # Check for unsubscribe keywords
+    for keyword in UNSUBSCRIBE_KEYWORDS_FOR_AI_HEURISTICS:
+        if keyword.lower() in combined_text:
+            result['has_unsubscribe_text'] = True
+            break
+    
+    # Check for promotional keywords
+    for keyword in PROMO_KEYWORDS_FOR_AI_HEURISTICS:
+        if keyword.lower() in combined_text:
+            result['has_promotional_keywords'] = True
+            break
+    
+    # Check for promotional formatting patterns
+    combined_text_original_case = f"{subject_text} {snippet_text}" if subject_text and snippet_text else ""
+    for pattern in FORMATTING_PATTERNS_FOR_AI_HEURISTICS:
+        if re.search(pattern, combined_text_original_case, re.IGNORECASE):
+            result['has_promotional_formatting'] = True
+            break
+    
+    # Check for List-Unsubscribe header
+    if list_unsubscribe_header:
+        result['has_list_unsubscribe_header'] = True
+    
+    # Overall assessment: likely unsubscribable if any of the criteria are met
+    # For training data preparation, we want to be somewhat inclusive in what we label as potentially unsubscribable
+    result['likely_unsubscribable'] = any([
+        result['has_unsubscribe_text'],
+        (result['has_promotional_keywords'] and result['has_promotional_formatting']),
+        result['has_list_unsubscribe_header']
+    ])
+    
+    return result
+
+
+# --- Text Cleaning Utilities ---
+
+def clean_html_text(html_content: str) -> str:
+    """
+    Clean HTML content and extract readable text.
+    
+    Args:
+        html_content: Raw HTML content string
+        
+    Returns:
+        Cleaned plain text extracted from HTML
+    """
+    if not html_content:
+        return ""
+    
+    try:
+        # Create BeautifulSoup object
+        soup = BeautifulSoup(html_content, 'html.parser')
+        
+        # Remove script and style elements
+        for script_or_style in soup(['script', 'style', 'head', 'title', 'meta', '[document]']):
+            script_or_style.decompose()
+        
+        # Get text content
+        text = soup.get_text()
+        
+        # Clean up text: replace multiple newlines, spaces, etc.
+        text = re.sub(r'\n+', '\n', text)
+        text = re.sub(r' +', ' ', text)
+        text = text.strip()
+        
+        return text
+    except Exception:
+        # If parsing fails, try to extract text with regex (fallback)
+        text = re.sub(r'<[^>]*>', ' ', html_content)
+        text = html.unescape(text)
+        text = re.sub(r'\s+', ' ', text)
+        return text.strip()
+
+
+def normalize_spaces(text: str) -> str:
+    """
+    Normalize whitespace in text.
+    
+    Args:
+        text: Input text
+        
+    Returns:
+        Text with normalized whitespace
+    """
+    if not text:
+        return ""
+    
+    # Replace newlines, tabs with spaces
+    text = re.sub(r'[\n\r\t]+', ' ', text)
+    # Collapse multiple spaces into one
+    text = re.sub(r' +', ' ', text)
+    return text.strip()
+
+
+def normalize_urls(text: str) -> str:
+    """
+    Replace URLs with a placeholder to reduce noise in training data.
+    
+    Args:
+        text: Input text
+        
+    Returns:
+        Text with URLs replaced by a placeholder
+    """
+    if not text:
+        return ""
+    
+    # URL regex pattern
+    url_pattern = r'(https?://[^\s]+)|(www\.[^\s]+\.[^\s]+)'
+    
+    # Replace URLs with placeholder
+    return re.sub(url_pattern, '[URL]', text)
+
+
+def clean_text_for_model(text: str, max_length: Optional[int] = None) -> str:
+    """
+    Clean and normalize text for model input.
+    
+    Args:
+        text: Input text (can be HTML or plain text)
+        max_length: Optional maximum length to truncate to
+        
+    Returns:
+        Cleaned text ready for model input
+    """
+    if not text:
+        return ""
+    
+    # Check if input is likely HTML
+    if re.search(r'<\w+[^>]*>.*?</\w+>', text, re.DOTALL):
+        text = clean_html_text(text)
+    
+    # Normalize whitespace
+    text = normalize_spaces(text)
+    
+    # Replace URLs with placeholder
+    text = normalize_urls(text)
+    
+    # Truncate if needed
+    if max_length and len(text) > max_length:
+        text = text[:max_length]
+    
+    return text
+
+
+# --- Timestamp and Path Utilities ---
+
+def get_current_timestamp() -> str:
+    """Returns ISO format timestamp for current time."""
+    return datetime.datetime.now().isoformat()
+
+
+def get_current_timestamp_log_prefix() -> str:
+    """Returns a formatted timestamp string for log entries."""
+    return f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
+
+
+def ensure_directory_exists(directory_path: str) -> bool:
+    """
+    Ensure that a directory exists, creating it if necessary.
+    
+    Args:
+        directory_path: Path to the directory
+        
+    Returns:
+        True if directory exists or was created, False on error
+    """
+    try:
+        os.makedirs(directory_path, exist_ok=True)
+        return True
+    except Exception:
+        return False
+
+
+# --- Email Header Analysis ---
+
+def extract_email_addresses(header_value: str) -> List[str]:
+    """
+    Extract email addresses from a header value.
+    
+    Args:
+        header_value: Raw header value containing email addresses
+        
+    Returns:
+        List of extracted email addresses
+    """
+    if not header_value:
+        return []
+    
+    # Basic email regex pattern
+    email_pattern = r'[\w.+-]+@[\w-]+\.[\w.-]+'
+    return re.findall(email_pattern, header_value)
+
+
+def parse_list_unsubscribe_header(header_value: str) -> Dict[str, Any]:
+    """
+    Parse the List-Unsubscribe header to extract URLs and email addresses.
+    
+    Args:
+        header_value: Raw List-Unsubscribe header value
+        
+    Returns:
+        Dict with extracted URLs and email addresses
+    """
+    if not header_value:
+        return {"urls": [], "emails": []}
+    
+    result = {"urls": [], "emails": []}
+    
+    # Split by comma and process each value
+    for item in header_value.split(','):
+        item = item.strip()
+        
+        # Handle <mailto:...> format
+        if item.startswith('<mailto:') and item.endswith('>'):
+            email = item[8:-1]  # Remove <mailto: and >
+            result["emails"].append(email)
+        
+        # Handle <http...> format
+        elif item.startswith('<http') and item.endswith('>'):
+            url = item[1:-1]  # Remove < and >
+            result["urls"].append(url)
+    
+    return result

prepare_deployment.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+# Script to prepare Hugging Face deployment without modifying the main app
+
+echo "Preparing Hugging Face deployment..."
+
+# Create symlinks to required directories (won't copy, just link)
+ln -sf ../ml_suite ml_suite
+ln -sf ../final_optimized_model final_optimized_model
+
+echo "Created symbolic links to ml_suite and model directories"
+echo ""
+echo "To deploy to Hugging Face Spaces:"
+echo "1. Create a new Space at https://huggingface.co/spaces"
+echo "2. Clone the space repository"
+echo "3. Copy the contents of this huggingface/ directory to the space repo"
+echo "4. Copy the ml_suite/ and final_optimized_model/ directories"
+echo "5. Git add, commit, and push"
+echo ""
+echo "Your main app remains completely untouched!"

requirements.txt

@@ -0,0 +1,11 @@
+# Core dependencies for Hugging Face deployment
+gradio>=4.0.0
+transformers==4.36.2
+torch==2.2.2
+accelerate==0.25.0
+scikit-learn==1.4.2
+pandas==2.2.1
+nltk==3.8.1
+beautifulsoup4>=4.12.0
+requests>=2.30.0
+joblib>=1.3.0