src/train.py

import glob
import logging
import os
import pickle
import json

import torch
from progress.bar import Bar
from sklearn.metrics import confusion_matrix, classification_report
from tabulate import tabulate
from torch.optim import AdamW
from tqdm.auto import tqdm
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from transformers import MarkupLMForTokenClassification
from transformers import MarkupLMProcessor
import evaluate
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

from timing import Timing
from consts import label2id, id2label

# pd.set_option('display.max_colwidth', 20)
# pd.set_option('display.max_columns', None)

MAX_LENGTH = 512
EPOCH_COUNT = 5
BATCH_SIZE = 25
SHUFFLE = True


class MarkupLMDataset(Dataset):
    """
    MarkupLM ile belirteç sınıflandırması için veri kümesi.

    Args:
        data (list): 'nodes', 'xpaths' ve 'node_labels' bilgisini içeren sözlük listesi.
        processor (MarkupLMProcessor, optional): Veriyi işlemek için MarkupLMProcessor örneği.
        max_length (int, optional): Kodlanmış dizilerin maksimum uzunluğu.

    Attributes:
        data (list): 'nodes', 'xpaths' ve 'node_labels' bilgisini içeren sözlük listesi.
        processor (MarkupLMProcessor): Veriyi işlemek için MarkupLMProcessor örneği.
        max_length (int): Kodlanmış dizilerin maksimum uzunluğu.

    Methods:
        __len__(): Veri kümesinin uzunluğunu döndürür.
        __getitem__(idx): Belirtilen dizindeki kodlanmış diziyi döndürür.

    """

    def __init__(self, data, processor: MarkupLMProcessor = None, max_length=MAX_LENGTH):
        self.data = data
        self.processor = processor
        self.max_length = max_length

    def __len__(self):
        """
        Veri kümesinin uzunluğunu döndürür.

        Returns:
            int: Veri kümesinin uzunluğu.

        """
        return len(self.data)

    def __getitem__(self, idx):
        """
        Belirtilen dizindeki kodlanmış diziyi döndürür.

        Args:
            idx (int): Alınacak öğenin dizini.

        Returns:
            dict: 'input_ids', 'attention_mask' ve 'labels' anahtarlarına sahip kodlanmış dizi.

        """
        # İlk olarak, 'nodes', 'xpaths' ve 'node_labels' bilgilerini alın
        item = self.data[idx]
        nodes, xpaths, node_labels = item['nodes'], item['xpaths'], item['node_labels']
        # İşlemciye iletilir
        encoding = self.processor(nodes=nodes, xpaths=xpaths, node_labels=node_labels, padding="max_length",
                                  max_length=self.max_length, return_tensors="pt", truncation=True)

        # Toplu boyutunu kaldır
        encoding = {k: v.squeeze() for k, v in encoding.items()}
        return encoding


class NewsTrainer:
    """
    NewsTrainer, MarkupLM ile haberler için belirteç sınıflandırması yapmak için kullanılan bir eğitim sınıfıdır.

    Methods:
        __init__(): NewsTrainer sınıfını başlatır.
        __get_labels(): Tahminleri ve referansları kullanarak etiketleri alır.
        __compute_metrics(): Metrikleri hesaplar.
        __load_train_data(): Eğitim verisini yükler.
        __get_dataset(): Veri kümesini oluşturur.
        __train(): Modeli eğitir.
        run(): Eğitimi çalıştırır.

    """

    def __init__(self):
        """
        NewsTrainer sınıfını başlatır.

        """
        logging.debug('NewsTrainer Sınıfı oluşturuldu')

    @staticmethod
    def __get_labels(predictions, references, label_list, device):
        """
        Tahminleri ve referansları kullanarak etiketleri alır.

        Args:
            predictions (torch.Tensor): Tahminler tensörü.
            references (torch.Tensor): Referanslar tensörü.
            label_list (list): Etiket listesi.
            device (torch.device): Cihaz türü.

        Returns:
            list, list: Gerçek tahminler ve gerçek etiketler listeleri.

        """
        # Tahminleri ve referansları numpy dizilerine dönüştürme
        if device.type == "cpu":
            y_pred = predictions.detach().clone().numpy()
            y_true = references.detach().clone().numpy()
        else:
            y_pred = predictions.detach().cpu().clone().numpy()
            y_true = references.detach().cpu().clone().numpy()

        # İgnor index'ini (özel belirteçler) kaldırma
        true_predictions = [
            [label_list[p] for (p, l) in zip(pred, gold_label) if l != -100]
            for pred, gold_label in zip(y_pred, y_true)
        ]
        true_labels = [
            [label_list[l] for (p, l) in zip(pred, gold_label) if l != -100]
            for pred, gold_label in zip(y_pred, y_true)
        ]
        return true_predictions, true_labels

    @staticmethod
    def __get_labels_2(predictions, references, label_list, device):
        """
        Tahminleri ve referansları kullanarak etiketleri alır.

        Args:
            predictions (torch.Tensor): Tahminler tensörü.
            references (torch.Tensor): Referanslar tensörü.
            label_list (list): Etiket listesi.
            device (torch.device): Cihaz türü.

        Returns:
            list, list: Gerçek tahminler ve gerçek etiketler listeleri.

        """
        # Tahminleri ve referansları numpy dizilerine dönüştürme
        if device.type == "cpu":
            y_pred = predictions.detach().clone().numpy()
            y_true = references.detach().clone().numpy()
        else:
            y_pred = predictions.detach().cpu().clone().numpy()
            y_true = references.detach().cpu().clone().numpy()

        # İgnor index'ini (özel belirteçler) kaldırma
        true_predictions = [
            [label_list[p] for (p, l) in zip(pred, gold_label)]
            for pred, gold_label in zip(y_pred, y_true)
        ]
        true_labels = [
            [label_list[l] for (p, l) in zip(pred, gold_label)]
            for pred, gold_label in zip(y_pred, y_true)
        ]
        return true_predictions, true_labels

    @staticmethod
    def __compute_metrics(metric, return_entity_level_metrics=True):
        """
        Metrikleri hesaplar.

        Args:
            metric: Metrik hesaplama nesnesi.
            return_entity_level_metrics (bool, optional): Bireysel düzeyde metrikleri döndürme.

        Returns:
            dict: Hesaplanan metrikler.

        """
        results = metric.compute()
        if return_entity_level_metrics:
            # İç içe geçmiş sözlükleri açma
            final_results = {}
            for key, value in results.items():
                if isinstance(value, dict):
                    for n, v in value.items():
                        final_results[f"{key}_{n}"] = v
                else:
                    final_results[key] = value
            return final_results
        else:
            return {
                "precision": results["overall_precision"],
                "recall": results["overall_recall"],
                "f1": results["overall_f1"],
                "accuracy": results["overall_accuracy"],
            }

    @staticmethod
    def __load_train_data(data_path):
        """
        Eğitim verisini yükler.

        Args:
            data_path (str): Veri dosyasının yolu.

        Returns:
            list: Yüklenen veri listesi.

        """
        # ./data/dataset/train
        file_dir = f"{data_path}"
        lfs = glob.glob(f"{file_dir}/*.pickle")
        _max = len(lfs)
        logging.info(f"Veri yükleme işlemi başlatıldı.")
        objects = []
        with Bar('Veri Kümeleri Birleştiriliyor', max=_max,
                 suffix='%(percent).1f%% | %(remaining)d | %(max)d | %(eta)ds') as bar:
            i = 0
            for lf in lfs:
                try:
                    with (open(lf, "rb")) as dataset_file:
                        while True:
                            try:
                                dataset = pickle.load(dataset_file)
                                for item in dataset:
                                    objects.append(item)
                            except EOFError:
                                break
                    bar.next()
                    i = i + 1
                except Exception as e:
                    logging.error(f"Bir hata oluştu. ID: {lf} Hata: {str(e)}")
        bar.finish()
        logging.info(f"Veri yükleme işlemi tamamlandı.\n")
        return objects

    def __get_dataset(self, data_path):
        """
        Veri kümesini oluşturur.

        Args:
            data_path (str): Veri dosyasının yolu.

        Returns:
            MarkupLMDataset: Oluşturulan veri kümesi.

        """
        _data = self.__load_train_data(data_path)
        processor = MarkupLMProcessor.from_pretrained("microsoft/markuplm-base")
        processor.parse_html = False
        dataset = MarkupLMDataset(data=_data, processor=processor, max_length=MAX_LENGTH)
        return dataset

    def __train(self, model_name, dataset, model_output_path):
        """
        Modeli eğitir.

        Args:
            model_name (str): Modelin adı.
            dataset (MarkupLMDataset): Eğitim veri kümesi.
            model_output_path (str): Model çıkış dizini.

        """
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=SHUFFLE)
        model = MarkupLMForTokenClassification.from_pretrained("microsoft/markuplm-base",
                                                               id2label=id2label,
                                                               label2id=label2id)

        label_list = ["B-" + x for x in list(id2label.values())]
        metric = evaluate.load("seqeval")

        optimizer = AdamW(model.parameters(), lr=5e-5)
        model.to(device)
        model.train()
        print("----------------------------")
        print("------- EĞİTİM BAŞLADI ----")
        print("----------------------------")
        timing = Timing(True)
        timing.start()
        eval_metric = None
        for epoch in range(EPOCH_COUNT):  # veri kümesi üzerinde birden fazla dön
            print(f"Epoch: {epoch} başladı.")
            i = 0
            for batch in tqdm(dataloader):
                i = i + 1
                # girdileri al;
                inputs = {k: v.to(device) for k, v in batch.items()}
                # parametre gradyanlarını sıfırla
                optimizer.zero_grad()
                # ileri yayılım + geri yayılım + optimize et
                outputs = model(**inputs)

                loss = outputs.loss
                loss.backward()  # gradyanı hesapla
                optimizer.step()  # ağırlıkları güncelle

                print(f"Epoch: {epoch} - Batch: {i} - Loss: {loss.item()}")

                predictions = outputs.logits.argmax(dim=-1)
                labels = batch["labels"]
                preds, refs = self.__get_labels(predictions, labels, label_list, device)
                metric.add_batch(
                    predictions=preds,
                    references=refs,
                )
            eval_metric = self.__compute_metrics(metric)
            df_eval_metric = pd.DataFrame(eval_metric, index=[0])
            print(f"Epoch {epoch}: ", eval_metric)
            print(tabulate(df_eval_metric.transpose(), headers='keys', tablefmt='psql'))
            # kontrol noktasını kaydet
            if not os.path.exists(model_output_path):
                os.makedirs(model_output_path)
            torch.save(model, f"{model_output_path}/{model_name}_{epoch}.pt")
            # kontrol noktası metriklerini kaydet
            with open(f"{model_output_path}/{model_name}_{epoch}_metrics.json", 'w', encoding='utf-8') as f:
                json.dump(eval_metric, f, default=str, ensure_ascii=False, indent=4)
            print(f"Epoch: {epoch} tamamlandı.")

        # final modeli kaydet
        torch.save(model, f"{model_output_path}/{model_name}.pth")
        # final metrikleri kaydet
        with open(f"{model_output_path}/{model_name}_metrics.json", 'w', encoding='utf-8') as f:
            json.dump(eval_metric, f, default=str, ensure_ascii=False, indent=4)
        timing.end()
        timing.print(f"Eğitim Tamamlandı.")
        print("----------------------------")
        print("------- EĞİTİM TAMAMLANDI ----")
        print("----------------------------")

    def evaluate(self):
        """
        Modelin performansını değerlendiren ve karmaşıklık matrisini çizen bir yöntem.

        Returns:
            None
        """
        train_data_path = "../data/dataset/test"
        model_path = "../model/model.pth"
        label_list = ["" + x for x in list(id2label.values())]
        label_list = label_list[:4]
        dataset = self.__get_dataset(train_data_path)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        dataloader = DataLoader(dataset, batch_size=10)
        model = torch.load(model_path, map_location=torch.device(device))
        i = 0
        y_pred = []
        y_true = []
        with torch.no_grad():
            for batch in tqdm(dataloader):
                i = i + 1
                inputs = {k: v.to(device) for k, v in batch.items()}
                outputs = model(**inputs)
                predictions = outputs.logits.argmax(dim=-1)
                labels = batch["labels"]
                preds, refs = self.__get_labels(predictions, labels, label_list, device)
                for pred in preds:
                    for pl in pred:
                        y_pred.append(pl)

                for ref in refs:
                    for rl in ref:
                        y_true.append(rl)
        print(classification_report(y_true, y_pred))
        cnf_matrix = confusion_matrix(y_true, y_pred, labels=label_list)
        ax = sns.heatmap(cnf_matrix, annot=True, fmt='', xticklabels=label_list,
                         yticklabels=label_list)

        plt.title('Karmaşıklık Matrisi', fontsize=20)  # title with fontsize 20
        plt.xlabel('Model tahminleri', fontsize=15)  # x-axis label with fontsize 15
        plt.ylabel('Gerçek Değerler', fontsize=15)  # x-axis label with fontsize 15
        plt.show()

    def run(self, model_name, train_data_path, model_output_path):
        """
        Eğitimi çalıştırır.

        Args:
            model_name (str): Modelin adı.
            train_data_path (str): Eğitim veri setinin dosya yolu.
            model_output_path (str): Model çıkış dizini.

        """
        dataset = self.__get_dataset(train_data_path)
        self.__train(model_name, dataset, model_output_path)


if __name__ == '__main__':
    trainer = NewsTrainer()
    # Eğitim
    # model_name = "model-10-1000"
    # _train_data_path = "./data/dataset/100"
    # _model_output_path = "./models"
    # trainer.run(model_name=model_name,
    #            train_data_path=_train_data_path,
    #            model_output_path=_model_output_path)

    # Değerlendirme
    trainer.evaluate()