Spaces:
Sleeping
Sleeping
File size: 4,050 Bytes
e3af320 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import os
import numpy as np
import json
import argparse
from utils_data import load_data, MyDataset
from modeling_bert import BertForSequenceClassification
from transformers import (
AutoConfig,
BertForSequenceClassification,
AutoTokenizer,
EvalPrediction,
Trainer,
TrainingArguments,
default_data_collator,
set_seed,
)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--data_root', type=str, default='data')
parser.add_argument('--model', type=str, default='bert-base-uncased')
parser.add_argument('--output_dir', type=str, default='experiments')
parser.add_argument('--epoch', type=int, default=2)
parser.add_argument('--lr', type=float, default=2e-5)
parser.add_argument('--bs', type=int, default=32)
parser.add_argument('--max_length', type=int, default=512)
parser.add_argument('--seed', type=int, default=666, help='random seed')
args = parser.parse_args()
return args
def main():
args = parse_args()
print("args",args)
print('====Input Arguments====')
print(json.dumps(vars(args), indent=2, sort_keys=False))
# Set seed before initializing model, for reproduction purpose.
set_seed(args.seed)
# Load pretrained model and tokenizer
config = AutoConfig.from_pretrained(args.model)
tokenizer = AutoTokenizer.from_pretrained(args.model)
model =BertForSequenceClassification.from_pretrained(args.model, config=config)
# Load data
train_data = load_data(args, "train")
train_dataset = MyDataset(train_data, tokenizer, args.max_length, is_test=False)
eval_data = load_data(args, "val")
eval_dataset = MyDataset(eval_data, tokenizer, args.max_length, is_test=False)
test_data = load_data(args, "test")
test_dataset = MyDataset(test_data, tokenizer, args.max_length, is_test=True)
def compute_metrics(p: EvalPrediction):
preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
preds = np.argmax(preds, axis=1)
correct = ((preds == p.label_ids).sum()).item()
return {'accuracy': 1.0*correct/len(preds)}
training_args = TrainingArguments(
output_dir = args.output_dir,
do_train=True,
do_eval=True,
do_predict=True,
logging_strategy="steps",
save_strategy="epoch",
learning_rate= args.lr,
per_device_train_batch_size=args.bs,
per_device_eval_batch_size=args.bs,
num_train_epochs=args.epoch,
report_to="none"
)
# Initialize our Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset if training_args.do_train else None,
eval_dataset=eval_dataset if training_args.do_eval else None,
compute_metrics=compute_metrics,
tokenizer=tokenizer,
data_collator=default_data_collator,
)
# Training
if training_args.do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.save_model() # Saves the tokenizer too for easy upload
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
# Evaluation
if training_args.do_eval:
metrics = trainer.evaluate(eval_dataset=eval_dataset)
trainer.log_metrics("eval", metrics)
trainer.save_metrics("eval", metrics)
if training_args.do_predict:
predictions = trainer.predict(test_dataset, metric_key_prefix="predict").predictions
predictions = np.argmax(predictions, axis=1)
output_predict_file = os.path.join(args.output_dir, "predict_results.txt")
if trainer.is_world_process_zero():
with open(output_predict_file, "w") as writer:
writer.write("index\tprediction\n")
for index, item in enumerate(predictions):
writer.write(f"{index}\t{item}\n")
if __name__ == "__main__":
main() |