added BERT training script

Author: valentinp72

evaluation/predict_bert models/BERT/predict

@@ -0,0 +1,189 @@
+#!/usr/bin/env python3
+
+import argparse
+import datetime
+
+parser = argparse.ArgumentParser(
+    description='Script to train (ie. finetune) the BERT models on a given ' \
+    'dataset. The dataset might be the validation set of the is24_news_topic' \
+    'dataset or a dataset of examples generated by Mixtral.'
+)
+parser.add_argument(
+    '--model_path', required=True, type=str,
+    help='Path or model name of the HuggingFace BERT model to finetune.'
+)
+parser.add_argument(
+    '--output_dir', required=True, type=str,
+    help='Output directory to save checkpoints and everything else.'
+)
+parser.add_argument('--device', default='cpu')
+parser.add_argument('--resume_from_checkpoint', default=None)
+parser.add_argument('--lowercase_text', action='store_true')
+
+parser.add_argument(
+    '--train_dataset', required=True, type=str,
+    help='Path to the training dataset to use.'
+)
+parser.add_argument('--train_subset', default='train')
+parser.add_argument(
+    '--validation_dataset', required=True, type=str,
+    help='Path to the validation dataset to use at each validation step'
+)
+parser.add_argument('--validation_subset', default='test')
+
+parser.add_argument('--validation_num_samples', default=0, type=int)
+parser.add_argument('--num_train_epochs', default=3, type=int)
+parser.add_argument('--validation_every_steps', default=500, type=int)
+
+parser.add_argument(
+    '--now', type=str,
+    default=datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')
+)
+
+args = parser.parse_args()
+print(args)
+
+################################################################################
+
+labels = [
+    'SPORT',
+    'ARTS/CULTURE/ENTERTAINMENT',
+    'EDUCATION',
+    'RELIGION/BELIEF',
+    'UNREST/CONFLICTS/WAR',
+    'CRIME/LAW/JUSTICE',
+    'HEALTH',
+    'LIFESTYLE/LEISURE',
+    'COMMERCIAL',
+    'SCIENCE/TECHNOLOGY',
+    'WEATHER',
+    'POLITICS',
+    'SOCIAL_ISSUE',
+    'OTHER',
+    'DISASTER/ACCIDENT',
+    'ECONOMY/BUSINESS/FINANCE',
+    'ENVIRONMENTAL_ISSUE',
+    'LABOUR'
+]
+id2label = {i:l for i,l in enumerate(labels)}
+label2id = {l:i for i,l in enumerate(labels)}
+
+################################################################################
+
+import json
+import torch
+import datasets
+import evaluate
+import transformers
+import numpy as np
+
+# preparing the models
+tokenizer = transformers.AutoTokenizer.from_pretrained(args.model_path)
+model = transformers.AutoModelForSequenceClassification.from_pretrained(
+    args.model_path,
+    num_labels=len(id2label),
+    id2label=id2label, label2id=label2id,
+    problem_type="multi_label_classification"
+).to(args.device)
+
+if 'flaubert' in args.model_path:
+    model.sequence_summary.summary_type = 'mean'
+model.config.max_length = 256
+
+clf_metrics = evaluate.combine(["accuracy", "f1", "precision", "recall"])
+
+def sigmoid(x):
+   return 1/(1 + np.exp(-x))
+
+def compute_metrics(eval_pred):
+   predictions, labels = eval_pred
+   predictions = sigmoid(predictions)
+   predictions = (predictions > 0.5).astype(int).reshape(-1)
+   return clf_metrics.compute(
+       predictions=predictions,
+       references=labels.astype(int).reshape(-1)
+   )
+
+# preparing the data
+if args.train_dataset == args.validation_dataset and \
+        args.train_subset == args.validation_subset:
+    print('WARNING: training and validation sets are equal. Will create a ' \
+          'random train/valid split with 0.80 ratio for train.')
+    hf_dataset = datasets \
+        .load_from_disk(args.train_dataset)[args.train_subset] \
+        .train_test_split(train_size=0.8, shuffle=True, seed=42)
+else:
+    hf_dataset = datasets.DatasetDict({
+        'train': datasets.load_from_disk(args.train_dataset)[args.train_subset],
+        'test': datasets.load_from_disk(args.validation_dataset)[args.validation_subset]
+    })
+
+if args.lowercase_text:
+    hf_dataset = hf_dataset.map(lambda x: {'text': x['text'].lower()})
+
+# prearing the labels
+def get_reference_labels(example):
+    if isinstance(example['classes'], list):
+        # loaded a mixtral annotation
+        labels_indices = [label2id[l] for l in example['classes']]
+    elif isinstance(example['classes'], dict):
+        # loaded human annotation
+        labels_indices = [label2id[l] for l,v in example['classes'].items() if v > 0]
+    labels = np \
+        .eye(len(label2id), dtype=float)[np.array(labels_indices, dtype=int)] \
+        .sum(axis=0)
+    return {'labels': labels}
+
+hf_dataset = hf_dataset.map(get_reference_labels)
+
+# tokenization
+def tokenize_function(examples):
+    return tokenizer(
+        examples["text"],
+        padding="max_length",
+        truncation=True,
+        max_length=model.config.max_length
+    )
+hf_dataset = hf_dataset.map(tokenize_function, batched=True, batch_size=1000)
+
+################################################################################
+
+# preparing for training
+output_dir = os.path.join(
+    args.output_dir,
+    args.now,
+    model.config._name_or_path.replace('/', '__')
+)
+training_args = transformers.TrainingArguments(
+    output_dir=output_dir,
+    evaluation_strategy="steps",
+    save_strategy="steps",
+    save_total_limit=3,
+    eval_steps=args.validation_every_steps,
+    save_steps=args.validation_every_steps,
+    weight_decay=0.01, learning_rate=2e-5,
+    num_train_epochs=args.num_train_epochs,
+    load_best_model_at_end=True, metric_for_best_model='f1'
+)
+
+print(hf_dataset)
+
+if args.validation_num_samples == 0:
+    eval_dataset = hf_dataset['test']
+else:
+    eval_dataset = hf_dataset['test'].select(range(args.validation_num_samples))
+
+trainer = transformers.Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=hf_dataset['train'],
+    eval_dataset=eval_dataset,
+    compute_metrics=compute_metrics,
+    data_collator=transformers.DefaultDataCollator()
+)
+
+# training!
+trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
+
+# saving the best model
+trainer.save_model()