task_selection_token.py

from datasets import  Dataset, DatasetDict
from transformers import AutoTokenizer, get_scheduler
from adapters import AutoAdapterModel, AdapterConfig
import torch
from torch.utils.data import DataLoader
from torch.optim import AdamW
from torch.nn.functional import l1_loss
import pandas as pd
import numpy as np
from tqdm import tqdm
import os
import json

import csv

import matplotlib.pyplot as plt

def load_dataframes():
    path = "/pvc/home/apprentissage-profond/data/"
    
    # Load train_test_df and validation
    train_test_df = pd.read_parquet(path+"train_test_df.parquet")
    validation_df = pd.read_parquet(path+"validation_df.parquet")
    # Load mean and std used for normalizing the data
    with open(path+"normalization_stats.json") as f:
        stats = json.load(f)
        mean = stats["mean"]
        std = stats["std"]

    print("Loading dataset")
    print(f"Mean : {mean}, std : {std}")
    print("Loading successful")

    return (train_test_df, validation_df, mean, std)



def tokenize_data(dataset, platform_tokens):
    tokenizer = AutoTokenizer.from_pretrained("roberta-base")
    tokenizer.add_special_tokens({'additional_special_tokens': ["<tw>", "<rd>"]})

    def add_platform_prefix(row):
        tok = platform_tokens[int(row["platform"])]   
        row["text"] = f"{tok} {row['text']}"     
        return row
    
    dataset_prefixed = dataset.map(add_platform_prefix)

    tokenized_dataset = dataset_prefixed.map(
        lambda df: tokenizer(df["text"], truncation=True, padding="max_length", max_length=512),
        batched=True,
        remove_columns=['text','platform', 'comment_id', '__index_level_0__']
    )
    return tokenized_dataset, tokenizer

def load_model(red_factor):
    print(f'Loading model with reduction factor = {red_factor}')
    model = AutoAdapterModel.from_pretrained("roberta-base")
    config = AdapterConfig.load("houlsby", reduction_factor=int(red_factor))
    model.add_adapter("linear_reg", config=config)
    model.add_classification_head("linear_reg", num_labels=1, layers=1, activation_function="identity")
    model.train_adapter(["linear_reg"])
    model.set_active_adapters("linear_reg")
    model.active_head = "linear_reg"
    model.get_input_embeddings().weight.requires_grad = True
    total_params     = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"{trainable_params:,} / {total_params:,} parameters are trainable "
        f"→ {trainable_params / total_params:.2%} of the model.")  
    

    return model

def train_epoch(model, train_dataloader, optimizer, scheduler, device, mean, std):
    model.train()
    epoch_loss = 0
    epoch_mae = 0
    n_samples = 0
    for batch in tqdm(train_dataloader):
        batch = {k: v.to(device) for k, v in batch.items()}

        optimizer.zero_grad()
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()

        #prevent update of the embeddings of the original vocab by cancelling their gradients
        with torch.no_grad():
            grad = model.roberta.embeddings.word_embeddings.weight.grad
            if grad is not None:
                grad[:-2] = 0

        optimizer.step()
        scheduler.step()


        epoch_loss += outputs.loss.item()

        labels = batch['labels']
        preds = outputs.logits.squeeze()

        mae = l1_loss(preds, labels, reduction='sum')

        n_samples += labels.size(0)
        epoch_mae += mae.item()

    epoch_loss /= len(train_dataloader)
    epoch_mae /= n_samples

    return epoch_loss, epoch_mae


def evaluate(model, test_dataloader, device, mean, std):
    model.eval()
    with torch.no_grad():
        epoch_loss = 0
        epoch_mae = 0
        n_samples = 0

        for batch in test_dataloader:
            batch = {k: v.to(device) for k, v in batch.items()}
            outputs = model(**batch)
            epoch_loss += outputs.loss.item()

            labels = batch['labels']
            preds = outputs.logits.squeeze()

            #de-normalize hate speech score if necessary
            #labels = std*labels + mean
            #preds = std*preds + mean

            mae = l1_loss(preds, labels, reduction='sum')

            n_samples += labels.size(0)
            epoch_mae += mae.item()

        epoch_loss /= len(test_dataloader)
        epoch_mae /= n_samples


    return epoch_loss, epoch_mae


def print_inference_samples(model, tokenizer, test_dataset, device):
    model.eval()
    sample_indices = np.random.choice(len(test_dataset), size=3, replace=False)
    samples = [test_dataset[int(i)] for i in sample_indices]

    for sample in samples:
        input_ids = sample["input_ids"].unsqueeze(0).to(device)
        attention_mask = sample["attention_mask"].unsqueeze(0).to(device)
        label = sample["labels"].item()

        with torch.no_grad():
            output = model(input_ids=input_ids, attention_mask=attention_mask)
            pred = output.logits.squeeze().item()

        input_ids_cpu = input_ids.squeeze().cpu()
        pad_token_id = tokenizer.pad_token_id
        if pad_token_id in input_ids_cpu:
            pad_index = (input_ids_cpu == pad_token_id).nonzero(as_tuple=True)[0][0].item()
            input_ids_cpu = input_ids_cpu[:pad_index]
        decoded = tokenizer.decode(input_ids_cpu, skip_special_tokens=False)

        print("\n---")
        print("Comment:", decoded)
        print(f"True score: {label:.3f} | Predicted: {pred:.3f}")



def train_cycle(epochs=10, batch_size=100, dataset_size=8000, learning_rate=5e-5, red_factor=12):
    platform_tokens = {
        0: '<rd>',
        2: '<tw>'
    }

    train_df, validation_df, mean, std = load_dataframes()

    train_df = train_df.sample(int(dataset_size*0.9), replace=False)
    validation_df = validation_df.sample(int(dataset_size*0.1), replace=False)

    train_dataset = Dataset.from_pandas(train_df)
    validation_dataset = Dataset.from_pandas(validation_df)
    dataset_dict = DatasetDict({
        "train": train_dataset,
        "test": validation_dataset,
    })




    tokenized_dataset, tokenizer = tokenize_data(dataset_dict, platform_tokens)

    tokenized_dataset = tokenized_dataset.rename_column("hate_speech_score", "labels") #compatibility with auto loss


    tokenized_dataset.set_format("torch")

    train_dataloader = DataLoader(tokenized_dataset['train'], shuffle=True, batch_size=batch_size)
    test_dataloader = DataLoader(tokenized_dataset['test'], batch_size=batch_size)

    model = load_model(red_factor)
    model.resize_token_embeddings(len(tokenizer)) #adds two new embeddings

    with torch.no_grad():
        before = model.roberta.embeddings.word_embeddings.weight.detach().cpu().clone()

    embed_weight = model.roberta.embeddings.word_embeddings.weight
    optimizer = AdamW(
    [
        {"params": [embed_weight], "weight_decay": 0.0},
        {"params": [p for p in model.parameters() if p is not embed_weight],
         "weight_decay": 0.01},
    ],
    lr=learning_rate,
    )
 

    num_training_steps = epochs*len(train_dataloader)
    lr_scheduler = get_scheduler(
        name="linear", optimizer=optimizer, num_warmup_steps=1, num_training_steps=num_training_steps
    )

    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    model.to(device)

    losses = np.zeros((2, epochs))

    metrics_df = pd.DataFrame(dtype=float, index=range(epochs), columns=['train_mae', 'validation_mae'])

    for i in range(epochs):
        train_loss, train_mae = train_epoch(model, train_dataloader, optimizer, lr_scheduler, device, mean, std)
        eval_loss, eval_mae = evaluate(model, test_dataloader, device, mean, std)

        loss_text_line = f'Epoch {i}   eval loss: {eval_loss}, train_loss: {train_loss}\neval MAE: {eval_mae}, train MAE: {train_mae}'

        print(loss_text_line)

        print_inference_samples(model, tokenizer, tokenized_dataset["test"], device)

        metrics_df.loc[i, ["train_mae", "validation_mae"]] = [train_mae, eval_mae]

        losses[0,i] = train_mae
        losses[1,i] = eval_mae

        plt.figure(figsize=(12, 5))

        plt.plot(losses[0, :i+1], c='blue', label='train')
        plt.plot(losses[1, :i+1], c='orange', label='test')
        plt.title('Mean Absolute Error') 
        plt.legend()



        plt.savefig(f"/pvc/home/apprentissage-profond/results/mae_plot_embeddings.png")
        plt.close()

    results_path = "/pvc/home/apprentissage-profond/results/models"

    metrics_df.to_csv(f"/pvc/home/apprentissage-profond/results/domain_tokens_output.csv", index=False)

    torch.save(model.state_dict(), results_path + "model.pt")

    model.save_adapter(results_path, "linear_reg", with_head=True)

    with torch.no_grad():
        after = model.roberta.embeddings.word_embeddings.weight.detach().cpu()
    unchanged = torch.allclose(before[:-2], after[:-2], atol=0, rtol=0) #verify that the embeddings of the original vocabulary are not trained

    print("Old-token embeddings exactly identical?   ", unchanged)




    
def lr_sweep(start, stop, n_steps):
    for lr in np.logspace(np.log10(start), np.log10(stop), n_steps):
        train_cycle(5, 50, 15000, (lr, 5e-5), 8)  

def reduction_sweep(start, stop, step):
    for red_factor in np.arange(start, stop+step, step):
        train_cycle(10, 100, 15000, red_factor=red_factor)


             
if __name__ == "__main__":

    train_cycle(epochs=10, batch_size=50, dataset_size=15000, learning_rate=5e-5, red_factor=8) # epoch, batch_size, dataset_size per platform