How to apply training ClassLabels on test / validation Dataset

I am trying to create a custom dataset using flat files. The data elements are 'Tokens" and ā€˜NER Tagsā€™ strings. I am able to create the ClassLabels for the Training dataset and now want to apply the same ClassLabels on the Test & Validation datasets. I have tried several approaches, but none have worked. Any help or insight will be greatly appreciated. The code segment is below :

def convert_to_featured_dataset(data):
formatted_data = {ā€œtokensā€: , ā€œner_tagsā€: }
features = Features(
{ā€œtokensā€: Value(dtype=ā€œstringā€), ā€œner_tagsā€: Value(dtype=ā€œstringā€)}
)
for sentence in data:
tokens = [token_data[1] for token_data in sentence]
ner_tags = [token_data[0] for token_data in sentence]

    formatted_data["tokens"].append(tokens)
    formatted_data["ner_tags"].append(ner_tags)

dataset = Dataset.from_dict(formatted_data, features=features)
return dataset

ā€¦
ā€¦
#Works
train_dataset = convert_to_featured_dataset(train_data)
train_dataset = train_dataset.class_encode_column(ā€œner_tagsā€)
label_feature = train_dataset.features[ā€œner_tagsā€]
test_features = train_dataset.features.copy()
ā€¦
ā€¦
#Does not work
def convert_ids(rec):
#Error : ValueError: Invalid string class label [ā€˜Oā€™, ā€˜Oā€™, ā€˜B-RATINGS_AVERAGEā€™, ā€˜I-RATINGS_AVERAGEā€™, ā€¦']
rec[ā€œner_tagsā€] = label_feature.str2int(rec[ā€œner_tagsā€])
return rec

test_dataset = convert_to_featured_dataset(test_data)
# Did not work : ValueError: Invalid string class label ['O'...
#test_dataset = test_dataset.map(convert_ids, batched=True)
#test_features["ner_tags"] = ClassLabel(label_feature.feature.names)
#test_dataset = test_dataset.cast(test_features)
 test_dataset = test_dataset.cast_column('ner_tags',label_feature)

Thank you for your time and help.

Tried encoding the ner_tags separately and then casting them to ClassLabel but got stuck on another point

# Failed  :  may have worked if ner_tags feature was like ['1','2','3']  but it is like [['1','2','3'] ,['1','2','3'] ]  
# Error : ValueError: Invalid string class label ['24', '24', '24', '0', '12']
train_dataset = train_dataset.cast_column(
    "ner_tags", ClassLabel(names=id_label_list)
)

You can pass Sequence(ClassLabel(...)) to .cast_column to cast a list of labels.