Обучение для задачи NER для русского языка с помощью GLiNER и корпуса Nerus
Распознавание именованного объекта или NER является одной из задач NLP. Целью ее решения является нахождение именованных объектов в тексте. Существуют множество методов и моделей, которые решают данную задачу, в том числе Большие Языковые Модели (LLMs).
GLiNER можно считать альтернативой LLMs, которую можно запускать на слабых GPU и на CPU. GLiNER - NER модель, способная идентифицировать любой тип объекта с использованием двунаправленного преобразователя-кодера.
Статья: https://arxiv.org/pdf/2311.08526
NERUS
Для обучения модели будет использовать корпус NERUS, а также Google Colab с картой T4:
https://github.com/natasha/nerus
Установка:
!pip install nerus
Загрузка корпуса:
!wget https://storage.yandexcloud.net/natasha-nerus/data/nerus_lenta.conllu.gz
Проверяем:
from nerus import load_nerus
docs = load_nerus("nerus_lenta.conllu.gz")
doc = next(docs)
doc.ner
Результат:
NERMarkup(
text='Вице-премьер по социальным вопросам Татьяна Голикова рассказала, в каких регионах России зафиксирована наиболее высокая смертность от рака, сообщает РИА Новости. По словам Голиковой, чаще всего онкологические заболевания становились причиной смерти в Псковской, Тверской, Тульской и Орловской областях, а также в Севастополе. Вице-премьер напомнила, что главные факторы смертности в России — рак и болезни системы кровообращения. В начале года стало известно, что смертность от онкологических заболеваний среди россиян снизилась впервые за три года. По данным Росстата, в 2017 году от рака умерли 289 тысяч человек. Это на 3,5 процента меньше, чем годом ранее.',
spans=[Span(
start=36,
stop=52,
type='PER'
),
Span(
start=82,
stop=88,
type='LOC'
),
Span(
start=149,
stop=160,
type='ORG'
),
Span(
start=172,
stop=181,
type='PER'
),
Span(
start=251,
stop=260,
type='LOC'
),
Span(
start=262,
stop=270,
type='LOC'
),
Span(
start=272,
stop=280,
type='LOC'
),
Span(
start=283,
stop=301,
type='LOC'
),
Span(
start=313,
stop=324,
type='LOC'
),
Span(
start=383,
stop=389,
type='LOC'
),
Span(
start=560,
stop=568,
type='ORG'
)]
)
Для токенизации будем использовать следующую простую пару функций:
import re
import json
def tokenize_text(text):
return re.findall(r'\w+(?:[-_]\w+)*|\S', text)
"""
docs - открытый load_nerus
max_seq - максимальное значение последовательностей. Их всего больше 700К,
поэтому для тестов можно задать значение конкретное. Иначе поставить: -1
"""
def process_entities(docs,max_seq):
all_data = [
doc = next(docs)
i = 0
while doc != None:
try:
tokenized_text = [
entity_spans = [
coord = 0
for s in doc.ner.spans:
if coord != s.start:
tokenized_text +=tokenize_text(doc.ner.text[coord:s.start])
a = len(tokenized_text)
tok = tokenize_text(doc.ner.text[s.start:s.stop])
tokenized_text += tok
entity_spans.append([a,a+len(tok)-1,s.type])
coord = s.stop
if len(doc.ner.spans) > 0 and doc.ner.spans[-1].stop < len(doc.ner.text):
tokenized_text +=tokenize_text(doc.ner.text[doc.ner.spans[-1].stop:])
except Exception as e:
continue
all_data.append({"tokenized_text": tokenized_text, "ner": entity_spans})
try:
doc = next(docs)
except Exception as e:
break
i+=1
if i == max_seq:
break
if i%1000 == 0:
print(str(int(i/1000)))
return all_data
Поскольку есть цель только попробовать обучить модель, то загрузим всего 2000 новостей из более 700 тыс. в корпусе.
def save_data_to_file(data, filepath):
with open(filepath, 'w') as f:
json.dump(data, f)
max_seq = 2000
docs = load_nerus("nerus_lenta.conllu.gz")
save_data_to_file(process_entities(docs,max_seq),"nerus_train.json")
Обучение
Инсталляция GLiNER:
!pip install gliner
Импорт:
from gliner import GLiNER
import torch
from tqdm import tqdm
from transformers import get_cosine_schedule_with_warmup
import os
Загружаем 2000 последовательностей токенов:
train_path = "nerus_train.json"
with open(train_path, "r") as f:
data = json.load(f)
Загружаем маленькую модель:
model = GLiNER.from_pretrained("urchade/gliner_small")
Настройки обучения:
from types import SimpleNamespace
# Define the hyperparameters in a config variable
config = SimpleNamespace(
num_steps=1000, # number of training iteration
train_batch_size=2,
eval_every=100, # evaluation/saving steps
save_directory="logs", # where to save checkpoints
warmup_ratio=0.1, # warmup steps
device='cuda',
lr_encoder=1e-5, # learning rate for the backbone
lr_others=5e-5, # learning rate for other parameters
freeze_token_rep=False, # freeze of not the backbone
# Parameters for set_sampling_params
max_types=25, # maximum number of entity types during training
shuffle_types=True, # if shuffle or not entity types
random_drop=True, # randomly drop entity types
max_neg_type_ratio=1, # ratio of positive/negative types, 1 mean 50%/50%, 2 mean 33%/66%, 3 mean 25%/75% ...
max_len=512 # maximum sentence length
)
Функция обучения (взято из репозитория GLiNER):
def train(model, config, train_data, eval_data=None):
model = model.to(config.device)
# Set sampling parameters from config
model.set_sampling_params(
max_types=config.max_types,
shuffle_types=config.shuffle_types,
random_drop=config.random_drop,
max_neg_type_ratio=config.max_neg_type_ratio,
max_len=config.max_len
)
model.train()
# Initialize data loaders
train_loader = model.create_dataloader(train_data, batch_size=config.train_batch_size, shuffle=True)
# Optimizer
optimizer = model.get_optimizer(config.lr_encoder, config.lr_others, config.freeze_token_rep)
pbar = tqdm(range(config.num_steps))
if config.warmup_ratio < 1:
num_warmup_steps = int(config.num_steps * config.warmup_ratio)
else:
num_warmup_steps = int(config.warmup_ratio)
scheduler = get_cosine_schedule_with_warmup(
optimizer,
num_warmup_steps=num_warmup_steps,
num_training_steps=config.num_steps
)
iter_train_loader = iter(train_loader)
for step in pbar:
try:
x = next(iter_train_loader)
except StopIteration:
iter_train_loader = iter(train_loader)
x = next(iter_train_loader)
for k, v in x.items():
if isinstance(v, torch.Tensor):
x[k] = v.to(config.device)
loss = model(x) # Forward pass
# Check if loss is nan
if torch.isnan(loss):
continue
loss.backward() # Compute gradients
optimizer.step() # Update parameters
scheduler.step() # Update learning rate schedule
optimizer.zero_grad() # Reset gradients
description = f"step: {step} | epoch: {step // len(train_loader)} | loss: {loss.item():.2f}"
pbar.set_description(description)
if (step + 1) % config.eval_every == 0:
model.eval()
if eval_data is not None:
results, f1 = model.evaluate(eval_data["samples"], flat_ner=True, threshold=0.5, batch_size=12,
entity_types=eval_data["entity_types"])
print(f"Step={step}\n{results}")
if not os.path.exists(config.save_directory):
os.makedirs(config.save_directory)
model.save_pretrained(f"{config.save_directory}/finetuned_{step}")
model.train()
Определения размера валидации:
valid_size = int(len(data)*.2)
Запуск обучения:
eval_data = {
"entity_types": ["PER","LOC","ORG"],
"samples": data[:valid_size]
}
train(model, config, data, eval_data)valid_size
Сохранение в файл:
model.save_pretrained("small")
Тестирование
Загрузка модели:
md = GLiNER.from_pretrained("small", local_files_only=True)
md.eval()
Проверяем работу модели:
text = """
Владимир Ильич Ленин купил шляпу, а она ему как раз. Россия - щедрая душа. Татьяна Борисовна ела кашу. Московский Государственный университет стал институтом.
"""
labels = ["PER","LOC","ORG"]
entities = md.predict_entities(text, labels, threshold=0.5)
for entity in entities:
print(entity["text"], "=>", entity["label"])
Результат:
Ильич Ленин => PER
Россия => LOC
Татьяна Борисовна => PER
Государственный университет => ORG
Качество не идеальное, для лучшего качества нужно использовать полный корпус для обучения и не менее 100000 итераций.