Условие

Дан CSV-файл с двумя колонками:

text — текст с подчёркнутыми тэгами <tag1> и <tag2>;
comment — комментарий с дополнительной информацией (может быть пустым).

Нужно для каждой строки определить семантическую связь между двумя сущностями (tag1, tag2). Возможные метки:

Cause-Effect (причина-следствие)
Component-Whole (компонент-целое)
Entity-Origin (сущность-источник)
Entity-Destination (сущность-получатель)
Instrument-Agency (инструмент-агент)
Member-Collection (член-коллекция)
Message-Topic (сообщение-тема)
Product-Producer (продукт-производитель)
Other — если ни одно не подходит

Порядок тэгов важен (Cause-Effect(tag1, tag2) ≠ Cause-Effect(tag2, tag1)).

Формат вывода

CSV с колонкой Relation_type → label, по одной метке на строку входа.

Решение

Подход

Это классическая задача Relation Classification в NLP. Данные — почти точно SemEval-2010 Task 8 (классическая benchmark). Сильный baseline — fine-tuning трансформера с маркерами сущностей.

1. Препроцессинг

Заменяем <tag1>...</tag1> и <tag2>...</tag2> на специальные токены:

def insert_markers(text):
    return (text
        .replace('<tag1>', ' [E1] ').replace('</tag1>', ' [/E1] ')
        .replace('<tag2>', ' [E2] ').replace('</tag2>', ' [/E2] '))

2. Модель

BERT (или xlm-roberta-base для мультиязычности) с классификацией по [CLS] или конкатенацией представлений [E1] и [E2]-токенов.

from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    Trainer, TrainingArguments
)
from datasets import Dataset
import pandas as pd
 
train = pd.read_csv("train.csv")
test = pd.read_csv("test.csv")
train['text'] = train['text'].apply(insert_markers)
test['text'] = test['text'].apply(insert_markers)
 
# 8 семантических связей × 2 направления (tag1→tag2 и tag2→tag1) + Other = 17 классов.
# Метка в data — это пара «(имя_связи, направление)»; если направление не выделено
# отдельной колонкой, оно зашито в `Relation_type` (например, `Cause-Effect(e1,e2)`).
LABELS = sorted(train['Relation_type'].unique())
label2id = {l: i for i, l in enumerate(LABELS)}
 
tok = AutoTokenizer.from_pretrained("xlm-roberta-base")
tok.add_special_tokens({"additional_special_tokens": ["[E1]", "[/E1]", "[E2]", "[/E2]"]})
 
def encode(batch):
    enc = tok(batch['text'], truncation=True, padding='max_length', max_length=128)
    enc['labels'] = [label2id[l] for l in batch['Relation_type']]
    return enc
 
ds_train = Dataset.from_pandas(train[['text', 'Relation_type']]).map(encode, batched=True)
 
model = AutoModelForSequenceClassification.from_pretrained(
    "xlm-roberta-base", num_labels=len(LABELS)
)
model.resize_token_embeddings(len(tok))
 
args = TrainingArguments(
    output_dir="out",
    num_train_epochs=4,
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    eval_strategy="epoch",          # transformers >= 4.41 (раньше: evaluation_strategy)
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="f1_macro",
    greater_is_better=True,
)
 
from sklearn.metrics import f1_score, accuracy_score
import numpy as np
 
def compute_metrics(p):
    preds = np.argmax(p.predictions, axis=1)
    return {
        "accuracy": accuracy_score(p.label_ids, preds),
        "f1_macro": f1_score(p.label_ids, preds, average='macro'),
    }
 
trainer = Trainer(
    model=model, args=args,
    train_dataset=ds_train,
    compute_metrics=compute_metrics,
)
trainer.train()

3. Использование комментария

Если comment непустой — конкатенируем:

text_with_comment = text + " [SEP] " + (comment if comment else "")

4. Предсказание

ds_test = Dataset.from_pandas(test[['text']]).map(
    lambda b: tok(b['text'], truncation=True, padding='max_length', max_length=128),
    batched=True,
)
preds = trainer.predict(ds_test).predictions.argmax(axis=1)
test['Relation_type'] = [LABELS[p] for p in preds]
test[['Relation_type']].to_csv('answers.csv', index=False)

5. Сильные приёмы

Entity Markers — обернуть сущности [E1]...[/E1] (как выше).
Использовать представление [E1] + [E2] вместо [CLS] (R-BERT).
Доменное предобучение на текстах из той же доменной области.
Аугментация: переставлять tag1 ↔ tag2 для соответствующего направления метки (учится симметрии).
Ensemble нескольких сидов.

6. Если без GPU

Базовый baseline без BERT:

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
 
# Признаки: tf-idf char-grams + слова между tag1 и tag2.
def extract_between(text):
    import re
    m = re.search(r'<tag1>.*?</tag1>(.*?)<tag2>.*?</tag2>', text, re.S)
    return m.group(1) if m else ""
 
train['between'] = train['text'].apply(extract_between)
vect = TfidfVectorizer(ngram_range=(1, 2), max_features=20000)
X = vect.fit_transform(train['text'] + ' ||| ' + train['between'])
clf = LogisticRegression(max_iter=1000, C=1.0)
clf.fit(X, train['Relation_type'])

Подводные камни

Направленность меток. Cause-Effect(A, B) ≠ Cause-Effect(B, A). Метки часто кодируются с направлением (Cause-Effect(e1,e2) и Cause-Effect(e2,e1)) → 17 классов вместо 9.
Other — большой класс. Он часто доминирует. Используйте f1_macro, не accuracy.
Длина текста. Тэги могут быть в длинных предложениях. max_length=128 мала; используйте 256.
Маркеры для tokenizer. После add_special_tokens обязательно model.resize_token_embeddings(len(tok)).
Утечка между train/test. Если в данных одни и те же entity-пары — модель запоминает, не учится связи.
Class imbalance. Внутри 9 классов один-два могут быть редкими (1-2%). Используйте class_weight или oversample.
comment иногда содержит ответ. Если в comment упомянута связь явно («Это причина») — модель может на это полагаться. Хорошая фича, но риск утечки.
Voted ensemble направления. Прогноз направления через два запуска (с переставленными tag1, tag2) и majority voting.

Эталонный ответ

Fine-tuning XLM-RoBERTa с entity markers [E1]/[E2] на задаче 9-классовой классификации семантических связей. CV с f1_macro, ансамбль 3-5 сидов. Базовый baseline без GPU — TF-IDF + LogReg на тексте + контексте между сущностями.

Стажировка ML — Инопланетный шифр: классификация семантических связей