Spaces:
Running on Zero
Running on Zero
File size: 1,238 Bytes
98d1e65 | 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 |
import torch
from src.infra.tokenizer import CharacterTokenizer
class TextClassificationDataset:
def __init__(
self,
texts,
labels,
tokenizer: CharacterTokenizer,
max_length, #we truncate anything after this length... we got a budget to keep ;-)
):
if len(texts) != len(labels):
raise ValueError(
"texts and labels must have the same length"
)
self.texts = texts
self.labels = torch.tensor(
labels,
dtype=torch.long,
)
encoded = [
tokenizer.encode(text)
for text in texts
]
self.tokens, self.padding_mask = tokenizer.pad(
encoded,
max_length=max_length,
)
self.tokens = torch.tensor(
self.tokens,
dtype=torch.long,
)
self.padding_mask = torch.tensor(
self.padding_mask,
dtype=torch.bool,
)
def __len__(self):
return len(self.texts)
def __getitem__(self, index):
return (
self.tokens[index],
self.padding_mask[index],
self.labels[index],
)
|