wehub-resource-sync
70 lines
2.9 KiB
Python
70 lines
2.9 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2020-06-12 18:00
|
|
from typing import Any
|
|
|
|
import torch
|
|
from hanlp_common.util import merge_locals_kwargs
|
|
|
|
import hanlp.utils.span_util
|
|
from hanlp.components.taggers.rnn_tagger import RNNTagger
|
|
from hanlp.metrics.chunking.conlleval import SpanF1
|
|
|
|
|
|
class RNNNamedEntityRecognizer(RNNTagger):
|
|
|
|
def __init__(self, **kwargs) -> None:
|
|
"""An old-school RNN tagger using word2vec or fasttext embeddings.
|
|
|
|
Args:
|
|
**kwargs: Predefined config.
|
|
"""
|
|
super().__init__(**kwargs)
|
|
|
|
def build_metric(self, **kwargs):
|
|
return SpanF1(self.tagging_scheme)
|
|
|
|
def evaluate_dataloader(self, data, criterion, logger=None, ratio_width=None, **kwargs):
|
|
loss, metric = super().evaluate_dataloader(data, criterion, logger, ratio_width, **kwargs)
|
|
if logger:
|
|
logger.info(metric.result(True, False)[-1])
|
|
return loss, metric
|
|
|
|
def fit(self, trn_data, dev_data, save_dir, batch_size=50, epochs=100, embed=100, rnn_input=None, rnn_hidden=256,
|
|
drop=0.5, lr=0.001, patience=10, crf=True, optimizer='adam', token_key='token', tagging_scheme=None,
|
|
anneal_factor: float = 0.5, delimiter=None, anneal_patience=2, devices=None,
|
|
token_delimiter=None,
|
|
logger=None,
|
|
verbose=True, **kwargs):
|
|
return super().fit(**merge_locals_kwargs(locals(), kwargs))
|
|
|
|
def update_metrics(self, metric, logits, y, mask, batch, prediction):
|
|
logits = self.decode_output(logits, mask, batch)
|
|
if isinstance(logits, torch.Tensor):
|
|
logits = logits.tolist()
|
|
metric(self._id_to_tags(logits), batch['tag'])
|
|
|
|
def predict(self, tokens: Any, batch_size: int = None, **kwargs):
|
|
return super().predict(tokens, batch_size, **kwargs)
|
|
|
|
def predict_data(self, data, batch_size, **kwargs):
|
|
outputs = super().predict_data(data, batch_size)
|
|
tagging_scheme = self.tagging_scheme
|
|
if tagging_scheme == 'IOBES':
|
|
entities = [hanlp.utils.span_util.iobes_tags_to_spans(y) for y in outputs]
|
|
elif tagging_scheme == 'BIO':
|
|
entities = [hanlp.utils.span_util.bio_tags_to_spans(y) for y in outputs]
|
|
elif tagging_scheme == 'BIOUL':
|
|
entities = [hanlp.utils.span_util.bioul_tags_to_spans(y) for y in outputs]
|
|
else:
|
|
raise ValueError(f'Unrecognized tag scheme {tagging_scheme}')
|
|
for i, (tokens, es) in enumerate(zip(data, entities)):
|
|
outputs[i] = [(self.config.token_delimiter.join(tokens[b:e + 1]), t, b, e + 1) for t, (b, e) in es]
|
|
return outputs
|
|
|
|
def save_config(self, save_dir, filename='config.json'):
|
|
if self.config.token_delimiter is None:
|
|
self.config.token_delimiter = '' if all(
|
|
[len(x) == 1 for x in self.vocabs[self.config.token_key].idx_to_token[-100:]]) else ' '
|
|
super().save_config(save_dir, filename)
|