File size: 3,554 Bytes
78a5823 |
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 85 86 87 88 89 |
from torch.utils.data import DataLoader, Dataset
import torch
from transformers import BertTokenizerFast, BertModel
from transformers import BertConfig, BertPreTrainedModel
import numpy as np
from typing import Dict, List, Union, Tuple
from utils import ner_labels_to_ids, intent_labels_to_ids, structure_data
class tokenized_dataset(Dataset):
"""
A Pytorch Dataset for tokenizing and encoding text data for a BERT-based model.
Args:
dataset (dict): A dictionary containing 'text', 'ner', and 'intent' keys.
tokenizer (BertTokenizerFast): A tokenizer for processing text input.
max_len (int, optionl): Maximum length of tokenized sequences (default: 128).
Attributes:
len (int): Number of samples in the dataset.
Methods:
__getitem__(self, index: int) -> Dict[str, torch.Tensor]:
Retrieve and preprocess a single sample from the dataset.
__len__(self) -> int:
Get the total number of samples int the dataset.
Returns:
Dict[str, torch.Tensor]: A dictionary containing tokenized and encoded text, NER and intent labels.
"""
def __init__(self, dataset: Dict[str, List[str]], tokenizer: BertTokenizerFast, max_len: int = 128):
self.len = len(dataset['text'])
self.ner_labels_to_ids = ner_labels_to_ids()
self.intent_labels_to_ids = intent_labels_to_ids()
self.text = dataset['text']
self.intent = dataset['intent']
self.ner = dataset['entities']
self.tokenizer = tokenizer()
self.max_len = max_len
def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
# step 1: get the sentence, ner label, and intent_label
sentence = self.text[index].strip()
intent_label = self.intent[index].strip()
ner_labels = self.ner[index]
# step 2: use tokenizer to encode a sentence (includes padding/truncation up to max length)
# BertTokenizerFast provides a handy "return_offsets_mapping" which highlights where each token starts and ends
encoding = self.tokenizer(
sentence,
return_offsets_mapping=True,
padding='max_length',
truncation=True,
max_length=self.max_len
)
# step 3: create ner token labels only for first word pieces of each tokenized word
tokenized_ner_labels = [self.ner_labels_to_ids[label] for label in ner_labels]
# create an empty array of -100 of length max_length
encoded_ner_labels = np.ones(len(encoding['offset_mapping']), dtype=int) * -100
# set only labels whose first offset position is 0 and the second is not 0
i = 0
prev = -1
for idx, mapping in enumerate(encoding['offset_mapping']):
if mapping[0] == mapping[1] == 0:
continue
if mapping[0] != prev:
# overwrite label
encoded_ner_labels[idx] = tokenized_ner_labels[i]
prev = mapping[1]
i += 1
else:
prev = mapping[1]
# create intent token labels
tokenized_intent_label = self.intent_labels_to_ids[intent_label]
# step 4: turn everything into Pytorch tensors
item = {key: torch.as_tensor(val) for key, val in encoding.items()}
item['ner_labels'] = torch.as_tensor(encoded_ner_labels)
item['intent_labels'] = torch.as_tensor(tokenized_intent_label)
return item
def __len__(self) -> int:
return self.len
|