modeling_empathy.py

"""
Custom Empathy Classification Model for HuggingFace

This file contains all necessary model architecture code to run independently.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch.nn import CrossEntropyLoss
from transformers import PreTrainedModel, PretrainedConfig, RobertaModel, RobertaConfig

# BertLayerNorm is deprecated, use nn.LayerNorm instead
try:
    from transformers.models.bert.modeling_bert import BertLayerNorm
except ImportError:
    # In newer versions, use LayerNorm directly
    BertLayerNorm = nn.LayerNorm


class Norm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.size = d_model
        self.alpha = nn.Parameter(torch.ones(self.size))
        self.bias = nn.Parameter(torch.zeros(self.size))
        self.eps = eps
    
    def forward(self, x):
        norm = self.alpha * (x - x.mean(dim=-1, keepdim=True)) / (x.std(dim=-1, keepdim=True) + self.eps) + self.bias
        return norm


class MultiHeadAttention(nn.Module):
    def __init__(self, heads, d_model, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.d_k = d_model // heads
        self.h = heads
        self.q_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
        self.out = nn.Linear(d_model, d_model)
    
    def forward(self, q, k, v, mask=None):
        bs = q.size(0)
        k = self.k_linear(k).view(bs, -1, self.h, self.d_k)
        q = self.q_linear(q).view(bs, -1, self.h, self.d_k)
        v = self.v_linear(v).view(bs, -1, self.h, self.d_k)
        k = k.transpose(1, 2)
        q = q.transpose(1, 2)
        v = v.transpose(1, 2)
        scores = self.attention(q, k, v, self.d_k, mask, self.dropout)
        concat = scores.transpose(1, 2).contiguous().view(bs, -1, self.d_model)
        output = self.out(concat)
        return output
    
    def attention(self, q, k, v, d_k, mask=None, dropout=None):
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
        if mask is not None:
            mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, -1e9)
        scores = F.softmax(scores, dim=-1)
        if dropout is not None:
            scores = dropout(scores)
        output = torch.matmul(scores, v)
        return output


class RobertaClassificationHead(nn.Module):
    """Head for sentence-level classification tasks."""
    def __init__(self, hidden_dropout_prob=0.1, hidden_size=768, empathy_num_labels=3):
        super().__init__()
        self.dense = nn.Linear(hidden_size, hidden_size)
        self.dropout = nn.Dropout(hidden_dropout_prob)
        self.out_proj = nn.Linear(hidden_size, empathy_num_labels)
    
    def forward(self, features, **kwargs):
        x = features[:, :]
        x = self.dropout(x)
        x = self.dense(x)
        x = torch.relu(x)
        x = self.dropout(x)
        x = self.out_proj(x)
        return x


class SeekerEncoder(nn.Module):
    """Wrapper for seeker encoder to match saved weight structure"""
    def __init__(self, config):
        super().__init__()
        self.roberta = RobertaModel(config, add_pooling_layer=False)

class ResponderEncoder(nn.Module):
    """Wrapper for responder encoder to match saved weight structure"""
    def __init__(self, config):
        super().__init__()
        self.roberta = RobertaModel(config, add_pooling_layer=False)

class BiEncoderAttentionWithRationaleClassification(nn.Module):
    def __init__(self, hidden_dropout_prob=0.2, rationale_num_labels=2, empathy_num_labels=3, hidden_size=768, attn_heads=1):
        super().__init__()
        self.dropout = nn.Dropout(hidden_dropout_prob)
        self.rationale_classifier = nn.Linear(hidden_size, rationale_num_labels)
        self.attn = MultiHeadAttention(attn_heads, hidden_size)
        self.norm = Norm(hidden_size)
        self.rationale_num_labels = rationale_num_labels
        self.empathy_num_labels = empathy_num_labels
        self.empathy_classifier = RobertaClassificationHead(hidden_size=768)
        
        self.apply(self._init_weights)
        
        # Initialize encoders with wrapper to match saved weight structure
        roberta_config = RobertaConfig.from_pretrained("roberta-base")
        self.seeker_encoder = SeekerEncoder(roberta_config)
        self.responder_encoder = ResponderEncoder(roberta_config)
    
    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, (nn.Linear, nn.Embedding)):
            initializer_range = 0.02
            module.weight.data.normal_(mean=0.0, std=initializer_range)
        elif isinstance(module, BertLayerNorm):
            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 forward(
        self,
        input_ids_SP=None,
        input_ids_RP=None,
        attention_mask_SP=None,
        attention_mask_RP=None,
        token_type_ids_SP=None,
        token_type_ids_RP=None,
        position_ids_SP=None,
        position_ids_RP=None,
        head_mask_SP=None,
        head_mask_RP=None,
        inputs_embeds_SP=None,
        inputs_embeds_RP=None,
        empathy_labels=None,
        rationale_labels=None,
        lambda_EI=1,
        lambda_RE=0.1
    ):
        outputs_SP = self.seeker_encoder.roberta(
            input_ids_SP,
            attention_mask=attention_mask_SP,
            token_type_ids=token_type_ids_SP,
            position_ids=position_ids_SP,
            head_mask=head_mask_SP,
            inputs_embeds=inputs_embeds_SP,
        )
        
        outputs_RP = self.responder_encoder.roberta(
            input_ids_RP,
            attention_mask=attention_mask_RP,
            token_type_ids=token_type_ids_RP,
            position_ids=position_ids_RP,
            head_mask=head_mask_RP,
            inputs_embeds=inputs_embeds_RP,
        )
        
        sequence_output_SP = outputs_SP[0]
        sequence_output_RP = outputs_RP[0]
        
        sequence_output_RP = sequence_output_RP + self.dropout(self.attn(sequence_output_RP, sequence_output_SP, sequence_output_SP))
        
        logits_empathy = self.empathy_classifier(sequence_output_RP[:, 0, :])
        sequence_output = self.dropout(sequence_output_RP)
        logits_rationales = self.rationale_classifier(sequence_output)
        outputs = (logits_empathy, logits_rationales) + outputs_RP[2:]
        
        loss_rationales = 0.0
        loss_empathy = 0.0
        
        if rationale_labels is not None:
            loss_fct = CrossEntropyLoss()
            if attention_mask_RP is not None:
                active_loss = attention_mask_RP.view(-1) == 1
                active_logits = logits_rationales.view(-1, self.rationale_num_labels)
                active_labels = torch.where(
                    active_loss, rationale_labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(rationale_labels)
                )
                loss_rationales = loss_fct(active_logits, active_labels)
            else:
                loss_rationales = loss_fct(logits_rationales.view(-1, self.rationale_num_labels), rationale_labels.view(-1))
        
        if empathy_labels is not None:
            loss_fct = CrossEntropyLoss()
            loss_empathy = loss_fct(logits_empathy.view(-1, self.empathy_num_labels), empathy_labels.view(-1))
            loss = lambda_EI * loss_empathy + lambda_RE * loss_rationales
            outputs = (loss, loss_empathy, loss_rationales) + outputs
        
        return outputs


class EmpathyModelConfig(PretrainedConfig):
    """Configuration for Empathy Model"""
    model_type = "empathy_classifier"
    
    def __init__(
        self,
        hidden_dropout_prob=0.2,
        rationale_num_labels=2,
        empathy_num_labels=3,
        hidden_size=768,
        attn_heads=1,
        max_length=64,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.hidden_dropout_prob = hidden_dropout_prob
        self.rationale_num_labels = rationale_num_labels
        self.empathy_num_labels = empathy_num_labels
        self.hidden_size = hidden_size
        self.attn_heads = attn_heads
        self.max_length = max_length


class EmpathyModel(PreTrainedModel):
    """HuggingFace wrapper for Empathy Model"""
    config_class = EmpathyModelConfig
    base_model_prefix = "model"  # Important: matches the saved state_dict structure
    
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.model = BiEncoderAttentionWithRationaleClassification(
            hidden_dropout_prob=config.hidden_dropout_prob,
            rationale_num_labels=config.rationale_num_labels,
            empathy_num_labels=config.empathy_num_labels,
            hidden_size=config.hidden_size,
            attn_heads=config.attn_heads
        )
    
    def forward(
        self,
        input_ids_SP=None,
        input_ids_RP=None,
        attention_mask_SP=None,
        attention_mask_RP=None,
        **kwargs
    ):
        return self.model(
            input_ids_SP=input_ids_SP,
            input_ids_RP=input_ids_RP,
            attention_mask_SP=attention_mask_SP,
            attention_mask_RP=attention_mask_RP,
            **kwargs
        )