Spaces:
Sleeping
Sleeping
File size: 1,695 Bytes
fa64206 |
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 |
from datasets import load_dataset
from transformers import BertTokenizer
def load_and_tokenize_data(config):
"""
Load and tokenize data based on the provided configuration.
Args:
config (dict): Configuration dictionary containing dataset and tokenizer details.
Returns:
tuple: A tuple containing the tokenized train and test datasets.
"""
# Load the dataset
dataset = load_dataset(config['dataset']['name'], split=config['dataset']['split'])
dataset = dataset.train_test_split(test_size=0.2)
train_dataset = dataset['train']
test_dataset = dataset['test']
# Initialize the tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
# Define the tokenization function
def tokenize_function(examples):
return tokenizer(examples['text'], padding='max_length', truncation=True)
# Apply tokenization to the train and test datasets
train_dataset = train_dataset.map(tokenize_function, batched=True)
test_dataset = test_dataset.map(tokenize_function, batched=True)
# Set the format to PyTorch tensors
train_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
test_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
return train_dataset, test_dataset
# Example usage
if __name__ == "__main__":
config = {
'dataset': {
'name': 'imdb',
'split': 'train[:10%]'
}
}
train_dataset, test_dataset = load_and_tokenize_data(config)
print("Train dataset and Test dataset have been loaded and tokenized successfully.")
|