nlu/DeBERTa/apps/tasks/ner_task.py

from collections import OrderedDict,defaultdict,Counter
from collections.abc import Sequence
import math
import numpy as np
import os
import pdb
import random
import torch
import ujson as json
from ...utils import xtqdm as tqdm
from ...utils import get_logger

from ..models import NERModel
from ...data import ExampleInstance, ExampleSet, DynamicDataset
from ...data.example import *
from ...data.example import _truncate_segments
from .task import EvalData, Task
from .task_registry import register_task

from seqeval import metrics as seq_metrics

__all__ = ['NERTask']

logger = get_logger()

@register_task(name="NER", desc="Named-entity recognition task")
class NERTask(Task):
  def __init__(self, data_dir, tokenizer, args, **kwargs):
    super().__init__(tokenizer, args, **kwargs)
    self.data_dir = data_dir

  def train_data(self, max_seq_len=512, dataset_size=None, epochs=1, mask_gen=None, **kwargs):
    train = self.load_data(os.path.join(self.data_dir, 'train.txt'), max_seq_len=max_seq_len)
    examples = ExampleSet(train)
    if dataset_size is None:
      dataset_size = len(examples)*epochs
    return DynamicDataset(examples, feature_fn = self.get_feature_fn(max_seq_len=max_seq_len, mask_gen=mask_gen), \
dataset_size = dataset_size, shuffle=True, **kwargs)

  def eval_data(self, max_seq_len=512, dataset_size=None, **kwargs):
    ds = [
        self._data('dev', 'valid.txt', 'dev', max_seq_len=max_seq_len),
        self._data('test', 'test.txt', 'test', max_seq_len=max_seq_len)
        ]
   
    for d in ds:
      if dataset_size is None:
        _size = len(d.data)
      d.data = DynamicDataset(d.data, feature_fn = self.get_feature_fn(max_seq_len=max_seq_len), dataset_size = _size, **kwargs)
    return ds

  def test_data(self,max_seq_len=512, dataset_size = None, **kwargs):
    """See base class."""
    ds = [
      self._data('test', 'test.txt', 'test', max_seq_len=max_seq_len)
      ]
    
    for d in ds:
      if dataset_size is None:
        _size = len(d.data)
      d.data = DynamicDataset(d.data, feature_fn = self.get_feature_fn(max_seq_len=max_seq_len), dataset_size = _size, **kwargs)
    return ds

  def _data(self, name, path, type_name = 'dev', ignore_metric=False, max_examples=None, shuffle=False, max_seq_len=512):
    input_src = os.path.join(self.data_dir, path)
    assert os.path.exists(input_src), f"{input_src} doesn't exists"
    data = self.load_data(input_src, max_seq_len=max_seq_len, max_examples=max_examples, shuffle=shuffle)
    examples = ExampleSet(data)
    predict_fn = self.get_predict_fn(examples)
    metrics_fn = self.get_metrics_fn()
    return EvalData(name, examples,
      metrics_fn = metrics_fn, predict_fn = predict_fn, ignore_metric=ignore_metric, critial_metrics=['f1'])

  def get_metrics_fn(self):
    """Calcuate metrics based on prediction results"""
    def metrics_fn(logits, labels):
      preds = np.argmax(logits, axis=-1)
      label_names = self.get_labels()
      y_true = []
      y_pred = []
      for pred,label in zip(preds, labels):
        y_true.append([label_names[l] for l in label if l>=0])
        y_pred.append([label_names[p] for p,l in zip(pred, label) if l>=0])
      return OrderedDict(
        accuracy = seq_metrics.accuracy_score(y_true, y_pred),
        f1 = seq_metrics.f1_score(y_true, y_pred),
        precision = seq_metrics.precision_score(y_true, y_pred),
        recall = seq_metrics.recall_score(y_true, y_pred)
        )
    return metrics_fn

  def get_predict_fn(self, examples):
    """Calcuate metrics based on prediction results"""
    def predict_fn(logits, output_dir, name, prefix):
      output=os.path.join(output_dir, 'submit-{}-{}.tsv'.format(name, prefix))
      preds = np.argmax(logits, axis=-1)
      labels = self.get_labels()
      with open(output, 'w', encoding='utf-8') as fs:
        fs.write('index\tpredictions\n')
        for i,(e,p) in enumerate(zip(examples,preds)):
          words = ''.join(e.sentence).split(' ')
          tokens = e.segments[0]
          bw = 0
          for w,t in zip(words,tokens):
            fs.write(f'{w} {labels[p[bw]]}\n')
            bw += len(t)
          fs.write('\n')

    return predict_fn

  def get_feature_fn(self, max_seq_len = 512, mask_gen = None):
    def _example_to_feature(example, rng=None, ext_params=None, **kwargs):
      return self.example_to_feature(self.tokenizer, example, max_seq_len = max_seq_len, \
        rng = rng, mask_generator = mask_gen, ext_params = ext_params, **kwargs)
    return _example_to_feature

  def get_model_class_fn(self):
    def partial_class(*wargs, **kwargs):
      return NERModel.load_model(*wargs, **kwargs)
    return partial_class

  def get_labels(self):
    """See base class."""
    return ["O", "B-MISC", "I-MISC",  "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]

  def load_data(self, path, max_seq_len=512, max_examples=None, shuffle=False):
    docs = self.extract_docs(path)
    examples=[]
    for doc in docs:
      merged_words = []
      merged_tokens = []
      merged_labels = []
      size = 0
      for sent in doc:
        words = [t[0] if i==0 else (' ' + t[0]) for i,t in enumerate(sent)]
        labels = [t[1] for t in sent]
        tokens = [self.tokenizer.tokenize(w) for w in words]
        l = sum(len(t) for t in tokens)
        if size+l > max_seq_len-2:
          examples.append(ExampleInstance(segments=[merged_tokens], label=merged_labels, sentence=merged_words))
          size = 0
          merged_words = []
          merged_tokens = []
          merged_labels = []
        size += l
        merged_words.extend(words)
        merged_tokens.extend(tokens)
        merged_labels.extend(labels)
      if size>0:
        examples.append(ExampleInstance(segments=[merged_tokens], label=merged_labels, sentence=merged_words))

    def get_stats(l):
      return f'Max={max(l)}, min={min(l)}, avg={np.mean(l)}'
    ctx_token_size = [sum(len(w) for w in  e.segments[0]) for e in examples]
    logger.info(f'Statistics: {get_stats(ctx_token_size)}, long={len([t for t in ctx_token_size if t > 500])}/{len(ctx_token_size)}')

    return examples

  def example_to_feature(self, tokenizer, example, max_seq_len=512, rng=None, mask_generator = None, ext_params=None, label_type='int', **kwargs):
    if not rng:
      rng = random
    max_num_tokens = max_seq_len - 2
    features = OrderedDict()
    tokens = ['[CLS]']
    target_labels = [-1]
    type_ids = [0]

    for i,w in enumerate(example.segments[0]):
      tokens.extend(w)
      type_ids.extend([0]*len(w))
      if example.label is not None:
        target_labels.append(self.label2id(example.label[i]))
        target_labels.extend([-1]*(len(w)-1))
    tokens.append('[SEP]')
    if example.label is not None:
      target_labels.extend([-1]*(max_seq_len-len(target_labels)))
    type_ids.append(0)
    token_ids = tokenizer.convert_tokens_to_ids(tokens)
    pos_ids = list(range(len(token_ids)))
    input_mask = [1]*len(token_ids)
    features['input_ids'] = token_ids
    features['type_ids'] = type_ids
    features['position_ids'] = pos_ids
    features['input_mask'] = input_mask
    padding_size = max(0, max_seq_len - len(token_ids))
    for f in features:
      features[f].extend([0]*padding_size)

    for f in features:
      features[f] = torch.tensor(features[f], dtype=torch.int)
    if example.label is not None: # and example.label[0]>=0 and example.label[1]>=0:
      features['labels'] = torch.tensor(target_labels, dtype=torch.int)
    return features

  def extract_docs(self, path):
    docs = []
    with open(path, 'r', encoding='utf-8') as fs:
      doc = []
      sent = []
      for line in fs:
        if line.startswith('-DOCSTART- '):
          if len(sent) > 0:
            doc.append(sent)
          sent = []
          if len(doc) > 0:
            docs.append(doc)
          doc = []
        elif line.strip() == '':
          if len(sent) > 0:
            doc.append(sent)
          sent = []
        else:
          tabs = line.split(' ')
          sent.append([tabs[0], tabs[-1].strip()])
    if len(sent) > 0:
      doc.append(sent)
      sent = []
    if len(doc) > 0:
      docs.append(doc)
      doc = []
    logger.info(f'Loaded {len(docs)} docs, {sum([len(d) for d in docs])} sentences.')
    return docs

def test_ner_load_data():
  tokenizer = GPT2Tokenizer()
  data='/mount/biglm/bert/NER/data/train.txt'
  task = NERTask(os.path.dirname(data), tokenizer)
  #docs = task.extract_docs(data)
  examples = task.load_data(data)
  feature = task.example_to_feature(tokenizer, examples[0], max_seq_len=512)
  pdb.set_trace()