Upload train.py with huggingface_hub

7817dee 9 months ago

No virus

6.25 kB

	from datasets import load_dataset
	from transformers import TrainingArguments

	from span_marker import SpanMarkerModel, Trainer


	def main() -> None:
	# Load the dataset, ensure "tokens" and "ner_tags" columns, and get a list of labels
	dataset = "Babelscape/multinerd"
	train_dataset = load_dataset(dataset, split="train")
	eval_dataset = load_dataset(dataset, split="validation").shuffle().select(range(3000))
	labels = [
	"O",
	"B-PER",
	"I-PER",
	"B-ORG",
	"I-ORG",
	"B-LOC",
	"I-LOC",
	"B-ANIM",
	"I-ANIM",
	"B-BIO",
	"I-BIO",
	"B-CEL",
	"I-CEL",
	"B-DIS",
	"I-DIS",
	"B-EVE",
	"I-EVE",
	"B-FOOD",
	"I-FOOD",
	"B-INST",
	"I-INST",
	"B-MEDIA",
	"I-MEDIA",
	"B-MYTH",
	"I-MYTH",
	"B-PLANT",
	"I-PLANT",
	"B-TIME",
	"I-TIME",
	"B-VEHI",
	"I-VEHI",
	]

	# Initialize a SpanMarker model using a pretrained BERT-style encoder
	model_name = "bert-base-multilingual-cased"
	model = SpanMarkerModel.from_pretrained(
	model_name,
	labels=labels,
	# SpanMarker hyperparameters:
	model_max_length=256,
	marker_max_length=128,
	entity_max_length=8,
	)

	# Prepare the 🤗 transformers training arguments
	args = TrainingArguments(
	output_dir="models/span_marker_mbert_base_multinerd",
	# Training Hyperparameters:
	learning_rate=5e-5,
	per_device_train_batch_size=32,
	per_device_eval_batch_size=32,
	# gradient_accumulation_steps=2,
	num_train_epochs=1,
	weight_decay=0.01,
	warmup_ratio=0.1,
	bf16=True, # Replace `bf16` with `fp16` if your hardware can't use bf16.
	# Other Training parameters
	logging_first_step=True,
	logging_steps=50,
	evaluation_strategy="steps",
	save_strategy="steps",
	eval_steps=1000,
	save_total_limit=2,
	dataloader_num_workers=2,
	)

	# Initialize the trainer using our model, training args & dataset, and train
	trainer = Trainer(
	model=model,
	args=args,
	train_dataset=train_dataset,
	eval_dataset=eval_dataset,
	)
	trainer.train()
	trainer.save_model("models/span_marker_mbert_base_multinerd/checkpoint-final")

	test_dataset = load_dataset(dataset, split="test")
	# Compute & save the metrics on the test set
	metrics = trainer.evaluate(test_dataset, metric_key_prefix="test")
	trainer.save_metrics("test", metrics)

	trainer.create_model_card(language="multilingual", license="apache-2.0")


	if __name__ == "__main__":
	main()

	"""
	Logs:

	Training set:
	This SpanMarker model will ignore 0.793782% of all annotated entities in the train dataset. This is caused by the SpanMarkerModel maximum entity length of 8 words and the maximum model input length of 256 tokens.
	These are the frequencies of the missed entities due to maximum entity length out of 4111958 total entities:
	- 12680 missed entities with 9 words (0.308369%)
	- 7308 missed entities with 10 words (0.177726%)
	- 4414 missed entities with 11 words (0.107345%)
	- 2474 missed entities with 12 words (0.060166%)
	- 1894 missed entities with 13 words (0.046061%)
	- 1130 missed entities with 14 words (0.027481%)
	- 744 missed entities with 15 words (0.018094%)
	- 582 missed entities with 16 words (0.014154%)
	- 344 missed entities with 17 words (0.008366%)
	- 226 missed entities with 18 words (0.005496%)
	- 84 missed entities with 19 words (0.002043%)
	- 46 missed entities with 20 words (0.001119%)
	- 20 missed entities with 21 words (0.000486%)
	- 20 missed entities with 22 words (0.000486%)
	- 12 missed entities with 23 words (0.000292%)
	- 18 missed entities with 24 words (0.000438%)
	- 2 missed entities with 25 words (0.000049%)
	- 4 missed entities with 26 words (0.000097%)
	- 4 missed entities with 27 words (0.000097%)
	- 2 missed entities with 31 words (0.000049%)
	- 8 missed entities with 32 words (0.000195%)
	- 6 missed entities with 33 words (0.000146%)
	- 2 missed entities with 34 words (0.000049%)
	- 4 missed entities with 36 words (0.000097%)
	- 8 missed entities with 37 words (0.000195%)
	- 2 missed entities with 38 words (0.000049%)
	- 2 missed entities with 41 words (0.000049%)
	- 2 missed entities with 72 words (0.000049%)
	Additionally, a total of 598 (0.014543%) entities were missed due to the maximum input length.

	Validation set:
	This SpanMarker model won't be able to predict 0.656224% of all annotated entities in the evaluation dataset. This is caused by the SpanMarkerModel maximum entity length of 8 words.
	These are the frequencies of the missed entities due to maximum entity length out of 4724 total entities:
	- 19 missed entities with 9 words (0.402202%)
	- 7 missed entities with 10 words (0.148180%)
	- 1 missed entities with 11 words (0.021169%)
	- 1 missed entities with 12 words (0.021169%)
	- 2 missed entities with 13 words (0.042337%)
	- 1 missed entities with 16 words (0.021169%)

	Testing set:
	This SpanMarker model won't be able to predict 0.794755% of all annotated entities in the evaluation dataset. This is caused by the SpanMarkerModel maximum entity length of 8 words and the maximum model input length of 256 tokens.
	These are the frequencies of the missed entities due to maximum entity length out of 511856 total entities:
	- 1520 missed entities with 9 words (0.296959%)
	- 870 missed entities with 10 words (0.169970%)
	- 640 missed entities with 11 words (0.125035%)
	- 346 missed entities with 12 words (0.067597%)
	- 224 missed entities with 13 words (0.043762%)
	- 172 missed entities with 14 words (0.033603%)
	- 72 missed entities with 15 words (0.014066%)
	- 66 missed entities with 16 words (0.012894%)
	- 16 missed entities with 17 words (0.003126%)
	- 14 missed entities with 18 words (0.002735%)
	- 14 missed entities with 19 words (0.002735%)
	- 2 missed entities with 20 words (0.000391%)
	- 12 missed entities with 21 words (0.002344%)
	- 2 missed entities with 24 words (0.000391%)
	- 2 missed entities with 25 words (0.000391%)
	Additionally, a total of 96 (0.018755%) entities were missed due to the maximum input length.
	"""