Что тут происходит?
# -*- coding: utf-8 -*-
"""Копия блокнота "bert_ner.ipynb"

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1FtLNOcOGE_aJq6IIxI17Na0DyFlNnUp_
"""

!python -m pip install spacy  # в данном примере нужен только displacy для красивых визуализаций

!python -m pip install datasets  # чтобы загружать датасеты с HF

!python -m pip install transformers  # чтобы работать с трансформерными нейросетями и загружать их с HF

!python -m pip install evaluate  # библиотека из экосистемы HF для интеграции метрик качества в Trainer

!python -m pip install seqeval  # нужно для расчёта F1 по сущностям, а не по токенам

!python -m pip install tensorboard  # чтобы выводить графики обучения

import os
import random
from typing import Dict, List, Tuple, Union  # эмуляция статической типизации

import matplotlib.colors as mcolors  # красиво раскрасим наши именованные сущности
from nltk.tokenize.treebank import TreebankWordDetokenizer
from spacy import displacy
import evaluate
import numpy as np
from datasets import load_dataset
import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification
from transformers import DataCollatorForTokenClassification, TrainingArguments, Trainer
from transformers import pipeline

# Хороший датасет для NER по текстам на тему астрофизики.
# Этой задаче в 2022 году была посвящена статья https://aclanthology.org/2022.wiesp-1.1.pdf
# и даже целый воркшоп https://ui.adsabs.harvard.edu/WIESP/
DATASET_NAME = 'adsabs/WIESP2022-NER'

# Фиксируем все рандом сиды, которые только можем
RANDOM_SEED = 42
random.seed(RANDOM_SEED)
torch.manual_seed(RANDOM_SEED)
torch.cuda.manual_seed(RANDOM_SEED)
np.random.seed(RANDOM_SEED)

trainset = load_dataset(DATASET_NAME, split='train')

for k in sorted(trainset[0].keys()):
    print(f'{k}\t{trainset[0][k]}')

label_set = set()
for it in trainset['ner_tags']:
    label_set |= set(it)
label_list = ['O'] + sorted(list(label_set - {'O'}))

for it in label_list: print(it)

print(f'Size of the entity tag dictionary is {len(label_list)}.')

entity_classes = sorted(list(set(
    map(
        lambda it2: it2[2:],
        filter(
            lambda it1: it1 != 'O',
            label_list
        )
    )
)))

for it in entity_classes: print(it)

print(f'Number of entity classes is {len(entity_classes)}.')

entity_colors = [mcolors.rgb2hex((0.5 + random.random() / 2, 0.5 +random.random() / 2, 0.5 +random.random() / 2))
                 for _ in range(len(entity_classes))]

len(entity_classes) == len(entity_colors)

def bio_to_spans(bio: List[str]) -> List[Dict[str, Union[int, str]]]:
    # Функция нужна для красивого изображения текста с именованными сущностями
    # с использованием displacy https://spacy.io/usage/visualizers в стиле "span".
    ne_tag = ''
    start_pos = -1
    bounds = []
    for idx, val in enumerate(bio):
        if val.upper() == 'O':
            if start_pos >= 0:
                bounds.append({
                    'start_token': start_pos,
                    'end_token': idx,
                    'label': ne_tag
                })
            start_pos = -1
            ne_tag = ''
        elif val.upper().startswith('B-'):
            if start_pos >= 0:
                bounds.append({
                    'start_token': start_pos,
                    'end_token': idx,
                    'label': ne_tag
                })
            start_pos = idx
            ne_tag = val[2:]
    if start_pos >= 0:
        bounds.append({
            'start_token': start_pos,
            'end_token': len(bio),
            'label': ne_tag
        })
    return bounds

sample_for_rendering = {
    'text': TreebankWordDetokenizer().detokenize(trainset[0]['tokens']),
    'spans': bio_to_spans(trainset[0]['ner_tags']),
    'tokens': trainset[0]['tokens'],
}

rendered = displacy.render(
    sample_for_rendering, style='span',
    options={'ents': entity_classes, 'colors': dict(zip(entity_classes, entity_colors))},
    manual=True, jupyter=True
)

def bio_to_ent(tokens: List[str], bio: List[str]) -> Tuple[str, List[Dict[str, Union[int, str]]]]:
    # Функция нужна для ещё более красивого изображения текста с именованными сущностями
    # с использованием displacy https://spacy.io/usage/visualizers в стиле "ent".
    if len(tokens) != len(bio):
        err_msg = f'Tokens do not correspond to their labels: {len(tokens)} != {len(bio)}!'
        raise RuntimeError(err_msg)
    full_text = TreebankWordDetokenizer().detokenize(tokens)
    token_bounds = []
    previous_pos = 0
    for cur in tokens:
        found_idx = full_text[previous_pos:].find(cur)
        if found_idx < 0:
            err_msg = f'The token {cur} is not found in the text "{full_text}".'
            raise RuntimeError(err_msg)
        token_start = found_idx + previous_pos
        token_end = token_start + len(cur)
        token_bounds.append((token_start, token_end))
        previous_pos = token_end
    entity_spans = bio_to_spans(bio)
    entity_bounds = []
    for cur in entity_spans:
        entity_class = cur['label']
        entity_start = token_bounds[cur['start_token']][0]
        entity_end = token_bounds[cur['end_token'] - 1][1]
        entity_bounds.append({
            'start': entity_start,
            'end': entity_end,
            'label': entity_class
        })
    del token_bounds, entity_spans
    return full_text, entity_bounds

sample_for_rendering_2 = dict(zip(
    ('text', 'ents'),
    bio_to_ent(trainset[0]['tokens'], trainset[0]['ner_tags'])
))

rendered_2 = displacy.render(
    sample_for_rendering_2, style='ent',
    options={'ents': entity_classes, 'colors': dict(zip(entity_classes, entity_colors))},
    manual=True, jupyter=True
)

MODEL_NAME = 'FacebookAI/xlm-roberta-base'  # бертоподобный языконезависимый трансформер-энкодер

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

example = trainset[0]
tokenized_input = tokenizer(example['tokens'], is_split_into_words=True)
tokens = tokenizer.convert_ids_to_tokens(tokenized_input['input_ids'])
print(tokens)

def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(examples['tokens'], truncation=True, is_split_into_words=True)
    labels = []
    for i, label in enumerate(examples['ner_tags']):
        word_ids = tokenized_inputs.word_ids(batch_index=i)  # Map tokens to their respective word.
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:  # Set the special tokens to -100.
            if word_idx is None:
                label_ids.append(-100)
            elif word_idx != previous_word_idx:  # Only label the first token of a given word.
                label_ids.append(label_list.index(label[word_idx]))
            else:
                label_ids.append(-100)
            previous_word_idx = word_idx
        labels.append(label_ids)

    tokenized_inputs['labels'] = labels
    return tokenized_inputs

tokenized_trainset = trainset.map(tokenize_and_align_labels, batched=True)

for k in sorted(tokenized_trainset[0].keys()):
    print(f'{k}\t{tokenized_trainset[0][k]}')

valset = load_dataset(DATASET_NAME, split='validation')

tokenized_valset = valset.map(tokenize_and_align_labels, batched=True)

data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

seqeval = evaluate.load('seqeval')

def compute_metrics(p):
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)

    true_predictions = [
        [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    true_labels = [
        [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]

    results = seqeval.compute(predictions=true_predictions, references=true_labels)
    return {
        'precision': results['overall_precision'],
        'recall': results['overall_recall'],
        'f1': results['overall_f1'],
        'accuracy': results['overall_accuracy'],
    }

id2label = dict(enumerate(label_list))

print(id2label)

label2id = dict((val, idx) for idx, val in enumerate(label_list))

print(label2id)

model = AutoModelForTokenClassification.from_pretrained(
    MODEL_NAME,
    num_labels=len(label_list),
    id2label=id2label,
    label2id=label2id
)

MODEL_NAME_ON_DISK = os.path.abspath('astro_ner')

training_args = TrainingArguments(
    output_dir=MODEL_NAME_ON_DISK,
    logging_dir=os.path.join(MODEL_NAME_ON_DISK, 'logs'),
    learning_rate=1e-4,
    warmup_ratio=0.5,  # делаем "прогрев": начинаем с околонулевого lr и до середины обучения (то есть до пятой эпохи, если у нас их 10) линейно увеличиваем до 1e-4
    per_device_train_batch_size=8,
    per_device_eval_batch_size=16,  # при тестировании не считаются градиенты, поэтому мини-батч можно и побольше
    num_train_epochs=10,
    weight_decay=0.01,  # для регуляризации обновлений весов
    eval_strategy='epoch',
    save_strategy='epoch',
    save_total_limit=2,  # чтобы не забивать жёсткий диск, будем сохранять только 2 чекпойнта: лучший и последний
    logging_strategy='epoch',
    report_to='tensorboard',  # хотим красиво нарисовать графики обучения в tensorboard
    metric_for_best_model='f1',
    greater_is_better=True,  # чем больше f1, тем лучше
    load_best_model_at_end=True,
    seed=RANDOM_SEED,
    data_seed=RANDOM_SEED
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_trainset,
    eval_dataset=tokenized_valset,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

# Commented out IPython magic to ensure Python compatibility.
# %load_ext tensorboard
# %tensorboard --logdir '{MODEL_NAME_ON_DISK}'/logs

for it in os.listdir(MODEL_NAME_ON_DISK): print(it)

possible_checkpoints = sorted(
    list(map(
        lambda it2: os.path.join(MODEL_NAME_ON_DISK, it2),
        filter(
            lambda it1: it1.startswith('checkpoint-'),
            os.listdir(MODEL_NAME_ON_DISK)
        )
    )),
    key=lambda it3: -len(os.listdir(it3))
)

for it in possible_checkpoints: print(it)

import gc
gc.collect()
torch.cuda.empty_cache()

classifier = pipeline('ner', model=possible_checkpoints[0], device=0)

original_text = 'The authors would like to thank Adam Burgasser, Brendan Bowler, Kelle Cruz, Mike Cushing, Michael Liu, and Emily Rice for useful discussions on benchmark systems, data treatment, and various data-model comparison approaches. The authors thank Richard Freedman and Roxana Lupu for providing gas opacities and Caroline Morley for radiative transfer code comparisons and helpful discussions. We thank Jacob Lustig-Yeager and Kyle Luther for rewriting portions of the code in python and C for significant speed improvements and also Dan Foreman-Mackey for making EMCEE available to the community. Finally, we thank the anonymous referee and statistics consultant for useful and insightful comments.'

original_res = classifier(original_text)
for it in original_res: print(it)

original_res_2 = classifier(original_text, aggregation_strategy='first')
for it in original_res_2: print(it)

sample_for_rendering_3 = {
    'text': original_text,
    'ents': [{'start': it['start'], 'end': it['end'], 'label': it['entity_group']}
             for it in original_res_2]
}

rendered_3 = displacy.render(
    sample_for_rendering_3, style='ent',
    options={'ents': entity_classes, 'colors': dict(zip(entity_classes, entity_colors))},
    manual=True, jupyter=True
)

ru_text = 'Авторы хотели бы поблагодарить Адама Бургассера, Брендана Боулера, Келли Круз, Майка Кушинга, Майкла Лю и Эмили Райс за полезные обсуждения систем бенчмарков, обработки данных и различных подходов к сравнению моделей данных. Авторы благодарят Ричарда Фридмана и Роксану Лупу за предоставление непрозрачности газа и Кэролайн Морли за сравнения кодов переноса излучения и полезные обсуждения. Мы благодарим Джейкоба Люстига-Йегера и Кайла Лютера за переписывание частей кода на Python и C для значительного улучшения скорости, а также Дэна Формана-Макки за предоставление EMCEE сообществу. Наконец, мы благодарим анонимного рецензента и консультанта по статистике за полезные и проницательные комментарии.'

ru_res = classifier(ru_text, aggregation_strategy='first')
for it in ru_res: print(it)

sample_for_rendering_4 = {
    'text': ru_text,
    'ents': [{'start': it['start'], 'end': it['end'], 'label': it['entity_group']}
             for it in ru_res]
}

rendered_4 = displacy.render(
    sample_for_rendering_4, style='ent',
    options={'ents': entity_classes, 'colors': dict(zip(entity_classes, entity_colors))},
    manual=True, jupyter=True
)

testset = load_dataset(DATASET_NAME, split='test')

def find_token(token_bounds: List[Tuple[int, int]], char_idx: int) -> int:
    res = -1
    for token_idx, (token_start, token_end) in enumerate(token_bounds):
        if (char_idx >= token_start) and (char_idx < token_end):
            res = token_idx
            break
    return res

def predictions_to_bio(text: str, tokens: List[str], predictions: List[Tuple[int, int, str]]) -> List[str]:
    token_bounds = []
    token_labels = []
    start_pos = 0
    for cur_token in tokens:
        found_idx = text[start_pos:].find(cur_token)
        if found_idx < 0:
            err_msg = f'The token {cur_token} is not found in the text {text}'
            raise RuntimeError(err_msg)
        token_start = found_idx + start_pos
        token_end = token_start + len(cur_token)
        start_pos = token_end
        token_bounds.append((token_start, token_end))
        token_labels.append('O')
    for span_start, span_end, span_label in predictions:
        start_token = find_token(token_bounds, span_start)
        end_token = find_token(token_bounds, span_end - 1)
        if (start_token >= 0) and (end_token >= 0):
            for token_idx in range(start_token, end_token + 1):
                token_labels[token_idx] = span_label
        elif start_token >= 0:
            token_labels[start_token] = span_label
        elif end_token >= 0:
            token_labels[end_token] = span_label
    corrected_token_labels = []
    previous_label = 'O'
    for cur_label in token_labels:
        if cur_label == previous_label:
            corrected_token_labels.append(cur_label)
        else:
            if (cur_label == 'O') or cur_label.startswith('B-'):
                corrected_token_labels.append(cur_label)
            else:
                if previous_label == 'O':
                    corrected_token_labels.append('B-' + cur_label[2:])
                elif previous_label[2:] != cur_label[2:]:
                    corrected_token_labels.append('B-' + cur_label[2:])
                else:
                    corrected_token_labels.append(cur_label)
        previous_label = cur_label
    return corrected_token_labels

from tqdm.notebook import tqdm

y_true = []
y_pred = []
for tokens, reference_tags in tqdm(zip(testset['tokens'], testset['ner_tags']), total=len(testset)):
    y_true.append(reference_tags)
    cur_text = TreebankWordDetokenizer().detokenize(tokens)
    cur_res = classifier(cur_text)
    y_pred.append(
        predictions_to_bio(
            cur_text,
            tokens,
            [(it['start'], it['end'], it['entity']) for it in cur_res]
        )
    )

for true_label, predicted_label in zip(y_true[0], y_pred[0]):
    print('{0:>25}   {1:>25}'.format(true_label, predicted_label))

from seqeval.scheme import IOB2
from seqeval.metrics import classification_report

print(classification_report(y_true, y_pred, digits=4))

print(classification_report(y_true, y_pred, digits=4, mode='strict', scheme=IOB2))