Huggingface TFBertForSequenceClassification всегда предсказывает одну и ту же метку

TL; DR: Моя модель всегда предсказывает одни и те же ярлыки, и я не знаю почему. Ниже приведен весь мой код для тонкой настройки в надежде, что кто-то укажет мне, в чем я ошибаюсь.

Я использую TFBertForSequenceClassification Huggingface для задачи классификации последовательностей, чтобы предсказать 4 метки предложений в немецком тексте.

Я использую модель bert-base-german-cased, поскольку я не использую только строчные буквы (поскольку немецкий более чувствителен к регистру, чем английский).

Я получаю данные из файла csv, который я создаю из аннотированного корпуса, который я получил. Вот пример этого:

0       Hier kommen wir ins Spiel Die App Cognitive At...
1       Doch wenn Athlet Lebron James jede einzelne Mu...
2       Wie kann ein Gehirn auf Hochleistung getrimmt ...
3       Wie schafft es Warren Buffett knapp 1000 Wörte...
4       Entfalte dein mentales Potenzial und werde ein...
Name: sentence_clean, Length: 3094, dtype: object

И это мои ярлыки из того же файла csv:

0       e_1
1       e_4
2       e_4
3       e_4
4       e_4

Различными ярлыками являются: e_1, e_2, e_3 и e_4.

Это код, который я использую для точной настройки моей модели:

import pandas as pd
import numpy as np
import os
    
# read in data
# sentences_df = pd.read_csv('path/file.csv')


X = sentences_df.sentence_clean
Y = sentences_df.classId

# =============================================================================
# One hot encode labels
# =============================================================================

# integer encode labels
from numpy import array
from numpy import argmax
from sklearn.preprocessing import LabelEncoder


label_encoder = LabelEncoder()
Y_integer_encoded = label_encoder.fit_transform(list(Y))


# one hot encode labels
from sklearn.preprocessing import OneHotEncoder

onehot_encoder = OneHotEncoder(sparse=False)
Y_integer_encoded_reshaped = Y_integer_encoded.reshape(len(Y_integer_encoded), 1)
Y_one_hot_encoded = onehot_encoder.fit_transform(Y_integer_encoded_reshaped)

# train test split
from sklearn.model_selection import train_test_split


X_train_raw, X_test_raw, y_train, y_test = train_test_split(X, Y_one_hot_encoded, test_size=0.20, random_state=42)


# =============================================================================
# Perpare datasets for finetuning
# =============================================================================
import tensorflow as tf
physical_devices = tf.config.list_physical_devices('GPU') 
tf.config.experimental.set_memory_growth(physical_devices[0], True)

from transformers import BertTokenizer, TFBertForSequenceClassification


tokenizer = BertTokenizer.from_pretrained('bert-base-german-cased') # initialize tokenizer


# tokenize trai and test sets
max_seq_length = 128

X_train_tokens = tokenizer(list(X_train_raw),
                            truncation=True,
                            padding=True)

X_test_tokens = tokenizer(list(X_test_raw),
                            truncation=True,
                            padding=True)


# create TF datasets as input for BERT model
bert_train_ds = tf.data.Dataset.from_tensor_slices((
    dict(X_train_tokens),
    y_train
))

bert_test_ds = tf.data.Dataset.from_tensor_slices((
    dict(X_test_tokens),
    y_test
))

# =============================================================================
# setup model and finetune
# =============================================================================

# define hyperparams
num_labels = 4
learninge_rate = 2e-5
epochs = 3
batch_size = 16

# create BERT model
bert_categorical_partial = TFBertForSequenceClassification.from_pretrained('bert-base-german-cased', num_labels=num_labels)

optimizer = tf.keras.optimizers.Adam(learning_rate=learninge_rate)
bert_categorical_partial.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy'])

history = bert_categorical_partial.fit(bert_train_ds.shuffle(100).batch(batch_size),
          epochs=epochs,
          # batch_size=batch_size,
          validation_data=bert_test_ds.shuffle(100).batch(batch_size))

А вот результат тонкой настройки:

Epoch 1/3
155/155 [==============================] - 31s 198ms/step - loss: 8.3038 - accuracy: 0.2990 - val_loss: 8.7751 - val_accuracy: 0.2811
Epoch 2/3
155/155 [==============================] - 30s 196ms/step - loss: 8.2451 - accuracy: 0.2913 - val_loss: 8.9314 - val_accuracy: 0.2779
Epoch 3/3
155/155 [==============================] - 30s 196ms/step - loss: 8.3101 - accuracy: 0.2913 - val_loss: 9.0355 - val_accuracy: 0.2746

Наконец, я пытаюсь предсказать метки тестового набора и проверить результаты с помощью матрицы путаницы:

X_test_tokens_new = {'input_ids': np.asarray(X_test_tokens['input_ids']),
                     'token_type_ids': np.asarray(X_test_tokens['token_type_ids']),
                     'attention_mask': np.asarray(X_test_tokens['attention_mask']),
                     }

pred_raw = bert_categorical_partial.predict(X_test_tokens_new)
pred_proba = tf.nn.softmax(pred_raw).numpy()
pred = pred_proba[0].argmax(axis = 1)
y_true = y_test.argmax(axis = 1)

cm = confusion_matrix(y_true, pred)

Выход отпечатка (см):

array([[  0,   0,   0,  41],
       [  2,   0,   0, 253],
       [  2,   0,   0, 219],
       [  6,   0,   0,  96]], dtype=int64)

Как видите, моя точность очень плохая, и когда я смотрю на сантиметры, я вижу, что моя модель в значительной степени просто предсказывает одну единственную метку. Я пробовал все и запускал модель несколько раз, но всегда получаю одни и те же результаты. Я знаю, что данные, с которыми я работаю, не очень хороши, и я тренируюсь только на двух тысячах предложений с ярлыками. Но я чувствую, что точность все равно должна быть выше и, что более важно, модель не должна предсказывать только одну единственную метку в 98% случаев, верно?

Я разместил все, что использую для запуска модели, в надежде, что кто-то сможет указать мне, где я ошибаюсь. Заранее большое спасибо за вашу помощь!