File size: 18,859 Bytes
b2d89cf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
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)
        }