```python
import tensorflow as tf
from datasets import load_dataset
from tran…sformers import AutoTokenizer, TFAutoModelForSequenceClassification, DataCollatorWithPadding
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.optimizers.schedules import PolynomialDecay
from tensorflow.keras.losses import SparseCategoricalCrossentropy
def prepare_imdb_dataset(tokenizer):
"""
Prepares the IMDB dataset for training and validation.
Args:
tokenizer: The tokenizer to use for text tokenization.
Returns:
A tuple containing the tokenized training and validation datasets.
"""
imdb = load_dataset("imdb")
train_set = imdb['train'].map(lambda x: tokenizer(x['text'], truncation=True), batched=True)
test_set = imdb['test'].map(lambda x: tokenizer(x['text'], truncation=True), batched=True)
return train_set, test_set
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = TFAutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)
train_set, test_set = prepare_imdb_dataset(tokenizer)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="tf")
tf_train_dataset = train_set.to_tf_dataset(
columns=["attention_mask", "input_ids"],
label_cols=["label"],
shuffle=True,
collate_fn=data_collator,
batch_size=8,
)
tf_validation_dataset = test_set.to_tf_dataset(
columns=["attention_mask", "input_ids"],
label_cols=["label"],
shuffle=False,
collate_fn=data_collator,
batch_size=8,
)
batch_size = 16
num_epochs = 1
num_train_steps = len(tf_train_dataset) * num_epochs
lr_scheduler = PolynomialDecay(
initial_learning_rate=5e-5, end_learning_rate=0.0, decay_steps=num_train_steps
)
optimizer = Adam(learning_rate=lr_scheduler)
loss = SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer=optimizer, loss=loss, metrics=["accuracy"])
model.fit(tf_train_dataset, validation_data=tf_validation_dataset, epochs=5)
```
```
Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFDistilBertForSequenceClassification: ['vocab_layer_norm.weight', 'vocab_transform.weight', 'vocab_projector.bias', 'vocab_transform.bias', 'vocab_layer_norm.bias']
- This IS expected if you are initializing TFDistilBertForSequenceClassification from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing TFDistilBertForSequenceClassification from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).
Some weights or buffers of the TF 2.0 model TFDistilBertForSequenceClassification were not initialized from the PyTorch model and are newly initialized: ['pre_classifier.weight', 'pre_classifier.bias', 'classifier.weight', 'classifier.bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Map: 100%
25000/25000 [00:23<00:00, 1086.84 examples/s]
Map: 100%
25000/25000 [00:20<00:00, 1304.86 examples/s]
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
[<ipython-input-17-ac80246ded67>](https://localhost:8080/#) in <cell line: 55>()
53 optimizer = Adam(learning_rate=lr_scheduler)
54 loss = SparseCategoricalCrossentropy(from_logits=True)
---> 55 model.compile(optimizer=optimizer, loss=loss, metrics=["accuracy"])
56
57 model.fit(tf_train_dataset, validation_data=tf_validation_dataset, epochs=5)
2 frames
[/usr/local/lib/python3.10/dist-packages/tf_keras/src/optimizers/__init__.py](https://localhost:8080/#) in get(identifier, **kwargs)
332 )
333 else:
--> 334 raise ValueError(
335 f"Could not interpret optimizer identifier: {identifier}"
336 )
ValueError: Could not interpret optimizer identifier: <keras.src.optimizers.adam.Adam object at 0x79d9071160e0>
```