|
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments |
|
from datasets import load_dataset, concatenate_datasets |
|
|
|
MODEL_NAME = "roberta-large" |
|
SAVE_MODEL_FOLDER = "img_intents_model" |
|
OUTPUT_DIR = "./results" |
|
output_dir = "/results" |
|
|
|
|
|
model = AutoModelForSequenceClassification.from_pretrained(SAVE_MODEL_FOLDER) |
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) |
|
|
|
|
|
positives_dataset = load_dataset('text', data_files='test_positives.txt') |
|
negatives_dataset = load_dataset('text', data_files='test_negatives.txt') |
|
|
|
|
|
positives_dataset = positives_dataset['train'].map(lambda example: {'label': 1}) |
|
negatives_dataset = negatives_dataset['train'].map(lambda example: {'label': 0}) |
|
|
|
|
|
train_dataset = concatenate_datasets([positives_dataset, negatives_dataset]) |
|
|
|
|
|
def preprocess_function(examples): |
|
|
|
return tokenizer(examples["text"], truncation=True, max_length=512, padding='max_length') |
|
|
|
train_dataset = train_dataset.map(preprocess_function, batched=True) |
|
|
|
|
|
train_dataset = train_dataset.remove_columns(["text"]).rename_column("label", "labels").with_format("torch") |
|
|
|
|
|
training_args = TrainingArguments( |
|
output_dir=OUTPUT_DIR, |
|
num_train_epochs=5, |
|
per_device_train_batch_size=16, |
|
per_device_eval_batch_size=64, |
|
warmup_steps=500, |
|
weight_decay=0.01, |
|
logging_dir=OUTPUT_DIR, |
|
logging_strategy='steps', |
|
logging_steps=10, |
|
evaluation_strategy='steps', |
|
eval_steps=100, |
|
save_strategy='steps', |
|
save_steps=500, |
|
no_cuda=False, |
|
gradient_accumulation_steps=2, |
|
fp16=True, |
|
report_to='tensorboard' |
|
) |
|
|
|
trainer = Trainer( |
|
model=model, |
|
args=training_args, |
|
train_dataset=train_dataset, |
|
) |
|
|
|
|
|
trainer.train() |
|
|
|
|
|
trainer.save_model(SAVE_MODEL_FOLDER) |
|
|
|
|
|
tokenizer.save_pretrained(OUTPUT_DIR) |
|
|