NER en documentos checos con XLM-RoBERTa usando 🤗 Accelerate | de Bohumir Buso | noviembre de 2024

🤗 Acelera
Habiendo comenzado en una época en la que los envoltorios eran menos comunes, Me acostumbré a escribir mis propios bucles de entrenamiento, que me resultan más fáciles de depurar. un enfoque que 🤗 Accelerate apoya eficazmente. Resultó beneficioso en este proyecto: no estaba completamente seguro de los datos requeridos y los formatos o formas de las etiquetas y mis datos no coincidían con los ejemplos bien organizados que a menudo se muestran en los tutoriales, pero tuve acceso completo a los cálculos intermedios durante el ciclo de entrenamiento. me permitió iterar rápidamente.

Longitud del contexto
La mayoría de los tutoriales sugieren utilizar cada oración como un único ejemplo de entrenamiento. Sin embargo, en este caso, decidí un contexto más largo sería más adecuado ya que los documentos normalmente contienen referencias a múltiples entidadesmuchos de los cuales son irrelevantes (por ejemplo, abogados, otros acreedores, números de casos). Este contexto más amplio permite que el modelo identifique mejor al cliente relevante. Utilicé 512 tokens de cada documento como ejemplo de formación. Este es un límite máximo común para los modelos, pero se adapta cómodamente a todas las entidades en la mayoría de mis documentos.

Etiquetado de subtokens
En el 🤗 tutorial de clasificación de tokens [1]el enfoque recomendado es:

Etiquetar únicamente el primer token de una palabra determinada. Asignar -100 a otras subtokens de la misma palabra.

Sin embargo, descubrí que el siguiente método sugerido en el tutorial de clasificación de tokens en su curso de PNL [2] funciona mucho mejor:

Cada token recibe la misma etiqueta que el token que inició la palabra que contiene, ya que son parte de la misma entidad. Para tokens dentro de una palabra pero no al principio, reemplazamos el B- con I-

La etiqueta “-100” es una etiqueta especial que se ignora por función de pérdida. Por lo tanto, utilicé sus funciones con cambios menores:

def align_labels_with_tokens(labels, word_ids):
new_labels = []
current_word = None
for word_id in word_ids:
if word_id != current_word:
# Start of a new word!
current_word = word_id
label = -100 if word_id is None else labels[word_id]
new_labels.append(label)
elif word_id is None:
# Special token
new_labels.append(-100)
else:
# Same word as previous token
label = labels[word_id]
# If the label is B-XXX we change it to I-XXX
if label % 2 == 1:
label += 1
new_labels.append(label)

return new_labels

def tokenize_and_align_labels(examples):
tokenizer = AutoTokenizer.from_pretrained("../model/xlm-roberta-large")
tokenized_inputs = tokenizer(
examples["tokens"], truncation=True, is_split_into_words=True,
padding="max_length", max_length=512)
all_labels = examples["ner_tags"]
new_labels = []
for i, labels in enumerate(all_labels):
word_ids = tokenized_inputs.word_ids(i)
new_labels.append(align_labels_with_tokens(labels, word_ids))

tokenized_inputs["labels"] = new_labels
return tokenized_inputs

También usé su postprocess()función:

Para simplificar su parte de evaluación, definimos esto postprocess() función que toma predicciones y etiquetas y las convierte en listas de cadenas.

def postprocess(predictions, labels):
predictions = predictions.detach().cpu().clone().numpy()
labels = labels.detach().cpu().clone().numpy()

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

Pesos de clase
La incorporación de ponderaciones de clases en la función de pérdida mejoró significativamente el rendimiento del modelo.
Si bien este ajuste puede parecer sencillo (sin él, el modelo enfatizaba demasiado la clase “O” mayoritaria), sorprendentemente está ausente en la mayoría de los tutoriales. Implementé una costumbre compute_weights() función para abordar este desequilibrio:

def compute_weights(trainset, num_labels):
c = Counter()
for t in trainset:
c += Counter(t['labels'].tolist())
weights = [sum(c.values())/(c[i]+1) for i in range(num_labels)]
return weights

Bucle de entrenamiento
Definí dos funciones adicionales: PyTorch DataLoader() para gestionar el procesamiento por lotes y un main() función para configurar objetos de entrenamiento distribuidos y ejecutar el ciclo de entrenamiento.

from accelerate import Accelerator, notebook_launcher
from collections import Counter
from datasets import Dataset
from datetime import datetime
import torch
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
from transformers import AutoModelForTokenClassification
from transformers import XLMRobertaConfig, XLMRobertaForTokenClassification
from seqeval.metrics import classification_report, f1_score

def create_dataloaders(trainset, evalset, batch_size, num_workers):
train_dataloader = DataLoader(trainset, shuffle=True,
batch_size=batch_size, num_workers=num_workers)
eval_dataloader = DataLoader(evalset, shuffle=False,
batch_size=batch_size, num_workers=num_workers)
return train_dataloader, eval_dataloader

def main(batch_size, num_workers, epochs, model_path, dataset_tr, dataset_ev, training_type, model_params, dt):
accelerator = Accelerator(split_batches=True)
num_labels = model_params['num_labels']

# Prepare data #
train_ds = Dataset.from_dict(
{"tokens": [d[2][:512] for d in dataset_tr],
"ner_tags": [d[1][:512] for d in dataset_tr]})
eval_ds = Dataset.from_dict(
{"tokens": [d[2][:512] for d in dataset_ev],
"ner_tags": [d[1][:512] for d in dataset_ev]})
trainset = train_ds.map(tokenize_and_align_labels, batched=True,
remove_columns=["tokens", "ner_tags"])
evalset = eval_ds.map(tokenize_and_align_labels, batched=True,
remove_columns=["tokens", "ner_tags"])
trainset.set_format("torch")
evalset.set_format("torch")
train_dataloader, eval_dataloader = create_dataloaders(trainset, evalset,
batch_size, num_workers)

# Type of training #
if training_type=='from_scratch':
config = XLMRobertaConfig.from_pretrained(model_path, **model_params)
model = XLMRobertaForTokenClassification(config)
elif training_type=='transfer_learning':
model = AutoModelForTokenClassification.from_pretrained(model_path,
ignore_mismatched_sizes=True, **model_params)
for param in model.parameters():
param.requires_grad=False
for param in model.classifier.parameters():
param.requires_grad=True
elif training_type=='fine_tuning':
model = AutoModelForTokenClassification.from_pretrained(model_path,
**model_params)
for param in model.parameters():
param.requires_grad=True
for param in model.classifier.parameters():
param.requires_grad=True

# Intantiate the optimizer #
optimizer = torch.optim.AdamW(params=model.parameters(), lr=2e-5)

# Instantiate the learning rate scheduler #
lr_scheduler = ReduceLROnPlateau(optimizer, patience=5)

# Define loss function #
weights = compute_weights(trainset, num_labels)
loss_fct = CrossEntropyLoss(weight=torch.tensor(weights))

# Prepare objects for distributed training #
loss_fct, train_dataloader, model, optimizer, eval_dataloader, lr_scheduler = accelerator.prepare(
loss_fct, train_dataloader, model, optimizer, eval_dataloader, lr_scheduler)

# Training loop #
max_f1 = 0 # for early stopping
for t in range(epochs):
# training
accelerator.print(f"\n\nEpoch {t+1}\n-------------------------------")
model.train()
tr_loss = 0
preds = list()
labs = list()
for batch in train_dataloader:
outputs = model(input_ids=batch['input_ids'],
attention_mask=batch['attention_mask'])
labels = batch["labels"]
loss = loss_fct(outputs.logits.view(-1, num_labels), labels.view(-1))
accelerator.backward(loss)
optimizer.step()
optimizer.zero_grad()
tr_loss += loss
predictions = outputs.logits.argmax(dim=-1)
predictions_gathered = accelerator.gather(predictions)
labels_gathered = accelerator.gather(labels)
true_predictions, true_labels = postprocess(predictions_gathered, labels_gathered)
preds.extend(true_predictions)
labs.extend(true_labels)

lr_scheduler.step(tr_loss)

accelerator.print(f"Train loss: {tr_loss/len(train_dataloader):>8f} \n")
accelerator.print(classification_report(labs, preds))

# evaluation
model.eval()
ev_loss = 0
preds = list()
labs = list()
for batch in eval_dataloader:
with torch.no_grad():
outputs = model(input_ids=batch['input_ids'],
attention_mask=batch['attention_mask'])
labels = batch["labels"]
loss = loss_fct(outputs.logits.view(-1, num_labels), labels.view(-1))

ev_loss += loss
predictions = outputs.logits.argmax(dim=-1)
predictions_gathered = accelerator.gather(predictions)
labels_gathered = accelerator.gather(labels)
true_predictions, true_labels = postprocess(predictions_gathered, labels_gathered)
preds.extend(true_predictions)
labs.extend(true_labels)

accelerator.print(f"Eval loss: {ev_loss/len(eval_dataloader):>8f} \n")
accelerator.print(classification_report(labs, preds))

accelerator.print(f"Current Learning Rate: {optimizer.param_groups[0]['lr']}")

# checkpoint best model
if f1_score(labs, preds) > max_f1:
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(f"../model/xlml_ner/{dt}/",
is_main_process=accelerator.is_main_process,
save_function=accelerator.save)
accelerator.print(f"Model saved during {t+1}. epoch.")
max_f1 = f1_score(labs, preds)
best_epoch = t

# early stopping
if (t - best_epoch) > 10:
accelerator.print(f"Early stopping after {t+1}. epoch.")
break

accelerator.print("Done!")

Con todo preparado, el modelo está listo para entrenar. Sólo necesito iniciar el proceso:

label_list = [
"O",
"B-evcu", "I-evcu", # variable symbol of creditor
"B-rc", "I-rc", # birth ID
"B-prijmeni", "I-prijmeni", # surname
"B-jmeno", "I-jmeno", # given name
"B-datum", "I-datum", # birth date
]
id2label = {a: b for a,b in enumerate(label_list)}
label2id = {b: a for a,b in enumerate(label_list)}

num_workers = 6 # number of GPUs
batch_size = num_workers*2
epochs = 100
model_path = "../model/xlm-roberta-large"
training_type = "fine_tuning" # from_scratch / transfer_learning / fine_tuning
model_params = {"id2label": id2label, "label2id": label2id, "num_labels": 11}
dt = datetime.now().strftime("%Y%m%d_%H%M%S")
os.mkdir(f"../model/xlml_ner/{dt}")

notebook_launcher(main, args=(batch_size, num_workers, epochs, model_path,
dataset_tr, dataset_ev, training_type, model_params, dt),
num_processes=num_workers, mixed_precision="fp16", use_port="29502")

encuentro usando notebook_launcher() conveniente, ya que me permite ejecutar el entrenamiento en la consola y luego trabajar fácilmente con los resultados.

Base XLM-RoBERTa vs grande vs Small-E-Czech
Experimenté ajustando tres modelos. El modelo básico XLM-RoBERTa [3] Obtuve un rendimiento satisfactorio, pero la capacidad del servidor también me permitió probar el modelo grande XLM-RoBERTa. [3]que tiene el doble de parámetros.

XLM-RoBERTa es una versión multilingüe de RoBERTa. Está previamente entrenado en 2,5 TB de datos filtrados de CommonCrawl que contienen 100 idiomas.

El modelo grande mostró una ligera mejora en los resultados, por lo que finalmente lo implementé. También probé Small-E-Czech [4]un modelo pequeño de Electra entrenado previamente con datos web checos, pero su rendimiento fue deficiente.

Ajuste, transferencia de aprendizaje y formación desde cero
Además del ajuste fino (actualizando todos los pesos del modelo), probé el aprendizaje por transferencia, ya que a veces se sugiere que entrenar solo la capa final (clasificación) puede ser suficiente. Sin embargo, la diferencia de rendimiento fue significativa, lo que favoreció el ajuste fino completo. También intenté entrenar desde cero importando solo la arquitectura del modelo, inicializando los pesos aleatoriamente y luego entrenando, pero como era de esperar, este enfoque no fue efectivo.

RoBERTa vs LLM (Claude 3.5 Soneto)
Exploré brevemente los LLM de tiro cero, aunque con una ingeniería mínima (así que 🥱). El modelo tuvo problemas incluso con solicitudes básicas, como (usé checo en el mensaje real):

Encuentre el símbolo variable del acreedor. Este número tiene exactamente 9 dígitos consecutivos del 0 al 9 sin letras ni otros caracteres especiales. Suele ir precedido de una de las siguientes abreviaturas: ‘ev.č.’, ‘zn. opr’, ‘VS. O’, ‘evidente. do. opr.’. Por el contrario, no me interesa un número de transacción con la abreviatura ‘č.j.’. Este número no aparece frecuentemente en los documentos, puede suceder que no puedas encontrarlo, entonces escribe ‘no lo encuentro’. Si no está seguro, escriba “no estoy seguro”.

En ocasiones, el modelo no lograba generar el formato de 9 dígitos con precisión. El posprocesamiento filtraría números más cortos, pero hubo muchos números de 9 dígitos falsos positivos.

Ocasionalmente el modelo infirió algo incorrecto. identificación de nacimientoSe basa únicamente en las fechas de nacimiento (incluso con la temperatura establecida en 0). Por otro lado, se destacó en la extracción nombres, apellidosy fechas de nacimiento.

En general, incluso en mis experimentos anteriores, Descubrí que los LLM (al momento de escribir) se desempeñan mejor en tareas generales, pero carecen de precisión y confiabilidad para tareas específicas o no convencionales. El desempeño en la identificación del cliente fue bastante similar para ambos enfoques. Por motivos internos, se implementó el modelo RoBERTa.