File size: 5,101 Bytes
1ae8986 | 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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import Dataset, DatasetDict, load_from_disk, concatenate_datasets
import torch
import os
import tensorflow as tf
OUTPUT_DIR = "./results"
DATASET_NAME = 'dataset_with_split.csv'
MODEL_NAME = "roberta-large"
LOG_DIR = "./logs"
SAVE_MODEL_FOLDER = "img_intents_model"
POS_NAME = "POSITIVE"
NEG_NAME = "NEGATIVE"
BATCH_SIZE_TRAIN = 16
BATCH_SIZE_EVAL = 64
EPOCS = 10
WARMUP_STEPS = 500
tf.debugging.experimental.enable_dump_debug_info(
LOG_DIR,
tensor_debug_mode="FULL_HEALTH",
circular_buffer_size=1000,
op_regex=None,
tensor_dtypes=None
)
# Load the dataset from the CSV file
dataset = Dataset.from_csv(DATASET_NAME)
# Create a DatasetDict object containing train, validation, and test datasets
datasets = DatasetDict({
'train': dataset.filter(lambda example: example['split'] == 'train'),
'validation': dataset.filter(lambda example: example['split'] == 'validation'),
'test': dataset.filter(lambda example: example['split'] == 'test'),
})
# Balance the datasets
for split in datasets.keys():
num_positive = len(datasets[split].filter(lambda example: example['label'] == POS_NAME))
num_negative = len(datasets[split].filter(lambda example: example['label'] == NEG_NAME))
if num_positive > num_negative:
# Downsample the positive examples
datasets[split] = concatenate_datasets([
datasets[split].filter(lambda example: example['label'] == POS_NAME).shuffle(seed=42).select(range(num_negative)),
datasets[split].filter(lambda example: example['label'] == NEG_NAME)
])
else:
# Downsample the negative examples
datasets[split] = concatenate_datasets([
datasets[split].filter(lambda example: example['label'] == POS_NAME),
datasets[split].filter(lambda example: example['label'] == NEG_NAME).shuffle(seed=42).select(range(num_positive))
])
# Shuffle the dataset to mix positive and negative examples
datasets[split] = datasets[split].shuffle(seed=42)
# Specify the model name
model_name = MODEL_NAME # Or whatever model you want to use
# Load the tokenizer associated with your model
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Load your datasets
train_dataset = datasets['train']
val_dataset = datasets['validation']
test_dataset = datasets['test']
# Preprocessing function
def preprocess_function(examples):
# Replace None in 'text' field with an empty string
examples["text"] = [text if text is not None else "" for text in examples["text"]]
# Convert labels from string to int
examples["label"] = [1 if label == POS_NAME else 0 for label in examples["label"]]
# Tokenize the texts
return tokenizer(examples["text"], truncation=True, max_length=512, padding='max_length')
train_dataset = train_dataset.map(preprocess_function, batched=True)
val_dataset = val_dataset.map(preprocess_function, batched=True)
test_dataset = test_dataset.map(preprocess_function, batched=True)
# Make sure all your tensors are the same size for batching together
train_dataset = train_dataset.remove_columns(["text"]).rename_column("label", "labels").with_format("torch")
val_dataset = val_dataset.remove_columns(["text"]).rename_column("label", "labels").with_format("torch")
test_dataset = test_dataset.remove_columns(["text"]).rename_column("label", "labels").with_format("torch")
# Load a pre-trained model for sequence classification
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2) # You have two labels: POSITIVE and NEGATIVE
# TrainingArguments
training_args = TrainingArguments(
output_dir=OUTPUT_DIR,
num_train_epochs=EPOCS,
per_device_train_batch_size=BATCH_SIZE_TRAIN, # decrease this if necessary
per_device_eval_batch_size=BATCH_SIZE_EVAL,
warmup_steps=WARMUP_STEPS,
weight_decay=0.01,
logging_dir=LOG_DIR,
logging_strategy='steps', # Log after every training step
logging_steps=10, # Adjust this to change how often logging occurs
evaluation_strategy='steps', # Evaluate after every training step
eval_steps=100, # Adjust this to change how often evaluation occurs
save_strategy='steps', # Save after every training step
save_steps=500, # Adjust this to change how often saving occurs
no_cuda=False, # use GPU
gradient_accumulation_steps=2, # if necessary
fp16=True, # use mixed precision training
report_to='tensorboard'
)
# Create a Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=val_dataset,
)
# Train the model
trainer.train()
# Save the model
trainer.save_model(SAVE_MODEL_FOLDER)
# Save the tokenizer
tokenizer.save_pretrained(OUTPUT_DIR)
# Save the training arguments
torch.save(training_args, os.path.join(OUTPUT_DIR, "training_args.bin"))
# Evaluate the model and print the results
eval_results = trainer.evaluate(test_dataset)
print(f"Test set evaluation results: {eval_results}") |