Longformer Tensorflow Int32 vs Int64 error

Models
1

I am trying an example of fine-tuning longformer and got the error of

TypeError: Input 'updates' of 'TensorScatterUpdate' Op has type int32 that does not match type int64 of argument 'tensor'.

Not sure what’s going on. Here is my code example. Any help would be great. Thank you very much!


from transformers import LongformerTokenizer, TFLongformerForSequenceClassification
from datasets import Dataset

import tensorflow as tf
import pickle
import numpy as np

tokenizer = LongformerTokenizerFast.from_pretrained('allenai/longformer-base-4096')
model = TFLongformerForSequenceClassification.from_pretrained('allenai/longformer-base-4096', num_labels=2, gradient_checkpointing=True)
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=5e-5),
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=tf.metrics.SparseCategoricalAccuracy(),
)

my_dict = {'text': ["random text 1", "random text 2", "random text 3"],
'label': np.array([0, 0, 1], dtype=np.int64)}

dataset = Dataset.from_dict(my_dict)

tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')

def tokenize_function(examples):
return tokenizer(examples["text"], padding="max_length", truncation=True)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

small_train_dataset = tokenized_datasets.shuffle(seed=42)

from transformers import DefaultDataCollator

data_collator = DefaultDataCollator(return_tensors="tf")

tf_train_dataset = small_train_dataset.to_tf_dataset(
columns=["attention_mask", "input_ids", "token_type_ids"],
label_cols=["labels"],
shuffle=True,
collate_fn=data_collator,
batch_size=8,
)

model.fit(tf_train_dataset, batch_size=1)
1 Like