[ehealth] draft version of ner and spo tasks

Author: LemonNoel

examples/biomedical/cblue/model.py

@@ -0,0 +1,118 @@
+import paddle
+import paddle.nn as nn
+from paddlenlp.transformers import ElectraPretrainedModel
+
+
+class ElectraForBinaryTokenClassification(ElectraPretrainedModel):
+    """
+    Electra Model with two linear layers on top of the hidden-states output layers,
+    designed for token classification tasks with nesting.
+
+    Args: 
+        electra (:class:`ElectraModel`):
+            An instance of ElectraModel.
+        num_classes (list):
+            The number of classes.
+        use_crf (bool, optional):
+            Use conditional random fields for named entity recognition.
+            Defaults to False.
+        dropout (float, optionl):
+            The dropout probability for output of Electra.
+            If None, use the same value as `hidden_dropout_prob' of 'ElectraModel`
+            instance `electra`. Defaults to None.
+    """
+
+    def __init__(self, electra, num_classes, dropout=None):
+        super(ElectraForBinaryTokenClassification, self).__init__()
+        assert (len(num_classes) == 2)
+        self.num_classes_oth = num_classes[0]
+        self.num_classes_sym = num_classes[1]
+        self.electra = electra
+        self.dropout = nn.Dropout(dropout if dropout is not None else
+                                  self.electra.config['hidden_dropout_prob'])
+        self.classifier_oth = nn.Linear(self.electra.config['hidden_size'],
+                                        self.num_classes_oth)
+        self.classifier_sym = nn.Linear(self.electra.config['hidden_size'],
+                                        self.num_classes_sym)
+        self.init_weights()
+
+    def forward(self,
+                input_ids=None,
+                token_type_ids=None,
+                position_ids=None,
+                attention_mask=None):
+        r"""
+        The ElectraForMedicalClassification forward method, overrides the __call__() special method.
+
+        TODO
+        """
+        sequence_output = self.electra(input_ids, token_type_ids, position_ids,
+                                       attention_mask)
+        sequence_output = self.dropout(sequence_output)
+
+        logits_sym = self.classifier_sym(sequence_output)
+        logits_oth = self.classifier_oth(sequence_output)
+        return logits_oth, logits_sym
+
+
+class MultiHeadAttentionForSPO(nn.Layer):
+    def __init__(self, embed_dim, num_heads, scale_value=768):
+        super(MultiHeadAttentionForSPO, self).__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.scale_value = scale_value**-0.5
+        self.q_proj = nn.Linear(embed_dim, embed_dim * num_heads)
+        self.k_proj = nn.Linear(embed_dim, embed_dim * num_heads)
+
+    def forward(self, query, key):
+        q = self.q_proj(query)
+        k = self.k_proj(query)
+        q = paddle.reshape(q, shape=[0, 0, self.num_heads, self.embed_dim])
+        k = paddle.reshape(k, shape=[0, 0, self.num_heads, self.embed_dim])
+        q = paddle.transpose(q, perm=[0, 2, 1, 3])
+        k = paddle.transpose(k, perm=[0, 2, 1, 3])
+        scores = paddle.matmul(q, k, transpose_y=True)
+        scores = paddle.scale(scores, scale=self.scale_value)
+        return scores
+
+
+class ElectraForSPO(ElectraPretrainedModel):
+    """
+    """
+
+    def __init__(self, electra, num_classes, dropout=None):
+        super(ElectraForSPO, self).__init__()
+        self.num_classes = num_classes
+        self.electra = electra
+        self.dropout = nn.Dropout(dropout if dropout is not None else
+                                  self.electra.config['hidden_dropout_prob'])
+        self.classifier = nn.Linear(self.electra.config['hidden_size'], 2)
+        self.span_attention = MultiHeadAttentionForSPO(
+            self.electra.config['hidden_size'], num_classes)
+        self.sigmoid = paddle.nn.Sigmoid()
+
+    def forward(self,
+                input_ids=None,
+                token_type_ids=None,
+                position_ids=None,
+                attention_mask=None):
+        sequence_outputs, _, all_hidden_states = self.electra(
+            input_ids,
+            token_type_ids,
+            position_ids,
+            attention_mask,
+            output_hidden_states=True)
+        sequence_outputs = self.dropout(sequence_outputs)
+        ent_logits = self.classifier(sequence_outputs)
+
+        subject_output = all_hidden_states[-2]
+        cls_output = paddle.unsqueeze(sequence_outputs[:, 0, :], axis=1)
+        subject_output = subject_output + cls_output
+
+        output_size = self.num_classes + self.electra.config['hidden_size']
+        rel_logits = self.span_attention(sequence_outputs, subject_output)
+
+        ent_logits = self.sigmoid(ent_logits)
+        rel_logits = self.sigmoid(rel_logits)
+
+        return ent_logits, rel_logits

examples/biomedical/cblue/train_classification.py

@@ -26,7 +26,7 @@
 import paddlenlp as ppnlp
 from paddlenlp.data import Stack, Tuple, Pad
 from paddlenlp.datasets import load_dataset
-from paddlenlp.transformers import ElectraForSequenceClassification, LinearDecayWithWarmup
+from paddlenlp.transformers import ElectraForSequenceClassification, ElectraTokenizer, LinearDecayWithWarmup
 from paddlenlp.metrics import MultiLabelsMetric, AccuracyAndF1
 from paddlenlp.ops.optimizer import ExponentialMovingAverage
 
@@ -117,16 +117,15 @@ def do_train():
 
     set_seed(args.seed)
 
-    train_ds, dev_ds, test_ds = load_dataset(
-        'cblue', args.dataset, splits=['train', 'dev', 'test'])
+    train_ds, dev_ds = load_dataset(
+        'cblue', args.dataset, splits=['train', 'dev'])
 
     model = ElectraForSequenceClassification.from_pretrained(
         'ehealth-chinese',
         num_classes=len(train_ds.label_list),
         activation='tanh',
         layer_norm_eps=1e-5)
-    tokenizer = ppnlp.transformers.ElectraTokenizer.from_pretrained(
-        'ehealth-chinese')
+    tokenizer = ElectraTokenizer.from_pretrained('ehealth-chinese')
 
     trans_func = partial(
         convert_example,

examples/biomedical/cblue/train_ner.py

@@ -0,0 +1,228 @@
+from functools import partial
+import argparse
+import os
+import random
+import time
+import distutils.util
+
+import numpy as np
+import paddle
+from paddlenlp.data import Pad, Dict
+from paddlenlp.datasets import load_dataset
+from paddlenlp.transformers import LinearDecayWithWarmup, ElectraTokenizer
+from paddlenlp.metrics import ChunkEvaluator
+
+from model import ElectraForBinaryTokenClassification
+from utils import create_dataloader, convert_example_ner
+
+# yapf: disable
+parser = argparse.ArgumentParser()
+parser.add_argument('--device', choices=['cpu', 'gpu', 'xpu', 'npu'], default='gpu', help='Select which device to train model, default to gpu.')
+parser.add_argument('--init_from_ckpt', default=None, type=str, help='The path of checkpoint to be loaded.')
+parser.add_argument('--batch_size', default=8, type=int, help='Batch size per GPU/CPU for training.')
+parser.add_argument('--learning_rate', default=6e-5, type=float, help='Learning rate for fine-tuning token classification task.')
+parser.add_argument('--max_seq_length', default=128, type=int, help='The maximum total input sequence length after tokenization.')
+parser.add_argument('--valid_steps', default=100, type=int, help='The interval steps to evaluate model performance.')
+parser.add_argument('--logging_steps', default=10, type=int, help='The interval steps to logging.')
+parser.add_argument('--save_steps', default=10000, type=int, help='The interval steps to save checkpoints.')
+parser.add_argument('--weight_decay', default=0.01, type=float, help='Weight decay if we apply some.')
+parser.add_argument('--warmup_proportion', default=0.1, type=float, help='Linear warmup proportion over the training process.')
+parser.add_argument('--use_amp', default=False, type=bool, help='Enable mixed precision training.')
+parser.add_argument('--epochs', default=1, type=int, help='Total number of training epochs.')
+parser.add_argument('--eval_mention', default=True, type=bool, help='.')
+parser.add_argument('--update_tokenizer', default=True, type=bool, help='Update the word tokenizer during training.')
+parser.add_argument('--seed', default=1000, type=int, help='Random seed.')
+parser.add_argument('--save_dir', default='./checkpoint', type=str, help='The output directory where the model checkpoints will be written.')
+
+args = parser.parse_args()
+# yapf: enable
+
+
+def set_seed(seed):
+    """set random seed"""
+    random.seed(seed)
+    np.random.seed(seed)
+    paddle.seed(seed)
+
+
+@paddle.no_grad()
+def evaluate(model, criterion, metrics, data_loader):
+    model.eval()
+    metrics[0].reset()
+    losses = []
+    for batch in data_loader:
+        input_ids, token_type_ids, position_ids, masks, label_oth, label_sym = batch
+        logits = model(input_ids, token_type_ids, position_ids)
+        loss_oth = criterion(logits[0], paddle.unsqueeze(label_oth, 2))
+        loss_oth = paddle.mean(loss_oth * paddle.unsqueeze(masks, 2))
+        loss_sym = criterion(logits[1], paddle.unsqueeze(label_sym, 2))
+        loss_sym = paddle.mean(loss_sym * paddle.unsqueeze(masks, 2))
+
+        losses.append([loss_oth.numpy(), loss_sym.numpy()])
+
+        lengths = paddle.sum(masks, axis=1)
+        pred_oth = paddle.argmax(logits[0], axis=2)
+        pred_sym = paddle.argmax(logits[1], axis=2)
+        correct_oth = metrics[0].compute(lengths, pred_oth, label_oth)
+        correct_sym = metrics[1].compute(lengths, pred_sym, label_sym)
+        correct_oth = [x.numpy() for x in correct_oth]
+        correct_sym = [x.numpy() for x in correct_sym]
+        metrics[0].update(*correct_oth)
+        metrics[0].update(*correct_sym)
+        _, _, result = metrics[0].accumulate()
+    loss = np.mean(losses, axis=0)
+    print('eval loss symptom: %.5f, loss others: %.5f, f1: %.5f' %
+          (loss[1], loss[0], result))
+    model.train()
+    metrics[0].reset()
+
+
+def do_train():
+    paddle.set_device(args.device)
+    rank = paddle.distributed.get_rank()
+    if paddle.distributed.get_world_size() > 1:
+        paddle.distributed.init_parallel_env()
+
+    set_seed(args.seed)
+
+    train_ds, dev_ds = load_dataset('cblue', 'CMeEE', splits=['train', 'dev'])
+
+    model = ElectraForBinaryTokenClassification.from_pretrained(
+        'ehealth-chinese', num_classes=[len(x) for x in train_ds.label_list])
+    tokenizer = ElectraTokenizer.from_pretrained('ehealth-chinese')
+
+    label_list = train_ds.label_list
+    pad_label_id = [len(label_list[0]) - 1, len(label_list[1]) - 1]
+    ignore_label_id = -100
+
+    trans_func = partial(
+        convert_example_ner,
+        tokenizer=tokenizer,
+        max_seq_length=args.max_seq_length,
+        pad_label_id=pad_label_id)
+
+    batchify_fn = lambda samples, fn=Dict({
+        'input_ids': Pad(axis=0, pad_val=tokenizer.pad_token_id, dtype='int64'),
+        'token_type_ids': Pad(axis=0, pad_val=tokenizer.pad_token_type_id, dtype='int64'),
+        'position_ids': Pad(axis=0, pad_val=tokenizer.pad_token_id, dtype='int64'),
+        'mask': Pad(axis=0, pad_val=0, dtype='float32'),
+        'label_oth': Pad(axis=0, pad_val=pad_label_id[0], dtype='int64'),
+        'label_sym': Pad(axis=0, pad_val=pad_label_id[1], dtype='int64')
+    }): fn(samples)
+
+    train_data_loader = create_dataloader(
+        train_ds,
+        mode='train',
+        batch_size=args.batch_size,
+        batchify_fn=batchify_fn,
+        trans_fn=trans_func)
+
+    dev_data_loader = create_dataloader(
+        dev_ds,
+        mode='dev',
+        batch_size=args.batch_size,
+        batchify_fn=batchify_fn,
+        trans_fn=trans_func)
+
+    if args.init_from_ckpt:
+        if not os.path.isfile(args.init_from_ckpt):
+            raise ValueError('init_from_ckpt is not a valid model filename.')
+        state_dict = paddle.load(args.init_from_ckpt)
+        model.set_dict(state_dict)
+    if paddle.distributed.get_world_size() > 1:
+        model = paddle.DataParallel(model)
+
+    num_training_steps = len(train_data_loader) * args.epochs
+
+    lr_scheduler = LinearDecayWithWarmup(args.learning_rate, num_training_steps,
+                                         args.warmup_proportion)
+
+    decay_params = [
+        p.name for n, p in model.named_parameters()
+        if not any(nd in n for nd in ['bias', 'norm'])
+    ]
+
+    optimizer = paddle.optimizer.AdamW(
+        learning_rate=lr_scheduler,
+        parameters=model.parameters(),
+        weight_decay=args.weight_decay,
+        apply_decay_param_fun=lambda x: x in decay_params)
+
+    criterion = paddle.nn.functional.softmax_with_cross_entropy
+
+    metrics = [ChunkEvaluator(label_list[0]), ChunkEvaluator(label_list[1])]
+
+    if args.use_amp:
+        scaler = paddle.amp.GradScaler(init_loss_scaling=args.scale_loss)
+
+    global_step = 0
+    tic_train = time.time()
+    total_train_time = 0
+    for epoch in range(1, args.epochs + 1):
+        for step, batch in enumerate(train_data_loader, start=1):
+            input_ids, token_type_ids, position_ids, masks, label_oth, label_sym = batch
+            with paddle.amp.auto_cast(
+                    args.use_amp,
+                    custom_white_list=['layer_norm', 'softmax', 'gelu'], ):
+                att_mask = paddle.unsqueeze(masks, axis=2)
+                att_mask = paddle.matmul(att_mask, att_mask, transpose_y=True)
+                logits = model(input_ids, token_type_ids, position_ids, masks)
+
+                loss_oth = criterion(logits[0], paddle.unsqueeze(label_oth, 2))
+                loss_sym = criterion(logits[1], paddle.unsqueeze(label_sym, 2))
+                loss_masks = paddle.unsqueeze(masks, 2)
+                loss_oth = paddle.mean(loss_oth * loss_masks)
+                loss_sym = paddle.mean(loss_sym * loss_masks)
+
+                loss = loss_oth + loss_sym
+
+                lengths = paddle.sum(masks, axis=1)
+                pred_oth = paddle.argmax(logits[0], axis=-1)
+                pred_sym = paddle.argmax(logits[1], axis=-1)
+                correct_oth = metrics[0].compute(lengths, pred_oth, label_oth)
+                correct_sym = metrics[1].compute(lengths, pred_sym, label_sym)
+                correct_oth = [x.numpy() for x in correct_oth]
+                correct_sym = [x.numpy() for x in correct_sym]
+                metrics[0].update(*correct_oth)
+                metrics[0].update(*correct_sym)
+                _, _, f1 = metrics[0].accumulate()
+
+                if args.use_amp:
+                    scaler.scale(loss).backward()
+                    scaler.minimize(optimizer, loss)
+                else:
+                    loss.backward()
+                    optimizer.step()
+                lr_scheduler.step()
+                optimizer.clear_grad()
+
+                global_step += 1
+                if global_step % args.logging_steps == 0 and rank == 0:
+                    time_diff = time.time() - tic_train
+                    total_train_time += time_diff
+                    print(
+                        'global step %d, epoch: %d, batch: %d, loss: %.5f, loss symptom: %.5f, loss others: %.5f, f1: %.5f, speed: %.2f step/s'
+                        % (global_step, epoch, step, loss, loss_sym, loss_oth,
+                           f1, args.logging_steps / time_diff))
+                    tic_train = time.time()
+
+                if global_step % args.valid_steps == 0 and rank == 0:
+                    evaluate(model, criterion, metrics, dev_data_loader)
+                    tic_train = time.time()
+
+                if global_step % args.save_steps == 0 and rank == 0:
+                    save_dir = os.patj.join(args.save_dir,
+                                            'model_%d' % global_step)
+                    if not os.path.exists(save_dir):
+                        os.makedirs(save_dir)
+                    if paddle.distributed.get_world_size() > 1:
+                        model._layers.save_pretrained(save_dir)
+                    else:
+                        model.save_pretrained(save_dir)
+                    tokenizer.save_pretrained(save_dir)
+                    tic_train = time.time()
+    print('Speed: %.2f steps/s' % (global_step / total_train_time))
+
+
+if __name__ == '__main__':
+    do_train()