File size: 3,604 Bytes
5207ca0 |
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 |
# Loading the Data.
import torch
import pandas as pd
from transformers import BartTokenizer, BartForSequenceClassification
# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load data from names.csv
# train_data_full = pd.read_csv('names_balanced_train.csv',header=None,names=["name","country"])
# test_data = pd.read_csv('names_balanced_test.csv',header=None,names=["name","country"])
joke_data = pd.read_csv('jokes.csv', sep='|', names=["joke", "label"], skiprows=1)
noJoke_data = pd.read_csv('not_jokes.csv', sep='|', names=["joke", "label"], skiprows=1)
frames = [joke_data, noJoke_data]
train_data = pd.concat(frames)
test_data = pd.read_csv('test_jokes.csv', sep='|', names=["joke", "label"], skiprows=1)
numCategories = 2
tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
model = BartForSequenceClassification.from_pretrained('facebook/bart-large', num_labels=numCategories)
model = model.to(device)
# Convert country column to one-hot encoding
one_hot_train = pd.get_dummies(train_data['label'])
one_hot_test = pd.get_dummies(test_data['label'])
# Tokenize names and convert to PyTorch dataset
inputs_train = tokenizer(list(train_data['joke']), return_tensors='pt', padding=True)
labels_train = torch.tensor(one_hot_train.values, dtype=torch.float32)
dataset_train = torch.utils.data.TensorDataset(inputs_train['input_ids'], inputs_train['attention_mask'], labels_train)
inputs_test = tokenizer(list(test_data['joke']), return_tensors='pt', padding=True)
labels_test = torch.tensor(one_hot_test.values, dtype=torch.float32)
dataset_test = torch.utils.data.TensorDataset(inputs_test['input_ids'], inputs_test['attention_mask'], labels_test)
# Define training parameters
epochs = 10
batch_size = 32
learning_rate = 1e-5
# Train model
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
data_loader_train = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True)
data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=batch_size)
print(f"\nTraining on {len(train_data)} examples\n")
print("Num. Parameters:", sum(p.numel() for p in model.parameters() if p.requires_grad))
for epoch in range(epochs):
# Compute average loss after 100 steps
avg_loss = 0
for step, batch in enumerate(data_loader_train):
input_ids, attention_mask, labels = batch
input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)
outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs[0]
avg_loss += loss.item()
if step % 100 == 0:
print(f"Step {step}/{len(data_loader_train)} Loss {loss} Avg Train Loss {avg_loss / (step + 1)}")
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss = avg_loss / len(data_loader_train)
# Print loss after every epoch
print(f"Epoch {epoch+1} Test Loss {loss}")
# Compute accuracy after every epoch
correct = 0
total = 0
for step, batch in enumerate(data_loader_test):
input_ids, attention_mask, labels = batch
input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)
outputs = model(input_ids, attention_mask=attention_mask)
predicted = torch.argmax(outputs[0], dim=1)
total += labels.size(0)
correct += (predicted == torch.argmax(labels, dim=1)).sum().item()
print(f"Test Accuracy {100*correct/total}%\n")
# Save model
model.save_pretrained('fine-tuned-bart_countries') |