84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
|
|
import torch
|
|
|
|
PATHS = ['train/in.tsv', 'train/expected.tsv', 'dev-0/in.tsv', 'test-A/in.tsv', './dev-0/out.tsv', './test-A/out.tsv']
|
|
OUTPUT_PATHS = ['dev-0/out.tsv', 'test-A/out.tsv']
|
|
PRE_TRAINED = ['roberta-base']
|
|
|
|
def get_data(path):
|
|
data = []
|
|
with open(path, encoding='utf-8') as f:
|
|
data = f.readlines()
|
|
|
|
return data
|
|
|
|
def generate_output(path, trainer, X_data):
|
|
data = []
|
|
with open(path, encoding='utf-8') as f:
|
|
for result in trainer.predict(X_data):
|
|
f.write(str(result) + '\n')
|
|
|
|
class IMDbDataset(torch.utils.data.Dataset):
|
|
def __init__(self, encodings, labels):
|
|
self.encodings = encodings
|
|
self.labels = labels
|
|
|
|
def __getitem__(self, idx):
|
|
item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
|
|
item['labels'] = torch.tensor(self.labels[idx])
|
|
return item
|
|
|
|
def __len__(self):
|
|
return len(self.labels)
|
|
|
|
def prepare(data_train_X, data_train_Y):
|
|
tokenizer = AutoTokenizer.from_pretrained(PRE_TRAINED[0])
|
|
model = AutoModelForSequenceClassification.from_pretrained(PRE_TRAINED[0], num_labels=2)
|
|
device = torch.device("cpu")
|
|
model.to(device)
|
|
encoded_input = tokenizer([text[0] for text in list(zip(data_train_X, data_train_Y))], truncation=True, padding=True)
|
|
train_dataset = IMDbDataset(encoded_input , [int(text[1]) for text in list(zip(data_train_X, data_train_Y))])
|
|
|
|
return train_dataset, model
|
|
|
|
def training(train_dataset, model):
|
|
training_args = TrainingArguments(
|
|
output_dir='./results', # output directory
|
|
num_train_epochs=3, # total number of training epochs
|
|
per_device_train_batch_size=16, # batch size per device during training
|
|
per_device_eval_batch_size=64, # batch size for evaluation
|
|
warmup_steps=500, # number of warmup steps for learning rate scheduler
|
|
weight_decay=0.01, # strength of weight decay
|
|
logging_dir='./logs', # directory for storing logs
|
|
logging_steps=10,
|
|
)
|
|
|
|
trainer = Trainer(
|
|
model=model, # the instantiated Transformers model to be trained
|
|
args=training_args, # training arguments, defined above
|
|
train_dataset=train_dataset, # training dataset
|
|
)
|
|
|
|
trainer.train()
|
|
|
|
return trainer
|
|
|
|
def main():
|
|
#data
|
|
X_train = get_data(PATHS[0])
|
|
y_train = get_data(PATHS[1])
|
|
X_dev = get_data(PATHS[2])
|
|
X_test = get_data(PATHS[3])
|
|
|
|
#prepare
|
|
train_dataset, model = prepare(X_train, y_train)
|
|
|
|
#trainer
|
|
trainer = training(train_dataset, model)
|
|
|
|
#output
|
|
generate_output(OUTPUT_PATHS[0], trainer, X_dev)
|
|
generate_output(OUTPUT_PATHS[1], trainer, X_test)
|
|
|
|
if __name__ == '__main__':
|
|
main() |