Обучение для задачи NER для русского языка с помощью GLiNER и корпуса Nerus

Распознавание именованного объекта или NER является одной из задач NLP. Целью ее решения является нахождение именованных объектов в тексте. Существуют множество методов и моделей, которые решают данную задачу, в том числе Большие Языковые Модели (LLMs).

GLiNER можно считать альтернативой LLMs, которую можно запускать на слабых GPU и на CPU. GLiNER - NER модель, способная идентифицировать любой тип объекта с использованием двунаправленного преобразователя-кодера.

Статья: https://arxiv.org/pdf/2311.08526

Git: https://github.com/urchade/GLiNER/

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

2024-05-14_16-03-34

Сохранение в файл:

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 итераций.