re_ranked_train.py

import argparse
import torch
import numpy as np
import pandas as pd
import random
import evaluate
from sklearn.metrics import classification_report
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import Trainer, TrainingArguments
from transformers import DataCollatorWithPadding
from sklearn.model_selection import train_test_split
from datasets import Dataset


# GPU device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Global variable
# 3 number of labels
TF_MAX_SIZE = 512

metric1 = evaluate.load("precision")
metric2 = evaluate.load("recall")
metric3 = evaluate.load("f1")
metric_name = "f_macro"


def preprocess_function(examples, tokenizer):
    """
    Tokeniza los textos y términos para la entrada del modelo.
    
    Args:
        examples (dict): Diccionario con las claves "text" y "term".
        tokenizer (AutoTokenizer): Tokenizador preentrenado.
    
    Returns:
        dict: Diccionario con las secuencias tokenizadas.
    """
    return tokenizer(examples["entity"], text_pair=examples["term"], max_length=TF_MAX_SIZE, padding="max_length", truncation=True)


def compute_metrics(eval_pred):
    """
    Calcula métricas de evaluación a partir de las predicciones.
    
    Args:
        eval_pred (tuple): Tupla con logits y etiquetas reales.
    
    Returns:
        dict: Diccionario con precisión, recall, f-score y f_macro.
    """
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    print(classification_report(labels, predictions, digits=4))
    precision = metric1.compute(predictions=predictions, references=labels, average='weighted')["precision"]
    recall = metric2.compute(predictions=predictions, references=labels, average='weighted')["recall"]
    f_score = metric3.compute(predictions=predictions, references=labels, average='weighted')["f1"]
    f_macro = metric3.compute(predictions=predictions, references=labels, average='macro')["f1"]
    return {"precision": precision, "recall": recall, "f_score": f_score, "f_macro": f_macro}


def train_model(transformer_model, train_df, val_df, model_name, num_label, label2id, id2label):
    """
    Entrena un modelo de clasificación de secuencias con Transformers.
    
    Args:
        transformer_model (str): Nombre del modelo preentrenado.
        train_df (pd.DataFrame): DataFrame con datos de entrenamiento.
        val_df (pd.DataFrame): DataFrame con datos de validación.
        model_name (str): Nombre del modelo.
        num_label (int): Número de etiquetas en la clasificación.
        label2id (dict): Mapeo de etiquetas a IDs.
        id2label (dict): Mapeo de IDs a etiquetas.
        save_path (str): Ruta donde guardar el modelo entrenado.
        dataset_name (str): Nombre del dataset.
    """
    tokenizer = AutoTokenizer.from_pretrained(transformer_model, max_length=TF_MAX_SIZE, padding="max_length", truncation=True)
    train_dataset = Dataset.from_pandas(train_df)
    val_dataset = Dataset.from_pandas(val_df)

    train_dataset = train_dataset.map(
            preprocess_function,
            fn_kwargs={"tokenizer": tokenizer}
    )
    
    val_dataset = val_dataset.map(
            preprocess_function,
            fn_kwargs={"tokenizer": tokenizer}
    )

    print(train_dataset[0])

    # For reproductivity 
    def model_init(trial):
        return AutoModelForSequenceClassification.from_pretrained(transformer_model, num_labels=num_label, id2label=id2label, label2id=label2id)
    
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    
    batch_train_size = 32
    batch_eval_size = 32
    training_args = TrainingArguments(
        output_dir="./reranked_log",
        # overwrite_output_dir=True,
        num_train_epochs=5,
        eval_strategy="epoch",
        save_strategy="epoch",
        per_device_train_batch_size=batch_train_size,
        per_device_eval_batch_size=batch_eval_size,
        metric_for_best_model=metric_name,
        save_total_limit=1,
        seed = 1,
        learning_rate = 2e-5,
        load_best_model_at_end=True,
        report_to="none"
    )
    trainer = Trainer(
        model_init=model_init,
        tokenizer=tokenizer,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
    )
    trainer.train()
    # Salvamos el modelo reentrenado
    trainer.save_model(f'models/{model_name}_reranked')
    tokenizer.save_pretrained(f'models/{model_name}_reranked')


def main():
    english_transformer_model = {'BERT-uncased': 'google-bert/bert-base-uncased', 
                                 'bio-bert': 'dmis-lab/biobert-v1.1',
                                 'XLM-RoBERTa': 'xlm-roberta-base',
                                }  
        
    # Read dataset     
    train_df = pd.read_csv("reranked_train.csv")
    train_df = train_df.dropna()
    val_df = pd.read_csv("reranked_valid.csv")
    val_df = val_df.dropna()

    label_list = sorted(train_df.label.unique().tolist())
    print(label_list)
    num_label = len(label_list)
    label2id = {label: i for i, label in enumerate(label_list)}
    id2label = {i: label for i, label in enumerate(label_list)}
    
    train_df['label'] = train_df['label'].apply(lambda x: label2id[x])
    val_df['label'] = val_df['label'].apply(lambda x: label2id[x])
       
    for k, v in english_transformer_model.items():

        train_model(v, train_df, val_df, k, num_label, label2id, id2label)


if __name__ == '__main__':
    main()