问无法让DataCollator准备tf数据集

EN

# load dataset, 
df_converters = {'tokens': ast.literal_eval, 'labels': ast.literal_eval}
train_df = pd.read_csv("train_df_pretokenization.csv", converters=df_converters)
train_df = train_df.head(10)

# model and pretrained tokenizer
model_ckpt = "indobenchmark/indobert-base-p2"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt) 

# tokenization, and align labels 
def tokenize_and_align_labels(batch): 

    tag2int = {'B-POI':0, 'B-STR':1, 'E-POI':2, 'E-STR':3, 'I-POI':4,
           'I-STR':5, 'S-POI':6, 'S-STR':7, 'O':8}
           
    #tokenized_inputs = tokenizer(batch['tokens'], is_split_into_words=True, truncation = True, padding = True)
    tokenized_inputs = tokenizer(batch['tokens'], is_split_into_words=True, truncation = True)
    labels=[]
    for idx, label in enumerate(batch['labels']):
        word_ids = tokenized_inputs.word_ids(batch_index = idx)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids: 
            if word_idx is None:
                label_ids.append(-100)
            elif word_idx != previous_word_idx:
                label_ids.append(tag2int[label[word_idx]])
            else: 
                label_ids.append(tag2int[label[word_idx]])
            previous_word_idx = word_idx
        labels.append(label_ids)
    tokenized_inputs['tags'] = labels

    return tokenized_inputs

def encode_dataset(ds):
    return ds.map(tokenize_and_align_labels, batched= True, batch_size=10, remove_columns=['labels','tokens', 'index'])
    
train_ds = Dataset.from_pandas(train_df)
train_ds_encoded = encode_dataset(train_ds)

# prepare model input 
data_collator = DataCollatorForTokenClassification(tokenizer, return_tensors="tf")

tf_train_dataset = train_ds_encoded.to_tf_dataset(
    columns= ['input_ids', 'token_type_ids', 'attention_mask', 'tags'],
    shuffle=False,
    batch_size=5,
    collate_fn=data_collator
)