Text summarization model training problem - trainer.train() [00:00<?, ?it/s]
HelloI am learning from NLP HF course, and after first part, I wanted to train a model for summarization task on this dataset:huggingface.coWiktorS/polish-news · Datasets at Hugging FaceWe’re on a journey to advance and democratize artificial intelligence through open source and open science.Dowloading and mapping works well, however the training bar do not even start progressing. I am working on SSH, with GPU, in conda env with:Name                    Version                   Build  Channel_libgcc_mutex             0.1                        main_openmp_mutex             5.1                       1_gnuabsl-py                   2.1.0                    pypi_0    pypiaccelerate                0.30.1                   pypi_0    pypiastunparse                1.6.3                    pypi_0    pypiattrs                     23.2.0                   pypi_0    pypibzip2                     1.0.8                h5eee18b_6ca-certificates           2024.3.11            h06a4308_0certifi                   2024.6.2                 pypi_0    pypiclick                     8.1.7                    pypi_0    pypiexpat                     2.6.2                h6a678d5_0flatbuffers               24.3.25                  pypi_0    pypigast                      0.5.4                    pypi_0    pypigoogle-pasta              0.2.0                    pypi_0    pypigrpcio                    1.64.1                   pypi_0    pypih5py                      3.11.0                   pypi_0    pypiidna                      3.7                      pypi_0    pypijinja2                    3.1.4                    pypi_0    pypijoblib                    1.4.2                    pypi_0    pypikeras                     3.3.3                    pypi_0    pypild_impl_linux-64          2.38                 h1181459_1libclang                  18.1.1                   pypi_0    pypilibffi                    3.4.4                h6a678d5_1libgcc-ng                 11.2.0               h1234567_1libgomp                   11.2.0               h1234567_1libstdcxx-ng              11.2.0               h1234567_1libuuid                   1.41.5               h5eee18b_0markdown                  3.6                      pypi_0    pypimarkdown-it-py            3.0.0                    pypi_0    pypimarkupsafe                2.1.5                    pypi_0    pypimdurl                     0.1.2                    pypi_0    pypiml-dtypes                 0.3.2                    pypi_0    pypimpmath                    1.3.0                    pypi_0    pypinamex                     0.0.8                    pypi_0    pypincurses                   6.4                  h6a678d5_0networkx                  3.3                      pypi_0    pypinltk                      3.8.1                    pypi_0    pypinvidia-cublas-cu12        12.1.3.1                 pypi_0    pypinvidia-cuda-cupti-cu12    12.1.105                 pypi_0    pypinvidia-cuda-nvrtc-cu12    12.1.105                 pypi_0    pypinvidia-cuda-runtime-cu12  12.1.105                 pypi_0    pypinvidia-cudnn-cu12         8.9.2.26                 pypi_0    pypinvidia-cufft-cu12         11.0.2.54                pypi_0    pypinvidia-curand-cu12        10.3.2.106               pypi_0    pypinvidia-cusolver-cu12      11.4.5.107               pypi_0    pypinvidia-cusparse-cu12      12.1.0.106               pypi_0    pypinvidia-nccl-cu12          2.20.5                   pypi_0    pypinvidia-nvjitlink-cu12     12.5.40                  pypi_0    pypinvidia-nvtx-cu12          12.1.105                 pypi_0    pypiopenssl                   3.0.13               h7f8727e_2opt-einsum                3.3.0                    pypi_0    pypioptree                    0.11.0                   pypi_0    pypipip                       24.0            py310h06a4308_0protobuf                  4.25.3                   pypi_0    pypipsutil                    5.9.8                    pypi_0    pypipygments                  2.18.0                   pypi_0    pypipython                    3.10.14              h955ad1f_1pytz                      2024.1                   pypi_0    pypipyyaml                    6.0.1                    pypi_0    pypireadline                  8.2                  h5eee18b_0rich                      13.7.1                   pypi_0    pypirouge-score               0.1.2                    pypi_0    pypiscikit-learn              1.5.0                    pypi_0    pypiscipy                     1.14.0                   pypi_0    pypisetuptools                69.5.1          py310h06a4308_0six                       1.16.0                   pypi_0    pypisqlite                    3.45.3               h5eee18b_0sympy                     1.12.1                   pypi_0    pypitensorboard               2.16.2                   pypi_0    pypitensorboard-data-server   0.7.2                    pypi_0    pypitensorflow                2.16.1                   pypi_0    pypitensorflow-io-gcs-filesystem 0.37.0                   pypi_0    pypitermcolor                 2.4.0                    pypi_0    pypitf-keras                  2.16.0                   pypi_0    pypithreadpoolctl             3.5.0                    pypi_0    pypitk                        8.6.14               h39e8969_0torch                     2.3.1                    pypi_0    pypitriton                    2.3.1                    pypi_0    pypitzdata                    2024a                h04d1e81_0urllib3                   2.2.1                    pypi_0    pypiwerkzeug                  3.0.3                    pypi_0    pypiwheel                     0.43.0          py310h06a4308_0wrapt                     1.16.0                   pypi_0    pypixz                        5.4.6                h5eee18b_1zlib                      1.2.13               h5eee18b_1# Text summarization on polish news dataset

from datasets import load_dataset
from transformers import AutoTokenizer    
import torch
from transformers import DataCollatorForSeq2Seq
from transformers import AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer
import evaluate
import numpy as np


import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
# choose GPU
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# load dataset
# dataset contains only a "train" split, so we need to split this
# train: Dataset({features: ['link', 'title', 'headline', 'content'], num_rows: 248123
raw_dataset = load_dataset("WiktorS/polish-news", split="train").shuffle(seed=42).select(range(100)) # choses which part of split to load #select 100 random
raw_dataset = raw_dataset.train_test_split(test_size=0.1) # split="train" - necessary to load in the purpose of future splitting

# dataset preprocessing: delete None Types in headline and content columns

def filter_none(example):
    return example['headline'] is not None and example['content'] is not None

raw_dataset = raw_dataset.filter(filter_none)

# choose model

checkpoint="facebook/bart-large-cnn"

# tokenize

tokenizer = AutoTokenizer.from_pretrained(checkpoint)

# tokenization function (mapping)

# Prefix the input with a prompt so model knows this is a summarization task. Some models capable of multiple NLP tasks require prompting for specific tasks.
prefix = "summarize: "

def tokenize_function(example):
    
    inputs = [prefix + doc for doc in example["content"]]
    model_inputs = tokenizer(inputs, max_length=1024, truncation=True) #dictionary with keys: input_ids, attention_mask
    #tokenize not only the input content but also labels
    labels = tokenizer(text_target=example["headline"], max_length=128, truncation=True)

    model_inputs["labels"] = labels["input_ids"] # add labels to the dictionary
    return model_inputs # return dictionary with keys: input_ids, attention_mask, labels


tokenize_dataset = raw_dataset.map(tokenize_function, batched=True) #features: ['link', 'title', 'headline', 'content', 'input_ids', 'attention_mask', 'labels']
print(tokenize_dataset["train"][0])

# create a batch of examples; dynamically pad the sentences to the longest length in a batch during collation, instead of padding the whole dataset to the maximum length
# data collator to load batch into tokanization function

data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=checkpoint)

# arguments for trainer

training_args = Seq2SeqTrainingArguments(
    output_dir="final-model",
    eval_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    weight_decay=0.01,
    save_total_limit=3,
    num_train_epochs=4,
    predict_with_generate=True,
    fp16=True,
)

# evaluation fuction included in the trainer. Including a metric during training is often helpful for evaluating your model’s performance. 

rouge = evaluate.load("rouge")

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True) #Converts tokens to human-readable text, skipping special tokens.
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id) #Replaces -100 values in labels with pad tokens so they can be decode
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True) 

    result = rouge.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True) #Computes the ROUGE metric for the model's predictions

    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions] #Counts the number of non-pad tokens in each prediction
    result["gen_len"] = np.mean(prediction_lens) #Calculates the average number of non-pad tokens in the predictions

    return {k: round(v, 4) for k, v in result.items()} #Returns the ROUGE metric results rounded to four decimal places

# initialize the model

model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)


# trainer

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenize_dataset["train"],
    eval_dataset=tokenize_dataset["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    #compute_metrics=compute_metrics,
)

# train model

trainer.train()