|
from transformers import BertModel |
|
import torch |
|
import onnx |
|
import pytorch_lightning as pl |
|
import wandb |
|
from metrics import MyAccuracy |
|
from utils import num_unique_labels |
|
from typing import Dict, Tuple, List, Optional |
|
|
|
class MultiTaskBertModel(pl.LightningModule): |
|
|
|
""" |
|
Multi-task Bert model for Named Entity Recognition (NER) and Intent Classification |
|
|
|
Args: |
|
config (BertConfig): Bert model configuration. |
|
dataset (Dict[str, Union[str, List[str]]]): A dictionary containing keys 'text', 'ner', and 'intent'. |
|
""" |
|
|
|
def __init__(self, config, dataset): |
|
super().__init__() |
|
|
|
self.num_ner_labels, self.num_intent_labels = num_unique_labels(dataset) |
|
|
|
self.dropout = torch.nn.Dropout(config.hidden_dropout_prob) |
|
|
|
self.model = BertModel(config=config) |
|
|
|
self.ner_classifier = torch.nn.Linear(config.hidden_size, self.num_ner_labels) |
|
self.intent_classifier = torch.nn.Linear(config.hidden_size, self.num_intent_labels) |
|
|
|
|
|
self.save_hyperparameters() |
|
|
|
self.accuracy = MyAccuracy() |
|
|
|
def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: |
|
|
|
""" |
|
Perform a forward pass through Multi-task Bert model. |
|
|
|
Args: |
|
input_ids (torch.Tensor, torch.shape: (batch, length_of_tokenized_sequences)): Input token IDs. |
|
attention_mask (Optional[torch.Tensor]): Attention mask for input tokens. |
|
|
|
Returns: |
|
Tuple[torch.Tensor,torch.Tensor]: NER logits, Intent logits. |
|
""" |
|
|
|
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask) |
|
|
|
sequence_output = outputs[0] |
|
sequence_output = self.dropout(sequence_output) |
|
ner_logits = self.ner_classifier(sequence_output) |
|
|
|
pooled_output = outputs[1] |
|
pooled_output = self.dropout(pooled_output) |
|
intent_logits = self.intent_classifier(pooled_output) |
|
|
|
return ner_logits, intent_logits |
|
|
|
def training_step(self: pl.LightningModule, batch, batch_idx: int) -> torch.Tensor: |
|
""" |
|
Perform a training step for the Multi-task BERT model. |
|
|
|
Args: |
|
batch: Input batch. |
|
batch_idx (int): Index of the batch. |
|
|
|
Returns: |
|
torch.Tensor: Loss value |
|
""" |
|
loss, ner_logits, intent_logits, ner_labels, intent_labels = self._common_step(batch, batch_idx) |
|
accuracy_ner = self.accuracy(ner_logits, ner_labels, self.num_ner_labels) |
|
accuracy_intent = self.accuracy(intent_logits, intent_labels, self.num_intent_labels) |
|
self.log_dict({'training_loss': loss, 'ner_accuracy': accuracy_ner, 'intent_accuracy': accuracy_intent}, |
|
on_step=False, on_epoch=True, prog_bar=True) |
|
return loss |
|
|
|
def on_validation_epoch_start(self): |
|
self.validation_step_outputs_ner = [] |
|
self.validation_step_outputs_intent = [] |
|
|
|
def validation_step(self, batch, batch_idx: int) -> torch.Tensor: |
|
""" |
|
Perform a validation step for the Multi-task BERT model. |
|
|
|
Args: |
|
batch: Input batch. |
|
batch_idx (int): Index of the batch. |
|
|
|
Returns: |
|
torch.Tensor: Loss value. |
|
""" |
|
loss, ner_logits, intent_logits, ner_labels, intent_labels = self._common_step(batch, batch_idx) |
|
|
|
accuracy_ner = self.accuracy(ner_logits, ner_labels, self.num_ner_labels) |
|
accuracy_intent = self.accuracy(intent_logits, intent_labels, self.num_intent_labels) |
|
self.log_dict({'validation_loss': loss, 'val_ner_accuracy': accuracy_ner, 'val_intent_accuracy': accuracy_intent}, |
|
on_step=False, on_epoch=True, prog_bar=True) |
|
|
|
self.validation_step_outputs_ner.append(ner_logits) |
|
self.validation_step_outputs_intent.append(intent_logits) |
|
return loss |
|
|
|
def on_validation_epoch_end(self): |
|
""" |
|
Perform actions at the end of validation epoch to track the training process in WandB. |
|
""" |
|
validation_step_outputs_ner = self.validation_step_outputs_ner |
|
validation_step_outputs_intent = self.validation_step_outputs_intent |
|
|
|
dummy_input = torch.zeros((1, 128), device=self.device, dtype=torch.long) |
|
model_filename = f"model_{str(self.global_step).zfill(5)}.onnx" |
|
torch.onnx.export(self, dummy_input, model_filename) |
|
artifact = wandb.Artifact(name="model.ckpt", type="model") |
|
artifact.add_file(model_filename) |
|
self.logger.experiment.log_artifact(artifact) |
|
|
|
flattened_logits_ner = torch.flatten(torch.cat(validation_step_outputs_ner)) |
|
flattened_logits_intent = torch.flatten(torch.cat(validation_step_outputs_intent)) |
|
self.logger.experiment.log( |
|
{"valid/ner_logits": wandb.Histogram(flattened_logits_ner.to('cpu')), |
|
"valid/intent_logits": wandb.Histogram(flattened_logits_intent.to('cpu')), |
|
"global_step": self.global_step} |
|
) |
|
|
|
def _common_step(self, batch, batch_idx): |
|
""" |
|
Common steps for both training and validation. Calculate loss for both NER and intent layer. |
|
|
|
Returns: |
|
Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
|
Combiner loss value, NER logits, intent logits, NER labels, intent labels. |
|
""" |
|
ids = batch['input_ids'] |
|
mask = batch['attention_mask'] |
|
ner_labels = batch['ner_labels'] |
|
intent_labels = batch['intent_labels'] |
|
|
|
ner_logits, intent_logits = self.forward(input_ids=ids, attention_mask=mask) |
|
|
|
criterion = torch.nn.CrossEntropyLoss() |
|
|
|
ner_loss = criterion(ner_logits.view(-1, self.num_ner_labels), ner_labels.view(-1).long()) |
|
intent_loss = criterion(intent_logits.view(-1, self.num_intent_labels), intent_labels.view(-1).long()) |
|
|
|
loss = ner_loss + intent_loss |
|
return loss, ner_logits, intent_logits, ner_labels, intent_labels |
|
|
|
def configure_optimizers(self): |
|
optimizer = torch.optim.Adam(self.parameters(), lr=1e-5) |
|
return optimizer |
