import tensorflow as tf
from transformers import TFAutoModel, AutoTokenizer
from datasets import load_dataset
model = TFAutoModel.from_pretrained("bert-base-uncased")
Downloading: 100%
570/570 [00:00<00:00, 22.2kB/s]
2023-01-21 14:49:15.558690: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-01-21 14:49:15.559793: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-01-21 14:49:15.560497: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-01-21 14:49:15.563043: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-01-21 14:49:15.563369: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-01-21 14:49:15.564106: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-01-21 14:49:15.564755: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-01-21 14:49:20.386731: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-01-21 14:49:20.387657: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-01-21 14:49:20.388329: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-01-21 14:49:20.388921: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1510] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 15043 MB memory:  -> device: 0, name: Tesla P100-PCIE-16GB, pci bus id: 0000:00:04.0, compute capability: 6.0
Downloading: 100%
511M/511M [00:23<00:00, 23.4MB/s]
Some layers from the model checkpoint at bert-base-uncased were not used when initializing TFBertModel: ['mlm___cls', 'nsp___cls']
- This IS expected if you are initializing TFBertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing TFBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
All the layers of TFBertModel were initialized from the model checkpoint at bert-base-uncased.
If your task is similar to the task the model of the checkpoint was trained on, you can already use TFBertModel for predictions without further training.
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
Downloading: 100%
28.0/28.0 [00:00<00:00, 1.10kB/s]
Downloading: 100%
226k/226k [00:00<00:00, 957kB/s]
Downloading: 100%
455k/455k [00:00<00:00, 919kB/s]
inputs = tokenizer(['Hello world', 'Hi how are you'], padding=True, truncation=True,
                  return_tensors='tf')
inputs
{'input_ids': <tf.Tensor: shape=(2, 6), dtype=int32, numpy=
array([[ 101, 7592, 2088,  102,    0,    0],
       [ 101, 7632, 2129, 2024, 2017,  102]], dtype=int32)>, 'token_type_ids': <tf.Tensor: shape=(2, 6), dtype=int32, numpy=
array([[0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0]], dtype=int32)>, 'attention_mask': <tf.Tensor: shape=(2, 6), dtype=int32, numpy=
array([[1, 1, 1, 1, 0, 0],
       [1, 1, 1, 1, 1, 1]], dtype=int32)>}
output = model(inputs)
output
TFBaseModelOutputWithPoolingAndCrossAttentions(last_hidden_state=<tf.Tensor: shape=(2, 6, 768), dtype=float32, numpy=
array([[[-0.16888332,  0.13606355, -0.13940018, ..., -0.6251125 ,
          0.05217262,  0.36714536],
        [-0.3632745 ,  0.14121903,  0.8799885 , ...,  0.10433032,
          0.2887578 ,  0.37267894],
        [-0.69859415, -0.69879794,  0.06450251, ..., -0.22103661,
          0.00986893, -0.5939796 ],
        [ 0.83098257,  0.12366717, -0.15119013, ...,  0.10309545,
         -0.67792666, -0.26285172],
        [-0.40266633, -0.01928236,  0.5732502 , ..., -0.20656842,
          0.02338582,  0.20126349],
        [-0.6228408 , -0.27453488,  0.1811763 , ..., -0.12944865,
         -0.03839079, -0.05733156]],

       [[ 0.09286558, -0.02636361, -0.12239343, ..., -0.21063566,
          0.17386371,  0.17250973],
        [ 0.40742022, -0.05930945,  0.55234593, ..., -0.6790563 ,
          0.6555748 , -0.2945646 ],
        [-0.21155298, -0.6858643 , -0.46280792, ...,  0.15278494,
          0.59774226, -0.9102001 ],
        [ 0.3992125 , -1.3207804 , -0.08008753, ..., -0.32125378,
          0.25572804, -0.57804394],
        [-0.07565154, -1.3393834 ,  0.18162948, ...,  0.07461128,
          0.40318406, -0.7079987 ],
        [ 0.5988934 , -0.28409335, -0.34899247, ...,  0.30420092,
         -0.4367556 , -0.20969652]]], dtype=float32)>, pooler_output=<tf.Tensor: shape=(2, 768), dtype=float32, numpy=
array([[-0.90615326, -0.31115308, -0.6216535 , ..., -0.30575195,
        -0.64009386,  0.91661745],
       [-0.930967  , -0.3380703 , -0.62161595, ..., -0.44018963,
        -0.68128854,  0.9348898 ]], dtype=float32)>, past_key_values=None, hidden_states=None, attentions=None, cross_attentions=None)
emotions = load_dataset('SetFit/emotion')
Downloading and preparing dataset json/SetFit--emotion to /root/.cache/huggingface/datasets/json/SetFit--emotion-ddb0f0516f8def09/0.0.0/ac0ca5f5289a6cf108e706efcf040422dbbfa8e658dee6a819f20d76bb84d26b...
Downloading data files: 100%
3/3 [00:06<00:00, 2.22s/it]
Downloading data: 100%
2.23M/2.23M [00:00<00:00, 3.59MB/s]
Downloading data: 100%
279k/279k [00:00<00:00, 588kB/s]
Downloading data: 100%
276k/276k [00:00<00:00, 622kB/s]
Extracting data files: 100%
3/3 [00:00<00:00, 112.52it/s]
Dataset json downloaded and prepared to /root/.cache/huggingface/datasets/json/SetFit--emotion-ddb0f0516f8def09/0.0.0/ac0ca5f5289a6cf108e706efcf040422dbbfa8e658dee6a819f20d76bb84d26b. Subsequent calls will reuse this data.
100%
3/3 [00:00<00:00, 96.61it/s]
emotions
DatasetDict({
    train: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 16000
    })
    test: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 2000
    })
    validation: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 2000
    })
})
def tokenize(batch):
    return tokenizer(batch["text"], padding=True, truncation=True)
emotions_encoded = emotions.map(tokenize, batched=True, batch_size=None)
100%
1/1 [00:02<00:00, 2.11s/ba]
100%
1/1 [00:00<00:00, 3.06ba/s]
100%
1/1 [00:00<00:00, 4.89ba/s]
emotions_encoded
DatasetDict({
    train: Dataset({
        features: ['text', 'label', 'label_text', 'input_ids', 'token_type_ids', 'attention_mask'],
        num_rows: 16000
    })
    test: Dataset({
        features: ['text', 'label', 'label_text', 'input_ids', 'token_type_ids', 'attention_mask'],
        num_rows: 2000
    })
    validation: Dataset({
        features: ['text', 'label', 'label_text', 'input_ids', 'token_type_ids', 'attention_mask'],
        num_rows: 2000
    })
})
# setting 'input_ids', 'attention_mask', 'token_type_ids', and 'label'
# to the tensorflow format. Now if you access this dataset you will get these
# columns in `tf.Tensor` format

emotions_encoded.set_format('tf', 
                            columns=['input_ids', 'attention_mask', 'token_type_ids', 'label'])

# setting BATCH_SIZE to 64.
BATCH_SIZE = 64

def order(inp):
    '''
    This function will group all the inputs of BERT
    into a single dictionary and then output it with
    labels.
    '''
    data = list(inp.values())
    return {
        'input_ids': data[1],
        'attention_mask': data[2],
        'token_type_ids': data[3]
    }, data[0]

# converting train split of `emotions_encoded` to tensorflow format
train_dataset = tf.data.Dataset.from_tensor_slices(emotions_encoded['train'][:])
# set batch_size and shuffle
train_dataset = train_dataset.batch(BATCH_SIZE).shuffle(1000)
# map the `order` function
train_dataset = train_dataset.map(order, num_parallel_calls=tf.data.AUTOTUNE)

# ... doing the same for test set ...
test_dataset = tf.data.Dataset.from_tensor_slices(emotions_encoded['test'][:])
test_dataset = test_dataset.batch(BATCH_SIZE)
test_dataset = test_dataset.map(order, num_parallel_calls=tf.data.AUTOTUNE)
inp, out = next(iter(train_dataset)) # a batch from train_dataset
print(inp, '\n\n', out)
{'input_ids': <tf.Tensor: shape=(64, 87), dtype=int64, numpy=
array([[  101,  1045,  2079, ...,     0,     0,     0],
       [  101,  1045,  2572, ...,     0,     0,     0],
       [  101,  1045,  2514, ...,     0,     0,     0],
       ...,
       [  101,  1045,  8343, ...,     0,     0,     0],
       [  101, 10047,  3110, ...,     0,     0,     0],
       [  101,  1045,  2514, ...,     0,     0,     0]])>, 'attention_mask': <tf.Tensor: shape=(64, 87), dtype=int64, numpy=
array([[0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       ...,
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0]])>, 'token_type_ids': <tf.Tensor: shape=(64, 87), dtype=int64, numpy=
array([[1, 1, 1, ..., 0, 0, 0],
       [1, 1, 1, ..., 0, 0, 0],
       [1, 1, 1, ..., 0, 0, 0],
       ...,
       [1, 1, 1, ..., 0, 0, 0],
       [1, 1, 1, ..., 0, 0, 0],
       [1, 1, 1, ..., 0, 0, 0]])>} 

 tf.Tensor(
[1 4 0 1 0 1 0 1 1 1 2 4 1 0 0 5 2 4 3 3 0 0 1 2 1 2 1 4 0 2 0 0 0 0 3 0 0
 1 0 0 0 2 3 3 0 0 1 4 3 0 1 1 5 2 1 3 1 1 0 5 0 1 5 3], shape=(64,), dtype=int64)
2023-01-21 14:50:04.788175: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:185] None of the MLIR Optimization Passes are enabled (registered 2)
class BERTForClassification(tf.keras.Model):
    
    def __init__(self, bert_model, num_classes):
        super().__init__()
        self.bert = bert_model
        self.fc = tf.keras.layers.Dense(num_classes, activation='softmax')
        
    def call(self, inputs):
        x = self.bert(inputs)[1]
        return self.fc(x)
classifier = BERTForClassification(model, num_classes=6)

classifier.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(),
    metrics=['accuracy']
)
history = classifier.fit(
    train_dataset,
    epochs=3
)
Epoch 1/3
250/250 [==============================] - 172s 617ms/step - loss: 1.0917 - accuracy: 0.5889
Epoch 2/3
250/250 [==============================] - 155s 618ms/step - loss: 0.2976 - accuracy: 0.8949
Epoch 3/3
250/250 [==============================] - 155s 619ms/step - loss: 0.1587 - accuracy: 0.9351
classifier.evaluate(test_dataset)

Give me an example of inference and outpot of this moddel