Как указать количество целевых классов для TFRobertaSequenceClassification?

Question

У меня есть задача классификации текста, и я хочу использовать предварительно обученную модель Роберты из библиотеки трансформаторов.

Вот документы: TFRobertaForSequenceClassification

Согласно документации для обучения мы должны использовать:

from transformers import RobertaTokenizer, TFRobertaForSequenceClassification

tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
model = TFRobertaForSequenceClassification.from_pretrained('roberta-base')

model.compile('adam', loss='sparse_categorical_crossentropy')
model.fit(x, y)

Итак, где мне указать количество целевых меток для классификации последовательностей?

Answer 1

Вы можете использовать num_labels параметр.

model = TFRobertaForSequenceClassification.from_pretrained('roberta-base', num_labels = 5)

ссылка: https://huggingface.co/transformers/main_classes/configuration.html