handler.py

import torch
import torch.nn as nn
from transformers import DebertaV2Tokenizer , DebertaV2Model
from typing import Dict, Any
import joblib
import os

# Define the EndpointHandler class
class EndpointHandler:
    def __init__(self, model_path =""):
        # Load tokenizer and model from the Hugging Face repository
        self.tokenizer = DebertaV2Tokenizer.from_pretrained(model_path)
        
        # Initialize the custom multitask DeBERTa model with pre-defined label counts
        self.model = MultitaskDebertaModel(num_emotion_labels=8, num_polarity_labels=4, num_hate_speech_labels=2)
        self.model.load_state_dict(torch.load(os.path.join(model_path, 'pytorch_model.bin')))
        
        # Use GPU if available
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model.to(self.device)
        self.model.eval()
        
        # Load the encoder files directly from the root of the repository
        self.emotion_encoder = joblib.load(os.path.join(model_path, 'emotion_encoder.pkl'))
        self.polarity_encoder = joblib.load(os.path.join(model_path, 'polarity_encoder.pkl'))
        self.hate_speech_encoder = joblib.load(os.path.join(model_path, 'hate_speech_encoder.pkl'))

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        # Extract the list of texts
        texts = data.get('inputs', [])

        # Batch processing for large inputs
        batch_size = 32
        results = {
            "emotions": [],
            "polarities": [],
            "hate_speech": []
        }

        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i+batch_size]
            
            # Tokenize the input text
            inputs = self.tokenizer(batch_texts, return_tensors='pt', max_length=256, truncation=True, padding=True)
            if 'token_type_ids' in inputs:
                del inputs['token_type_ids']
            inputs = {key: val.to(self.device) for key, val in inputs.items()}

            # Run the input through the model
            with torch.no_grad():
                outputs = self.model(**inputs)
                emotion_logits = outputs.get('emotion')
                polarity_logits = outputs.get('polarity')
                hate_speech_logits = outputs.get('hate_speech')
    
            # Decode predictions from logits using argmax and the label encoders
            emotion_preds = torch.argmax(emotion_logits, dim=1).cpu().numpy().tolist()  
            polarity_preds = torch.argmax(polarity_logits, dim=1).cpu().numpy().tolist()  
            hate_speech_preds = torch.argmax(hate_speech_logits, dim=1).cpu().numpy().tolist()  
    
            # Inverse transform the predictions to get human-readable labels
            decoded_emotions = self.emotion_encoder.inverse_transform(emotion_preds).tolist()
            decoded_polarities = self.polarity_encoder.inverse_transform(polarity_preds).tolist()
            decoded_hate_speech = self.hate_speech_encoder.inverse_transform(hate_speech_preds).tolist()

            results["emotions"].extend(decoded_emotions)
            results["polarities"].extend(decoded_polarities)
            results["hate_speech"].extend(decoded_hate_speech)

        return results

    def load_model(self, model_path):
        #Load model weights from the specified path
        self.load_state_dict(torch.load(model_path))

# Define your custom multitask model architecture here
class MultitaskDebertaModel(nn.Module):
    def __init__(self, num_emotion_labels, num_polarity_labels, num_hate_speech_labels):
        super(MultitaskDebertaModel, self).__init__()
        self.deberta = DebertaV2Model.from_pretrained('microsoft/deberta-v3-base')

        # Freeze the first 5 layers of DeBERTa to speed up training and inference
        for param in self.deberta.encoder.layer[:5]:
            for p in param.parameters():
                p.requires_grad = False

        # LSTM layers for each task
        self.emotion_lstm = nn.LSTM(768, 128, bidirectional=True, batch_first=True)
        self.polarity_lstm = nn.LSTM(768, 128, bidirectional=True, batch_first=True)
        self.hate_speech_lstm = nn.LSTM(768, 128, bidirectional=True, batch_first=True)

        # Attention layers for each task
        self.emotion_attention = nn.MultiheadAttention(embed_dim=256, num_heads=8, batch_first=True)
        self.polarity_attention = nn.MultiheadAttention(embed_dim=256, num_heads=8, batch_first=True)
        self.hate_speech_attention = nn.MultiheadAttention(embed_dim=256, num_heads=8, batch_first=True)

        # Dense layers for each task after attention
        self.emotion_dense = nn.Linear(256, 128)
        self.polarity_dense = nn.Linear(256, 128)
        self.hate_speech_dense = nn.Linear(256, 128)

        # Fusion layer that combines the task-specific features and the CLS token
        self.fusion_dense = nn.Linear(128 + 128 + 128 + 768, 128)

        # Task-specific classifiers
        self.emotion_classifier = nn.Linear(128, num_emotion_labels)
        self.polarity_classifier = nn.Linear(128, num_polarity_labels)
        self.hate_speech_classifier = nn.Linear(128, num_hate_speech_labels)

        # Regularization layers: layer normalization and dropout
        self.layer_norm = nn.LayerNorm(128)
        self.dropout = nn.Dropout(p=0.3)
        self.relu = nn.ReLU()

    def forward(self, input_ids, attention_mask):
        # Extract DeBERTa outputs
        deberta_outputs = self.deberta(input_ids, attention_mask=attention_mask)
        sequence_output = deberta_outputs.last_hidden_state
        cls_output = sequence_output[:, 0, :]  # CLS token output

        # Task-specific LSTM outputs
        emotion_lstm_output, _ = self.emotion_lstm(sequence_output)
        polarity_lstm_output, _ = self.polarity_lstm(sequence_output)
        hate_speech_lstm_output, _ = self.hate_speech_lstm(sequence_output)

        # Task-specific attention outputs
        emotion_attention_output, _ = self.emotion_attention(emotion_lstm_output, emotion_lstm_output, emotion_lstm_output)
        polarity_attention_output, _ = self.polarity_attention(polarity_lstm_output, polarity_lstm_output, polarity_lstm_output)
        hate_speech_attention_output, _ = self.hate_speech_attention(hate_speech_lstm_output, hate_speech_lstm_output, hate_speech_lstm_output)

        # Pool the attention outputs
        emotion_features = torch.mean(emotion_attention_output, dim=1)
        polarity_features = torch.mean(polarity_attention_output, dim=1)
        hate_speech_features = torch.mean(hate_speech_attention_output, dim=1)

        # Apply dense layers after pooling
        emotion_features = self.relu(self.emotion_dense(emotion_features))
        polarity_features = self.relu(self.polarity_dense(polarity_features))
        hate_speech_features = self.relu(self.hate_speech_dense(hate_speech_features))

        # Combine all features (task-specific features + CLS token)
        combined_features = torch.cat([emotion_features, polarity_features, hate_speech_features, cls_output], dim=-1)
        combined_features = self.relu(self.fusion_dense(combined_features))

        # Apply layer normalization and dropout
        combined_features = self.layer_norm(combined_features)
        combined_features = self.dropout(combined_features)

        # Task-specific logits
        emotion_logits = self.emotion_classifier(combined_features)
        polarity_logits = self.polarity_classifier(combined_features)
        hate_speech_logits = self.hate_speech_classifier(combined_features)

        return {
            'emotion': emotion_logits,
            'polarity': polarity_logits,
            'hate_speech': hate_speech_logits
        }