Upload modeling_conceptframemet.py with huggingface_hub

modeling_conceptframemet.py

@@ -1,102 +1,96 @@
 """
-ConceptFrameMet: Metaphor Detection with Frame and Source Domain Prediction
+Adaptive Source QA MelBERT with Configurable Blending Strategies
 
-This model detects metaphors and predicts their semantic frames and source domains.
-Based on AdaptiveSourceQAMelBert architecture.
+This model provides configurable approaches to incorporating source domain information:
+
+FLAGS:
+1. --source_blend_mode: 'additive' or 'replacement' (default: 'replacement')
+   - additive: enhanced = target + alpha * source (keeps target strength)
+   - replacement: blended = conf * source + (1-conf) * target (original approach)
+
+2. --source_use_mode: 'metaphor_only' or 'all' (default: 'all')
+   - metaphor_only: Only use source for samples with high metaphor probability
+   - all: Use source for all samples
+
+3. --source_alpha: float (default: 0.3) - scaling factor for additive mode
+
+4. --metaphor_threshold: float (default: 0.5) - threshold for metaphor-only mode
+
+Architecture:
+- CONTEXT: target_word in full sentence → encoder 1 → target_context_embedding  
+- SOURCE: [SEP] sentence [SEP] target [SEP] → QA model → predict source + confidence
+- ISOLATED: isolated target → encoder 2 → target_embedding
+- BLEND: Configurable (additive or replacement)
+- FILTER: Configurable (metaphor-only or all)
+- MIP: [enhanced_embedding, target_context_embedding]
+- SPV: [pooled, enhanced_embedding] or [pooled, target_context_embedding]
 """
 
 import torch
 import torch.nn as nn
-from transformers import RobertaModel, RobertaTokenizer, AutoModelForQuestionAnswering, AutoTokenizer
-from typing import Dict, List, Tuple, Optional
-import json
-import os
+import torch.nn.functional as F
 
 
-class ConceptFrameMetForMetaphorDetection(nn.Module):
-    """
-    Metaphor detection model with semantic frame and source domain prediction capabilities.
-    
-    This model:
-    - Detects metaphors in text
-    - Predicts semantic frames for target words
-    - Predicts source domains for metaphors
-    """
-    
-    def __init__(
-        self,
-        encoder_model_name="roberta-base",
-        frame_qa_model_name="nixie1981/sem_frames",
-        source_qa_model_name=None,
-        classifier_hidden=768,
-        drop_ratio=0.2,
-        num_labels=2,
-        source_blend_mode='replacement',
-        source_use_mode='metaphor_only',
-        source_alpha=0.3,
-        metaphor_threshold=0.5,
-    ):
-        super().__init__()
+class AdaptiveSourceQAMelBert(nn.Module):
+    """MelBERT with configurable source domain blending strategies"""
+
+    def __init__(self, args, Model, config, Source_QA_Model, 
+                 source_qa_tokenizer, melbert_tokenizer, num_labels=2):
+        """
+        Initialize the model with configurable flags
         
+        Args:
+            args: Configuration arguments with:
+                - source_blend_mode: 'additive' or 'replacement'
+                - source_use_mode: 'metaphor_only' or 'all'
+                - source_alpha: scaling factor for additive mode
+                - metaphor_threshold: threshold for metaphor-only mode
+            Model: MelBert encoder (RoBERTa/BERT)
+            config: Model configuration
+            Source_QA_Model: QA-style model to predict source domain
+            source_qa_tokenizer: Tokenizer for QA model
+            melbert_tokenizer: Tokenizer for MelBert
+            num_labels: Number of metaphor classes (2: literal/metaphorical)
+        """
+        super(AdaptiveSourceQAMelBert, self).__init__()
         self.num_labels = num_labels
-        self.classifier_hidden = classifier_hidden
-        self.drop_ratio = drop_ratio
-        
-        # Configuration
-        self.source_blend_mode = source_blend_mode
-        self.source_use_mode = source_use_mode
-        self.source_alpha = source_alpha
-        self.metaphor_threshold = metaphor_threshold
-        
-        # Load encoder (RoBERTa) with correct type_vocab_size
-        from transformers import RobertaConfig
-        
-        # Load base model first
-        self.encoder = RobertaModel.from_pretrained(encoder_model_name)
-        
-        # Resize token_type_embeddings to match training (type_vocab_size=4)
-        # This is needed because the model was trained with 4 token types
-        if self.encoder.embeddings.token_type_embeddings.weight.shape[0] != 4:
-            old_embeddings = self.encoder.embeddings.token_type_embeddings
-            new_embeddings = nn.Embedding(4, old_embeddings.embedding_dim)
-            # Copy the original embedding (for type 0)
-            new_embeddings.weight.data[0] = old_embeddings.weight.data[0]
-            # Initialize the rest
-            new_embeddings.weight.data[1:].normal_(mean=0.0, std=self.encoder.config.initializer_range)
-            self.encoder.embeddings.token_type_embeddings = new_embeddings
-            self.encoder.config.type_vocab_size = 4
-        
-        self.tokenizer = RobertaTokenizer.from_pretrained(encoder_model_name)
-        self.config = self.encoder.config
-        
-        # Load frame QA model
-        try:
-            self.frame_qa_model = AutoModelForQuestionAnswering.from_pretrained(frame_qa_model_name)
-            self.frame_qa_tokenizer = AutoTokenizer.from_pretrained(frame_qa_model_name)
-            self.has_frame_predictor = True
-        except:
-            print("Warning: Frame QA model not available")
-            self.has_frame_predictor = False
-        
-        # Load source QA model (if available)
-        if source_qa_model_name:
-            try:
-                self.source_qa_model = AutoModelForQuestionAnswering.from_pretrained(source_qa_model_name)
-                self.source_qa_tokenizer = AutoTokenizer.from_pretrained(source_qa_model_name)
-                self.has_source_predictor = True
-            except:
-                print("Warning: Source QA model not available")
-                self.has_source_predictor = False
+        self.encoder = Model
+        self.source_qa_model = Source_QA_Model
+        self.source_qa_tokenizer = source_qa_tokenizer
+        self.melbert_tokenizer = melbert_tokenizer
+        self.config = config
+        self.dropout = nn.Dropout(args.drop_ratio)
+        self.args = args
+
+        # Configuration flags with defaults
+        self.source_blend_mode = getattr(args, 'source_blend_mode', 'replacement')
+        self.source_use_mode = getattr(args, 'source_use_mode', 'all')
+        self.source_alpha = getattr(args, 'source_alpha', 0.3)
+        self.metaphor_threshold = getattr(args, 'metaphor_threshold', 0.5)
+
+        # Freeze or unfreeze source QA model
+        if not getattr(args, 'unfreeze_source_qa', False):
+            for param in self.source_qa_model.parameters():
+                param.requires_grad = False
         else:
-            self.has_source_predictor = False
-        
-        # Dropout
-        self.dropout = nn.Dropout(drop_ratio)
-        
-        # Classification layers
-        self.SPV_linear = nn.Linear(self.config.hidden_size * 2, classifier_hidden)
-        self.MIP_linear = nn.Linear(self.config.hidden_size * 2, classifier_hidden)
-        self.classifier = nn.Linear(classifier_hidden * 2, num_labels)
+            for param in self.source_qa_model.parameters():
+                param.requires_grad = True
+
+        # Load source labels
+        self.source_id2label = {}
+        try:
+            import json
+            with open('source_finder/source_labels.json', 'r') as f:
+                source_label2id = json.load(f)
+                self.source_id2label = {v: k for k, v in source_label2id.items()}
+                print(f"✓ Loaded {len(self.source_id2label)} source domain labels")
+        except Exception as e:
+            print(f"❌ Warning: Could not load source labels: {e}")
+
+        # SPV and MIP linear layers
+        self.SPV_linear = nn.Linear(config.hidden_size * 2, args.classifier_hidden)
+        self.MIP_linear = nn.Linear(config.hidden_size * 2, args.classifier_hidden)
+        self.classifier = nn.Linear(args.classifier_hidden * 2, num_labels)
         
         self._init_weights(self.SPV_linear)
         self._init_weights(self.MIP_linear)
@@ -104,245 +98,237 @@ class ConceptFrameMetForMetaphorDetection(nn.Module):
         
         self.logsoftmax = nn.LogSoftmax(dim=1)
         
-        # Load source and frame labels
-        self.source_id2label = {}
-        self.frame_id2label = {}
-        
+        # Print configuration
+        print(f"\n{'='*80}")
+        print(f"✓ AdaptiveSourceQAMelBert initialized")
+        print(f"  - Blend Mode: {self.source_blend_mode.upper()}")
+        if self.source_blend_mode == 'additive':
+            print(f"    → enhanced = target + {self.source_alpha} * source")
+        else:
+            print(f"    → blended = conf * source + (1-conf) * target")
+        print(f"  - Use Mode: {self.source_use_mode.upper()}")
+        if self.source_use_mode == 'metaphor_only':
+            print(f"    → Only use source when metaphor_score > {self.metaphor_threshold}")
+        else:
+            print(f"    → Use source for all samples")
+        print(f"{'='*80}\n")
+
     def _init_weights(self, module):
         """Initialize the weights"""
         if isinstance(module, (nn.Linear, nn.Embedding)):
             module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
         if isinstance(module, nn.Linear) and module.bias is not None:
             module.bias.data.zero_()
-    
-    def predict_frames(self, sentence: str, target_word: str) -> Dict[str, any]:
+
+    def predict_source_and_embeddings(self, input_ids, target_mask, attention_mask, 
+                                       input_ids_2, target_mask_2, attention_mask_2):
         """
-        Predict semantic frame for a target word in context
+        Predict source domain and get source/target embeddings
         
-        Args:
-            sentence: Input sentence
-            target_word: Target word to analyze
-            
         Returns:
-            Dictionary with frame prediction and confidence
+            source_embeddings: [batch_size, hidden_size]
+            target_embeddings: [batch_size, hidden_size]
+            confidences: [batch_size] - confidence scores
         """
-        if not self.has_frame_predictor:
-            return {"frame": "UNKNOWN", "confidence": 0.0}
+        batch_size = input_ids.size(0)
         
-        try:
-            inputs = self.frame_qa_tokenizer(
-                sentence,
-                target_word,
-                max_length=150,
+        # 1. Decode sentences and extract target words
+        sentences = []
+        target_words = []
+        
+        for i in range(batch_size):
+            sentence = self.melbert_tokenizer.decode(input_ids[i], skip_special_tokens=True)
+            target_positions = target_mask[i].nonzero(as_tuple=True)[0]
+            
+            if len(target_positions) > 0:
+                target_tokens = input_ids[i][target_positions]
+                target_word = self.melbert_tokenizer.decode(target_tokens, skip_special_tokens=True)
+            else:
+                target_word = "unknown"
+            
+            sentences.append(sentence)
+            target_words.append(target_word)
+        
+        # 2. Format QA input and predict source
+        with torch.no_grad():
+            qa_inputs = self.source_qa_tokenizer(
+                sentences,
+                target_words,
+                max_length=self.args.max_seq_length,
                 padding='max_length',
                 truncation=True,
                 return_tensors='pt'
             )
+            qa_inputs = {k: v.to(input_ids.device) for k, v in qa_inputs.items()}
             
-            with torch.no_grad():
-                outputs = self.frame_qa_model(**inputs)
-                
-                # Check if it has start/end logits
-                if hasattr(outputs, 'start_logits') and hasattr(outputs, 'end_logits'):
-                    start_logits = outputs.start_logits
-                    end_logits = outputs.end_logits
-                    
-                    start_idx = torch.argmax(start_logits)
-                    end_idx = torch.argmax(end_logits)
-                    
-                    confidence = (torch.max(torch.softmax(start_logits, dim=-1)) + 
-                                 torch.max(torch.softmax(end_logits, dim=-1))) / 2.0
-                    
-                    frame_tokens = inputs['input_ids'][0][start_idx:end_idx+1]
-                    frame = self.frame_qa_tokenizer.decode(frame_tokens, skip_special_tokens=True)
-                else:
-                    # Fallback if model structure is different
-                    frame = "Self_motion"
-                    confidence = 0.5
-            
-            return {
-                "frame": frame if frame else "UNKNOWN",
-                "confidence": confidence.item() if isinstance(confidence, torch.Tensor) else confidence
-            }
-        except Exception as e:
-            # If frame prediction fails, return a default
-            print(f"Frame prediction warning: {e}")
-            return {"frame": "UNKNOWN", "confidence": 0.0}
-    
-    def predict_source(self, sentence: str, target_word: str) -> Dict[str, any]:
-        """
-        Predict source domain for a metaphor
-        
-        Args:
-            sentence: Input sentence
-            target_word: Target word to analyze
-            
-        Returns:
-            Dictionary with source prediction and confidence
-        """
-        if not self.has_source_predictor:
-            return {"source": "UNKNOWN", "confidence": 0.0}
-        
-        inputs = self.source_qa_tokenizer(
-            sentence,
-            target_word,
-            max_length=150,
-            padding='max_length',
-            truncation=True,
-            return_tensors='pt'
-        )
-        
+            # If source model is FrameAwareSourcePredictor, also pass frame inputs
+            # (frame inputs are the same as source inputs for this use case)
+            if hasattr(self.source_qa_model, 'frame_finder'):
+                qa_inputs['frame_input_ids'] = qa_inputs['input_ids']
+                qa_inputs['frame_attention_mask'] = qa_inputs['attention_mask']
+        
+        # 3. Get source predictions with confidence
+        qa_outputs = self.source_qa_model(**qa_inputs)
+        source_logits = qa_outputs.logits
+        source_probs = torch.softmax(source_logits, dim=-1)
+        predicted_source_ids = torch.argmax(source_logits, dim=-1)
+        
+        # Get confidence scores
+        confidences = source_probs.gather(1, predicted_source_ids.unsqueeze(1)).squeeze(1)
+        
+        # Map to source words
         with torch.no_grad():
-            outputs = self.source_qa_model(**inputs)
-            logits = outputs.logits if hasattr(outputs, 'logits') else outputs.start_logits
-            
-            probs = torch.softmax(logits, dim=-1)
-            predicted_id = torch.argmax(probs, dim=-1)
-            confidence = probs.gather(-1, predicted_id.unsqueeze(-1)).squeeze(-1)
-            
-            source = self.source_id2label.get(predicted_id.item(), "UNKNOWN")
-        
-        return {
-            "source": source,
-            "confidence": confidence.item()
-        }
-    
-    def predict_metaphor(
-        self,
-        sentence: str,
-        target_word: str,
-        target_positions: Optional[List[int]] = None
-    ) -> Dict[str, any]:
-        """
-        Predict if target word is metaphorical in context
+            predicted_sources = [self.source_id2label.get(sid.item(), "UNKNOWN") 
+                                for sid in predicted_source_ids]
         
-        Args:
-            sentence: Input sentence
-            target_word: Target word to analyze
-            target_positions: Token positions of target word (optional)
-            
-        Returns:
-            Dictionary with metaphor prediction, frame, and source
-        """
-        # Tokenize input
-        inputs = self.tokenizer(
-            sentence,
-            max_length=150,
+        # 4. Encode predicted source words
+        source_inputs = self.melbert_tokenizer(
+            predicted_sources,
+            max_length=self.args.max_seq_length,
             padding='max_length',
             truncation=True,
             return_tensors='pt'
         )
+        source_inputs = {k: v.to(input_ids.device) for k, v in source_inputs.items()}
+        source_target_mask = (source_inputs['input_ids'] != self.melbert_tokenizer.pad_token_id).float()
         
-        # Create target mask
-        if target_positions is None:
-            # Find target word positions
-            target_tokens = self.tokenizer.tokenize(target_word)
-            sentence_tokens = self.tokenizer.tokenize(sentence)
-            target_positions = []
-            for i in range(len(sentence_tokens) - len(target_tokens) + 1):
-                if sentence_tokens[i:i+len(target_tokens)] == target_tokens:
-                    target_positions = list(range(i+1, i+1+len(target_tokens)))  # +1 for CLS token
-                    break
-        
-        target_mask = torch.zeros_like(inputs['input_ids'], dtype=torch.float)
-        if target_positions:
-            for pos in target_positions:
-                if pos < target_mask.size(1):
-                    target_mask[0, pos] = 1.0
-        
-        # Forward pass for metaphor detection
-        with torch.no_grad():
-            outputs = self.encoder(**inputs)
-            sequence_output = outputs[0]
-            pooled_output = outputs[1]
-            
-            # Get target output
-            target_output = sequence_output * target_mask.unsqueeze(2)
-            target_output = target_output.sum(dim=1) / (target_mask.sum(-1, keepdim=True) + 1e-10)
-            target_output = self.dropout(target_output)
-            pooled_output = self.dropout(pooled_output)
-            
-            # SPV and MIP
-            SPV_hidden = self.SPV_linear(torch.cat([pooled_output, target_output], dim=1))
-            MIP_hidden = self.MIP_linear(torch.cat([target_output, target_output], dim=1))
-            
-            # Classification
-            logits = self.classifier(torch.cat([SPV_hidden, MIP_hidden], dim=1))
-            logits = self.logsoftmax(logits)
-            probs = torch.exp(logits)
-            
-            is_metaphor = torch.argmax(probs, dim=1).item() == 1
-            metaphor_confidence = probs[0, 1].item()
-        
-        # Predict frame and source
-        frame_result = self.predict_frames(sentence, target_word)
-        source_result = self.predict_source(sentence, target_word) if is_metaphor else {"source": "N/A", "confidence": 0.0}
+        source_outputs = self.encoder(
+            source_inputs['input_ids'],
+            attention_mask=source_inputs['attention_mask']
+        )
         
-        return {
-            "is_metaphor": is_metaphor,
-            "metaphor_confidence": metaphor_confidence,
-            "frame": frame_result["frame"],
-            "frame_confidence": frame_result["confidence"],
-            "source": source_result["source"],
-            "source_confidence": source_result["confidence"]
-        }
-    
-    @classmethod
-    def from_pretrained(cls, model_path, **kwargs):
-        """Load model from pretrained checkpoint"""
-        # Load weights first to check what's in checkpoint
-        weights_path = os.path.join(model_path, "pytorch_model.bin")
-        state_dict = torch.load(weights_path, map_location='cpu')
+        source_sequence_output = source_outputs[0]
+        source_target_output = source_sequence_output * source_target_mask.unsqueeze(2)
         
-        # Check what's in the checkpoint
-        has_source_in_checkpoint = any(k.startswith('source_qa_model.') for k in state_dict.keys())
-        has_frame_in_checkpoint = any(k.startswith('frame_qa_model.') for k in state_dict.keys())
+        if self.args.small_mean:
+            source_embeddings = source_target_output.mean(1)
+        else:
+            source_embeddings = source_target_output.sum(dim=1) / source_target_mask.sum(-1, keepdim=True)
         
-        # Initialize model:
-        # - Download frame_qa_model (nixie1981/sem_frames) - NOT in checkpoint
-        # - Don't download source_qa_model - IS in checkpoint
-        model = cls(
-            frame_qa_model_name="nixie1981/sem_frames",  # Download - needed for frames!
-            source_qa_model_name=None,  # Don't download - in checkpoint
-            **kwargs
+        # 5. Encode original isolated target words
+        target_outputs_2 = self.encoder(
+            input_ids_2,
+            attention_mask=attention_mask_2
         )
         
-        # Manually set source flag since weights are in checkpoint
-        if has_source_in_checkpoint:
-            model.has_source_predictor = True
-        
-        # Load ALL weights from checkpoint (including source_qa_model)
-        missing, unexpected = model.load_state_dict(state_dict, strict=False)
+        target_sequence_output_2 = target_outputs_2[0]
+        target_output_2 = target_sequence_output_2 * target_mask_2.unsqueeze(2)
         
-        print(f"Loaded {len(state_dict)} weights from checkpoint")
-        if missing:
-            print(f"Missing {len(missing)} keys")
+        if self.args.small_mean:
+            target_embeddings_2 = target_output_2.mean(1)
+        else:
+            target_embeddings_2 = target_output_2.sum(dim=1) / target_mask_2.sum(-1, keepdim=True)
         
-        return model
-    
-    def save_pretrained(self, save_directory):
-        """Save model to directory"""
-        os.makedirs(save_directory, exist_ok=True)
+        return source_embeddings, target_embeddings_2, confidences
+
+    def blend_embeddings(self, source_embeddings, target_embeddings, confidences):
+        """
+        Blend source and target embeddings based on configuration
         
-        # Save weights
-        torch.save(self.state_dict(), os.path.join(save_directory, "pytorch_model.bin"))
+        Args:
+            source_embeddings: [batch_size, hidden_size]
+            target_embeddings: [batch_size, hidden_size]
+            confidences: [batch_size]
+            
+        Returns:
+            blended_embeddings: [batch_size, hidden_size]
+        """
+        confidence_weights = confidences.unsqueeze(1)
         
-        # Save config
-        config = {
-            "_name_or_path": "ConceptFrameMet",
-            "architectures": ["ConceptFrameMetForMetaphorDetection"],
-            "model_type": "conceptframemet",
-            "num_labels": self.num_labels,
-            "classifier_hidden": self.classifier_hidden,
-            "drop_ratio": self.drop_ratio,
-            "source_blend_mode": self.source_blend_mode,
-            "source_use_mode": self.source_use_mode,
-            "source_alpha": self.source_alpha,
-            "metaphor_threshold": self.metaphor_threshold,
-        }
+        if self.source_blend_mode == 'additive':
+            # ADDITIVE: enhanced = target + alpha * source
+            # Keeps target strength, adds source as enhancement
+            enhanced = target_embeddings + self.source_alpha * confidence_weights * source_embeddings
+            return enhanced
+        else:
+            # REPLACEMENT: blended = conf * source + (1-conf) * target
+            # Original soft confidence approach
+            blended = confidence_weights * source_embeddings + (1 - confidence_weights) * target_embeddings
+            return blended
+
+    def forward(
+        self,
+        input_ids,
+        input_ids_2,
+        target_mask,
+        target_mask_2,
+        attention_mask_2,
+        token_type_ids=None,
+        attention_mask=None,
+        labels=None,
+        head_mask=None,
+        input_with_mask_ids=None
+    ):
+        """
+        Forward pass with configurable source blending
+        """
+        # ===== ENCODER 1: Target in context =====
+        outputs = self.encoder(
+            input_ids,
+            token_type_ids=token_type_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+        )
+
+        sequence_output = outputs[0]
+        pooled_output = outputs[1]
+
+        # Get target output with target mask
+        target_output = sequence_output * target_mask.unsqueeze(2)
+        target_output = self.dropout(target_output)
+        pooled_output = self.dropout(pooled_output)
+
+        if self.args.small_mean:
+            target_output = target_output.mean(1)
+        else:
+            target_output = target_output.sum(dim=1) / target_mask.sum(-1, keepdim=True)
+
+        # ===== ENCODER 2: Get source and target embeddings =====
+        source_embeddings, target_embeddings_2, confidences = self.predict_source_and_embeddings(
+            input_ids, target_mask, attention_mask,
+            input_ids_2, target_mask_2, attention_mask_2
+        )
+
+        # ===== METAPHOR-ONLY FILTERING (if enabled) =====
+        if self.source_use_mode == 'metaphor_only':
+            # Get preliminary metaphor score
+            # Use simple heuristic based on target context
+            prelim_features = torch.cat([pooled_output, target_output], dim=1)
+            prelim_hidden = self.SPV_linear(prelim_features)
+            prelim_logits = self.classifier(torch.cat([prelim_hidden, prelim_hidden], dim=1))
+            prelim_probs = torch.exp(self.logsoftmax(prelim_logits))
+            metaphor_scores = prelim_probs[:, 1]  # Probability of metaphor class
+            
+            # Only use source for samples with high metaphor probability
+            use_source_mask = (metaphor_scores > self.metaphor_threshold).float().unsqueeze(1)
+        else:
+            # Use source for all samples
+            use_source_mask = torch.ones(source_embeddings.size(0), 1).to(source_embeddings.device)
+
+        # ===== BLEND: Apply configured blending strategy =====
+        blended_embedding = self.blend_embeddings(source_embeddings, target_embeddings_2, confidences)
         
-        with open(os.path.join(save_directory, "config.json"), 'w') as f:
-            json.dump(config, f, indent=2)
+        # Apply metaphor-only mask
+        final_embedding = use_source_mask * blended_embedding + (1 - use_source_mask) * target_embeddings_2
+        final_embedding = self.dropout(final_embedding)
+
+        # ===== SPV and MIP =====
+        if self.args.spv_isolate:
+            SPV_hidden = self.SPV_linear(torch.cat([pooled_output, final_embedding], dim=1))
+        else:
+            SPV_hidden = self.SPV_linear(torch.cat([pooled_output, target_output], dim=1))
         
-        # Save tokenizer
-        self.tokenizer.save_pretrained(save_directory)
+        MIP_hidden = self.MIP_linear(torch.cat([final_embedding, target_output], dim=1))
+
+        # Final classification
+        logits = self.classifier(self.dropout(torch.cat([SPV_hidden, MIP_hidden], dim=1)))
+        logits = self.logsoftmax(logits)
+
+        if labels is not None:
+            loss_fct = nn.NLLLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            return loss
+        return logits