Added pretrained BART wights for the encoder/decoder while keeping custom architectiure and task heads

configs/model/base.yaml

@@ -1,6 +1,8 @@
-d_model: 512
+d_model: 768
 num_encoder_layers: 6
 num_decoder_layers: 6
-num_attention_heads: 8
-ffn_dim: 2048
+num_attention_heads: 12
+ffn_dim: 3072
 dropout: 0.1
+use_pretrained: true
+pretrained_model_name: facebook/bart-base

configs/training/default.yaml

@@ -10,3 +10,5 @@ scheduler:
 trainer:
   max_epochs: 5
   gradient_clip_norm: 1.0
+  validation_samples: 3
+  validation_max_length: 128

scripts/inference.py

@@ -3,9 +3,14 @@ from __future__ import annotations
 
 import argparse
 import json
+import sys
 from pathlib import Path
 from typing import List, cast
 
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
 from src.data.tokenization import TokenizerConfig
 from src.inference import EmotionPrediction, TopicPrediction, create_inference_pipeline
 

src/models/attention.py

@@ -83,8 +83,6 @@ class ScaledDotProductAttention(nn.Module):
             mask_bool = mask.to(dtype=torch.bool, device=scores.device)
             # masked_fill expects broadcastable mask: True means keep, False means mask out
             scores = scores.masked_fill(~mask_bool, float("-1e9"))
-        # Applying Softmax to get attention weights
-        attention_weights = F.softmax(scores, dim=-1)
         
         # Softmax to get attention probabilities
         p_attn = F.softmax(scores, dim=-1)

src/models/factory.py

@@ -5,6 +5,9 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 
+import torch
+from transformers import BartModel
+
 from ..data.tokenization import Tokenizer
 from ..utils.config import load_yaml
 from .decoder import TransformerDecoder
@@ -23,6 +26,8 @@ class ModelConfig:
     num_attention_heads: int = 8
     ffn_dim: int = 2048
     dropout: float = 0.1
+    use_pretrained: bool = False
+    pretrained_model_name: str = "facebook/bart-base"
 
     def __post_init__(self):
         if self.d_model % self.num_attention_heads != 0:
@@ -51,9 +56,93 @@ def load_model_config(path: Optional[str | Path]) -> ModelConfig:
         num_attention_heads=int(data.get("num_attention_heads", 8)),
         ffn_dim=int(data.get("ffn_dim", 2048)),
         dropout=float(data.get("dropout", 0.1)),
+        use_pretrained=bool(data.get("use_pretrained", False)),
+        pretrained_model_name=str(data.get("pretrained_model_name", "facebook/bart-base")),
     )
 
 
+def _load_pretrained_weights(encoder: TransformerEncoder, decoder: TransformerDecoder, model_name: str) -> None:
+    """Load pretrained BART weights into custom encoder/decoder."""
+    print(f"Loading pretrained weights from {model_name}...")
+    bart = BartModel.from_pretrained(model_name)
+    
+    # Load encoder weights
+    print("Transferring encoder weights...")
+    encoder.embedding.weight.data.copy_(bart.encoder.embed_tokens.weight.data)
+    encoder.pos_encoder.pe.data.copy_(bart.encoder.embed_positions.weight.data.unsqueeze(0))
+    
+    for i, (custom_layer, bart_layer) in enumerate(zip(encoder.layers, bart.encoder.layers)):
+        # Self-attention
+        custom_layer.self_attn.W_Q.weight.data.copy_(bart_layer.self_attn.q_proj.weight.data)
+        custom_layer.self_attn.W_Q.bias.data.copy_(bart_layer.self_attn.q_proj.bias.data)
+        custom_layer.self_attn.W_K.weight.data.copy_(bart_layer.self_attn.k_proj.weight.data)
+        custom_layer.self_attn.W_K.bias.data.copy_(bart_layer.self_attn.k_proj.bias.data)
+        custom_layer.self_attn.W_V.weight.data.copy_(bart_layer.self_attn.v_proj.weight.data)
+        custom_layer.self_attn.W_V.bias.data.copy_(bart_layer.self_attn.v_proj.bias.data)
+        custom_layer.self_attn.W_O.weight.data.copy_(bart_layer.self_attn.out_proj.weight.data)
+        custom_layer.self_attn.W_O.bias.data.copy_(bart_layer.self_attn.out_proj.bias.data)
+        
+        # Layer norms
+        custom_layer.norm1.weight.data.copy_(bart_layer.self_attn_layer_norm.weight.data)
+        custom_layer.norm1.bias.data.copy_(bart_layer.self_attn_layer_norm.bias.data)
+        custom_layer.norm2.weight.data.copy_(bart_layer.final_layer_norm.weight.data)
+        custom_layer.norm2.bias.data.copy_(bart_layer.final_layer_norm.bias.data)
+        
+        # FFN
+        custom_layer.ffn.fc1.weight.data.copy_(bart_layer.fc1.weight.data)
+        custom_layer.ffn.fc1.bias.data.copy_(bart_layer.fc1.bias.data)
+        custom_layer.ffn.fc2.weight.data.copy_(bart_layer.fc2.weight.data)
+        custom_layer.ffn.fc2.bias.data.copy_(bart_layer.fc2.bias.data)
+    
+    encoder.final_norm.weight.data.copy_(bart.encoder.layernorm_embedding.weight.data)
+    encoder.final_norm.bias.data.copy_(bart.encoder.layernorm_embedding.bias.data)
+    
+    # Load decoder weights
+    print("Transferring decoder weights...")
+    decoder.embedding.weight.data.copy_(bart.decoder.embed_tokens.weight.data)
+    decoder.pos_encoder.pe.data.copy_(bart.decoder.embed_positions.weight.data.unsqueeze(0))
+    
+    for i, (custom_layer, bart_layer) in enumerate(zip(decoder.layers, bart.decoder.layers)):
+        # Self-attention
+        custom_layer.self_attn.W_Q.weight.data.copy_(bart_layer.self_attn.q_proj.weight.data)
+        custom_layer.self_attn.W_Q.bias.data.copy_(bart_layer.self_attn.q_proj.bias.data)
+        custom_layer.self_attn.W_K.weight.data.copy_(bart_layer.self_attn.k_proj.weight.data)
+        custom_layer.self_attn.W_K.bias.data.copy_(bart_layer.self_attn.k_proj.bias.data)
+        custom_layer.self_attn.W_V.weight.data.copy_(bart_layer.self_attn.v_proj.weight.data)
+        custom_layer.self_attn.W_V.bias.data.copy_(bart_layer.self_attn.v_proj.bias.data)
+        custom_layer.self_attn.W_O.weight.data.copy_(bart_layer.self_attn.out_proj.weight.data)
+        custom_layer.self_attn.W_O.bias.data.copy_(bart_layer.self_attn.out_proj.bias.data)
+        
+        # Cross-attention
+        custom_layer.cross_attn.W_Q.weight.data.copy_(bart_layer.encoder_attn.q_proj.weight.data)
+        custom_layer.cross_attn.W_Q.bias.data.copy_(bart_layer.encoder_attn.q_proj.bias.data)
+        custom_layer.cross_attn.W_K.weight.data.copy_(bart_layer.encoder_attn.k_proj.weight.data)
+        custom_layer.cross_attn.W_K.bias.data.copy_(bart_layer.encoder_attn.k_proj.bias.data)
+        custom_layer.cross_attn.W_V.weight.data.copy_(bart_layer.encoder_attn.v_proj.weight.data)
+        custom_layer.cross_attn.W_V.bias.data.copy_(bart_layer.encoder_attn.v_proj.bias.data)
+        custom_layer.cross_attn.W_O.weight.data.copy_(bart_layer.encoder_attn.out_proj.weight.data)
+        custom_layer.cross_attn.W_O.bias.data.copy_(bart_layer.encoder_attn.out_proj.bias.data)
+        
+        # Layer norms
+        custom_layer.norm1.weight.data.copy_(bart_layer.self_attn_layer_norm.weight.data)
+        custom_layer.norm1.bias.data.copy_(bart_layer.self_attn_layer_norm.bias.data)
+        custom_layer.norm2.weight.data.copy_(bart_layer.encoder_attn_layer_norm.weight.data)
+        custom_layer.norm2.bias.data.copy_(bart_layer.encoder_attn_layer_norm.bias.data)
+        custom_layer.norm3.weight.data.copy_(bart_layer.final_layer_norm.weight.data)
+        custom_layer.norm3.bias.data.copy_(bart_layer.final_layer_norm.bias.data)
+        
+        # FFN
+        custom_layer.ffn.fc1.weight.data.copy_(bart_layer.fc1.weight.data)
+        custom_layer.ffn.fc1.bias.data.copy_(bart_layer.fc1.bias.data)
+        custom_layer.ffn.fc2.weight.data.copy_(bart_layer.fc2.weight.data)
+        custom_layer.ffn.fc2.bias.data.copy_(bart_layer.fc2.bias.data)
+    
+    decoder.final_norm.weight.data.copy_(bart.decoder.layernorm_embedding.weight.data)
+    decoder.final_norm.bias.data.copy_(bart.decoder.layernorm_embedding.bias.data)
+    
+    print("Pretrained weights loaded successfully!")
+
+
 def build_multitask_model(
     tokenizer: Tokenizer,
     *,
@@ -68,6 +157,7 @@ def build_multitask_model(
         raise ValueError("num_emotions must be a positive integer")
     if not isinstance(num_topics, int) or num_topics <= 0:
         raise ValueError("num_topics must be a positive integer")
+    
     encoder = TransformerEncoder(
         vocab_size=tokenizer.vocab_size,
         d_model=cfg.d_model,
@@ -88,7 +178,12 @@ def build_multitask_model(
         max_len=tokenizer.config.max_length,
         pad_token_id=tokenizer.pad_token_id,
     )
+    
+    # Load pretrained weights if requested
+    if cfg.use_pretrained:
+        _load_pretrained_weights(encoder, decoder, cfg.pretrained_model_name)
 
+    
     model = MultiTaskModel(encoder=encoder, decoder=decoder, decoder_outputs_logits=True)
     model.add_head(
         "summarization",

src/training/trainer.py

@@ -20,6 +20,8 @@ class TrainerConfig:
     gradient_clip_norm: float = 1.0
     logging_interval: int = 50
     task_weights: Dict[str, float] | None = None
+    validation_samples: int = 3
+    validation_max_length: int = 128
 
 
 class Trainer:
@@ -63,6 +65,9 @@ class Trainer:
             if val_loaders:
                 val_metrics = self._run_epoch(val_loaders, train=False, epoch=epoch)
                 history[f"val_epoch_{epoch}"] = val_metrics
+                # Generate sample summaries for validation
+                if "summarization" in val_loaders:
+                    self._validate_generation(val_loaders["summarization"], epoch)
             epoch_duration = time.perf_counter() - epoch_start
             total_elapsed = time.perf_counter() - start_time
             self._print_epoch_progress(epoch, total_epochs, epoch_duration, total_elapsed)
@@ -223,6 +228,72 @@ class Trainer:
         valid[valid == -100] = self.tokenizer.pad_token_id
         return self.tokenizer.decode_batch(valid.tolist())
 
+    def _validate_generation(self, val_loader: DataLoader, epoch: int) -> None:
+        """Generate and print sample summaries to monitor quality during training."""
+        self.model.eval()
+        samples_generated = 0
+        print(f"\n{'='*80}")
+        print(f"[Validation Generation - Epoch {epoch}]")
+        print(f"{'='*80}")
+        
+        with torch.no_grad():
+            for batch in val_loader:
+                if samples_generated >= self.config.validation_samples:
+                    break
+                    
+                batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+                src_ids = batch["src_ids"]
+                src_mask = batch.get("src_mask")
+                labels = batch["labels"]
+                
+                # Only process first item from batch
+                src_ids = src_ids[:1]
+                if src_mask is not None:
+                    src_mask = src_mask[:1]
+                labels = labels[:1]
+                
+                # Encode source
+                encoder_mask = None
+                if src_mask is not None:
+                    encoder_mask = src_mask.unsqueeze(1) & src_mask.unsqueeze(2)
+                memory = self.model.encoder(src_ids, mask=encoder_mask)
+                
+                # Ban special tokens from generation
+                ban_token_ids = [self.tokenizer.bos_token_id, self.tokenizer.pad_token_id]
+                unk_id = getattr(self.tokenizer._tokenizer, 'unk_token_id', None)
+                if isinstance(unk_id, int):
+                    ban_token_ids.append(unk_id)
+                ban_token_ids = [tid for tid in ban_token_ids if tid is not None]
+                
+                # Generate
+                generated = self.model.decoder.greedy_decode(
+                    memory=memory,
+                    max_len=self.config.validation_max_length,
+                    start_token_id=self.tokenizer.bos_token_id,
+                    end_token_id=self.tokenizer.eos_token_id,
+                    device=self.device,
+                    min_len=10,
+                    ban_token_ids=ban_token_ids,
+                    no_repeat_ngram_size=3,
+                    memory_mask=src_mask,
+                )
+                
+                # Decode
+                source_text = self.tokenizer.decode(src_ids[0].tolist())
+                generated_text = self.tokenizer.decode(generated[0].tolist())
+                reference_text = self._decode_labels(labels)[0]
+                
+                print(f"\nSample {samples_generated + 1}:")
+                print(f"Source: {source_text[:200]}..." if len(source_text) > 200 else f"Source: {source_text}")
+                print(f"Generated: {generated_text}")
+                print(f"Reference: {reference_text[:200]}..." if len(reference_text) > 200 else f"Reference: {reference_text}")
+                print("-" * 80)
+                
+                samples_generated += 1
+        
+        print(f"{'='*80}\n")
+        self.model.train()
+
     def _print_epoch_progress(
         self,
         epoch: int,