from vncorenlp import VnCoreNLP
from typing import Union
from transformers import AutoConfig, AutoTokenizer
from Model.NER.VLSP2021.Ner_CRF import PhoBertCrf,PhoBertSoftmax,PhoBertLstmCrf
import re
import os
import torch
import itertools
import numpy as np
MODEL_MAPPING = {
'vinai/phobert-base': {
'softmax': PhoBertSoftmax,
'crf': PhoBertCrf,
'lstm_crf': PhoBertLstmCrf
},
}
def normalize_text(txt: str) -> str:
# Remove special character
txt = re.sub("\xad|\u200b|\ufeff", "", txt)
# Normalize vietnamese accents
txt = re.sub(r"òa", "oà", txt)
txt = re.sub(r"óa", "oá", txt)
txt = re.sub(r"ỏa", "oả", txt)
txt = re.sub(r"õa", "oã", txt)
txt = re.sub(r"ọa", "oạ", txt)
txt = re.sub(r"òe", "oè", txt)
txt = re.sub(r"óe", "oé", txt)
txt = re.sub(r"ỏe", "oẻ", txt)
txt = re.sub(r"õe", "oẽ", txt)
txt = re.sub(r"ọe", "oẹ", txt)
txt = re.sub(r"ùy", "uỳ", txt)
txt = re.sub(r"úy", "uý", txt)
txt = re.sub(r"ủy", "uỷ", txt)
txt = re.sub(r"ũy", "uỹ", txt)
txt = re.sub(r"ụy", "uỵ", txt)
txt = re.sub(r"Ủy", "Uỷ", txt)
txt = re.sub(r'"', '”', txt)
# Remove multi-space
txt = re.sub(" +", " ", txt)
return txt.strip()
class ViTagger(object):
def __init__(self, model_path: Union[str or os.PathLike], no_cuda=False):
self.device = 'cuda' if not no_cuda and torch.cuda.is_available() else 'cpu'
print("[ViTagger] VnCoreNLP loading ...")
self.rdrsegmenter = VnCoreNLP("VnCoreNLP/VnCoreNLP-1.1.1.jar", annotators="wseg", max_heap_size='-Xmx500m')
print("[ViTagger] Model loading ...")
self.model, self.tokenizer, self.max_seq_len, self.label2id, self.use_crf = self.load_model(model_path, device=self.device)
self.id2label = {idx: label for idx, label in enumerate(self.label2id)}
print("[ViTagger] All ready!")
@staticmethod
def load_model(model_path: Union[str or os.PathLike], device='cpu'):
if device == 'cpu':
checkpoint_data = torch.load(model_path, map_location='cpu')
else:
checkpoint_data = torch.load(model_path)
args = checkpoint_data["args"]
max_seq_len = args.max_seq_length
use_crf = True if 'crf' in args.model_arch else False
tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=False)
config = AutoConfig.from_pretrained(args.model_name_or_path, num_labels=len(args.label2id))
model_clss = MODEL_MAPPING[args.model_name_or_path][args.model_arch]
model = model_clss(config=config)
model.load_state_dict(checkpoint_data['model'],strict=False)
model.to(device)
model.eval()
return model, tokenizer, max_seq_len, args.label2id, use_crf
def preprocess(self, in_raw: str):
norm_text = normalize_text(in_raw)
sents = []
sentences = self.rdrsegmenter.tokenize(norm_text)
for sentence in sentences:
sents.append(sentence)
return sents
def convert_tensor(self, tokens):
seq_len = len(tokens)
encoding = self.tokenizer(tokens,
padding='max_length',
truncation=True,
is_split_into_words=True,
max_length=self.max_seq_len)
if 'vinai/phobert' in self.tokenizer.name_or_path:
print(' '.join(tokens))
subwords = self.tokenizer.tokenize(' '.join(tokens))
valid_ids = np.zeros(len(encoding.input_ids), dtype=int)
label_marks = np.zeros(len(encoding.input_ids), dtype=int)
i = 1
for idx, subword in enumerate(subwords[:self.max_seq_len - 2]):
if idx != 0 and subwords[idx - 1].endswith("@@"):
continue
if self.use_crf:
valid_ids[i - 1] = idx + 1
else:
valid_ids[idx + 1] = 1
i += 1
else:
valid_ids = np.zeros(len(encoding.input_ids), dtype=int)
label_marks = np.zeros(len(encoding.input_ids), dtype=int)
i = 1
word_ids = encoding.word_ids()
for idx in range(1, len(word_ids)):
if word_ids[idx] is not None and word_ids[idx] != word_ids[idx - 1]:
if self.use_crf:
valid_ids[i - 1] = idx
else:
valid_ids[idx] = 1
i += 1
if self.max_seq_len >= seq_len + 2:
label_marks[:seq_len] = [1] * seq_len
else:
label_marks[:-2] = [1] * (self.max_seq_len - 2)
if self.use_crf and label_marks[0] == 0:
raise f"{tokens} have mark == 0 at index 0!"
item = {key: torch.as_tensor([val]).to(self.device, dtype=torch.long) for key, val in encoding.items()}
item['valid_ids'] = torch.as_tensor([valid_ids]).to(self.device, dtype=torch.long)
item['label_masks'] = torch.as_tensor([valid_ids]).to(self.device, dtype=torch.long)
return item
def extract_entity_doc(self, in_raw: str):
sents = self.preprocess(in_raw)
print(sents)
entities_doc = []
for sent in sents:
item = self.convert_tensor(sent)
with torch.no_grad():
outputs = self.model(**item)
entity = None
if isinstance(outputs.tags[0], list):
tags = list(itertools.chain(*outputs.tags))
else:
tags = outputs.tags
for w, l in list(zip(sent, tags)):
w = w.replace("_", " ")
tag = self.id2label[l]
if not tag == 'O':
parts = tag.split('-', 1)
prefix = parts[0]
tag = parts[1] if len(parts) > 1 else ""
if entity is None:
entity = (w, tag)
else:
if entity[-1] == tag:
if prefix == 'I':
entity = (entity[0] + f' {w}', tag)
else:
entities_doc.append(entity)
entity = (w, tag)
else:
entities_doc.append(entity)
entity = (w, tag)
elif entity is not None:
entities_doc.append(entity)
if w != ' ':
entities_doc.append((w, 'O'))
entity = None
elif w != ' ':
entities_doc.append((w, 'O'))
entity = None
return entities_doc
def __call__(self, in_raw: str):
sents = self.preprocess(in_raw)
entites = []
for sent in sents:
item = self.convert_tensor(sent)
with torch.no_grad():
outputs = self.model(**item)
entity = None
if isinstance(outputs.tags[0], list):
tags = list(itertools.chain(*outputs.tags))
else:
tags = outputs.tags
for w, l in list(zip(sent, tags)):
w = w.replace("_", " ")
tag = self.id2label[l]
if not tag == 'O':
prefix, tag = tag.split('-')
if entity is None:
entity = (w, tag)
else:
if entity[-1] == tag:
if prefix == 'I':
entity = (entity[0] + f' {w}', tag)
else:
entites.append(entity)
entity = (w, tag)
else:
entites.append(entity)
entity = (w, tag)
elif entity is not None:
entites.append(entity)
entity = None
else:
entity = None
return entites