File size: 7,742 Bytes
863b0e4 b9bce35 863b0e4 b9bce35 3667632 b9bce35 863b0e4 b9bce35 863b0e4 b9bce35 863b0e4 b9bce35 863b0e4 b64c23f 863b0e4 |
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 |
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
}
|