Помоги мне поменять модель в данном скрипте, а то у меня train() происходит очень долгое время
!pip install datasets
!pip install transformers
!pip install evaluate


from typing import List, Tuple
from datasets import load_dataset, Dataset
import evaluate
import numpy as np
from sklearn.metrics import classification_report
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers import DataCollatorWithPadding
from transformers import TrainingArguments, Trainer
from transformers import pipeline


MODEL_NAME = 'DeepPavlov/rubert-base-cased'
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)


DATASET_NAME = 'Davlan/sib200'
DATASET_LANGUAGE = 'rus_Cyrl'
train_set = load_dataset(DATASET_NAME, DATASET_LANGUAGE, split='train')
validation_set = load_dataset(DATASET_NAME, DATASET_LANGUAGE, split='validation')
test_set = load_dataset(DATASET_NAME, DATASET_LANGUAGE, split='test')

print(train_set)


MINIBATCH_SIZE = 64
tokenized_train_set = train_set.map(
    lambda it: tokenizer(it['text'], truncation=True),
    batched=True, batch_size=MINIBATCH_SIZE
)
tokenized_validation_set = validation_set.map(
    lambda it: tokenizer(it['text'], truncation=True),
    batched=True, batch_size=MINIBATCH_SIZE
)

print(tokenized_train_set)


cls_metric = evaluate.load('f1')

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return cls_metric.compute(predictions=predictions, references=labels, average='macro')


data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

list_of_categories = sorted(list(
    set(train_set['category']) | set(validation_set['category']) | set(test_set['category'])
))
indices_of_categories = list(range(len(list_of_categories)))
n_categories = len(list_of_categories)
print(f'Categories for classification are: {list_of_categories}')
id2label = dict(zip(indices_of_categories, list_of_categories))
label2id = dict(zip(list_of_categories, indices_of_categories))

labeled_train_set = tokenized_train_set.add_column(
    'label',
    [label2id[val] for val in tokenized_train_set['category']]
)
labeled_validation_set = tokenized_validation_set.add_column(
    'label',
    [label2id[val] for val in tokenized_validation_set['category']]
)


print(labeled_train_set)


classifier = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME,
    num_labels=n_categories, id2label=id2label, label2id=label2id
).cuda()
for param in classifier.parameters(): param.data = param.data.contiguous()



training_args = TrainingArguments(
    output_dir='rubert_sib200',
    learning_rate=2e-5,
    per_device_train_batch_size=MINIBATCH_SIZE,
    per_device_eval_batch_size=MINIBATCH_SIZE,
    num_train_epochs=20,
    weight_decay=1e-3,
    eval_strategy='epoch',
    save_strategy='epoch',
    load_best_model_at_end=True
)

trainer = Trainer(
    model=classifier,
    args=training_args,
    train_dataset=labeled_train_set,
    eval_dataset=labeled_validation_set,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()


classifiсation_pipeline = pipeline('text-classification', model=classifier, tokenizer=tokenizer, device=0)

y_pred = list(map(lambda x: x['label'], classifiсation_pipeline(validation_set['text'])))
y_true = validation_set['category']
print(classification_report(y_true=y_true, y_pred=y_pred))

y_pred = list(map(lambda x: x['label'], classifiсation_pipeline(test_set['text'])))
y_true = test_set['category']
print(classification_report(y_true=y_true, y_pred=y_pred))