kazu.training.train_multilabel_ner

Functions

calculate_metrics(epoch_loss, test_docs, ...)
move_entities_to_metadata(docs)

Classes

KazuMultiHotNerMultiLabelTrainingDataset
ModelSaver
SavedModel	SavedModel(path: pathlib.Path, step: int, metrics: dict[str, typing.Any] = <factory>)
Trainer

class kazu.training.train_multilabel_ner.KazuMultiHotNerMultiLabelTrainingDataset

Bases: Dataset[dict[str, Tensor]]

__init__(docs_iter, model_tokenizer, labels, tmp_dir, use_cache=True, max_length=128, stride=64, max_docs=None, keep_doc_reference=False)

Parameters:

• docs_iter (Iterable[Document])

• model_tokenizer (PreTrainedTokenizerFast)

• labels (list[str])

• tmp_dir (Path)

• use_cache (bool)

• max_length (int)

• stride (int)

• max_docs (int | None)

• keep_doc_reference (bool)

get_docs_copy()

Return type:

list[Document]

tokenize_and_align(section)

Parameters:

section (Section)

Return type:

dict[str, Tensor]

class kazu.training.train_multilabel_ner.ModelSaver

Bases: object

__init__(save_dir, max_to_keep=5, patience=5)

Parameters:

• save_dir (Path)

• max_to_keep (int)

• patience (int)

save(model, step, tokenizer, metrics, stopping_metric, test_docs=None)

Parameters:

• model (PreTrainedModel)

• step (int)

• tokenizer (PreTrainedTokenizerFast)

• metrics (dict[str, float])

• stopping_metric (str)

• test_docs (list[Document] | None)

Return type:

None

static save_model(tokenizer, model, path)

Parameters:

• tokenizer (PreTrainedTokenizerFast)

• model (PreTrainedModel)

• path (Path)

Return type:

None

class kazu.training.train_multilabel_ner.SavedModel

Bases: object

SavedModel(path: pathlib.Path, step: int, metrics: dict[str, typing.Any] = <factory>)

__init__(path, step, metrics=<factory>)

Parameters:

• path (Path)

• step (int)

• metrics (dict[str, Any])

Return type:

None

metrics: dict[str, Any]

path: Path

step: int

class kazu.training.train_multilabel_ner.Trainer

Bases: object

__init__(training_config, pretrained_model_name_or_path, label_list, train_dataset, test_dataset, working_dir, summary_writer=None, ls_wrapper=None)

Parameters:

• training_config (TrainingConfig)

• pretrained_model_name_or_path (str)

• label_list (list[str])

• train_dataset (KazuMultiHotNerMultiLabelTrainingDataset)

• test_dataset (KazuMultiHotNerMultiLabelTrainingDataset)

• working_dir (Path)

• summary_writer (SummaryWriter | None)

• ls_wrapper (LSManagerViewWrapper | None)

evaluate_model(model, global_step, save_model=True)

Parameters:

• model (PreTrainedModel)

• global_step (int)

• save_model (bool)

Return type:

None

log_metrics(tensorboard_loggables, global_step)

Parameters:

• tensorboard_loggables (dict[str, dict[str, Any]])

• global_step (int)

Return type:

None

train_model()

Return type:

None

kazu.training.train_multilabel_ner.calculate_metrics(epoch_loss, test_docs, label_list)

Parameters:

• epoch_loss (float)

• test_docs (list[Document])

• label_list (list[str])

Return type:

tuple[dict[str, Any], dict[str, dict[str, Any]]]

kazu.training.train_multilabel_ner.move_entities_to_metadata(docs)

Parameters:

docs (list[Document])

Return type:

list[Document]