Solution author: Walnit
If you're paying attention to the problem, data and code, you'll stumble across the three biggest pitfalls:
- More often than not, error correction is a local task that does not require much global context - mostly just the surrounding characters.
- Following from that, the corrected text always has the same number of characters as the text with errors. This should be a hint that you won't need an encoder-decoder model, and an encoder-only model will be sufficient! So please don't blindly attempt to throw a Seq2Seq model at it...
- A decent proportion (>25%!) have more characters than the default truncation length of 768 characters, so you lose that training data. (Please read the code properly 😈 we put a lot of hard work into writing questions! This 'pitfall' doesn't really matter that much though, so you'll be fine for now.)
Thus, the question is "how can we train a model that learns as much as possible, within the possible training time?"
Loading the dataset
import datasets
import pandas as pd
from tqdm.auto import tqdm
ds = datasets.load_dataset("csv", data_files="train.csv")
ds = ds["train"].train_test_split(seed=42)
ds
DatasetDict({
train: Dataset({
features: ['text', 'misspell'],
num_rows: 7928
})
test: Dataset({
features: ['text', 'misspell'],
num_rows: 2643
})
})
train_ds = ds["train"]
val_ds = ds["test"]
# for later evaluation
val_ds_input = val_ds.select_columns("misspell")
val_ds_solution = val_ds["text"]
By the way, the purpose of the analysis below was to show you that the number of characters in both datasets are identical - to clue you into realizing point 2 above.
pd.DataFrame(val_ds)["text"].str.len().describe()
count 2643.000000
mean 628.893303
std 336.841864
min 102.000000
25% 379.000000
50% 581.000000
75% 819.500000
max 2291.000000
Name: text, dtype: float64
pd.DataFrame(val_ds)["misspell"].str.len().describe()
count 2643.000000
mean 628.893303
std 336.841864
min 102.000000
25% 379.000000
50% 581.000000
75% 819.500000
max 2291.000000
Name: misspell, dtype: float64
Preparing the data
from tokenizers import Tokenizer, models, pre_tokenizers, processors
tokenizer = Tokenizer(models.WordLevel(unk_token=""))
tokenizer.pre_tokenizer = pre_tokenizers.Split("", "isolated")
tokenizer.enable_padding(pad_token="")
trainer = tokenizer.model.get_trainer()
trainer.vocab_size = 1000
trainer.special_tokens = ["", "", "<SOS>", "<EOS>"]
def ds_iterator():
for row in train_ds:
yield row["text"]
yield row["misspell"]
tokenizer.train_from_iterator(ds_iterator(), trainer=trainer)
tokenizer.post_processor = processors.TemplateProcessing(
single="<SOS> $0 <EOS>",
special_tokens=[("<SOS>", 2), ("<EOS>", 3)]
)
from transformers import PreTrainedTokenizerFast
tokenizer = PreTrainedTokenizerFast(tokenizer_object=tokenizer)
tokenizer.add_special_tokens({"pad_token": "", "unk_token": "", "cls_token": "<SOS>", "eos_token": "<EOS>"})
tokenizer
PreTrainedTokenizerFast(name_or_path='', vocab_size=96, model_max_length=1000000000000000019884624838656, is_fast=True, padding_side='right', truncation_side='right', special_tokens={'eos_token': '<EOS>', 'unk_token': '', 'pad_token': '', 'cls_token': '<SOS>'}, clean_up_tokenization_spaces=False, added_tokens_decoder={
0: AddedToken("", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
1: AddedToken("", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
2: AddedToken("<SOS>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
3: AddedToken("<EOS>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
}
)
Here is where the solution starts. In order to make use of the local nature of the task, as well as the fact that shorter sentences train faster, we will chunk all sentences into chunks of 256 characters.
def tokenize_fn(examples):
input_encodings = tokenizer(
examples['misspell'],
padding='max_length',
truncation=True,
max_length=256,
stride=128,
return_overflowing_tokens=True,
return_attention_mask=False,
return_tensors='pt'
)
label_encodings = tokenizer(
examples['text'],
padding='max_length',
truncation=True,
max_length=256,
stride=128,
return_overflowing_tokens=True,
return_attention_mask=False,
return_tensors='pt'
)
input_lengths = (input_encodings['input_ids'] != tokenizer.pad_token_id).sum(dim=1)
label_lengths = (label_encodings['input_ids'] != tokenizer.pad_token_id).sum(dim=1)
return {
'input_ids': input_encodings['input_ids'],
'labels': label_encodings['input_ids'],
'input_lengths': input_lengths,
'label_lengths': label_lengths,
}
from torch.utils.data import DataLoader
train_ds = train_ds.map(tokenize_fn, batched=True, remove_columns=['text', 'misspell'])
train_ds.set_format("torch")
train_loader = DataLoader(train_ds, batch_size=1024, shuffle=True)
Since for the actual task we won't have the luxury of truncating however we want, we must tokenize the validation and test sets as per the orignal method, without chunking.
Educational Tidbit
There's also a sublety about using pack_padded_sequence that you should know about: it's rather computationally expensive, doubling your execution time per batch.
Why is this so, and why is it recommended in most cases but not here? In most cases, like generation or classification tasks, global context is needed. Thus, it is important to process the entire sentence at once, and sentences of varying lengths becomes an issue. So instead of padding and processing to the length of the longest sentence (which is very computationally expensive!) it's much better to use pack_padded_sequence to pack sequences and reduce the wasted computational power on lengthy sequences of pad tokens.
But in this case, most sentences are longer than 256 characters. And, for the validation set, the median length of 581 isn't that far off from the original truncation length of 768. Thus, padding takes up a minimal proportion of our sequence, and it simply isn't worth using it.
Yet, during the final evaluation on the test set, you're going to want to process all the characters. For that case, we're padding all the way out to 4096 characters, so it's a must that we use pack_padded_sequence there.
This may not be super useful, but if you knew about this, you'd halve the time spent for each iteration of your model in this challenge!
Now, back to processing the data.
val_ds = ds["test"]
val_ds_sol = val_ds["text"]
def tokenize_val_fn(examples):
input_tokens = tokenizer(examples['misspell'], padding='max_length', truncation=True, max_length=768, return_tensors='pt')['input_ids']
label_tokens = tokenizer(examples['text'], padding='max_length', truncation=True, max_length=768, return_tensors='pt')['input_ids']
input_lengths = (input_tokens != tokenizer.pad_token_id).sum(dim=1)
label_lengths = (label_tokens != tokenizer.pad_token_id).sum(dim=1)
return {
'input_ids': input_tokens,
'labels': label_tokens,
'input_lengths': input_lengths,
'label_lengths': label_lengths,
}
val_ds = val_ds.map(tokenize_val_fn, batched=True, remove_columns=['text', 'misspell'])
val_ds.set_format("torch")
Map: 0%| | 0/2643 [00:00<?, ? examples/s]
val_loader = DataLoader(val_ds, batch_size=256, shuffle=False)
Building the model
For this section, we want a good balance of architectural capability and performance.
You should quickly realize that using a transformer is not ideal for this task, due to the short training time allowed, tight execution budget, and low need for global context. So in the end we won't need to go far from the original LSTM!
However, you should not take this as an opportunity to mindlessly tune hyperparameters. Let's break it down.
There are three main hyperparameters to tune:
- Embedding dimension: This is only used by the initial tokens. And since our tokens are characters, we don't need a lot of dimensions for this, as characters are unlikely to carry lots of meaning alone.
- Hidden dimension: Most operations happen on this dimension. The recurrence 'stores' its information here, and outputs for future layers also take in the hidden state of the previous layer. However, even though it's important, number of parameters grows quadratically with dimension size! So we can't recklessly increase the hidden dimension.
- Number of Layers: Offers a greater temporal (or in our case, character) range, as well as a deeper network for greater representational ability. Number of paramters grows linearly, but remember each layer must be calculated one after another and is not parallelizable.
Given the local nature of the task, a balance needs to be struck between the hidden layer and the number of layers. The embedding dimension isn't too important, and doesn't need to be huge.
import torch
import torch.nn as nn
import torch.optim as optim
class BiGRUAutocorrect(nn.Module):
def __init__(self, vocab_size, embed_size, hidden_size, num_layers, dropout):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_size)
self.gru = nn.GRU(
embed_size,
hidden_size,
num_layers=num_layers,
batch_first=True,
dropout=dropout,
bidirectional=True
)
# bidirectional doubles hidden size
self.fc = nn.Linear(hidden_size * 2, vocab_size)
def forward(self, x, lens=None, hidden=None):
x = self.embedding(x)
if lens is None:
out, hidden = self.gru(x, hidden)
else:
packed = nn.utils.rnn.pack_padded_sequence(x, lens.cpu(), batch_first=True, enforce_sorted=False)
packed_outputs, hidden = self.gru(packed, hidden)
out, _ = nn.utils.rnn.pad_packed_sequence(packed_outputs, batch_first=True)
out = self.fc(out)
return out, hidden
In this reference solution, we use a GRU with the following hyperparameters:
- Embedding dimension: 128. This is one dimension per ASCII character, which is more than enough expressiveness for the embeddings.
- Hidden dimension: 256. This was decided after the number of layers. It just so happens this is the a feasible nice number given the limitations of the challenge.
- Number of Layers: 2. Having more layers would detriment performance, and we don't need the temporal range - so 2 is a reasonable value.
Unfortunately, some amount of hyperparamter tuning is required. However, a tip you should know, in the rare occasion you again need to consider the final speed of a model: an untrained model takes the same time as a trained model to predict on test data! So run a few epochs of the final evaluation code using a randomly initialized model with your chosen hyperparameters and see if the approximate time to finish looks good to you. (tqdm's expected remaining time stabilizes after just a few batches, so it's a good estimate!)
vocab_size = tokenizer.vocab_size
embed_size = 128
hidden_size = 256
num_layers = 2
dropout = 0.2
num_epochs = 15
learning_rate = 4e-3
device = "cuda" if torch.cuda.is_available() else "cpu"
# Initialize model, loss function, and optimizer
model = BiGRUAutocorrect(vocab_size, embed_size, hidden_size, num_layers, dropout).to(device)
criterion = nn.CrossEntropyLoss() # you should probably use ignore_index=tokenizer.pad_token_type_id
optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=1e-4)
import gc
gc.collect()
torch.cuda.empty_cache()
We first pre-train the model on an autoencoding task, by passing in the input ids as the target labels. 3 epochs is an arbitrary value.
for epoch in range(3):
# Training
model.train()
total_loss = 0
scaler = torch.amp.GradScaler("cuda")
for batch in tqdm(train_loader):
optimizer.zero_grad()
input_ids = batch['input_ids'].to(device)
labels = batch['labels'].to(device)
with torch.autocast(device_type="cuda"):
outputs, _ = model(labels)
loss = criterion(outputs.view(-1, vocab_size), input_ids.view(-1))
scaler.scale(loss).backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
total_loss += loss.item()
avg_train_loss = total_loss / (len(train_loader))
# Testing
model.eval()
val_loss = 0
with torch.no_grad():
for batch in val_loader:
input_ids = batch['input_ids'].to(device)
labels = batch['labels'].to(device)
with torch.autocast(device_type="cuda"):
outputs, _ = model(input_ids)
loss = criterion(outputs.view(-1, vocab_size), input_ids.view(-1))
val_loss += loss.item()
avg_val_loss = val_loss / len(val_loader)
print(f"Epoch [{epoch+1}/{num_epochs}] | Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
0%| | 0/37 [00:00<?, ?it/s]Epoch [1/15] | Train Loss: 1.2227 | Val Loss: 0.0972
0%| | 0/37 [00:00<?, ?it/s]Epoch [2/15] | Train Loss: 0.6141 | Val Loss: 0.0948
0%| | 0/37 [00:00<?, ?it/s]Epoch [3/15] | Train Loss: 0.5953 | Val Loss: 0.0971
Now that the model is sufficiently initialized, we now train with the training data.
for epoch in range(num_epochs):
# Training
model.train()
total_loss = 0
scaler = torch.amp.GradScaler("cuda")
for batch in tqdm(train_loader):
optimizer.zero_grad()
input_ids = batch['input_ids'].to(device)
labels = batch['labels'].to(device)
with torch.autocast(device_type="cuda"):
outputs, _ = model(input_ids)
loss = criterion(outputs.view(-1, vocab_size), labels.view(-1))
scaler.scale(loss).backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
total_loss += loss.item()
avg_train_loss = total_loss / (len(train_loader))
# Testing
model.eval()
val_loss = 0
with torch.no_grad():
for batch in val_loader:
input_ids = batch['input_ids'].to(device)
labels = batch['labels'].to(device)
with torch.autocast(device_type="cuda"):
outputs, _ = model(input_ids)
loss = criterion(outputs.view(-1, vocab_size), labels.view(-1))
val_loss += loss.item()
avg_val_loss = val_loss / len(val_loader)
print(f"Epoch [{epoch+1}/{num_epochs}] | Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
0%| | 0/37 [00:00<?, ?it/s]Epoch [1/15] | Train Loss: 0.4792 | Val Loss: 0.2818
0%| | 0/37 [00:00<?, ?it/s]Epoch [2/15] | Train Loss: 0.3260 | Val Loss: 0.2069
0%| | 0/37 [00:00<?, ?it/s]Epoch [3/15] | Train Loss: 0.2567 | Val Loss: 0.1724
0%| | 0/37 [00:00<?, ?it/s]Epoch [4/15] | Train Loss: 0.2178 | Val Loss: 0.1502
0%| | 0/37 [00:00<?, ?it/s]Epoch [5/15] | Train Loss: 0.1931 | Val Loss: 0.1363
0%| | 0/37 [00:00<?, ?it/s]Epoch [6/15] | Train Loss: 0.1746 | Val Loss: 0.1276
0%| | 0/37 [00:00<?, ?it/s]Epoch [7/15] | Train Loss: 0.1601 | Val Loss: 0.1194
0%| | 0/37 [00:00<?, ?it/s]Epoch [8/15] | Train Loss: 0.1487 | Val Loss: 0.1154
0%| | 0/37 [00:00<?, ?it/s]Epoch [9/15] | Train Loss: 0.1386 | Val Loss: 0.1112
0%| | 0/37 [00:00<?, ?it/s]Epoch [10/15] | Train Loss: 0.1294 | Val Loss: 0.1073
0%| | 0/37 [00:00<?, ?it/s]Epoch [11/15] | Train Loss: 0.1214 | Val Loss: 0.1047
0%| | 0/37 [00:00<?, ?it/s]Epoch [12/15] | Train Loss: 0.1141 | Val Loss: 0.1036
0%| | 0/37 [00:00<?, ?it/s]Epoch [13/15] | Train Loss: 0.1075 | Val Loss: 0.1009
0%| | 0/37 [00:00<?, ?it/s]Epoch [14/15] | Train Loss: 0.1015 | Val Loss: 0.1006
0%| | 0/37 [00:00<?, ?it/s]Epoch [15/15] | Train Loss: 0.0965 | Val Loss: 0.1000
To ensure the model is exposed to sequences of longer lengths, we reduce the learning rate and train it on full-length sequences.
train_ds = ds["train"]
train_ds = train_ds.map(tokenize_val_fn, batched=True, remove_columns=['text', 'misspell'])
train_ds.set_format("torch")
train_loader = DataLoader(train_ds, batch_size=256, shuffle=True)
learning_rate = 1e-3
optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=1e-4)
import gc
gc.collect()
torch.cuda.empty_cache()
for epoch in range(5):
# Training
model.train()
total_loss = 0
scaler = torch.amp.GradScaler("cuda")
for batch in tqdm(train_loader):
optimizer.zero_grad()
input_ids = batch['input_ids'].to(device)
labels = batch['labels'].to(device)
with torch.autocast(device_type="cuda"):
outputs, _ = model(input_ids)
loss = criterion(outputs.view(-1, vocab_size), labels.view(-1))
scaler.scale(loss).backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
total_loss += loss.item()
avg_train_loss = total_loss / (len(train_loader))
# Testing
model.eval()
val_loss = 0
with torch.no_grad():
for batch in val_loader:
input_ids = batch['input_ids'].to(device)
labels = batch['labels'].to(device)
with torch.autocast(device_type="cuda"):
outputs, _ = model(input_ids)
loss = criterion(outputs.view(-1, vocab_size), labels.view(-1))
val_loss += loss.item()
avg_val_loss = val_loss / len(val_loader)
print(f"Epoch [{epoch+1}/{num_epochs}] | Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
0%| | 0/31 [00:00<?, ?it/s]Epoch [1/15] | Train Loss: 0.0690 | Val Loss: 0.0972
0%| | 0/31 [00:00<?, ?it/s]Epoch [2/15] | Train Loss: 0.0660 | Val Loss: 0.0971
0%| | 0/31 [00:00<?, ?it/s]Epoch [3/15] | Train Loss: 0.0645 | Val Loss: 0.0970
0%| | 0/31 [00:00<?, ?it/s]Epoch [4/15] | Train Loss: 0.0632 | Val Loss: 0.0973
0%| | 0/31 [00:00<?, ?it/s]Epoch [5/15] | Train Loss: 0.0620 | Val Loss: 0.0973
And we save the model! Thatt's all the training we're doing.
torch.save(model.state_dict(), "model_weights.pth")
Here's a toy function to see the output of our autocorrect system. It's not the best, makes additional mistakes, but the real test is the evaluation later.
def predict_sentence(model, tokenizer, input_ids, device="cpu"):
"""
input_ids: tensor of shape (1, seq_len) with tokenized input
"""
model.eval()
with torch.no_grad():
input_ids = input_ids.to(device)
outputs = model(input_ids)[0] # (1, seq_len, vocab_size)
predictions = outputs.argmax(dim=-1) # (1, seq_len)
# Convert ids back to tokens / string
predicted_tokens = tokenizer.convert_ids_to_tokens(predictions[0][1:-1].tolist())
return "".join(predicted_tokens)
# Suppose you have a test sentence
test_text = " Any Athlete wha josfkes cy obstructs another athlete , Inf wat thatimpedes his progress , should be dis qualitierfrom than Evemt . howerer , yf yn athlete re pushed or forced by anotherp erson rk rin outside hhs labe , and id no material advantage th gained , the athlete should aot se diswualified . " # misspelled input
# Tokenize and convert to tensor
input_ids = tokenizer.encode(test_text, return_tensors="pt")
# Get prediction
corrected = predict_sentence(model, tokenizer, input_ids, device=device)
print("Input:", test_text)
print("Corrected:", corrected)
Input: Any Athlete wha josfkes cy obstructs another athlete , Inf wat thatimpedes his progress , should be dis qualitierfrom than Evemt . howerer , yf yn athlete re pushed or forced by anotherp erson rk rin outside hhs labe , and id no material advantage th gained , the athlete should aot se diswualified .
Corrected: Any athlete who jostles or obstructs another athlete , In a was that impedes his progress , should be disqqulitier from than event . However , if an athle tere pushed or forced by another person to run outside his lage , and it no material advantage oo gained , the athlete should not be disqualified .
Evaluation
Predicting on Test
test_ds = datasets.load_dataset("csv", data_files="test.csv")
# test_ds = val_ds_input # for validation
test_ds = test_ds["train"] # for test
test_ds
Dataset({
features: ['misspell'],
num_rows: 2643
})
def tokenize_test_fn(examples):
input_tokens = tokenizer(examples['misspell'], padding='max_length', truncation=True, max_length=4096, return_tensors='pt')['input_ids']
input_lengths = (input_tokens != tokenizer.pad_token_id).sum(dim=1)
return {
'input_ids': input_tokens,
'input_lengths': input_lengths
}
test_ds = test_ds.map(tokenize_test_fn, batched=True, remove_columns=['misspell'])
test_ds.set_format("torch")
test_loader = DataLoader(test_ds, batch_size=32, shuffle=False)
import gc
gc.collect()
torch.cuda.empty_cache()
As mentioned above, we want to pack the sequences here.
def wrap_model_prediction(model, batch, device):
input_ids = batch['input_ids'].to(device)
lengths = batch['input_lengths'].to(device)
output = model(input_ids, lengths)[0]
for i, length in enumerate(lengths):
output[i, length:] = 0
return output
Running on Google Colab, it comfortably breezes past the requirement, taking only 149 seconds!
%%time
device = "cpu" # change to cuda for validation
preds_all = []
model.eval()
model.to(device)
with torch.no_grad():
for batch in tqdm(test_loader):
outputs = wrap_model_prediction(model, batch, device)
predictions = outputs.argmax(dim=-1)
preds_all.append(predictions.cpu())
0%| | 0/83 [00:00<?, ?it/s]CPU times: user 2min 6s, sys: 22.5 s, total: 2min 29s
Wall time: 2min 30s
results = [] # convert tokens back to string, excluded from loop as this takes quite a while
for pred in tqdm(preds_all):
results += [ "".join(tokenizer.convert_ids_to_tokens(x, skip_special_tokens=True)) for x in pred ]
0%| | 0/83 [00:00<?, ?it/s]
import csv
with open('test_submission.csv', 'w', newline='') as f:
writer = csv.writer(f)
writer.writerows([[l] for l in results])
print("test_submission.csv generated!")
test_submission.csv generated!
Code for Evaluation (for Validation)
! pip install jiwer evaluate --quiet
import evaluate
cer = evaluate.load("cer")
import pandas as pd
submission = pd.read_csv("test_submission.csv", header=None)
# solution = pd.DataFrame(pd.Series(val_ds_solution)) # for validation
solution = pd.read_csv("test_sol.csv") # for testing
cer.compute(
predictions=submission.iloc[:, 0],
references=solution.iloc[:, 0]
)
0.03616788264131573
And tada! A final score of 0.0362.
This is NOT the lowest possible score. Further tuning of hyperparameters and techniques, or training for longer, will definitely yield better results. But this score is a semi-realistic standard you should hopefully be able to reach within the time limit.
If you have any questions, concerns, complaints, confessions, hate mail, or extra GPU credits for me, feel free to message me on Discord (@walnit). 😛
See you in the next task!! 👋