When I freshly train the Token Classification model (DistilBertForTokenClassification) and run a prediction for a single sentence that I manually type out, it runs fine, but when I try to run it on my dataset, it fails with RuntimeError: CUDA error: device-side assert triggered when using.
This is how I’m performing the training:
if __name__ == '__main__':
import torch
class WNUTDataset(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)
train_encodings.pop("offset_mapping") # we don't want to pass this to the model
val_encodings.pop("offset_mapping")
train_dataset = WNUTDataset(train_encodings, train_labels)
val_dataset = WNUTDataset(val_encodings, val_labels)
from transformers import DistilBertForTokenClassification
model = DistilBertForTokenClassification.from_pretrained('distilbert-base-cased', num_labels=len(unique_tags))
model.resize_token_embeddings(len(tokenizer))
from transformers import Trainer, TrainingArguments
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
eval_dataset=val_dataset # evaluation dataset
)
trainer.train()
This is the function that I’m using for prediction.
def ner_normalize(text, clean=False):
if clean:
text = clean_posts(text)
inputs = tokenizer(text, return_tensors="pt")
labels = torch.tensor([1] * inputs["input_ids"].size(1)).unsqueeze(0)
# Running the prediction
outputs = model(**inputs, token_type_ids=None)
text = [inverse_vocab[i] for i in inputs['input_ids'][0].numpy()]
preds = [id2tag[i.item()] for i in torch.argmax(outputs['logits'], axis=-1)[0]]
collapsed_tokens = []
collapsed_preds = []
curr_token = ''
curr_pred = None
for i, subword_pred_pair in enumerate(zip(text, preds)):
subword, pred = subword_pred_pair
if subword.startswith('##'):
curr_token += subword[2:]
if i == len(text) - 1:
collapsed_tokens.append(curr_token)
collapsed_preds.append(curr_pred)
else:
collapsed_tokens.append(curr_token)
collapsed_preds.append(curr_pred)
curr_token = subword
curr_pred = pred
outputs = zip(collapsed_tokens[2:], collapsed_preds[2:])
normalized = ' '.join(
[x[0] if x[1] == "O" else "CODE" for x in outputs]
)
return normalized