File size: 3,322 Bytes
a0e0ff1 | 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 | import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
from model import CustomBERTModel
from config import Config
import pandas as pd
from tqdm import tqdm
def load_data(file_path):
df = pd.read_csv(file_path, header=None)
return torch.tensor(df.values, dtype=torch.float32)
def create_mlm_data(data, mlm_probability):
labels = data.clone()
probability_matrix = torch.full(labels.shape, mlm_probability)
masked_indices = torch.bernoulli(probability_matrix).bool()
labels[~masked_indices] = -100 # We only compute loss on masked tokens
# 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
data[indices_replaced] = 0 # Assume 0 is the representation of [MASK]
# 10% of the time, we replace masked input tokens with random word
indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
random_words = torch.randint(Config.vocab_size, labels.shape, dtype=torch.long)
data[indices_random] = random_words[indices_random].float()
return data, labels
def train():
config = Config()
model = CustomBERTModel(config).to(config.device)
optimizer = optim.AdamW(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay)
train_data = load_data(config.train_file)
val_data = load_data(config.val_file)
train_dataset = TensorDataset(train_data)
val_dataset = TensorDataset(val_data)
train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=config.batch_size)
for epoch in range(config.num_train_epochs):
model.train()
total_loss = 0
for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{config.num_train_epochs}"):
inputs = batch[0].to(config.device)
masked_inputs, labels = create_mlm_data(inputs, config.mlm_probability)
optimizer.zero_grad()
outputs = model(masked_inputs, labels=labels)
loss = outputs.loss
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
optimizer.step()
total_loss += loss.item()
avg_train_loss = total_loss / len(train_loader)
print(f"Epoch {epoch+1}/{config.num_train_epochs}, Average training loss: {avg_train_loss:.4f}")
# Validation
model.eval()
total_val_loss = 0
with torch.no_grad():
for batch in val_loader:
inputs = batch[0].to(config.device)
masked_inputs, labels = create_mlm_data(inputs, config.mlm_probability)
outputs = model(masked_inputs, labels=labels)
total_val_loss += outputs.loss.item()
avg_val_loss = total_val_loss / len(val_loader)
print(f"Validation loss: {avg_val_loss:.4f}")
# Save the model
torch.save(model.state_dict(), "bert_mlm_model.pth")
print("Model saved as bert_mlm_model.pth")
if __name__ == "__main__":
train()
|