BERT with lightautoml

# Решение задачи NLP с помощью lightautoml (включить GPU!!!)
# Сейчас настроен самый "тяжелый" вариант, когда дообучается BERT, можно настроить более быстрые и простые варианты. См.док.:
# https://github.com/sberbank-ai-lab/LightAutoML/blob/master/lightautoml/automl/presets/text_config.yml
# https://github.com/sberbank-ai-lab/LightAutoML

import pandas as pd
import numpy as np

from lightautoml.automl.presets.text_presets import TabularNLPAutoML
from lightautoml.tasks import Task

from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score

bert_models = {
    'en': 'bert-base-cased',
    'ru': 'DeepPavlov/rubert-base-cased-conversational',
    'multi': 'bert-base-multilingual-cased'
}

class params:
    random_state = 12345
    timeout = 30_000 # сколько секунд на выполнение алгоритма
    col_target = 'toxic'
    cols_text = ['text'] # колонки с текстом, которые надо обрабатывать
    n_threads = 4
    test_size = 0.2
    epochs = 3 # кол-во эпох обучения (но вроде важнее timeout)
    batch_size = 24 # если позволяет память GPU, то увеличить до 32 и выше. Что ускорит обучение
    cv = 3 # кол-во folds для кросс-валидации (лучше 5, но будет дольше обучаться)
    lang = 'en'

df = pd.read_csv('toxic_comments.csv')
print(df.shape)

train, test = train_test_split(df,
                               test_size=params.test_size,
                               random_state=params.random_state)
roles = {
    'target': params.col_target,
    'text': params.cols_text,
}

task = Task('binary')

automl = TabularNLPAutoML(
    task = task,
    timeout = params.timeout,
    cpu_limit = params.n_threads,
    general_params = {'nested_cv': False, # не срабатывает, всё равно исп.кросс-валидация
                      'use_algos': [['nn']]},
    text_params = {'lang': params.lang, 'bert_model': bert_models[params.lang]},
    nn_params = {'opt_params': { 'lr': 1e-5},
                 'max_length': 150, 'bs': params.batch_size,
                 'epoch': params.epochs # не срабатывает для обучения Bert. Видимо, только для посл.сети.
                },
    nested_cv_params = {'cv': params.cv}, # не срабатывает 1, всё равно исп. кросс-валидация
)

oof_pred = automl.fit_predict(train, roles=roles)
test_pred = automl.predict(test)

treshold = 0.56 # порог для предсказания класса 1 (нужно подбирать)
f1_val = f1_score(train[params.col_target].values,
                  np.where(oof_pred.data[:, 0] > treshold, 1, 0))
f1_test = f1_score(test[params.col_target].values,
                   np.where(test_pred.data[:, 0] > treshold, 1, 0))
print(f'VAL score: {f1_val}')
print(f'TEST score: {f1_test}')

# F1=0.844 - TEST получается с параметрами:timeout = 30_000; epochs = 3; treshold = 0.56