|
|
|
|
|
"""ZGBot Training Script |
|
|
File can be executed here (Link to Google Colab): |
|
|
https://colab.research.google.com/drive/1Dyn37CljZnYaQ1dXs3rCOQmdpioKB6-t |
|
|
""" |
|
|
|
|
|
!pip install datasets wandb evaluate accelerate -qU |
|
|
!pip install transformers |
|
|
|
|
|
from huggingface_hub import notebook_login |
|
|
notebook_login() |
|
|
import wandb |
|
|
wandb.login() |
|
|
from datasets import load_dataset |
|
|
squad = load_dataset("squad", split="train[:5000]") |
|
|
squad = squad.train_test_split(test_size=0.2) |
|
|
from transformers import AutoTokenizer |
|
|
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased") |
|
|
|
|
|
def preprocess_function(examples): |
|
|
questions = [q.strip() for q in examples["question"]] |
|
|
inputs = tokenizer( |
|
|
questions, |
|
|
examples["context"], |
|
|
max_length=384, |
|
|
truncation="only_second", |
|
|
return_offsets_mapping=True, |
|
|
padding="max_length", |
|
|
) |
|
|
|
|
|
offset_mapping = inputs.pop("offset_mapping") |
|
|
answers = examples["answers"] |
|
|
start_positions = [] |
|
|
end_positions = [] |
|
|
|
|
|
for i, offset in enumerate(offset_mapping): |
|
|
answer = answers[i] |
|
|
start_char = answer["answer_start"][0] |
|
|
end_char = answer["answer_start"][0] + len(answer["text"][0]) |
|
|
sequence_ids = inputs.sequence_ids(i) |
|
|
|
|
|
|
|
|
idx = 0 |
|
|
while sequence_ids[idx] != 1: |
|
|
idx += 1 |
|
|
context_start = idx |
|
|
while sequence_ids[idx] == 1: |
|
|
idx += 1 |
|
|
context_end = idx - 1 |
|
|
|
|
|
|
|
|
if offset[context_start][0] > end_char or offset[context_end][1] < start_char: |
|
|
start_positions.append(0) |
|
|
end_positions.append(0) |
|
|
else: |
|
|
|
|
|
idx = context_start |
|
|
while idx <= context_end and offset[idx][0] <= start_char: |
|
|
idx += 1 |
|
|
start_positions.append(idx - 1) |
|
|
|
|
|
idx = context_end |
|
|
while idx >= context_start and offset[idx][1] >= end_char: |
|
|
idx -= 1 |
|
|
end_positions.append(idx + 1) |
|
|
|
|
|
inputs["start_positions"] = start_positions |
|
|
inputs["end_positions"] = end_positions |
|
|
return inputs |
|
|
|
|
|
dataset = squad.map(preprocess_function, batched=True, remove_columns=squad["train"].column_names) |
|
|
from transformers import DefaultDataCollator |
|
|
data_collator = DefaultDataCollator() |
|
|
from transformers import AutoModelForQuestionAnswering, TrainingArguments, Trainer |
|
|
model = AutoModelForQuestionAnswering.from_pretrained("distilbert-base-uncased") |
|
|
|
|
|
!pip install transformers[torch] |
|
|
|
|
|
import evaluate |
|
|
import numpy as np |
|
|
metric=evaluate.load("accuracy") |
|
|
def compute_metrics(eval_pred): |
|
|
logits, labels = eval_pred |
|
|
predictions = np.argmax(logits, axis=-1) |
|
|
return metric.compute(predictions=predictions, references=labels) |
|
|
|
|
|
from transformers import Trainer, TrainingArguments |
|
|
|
|
|
args = TrainingArguments( |
|
|
output_dir = "MA-saemi-5", |
|
|
report_to = 'wandb', |
|
|
evaluation_strategy = 'steps', |
|
|
learning_rate = 3e-5, |
|
|
max_steps = 3000, |
|
|
logging_steps = 100, |
|
|
eval_steps = 250, |
|
|
save_steps = 10000, |
|
|
load_best_model_at_end = True, |
|
|
metric_for_best_model = 'accuracy', |
|
|
run_name = 'training5', |
|
|
per_device_train_batch_size=16, |
|
|
per_device_eval_batch_size=16, |
|
|
push_to_hub=True, |
|
|
|
|
|
) |
|
|
|
|
|
trainer = Trainer( |
|
|
model = model, |
|
|
args = args, |
|
|
train_dataset=dataset['train'], |
|
|
eval_dataset=dataset['test'], |
|
|
tokenizer=tokenizer, |
|
|
compute_metrics=compute_metrics, |
|
|
) |
|
|
|
|
|
trainer.train() |
|
|
wandb.finish() |
|
|
trainer.push_to_hub() |
|
|
|
