Reproducibility/Finetune_sum.py

import os
# Set the HF_HOME environment variable to point to the desired cache location
# os.environ["HF_TOKEN"] = "your_hugging_face_token_here"  # Replace with your Hugging Face token
# Specify the directory path 
cache_dir = '/network/rit/lab/Lai_ReSecureAI/kiel/wmm'
# Set the HF_HOME environment variable
os.environ['HF_HOME'] = cache_dir

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

import matplotlib.pyplot as plt
import logging
import time
import torch
import json
import torch.nn as nn
from typing import Optional
import pandas as pd
from datasets import Dataset
from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training
from dataclasses import dataclass, field
from transformers import (
    HfArgumentParser,
    AutoTokenizer,
    TrainingArguments,
    BitsAndBytesConfig,
    TrainerCallback,
    AutoModelForCausalLM
)
from trl import SFTTrainer
import warnings


# Ignore all warnings
warnings.filterwarnings("ignore")

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Clear cache
torch.cuda.empty_cache()
device = "cuda" if torch.cuda.is_available() else "cpu"

# Initialize parameters
model_name =  "Llama2" #'DeepSeek' #'Llama3' #
WM = "TW"
num_data = 10000
num_epochs = 5
learning_rate_ = 1e-5


# Print parameters
print(f'Device: {device}')
print(f'Model: {model_name}')
print(f'WM: {WM}')
print(f'Number of data: {num_data}')
print(f'Number of epochs: {num_epochs}')
print(f'Learning rate: {learning_rate_}')

start_time = time.time()
# Load data
def load_data(file_path, num_data):
    with open(file_path, 'r') as f:
        data = json.load(f)
        return [
            {
                "text": "Now summarize the following text with maximum 60 words: " +
                        item["article"] + 
                        "\nThe summary is: " +
                        item['Watermarked_summary']
            }
            for item in data[:num_data]
        ]


# Create dataset
def create_dataset(data):
    """
    Convert the concatenated data into a Hugging Face Dataset format.
    """
    df = pd.DataFrame(data)  # Each element in 'data' is a dictionary with 'text' as the key
    return Dataset.from_pandas(df)

def get_file_paths(model_name,WM):
    base_path = '/network/rit/lab/Lai_ReSecureAI/kiel/Website/Stealing/'
    if WM == "SafeSeal":
        paths = {
            'DeepSeek': ('DeepSeek_train_Summarization_Safeseal_top_3_threshold_0.8_Uniform_0_20000_20k.json', 'DeepSeek_test_Summarization_Safeseal_top_3_threshold_0.8_Uniform_0_1000_1000.json'),
            'Llama3': ('Llama3_train_Summarization_Safeseal_top_3_threshold_0.8_Uniform_0_20000_20k.json', 'Llama3_test_Summarization_Safeseal_top_3_threshold_0.8_Uniform_0_1000_1000.json')
        }
    elif WM == "DTM":
        paths = {
            'Llama3': ('Llama3_DTM_Summarization_train__20000.json', 'Llama3_DTM_Summarization_test__1000.json'),
            'DeepSeek': ('DeepSeek_DTM_Summarization_train__20000.json', 'DeepSeek_DTM_Summarization_test__1000.json'),
            'Llama2': ('Llama2_DTM_Summarization_train_20k.json', 'Llama2_DTM_Summary_test_1000.json'),
            'Mistral': ('Mistral_DTM_Summarization_train_20k.json', 'Mistral_DTM_Summary_test_1000.json') 
        }
    elif WM == "KGW":
        paths = {
            'Llama3': ('Llama3_KGW_Summarization_train_0_20000_20000.json', 'Llama3_KGW_Summarization_test_0_1000_1000.json'),
            'DeepSeek': ('DeepSeek_KGW_Summarization_train_0_20000_20000.json', 'DeepSeek_KGW_Summarization_test_0_1000_1000.json')
        }
    elif WM == "SIR":
        paths = {
            'DeepSeek': ('DeepSeek_SIR_Summarization_train_0_20000_20000.json', 'DeepSeek_SIR_Summarization_test_0_1000_1000.json'),
            'Llama3': ('Llama3_SIR_Summarization_train_0_20000_20000.json', 'Llama3_SIR_Summarization_test_0_1000_1000.json')
        }
    elif WM == "SynthID":
        paths = {
            'DeepSeek': ('DeepSeek_SynthID_Summarization_train_0_20000_20000.json', 'DeepSeek_SynthID_Summarization_test_0_1000_1000.json'),
            'Llama3': ('Llama3_SynthID_Summarization_train_0_20000_20000.json', 'Llama3_SynthID_Summarization_test_0_1000_1000.json')
        }
    elif WM == "TW":
        paths = {
            'DeepSeek': ('DeepSeek_TW_Summarization_train_20000.json', 'DeepSeek_TW_Summarization_test__1000.json'),
            'Llama3': ('Llama3_TW_Summarization_train__20000.json', 'Llama3_TW_Summarization_test__1000.json'),
            'Llama2': ('Llama2_TW_Summarization_train_20k.json', 'Llama2_TW_Summary_test_1000.json'),
            'Mistral': ('Mistral_TW_Summarization_train_20k.json', 'Mistral_TW_Summary_Test_1000.json') 
        }

    return base_path + paths[model_name][0], base_path + paths[model_name][1]

def get_new_model_path(model_name,WM, num_epochs, learning_rate_, num_data):
    #return f"/network/rit/lab/Lai_ReSecureAI/phung/adversary_models/{model_name}_epoch{num_epochs}_lr{learning_rate_}_K{K}_Threshold{Threshold}_data{num_data}_testing_batch{batch_no}_"
    return f"./adversary_models/{model_name}_{WM}_epoch{num_epochs}_lr{learning_rate_}_data{num_data}_"
    #return f"/network/rit/lab/Lai_ReSecureAI/phung/adversary_models/{model_name}_{WM}_epoch{num_epochs}_lr{learning_rate_}_data{num_data}_"

train_file, test_file = get_file_paths(model_name, WM)
train_data = load_data(train_file, num_data)
test_data = load_data(test_file, num_data)

train_dataset = create_dataset(train_data)
test_dataset = create_dataset(test_data)

new_model = get_new_model_path(model_name, WM, num_epochs, learning_rate_, num_data)
print(f'New model path: {new_model}')

# Load parameters
@dataclass
class ScriptArguments:
    use_8_bit: Optional[bool] = field(default=False, metadata={"help": "use 8 bit precision"})
    use_4_bit: Optional[bool] = field(default=False, metadata={"help": "use 4 bit precision"})
    bnb_4bit_quant_type: Optional[str] = field(default="nf4", metadata={"help": "precise the quantization type (fp4 or nf4)"})
    use_bnb_nested_quant: Optional[bool] = field(default=False, metadata={"help": "use nested quantization"})
    use_multi_gpu: Optional[bool] = field(default=True, metadata={"help": "use multi GPU"})
    use_adapters: Optional[bool] = field(default=True, metadata={"help": "use adapters"})
    batch_size: Optional[int] = field(default=8, metadata={"help": "input batch size"})
    max_seq_length: Optional[int] = field(default=400, metadata={"help": "max sequence length"})
    optimizer_name: Optional[str] = field(default="adamw_hf", metadata={"help": "Optimizer name"})

parser = HfArgumentParser(ScriptArguments)
script_args = parser.parse_args_into_dataclasses()[0]

# Device map
device_map = "auto" if script_args.use_multi_gpu else "cpu"

# Check precision settings
if script_args.use_8_bit and script_args.use_4_bit:
    raise ValueError("You can't use 8 bit and 4 bit precision at the same time")

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type=script_args.bnb_4bit_quant_type,
    bnb_4bit_use_double_quant=script_args.use_bnb_nested_quant,
) if script_args.use_4_bit else None

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B" if model_name == 'Llama3' 
    else "meta-llama/Llama-2-7b-chat-hf" if model_name == 'Llama2'
    else "mistralai/Mistral-7B-Instruct-v0.2" if model_name == 'Mistral'
    else "deepseek-ai/deepseek-llm-7b-base",
    cache_dir=cache_dir,
    quantization_config=bnb_config,
    device_map={"": 0}
)

model.config.use_cache = False
model.config.pretraining_tp = 1
model = prepare_model_for_kbit_training(model)

tokenizer = AutoTokenizer.from_pretrained(
    "meta-llama/Meta-Llama-3-8B" if model_name == 'Llama3' 
    else "meta-llama/Llama-2-7b-chat-hf" if model_name == 'Llama2'
    else "mistralai/Mistral-7B-Instruct-v0.2" if model_name == 'Mistral'
    else "deepseek-ai/deepseek-llm-7b-base",
    use_fast=False
)

tokenizer.add_special_tokens({'pad_token': '[PAD]'})
tokenizer.pad_token = tokenizer.eos_token

# LoRA Config
peft_config = LoraConfig(
        lora_alpha=32, # Alpha value for LoRA, the higher the value, the more aggressive the sparsity
        lora_dropout=0.05,
        r=16, # Rank of the LoRA decomposition
        target_modules= ['q_proj','k_proj','v_proj','o_proj','gate_proj','down_proj','up_proj','lm_head'],
        bias="none",
        task_type="CAUSAL_LM",
    )

# Create adapter directory
os.makedirs(new_model, exist_ok=True)

# Store loss for visualization
class LoggingCallback(TrainerCallback):
    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs:
            output_log_file = os.path.join(args.output_dir, "train_results.json")
            with open(output_log_file, "a") as writer:
                writer.write(json.dumps(logs) + "\n")

# Training arguments
training_arguments = TrainingArguments(
    num_train_epochs=num_epochs,
    evaluation_strategy="steps",
    save_steps=-1,
    save_total_limit=1,
    logging_steps=500,
    eval_steps=500,
    learning_rate=learning_rate_,
    weight_decay=0.001,
    per_device_train_batch_size=script_args.batch_size,
    max_steps=-1,
    gradient_accumulation_steps=4,
    per_device_eval_batch_size=script_args.batch_size,
    output_dir=new_model,
    max_grad_norm=0.3,
    warmup_ratio=0.03,
    lr_scheduler_type="constant",
    optim=script_args.optimizer_name,
    fp16=True,
    logging_strategy="steps",
    log_level='info'
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    dataset_text_field="text",
    peft_config=peft_config,
    max_seq_length=script_args.max_seq_length,
    args=training_arguments,
    callbacks=[LoggingCallback()]
)

trainer.train()
trainer.model.save_pretrained(new_model)
trainer.tokenizer.save_pretrained(new_model)
print('Done in ', time.time() - start_time)

# Save plots
epochs, train_losses, eval_losses = [], [], []

# Load evaluation results
eval_results_file = os.path.join(new_model, "train_results.json")
with open(eval_results_file, "r") as f:
    for line in f:
        data = json.loads(line)
        if 'epoch' in data:
            epoch = data['epoch']
            if 'loss' in data:
                train_losses.append(data['loss'])
                epochs.append(epoch)
            if 'eval_loss' in data:
                eval_losses.append(data['eval_loss'])
                if epoch not in epochs:
                    epochs.append(epoch)

# Plotting
plt.figure(figsize=(10, 5))
plt.plot(epochs[:len(train_losses)], train_losses, label='Train Loss', color='blue')
plt.plot(epochs[:len(eval_losses)], eval_losses, label='Eval Loss', color='red')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training and Evaluation Loss', fontsize=10)
plt.legend()
plt.tight_layout()

# Save the plot
plot_path = os.path.join(new_model, 'training_evaluation_loss_plot.png')
plt.savefig(plot_path)
plt.close()

print(f"Plot saved in the current directory as 'training_evaluation_loss_plot.png'.")