Planning/data/train_eval.py

"""
Description: This file is used to train and evaluate the llama model.

"""
import os
from unsloth import FastLanguageModel
from trl import SFTTrainer
from transformers import TrainingArguments, DataCollatorForSeq2Seq
from unsloth import is_bfloat16_supported
from datasets import load_dataset
from unsloth.chat_templates import train_on_responses_only
from unsloth.chat_templates import get_chat_template
from sklearn.model_selection import train_test_split
from transformers import TextStreamer
import ast
import contractions
import re
from utils import remove_inner_single_quotes

# *****load model and tokenizer*****
max_seq_length = 2048
dtype = None    # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
load_in_4bit = True # Use 4bit quantization to reduce memory usage.

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Llama-3.2-3B-Instruct",
    max_seq_length = max_seq_length,    # specify the maximum length of input the model can accept
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

# add LoRA adapters so we only need to update 1 to 10% of all parameters!
model = FastLanguageModel.get_peft_model(
    model,
    r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",# Reference: https://colab.research.google.com/drive/1Ys44kVvmeZtnICzWz0xgpRnrIOjZAuxp?usp=sharing#scrollTo=2eSvM9zX_2d3
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 16,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
    random_state = 3407,
    use_rslora = False,  # We support rank stabilized LoRA
    loftq_config = None, # And LoftQ
)

tokenizer = get_chat_template(
    tokenizer,
    chat_template = "llama-3.1",
)

def formatting_prompts_func(data):
    convos = data['conversations']
    texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) for convo in convos]
    
    return { "text" : texts, }

def train(train_dataset, val_dataset):
    # format the dataset
    texts_train = train_dataset.map(formatting_prompts_func, batched=True)
    print(texts_train)
    texts_val = val_dataset.map(formatting_prompts_func, batched=True)
    print(texts_val)
    
    # load trainer
    trainer = SFTTrainer(
        model = model,
        tokenizer = tokenizer,
        train_dataset = texts_train,
        eval_dataset= texts_val,
        dataset_text_field = "text",
        max_seq_length = max_seq_length,
        data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),  # A function or mechanism to batch and prepare data during training and evaluation
        dataset_num_proc = 2,   # number of processes to use for data preprocessing
        packing = False, # Can make training 5x faster for short sequences.
        args = TrainingArguments(
            per_device_train_batch_size = 4, 
            gradient_accumulation_steps = 8,
            warmup_steps = 5,
            num_train_epochs = 100, # Set this for 1 full training run. # TODO: increase the value
            max_steps = 1000, # TODO: increase the value
            learning_rate = 2e-4,
            fp16 = not is_bfloat16_supported(),
            bf16 = is_bfloat16_supported(),
            logging_steps = 1,
            optim = "adamw_8bit",
            weight_decay = 0.01,
            lr_scheduler_type = "linear",
            seed = 3407,
            output_dir = "outputs",
            do_eval=True,
            report_to='wandb',
            evaluation_strategy="epoch",    # Specifies that evaluations will happen at the end of each epoch, using texts_val for metrics calculation. If we set the evaluation strategy without passing evaluation dataset, there is an error.
            save_strategy="epoch",      # Save checkpoints based on epoch.
            load_best_model_at_end=True,    # Load the best one from the disk to memory. Otherwise, the model is still the one trained after last epoch.
            metric_for_best_model="loss",       # Evaluation metric on evaluation set. Make sure the metric is an option included in TrainingArguments.
            greater_is_better=False,     # The metric is "greater_is_better".
            save_total_limit=1      # How many checkpoints will be saved.
        ),
    )
    
    # train only on outputs and ignore the loss of user's inputs
    trainer = train_on_responses_only(
        trainer,
        instruction_part = "<|start_header_id|>user<|end_header_id|>\n\n",
        response_part = "<|start_header_id|>assistant<|end_header_id|>\n\n",
    )

    trainer_stats = trainer.train()
    
    # save the model
    checkpoint_dir = f"./checkpoints/{dataset_name}"
    os.makedirs(checkpoint_dir, exist_ok=True)
    checkpoint_path = os.path.join(checkpoint_dir, f"lora_model")
    model.save_pretrained(checkpoint_path) # Local saving 
    tokenizer.save_pretrained(checkpoint_path)
    
    print("Training completed.")
    
    return checkpoint_path

def evaluate(test_dataset, checkpoint_path):
    
    # load model
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name = checkpoint_path,
        max_seq_length = max_seq_length,    # specify the maximum length of input the model can accept
        dtype = dtype,
        load_in_4bit = load_in_4bit,
    )
    
    FastLanguageModel.for_inference(model)
    
    # evaluate
    acc = 0
    for idx, dp in enumerate(test_dataset): # TODO: batch
        message = dp['conversations'][0]
        label = dp['conversations'][1]
        assert label['role'] == "assistant"
        
        # format the input
        inputs = tokenizer.apply_chat_template([message], tokenize = True, add_generation_prompt = True, return_tensors = 'pt').to('cuda')
        print(f"222, {inputs}")
        
        text_streamer = TextStreamer(tokenizer, skip_prompt = True)
        outputs = model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 128,
                        use_cache = True, temperature = 1.5, min_p = 0.1)
        outputs = tokenizer.batch_decode(outputs)
        parts = outputs[0].split("<|start_header_id|>assistant<|end_header_id|>\n\n")
        # print(f"111, {parts}")
        results = parts[1].strip("<|eot_id|>")
        results = contractions.fix(results)
        try:
            results = ast.literal_eval(results)
        except:
            try:
                results = re.sub(r"\['(.*?)'", remove_inner_single_quotes, results) 
                results = ast.literal_eval(results)
            except:
                results = {
                    "Metapath":"",
                    "Restriction":{},
                }

        
        
        pred_metapath = results['Metapath']
        
        ground_metapath = ast.literal_eval(label['content'])['Metapath']
        if pred_metapath == ground_metapath:
            acc += 1
        print(f"Prediction: {pred_metapath}")
        print(f"Ground truth: {ground_metapath}")
        
        
    
    print(f"Accuracy: {acc / len(test_dataset)}")
    

def main(dataset_name, model_name):
    # *****load dataset*****
    data_dir = f"./data/finetune"
    data_path = os.path.join(data_dir, f"{dataset_name}/llama_ft.jsonl")
    # data_path = os.path.join(data_dir, f"{dataset_name}/llama_ft_{model_name}.jsonl")
    dataset = load_dataset("json", data_files=data_path)
    dataset = dataset['train']

    # Add an index column to keep track of original indices
    dataset = dataset.add_column("index", list(range(len(dataset))))

    # Perform train-test split
    train_test = dataset.train_test_split(test_size=0.2, seed=42)
    val_test = train_test['test'].train_test_split(test_size=0.5, seed=42)

    # Extract the datasets
    train_dataset = train_test['train']
    val_dataset = val_test['train']
    test_dataset = val_test['test']

    # Remove the index column if not needed
    train_dataset = train_dataset.remove_columns(['index'])
    val_dataset = val_dataset.remove_columns(['index'])
    test_dataset = test_dataset.remove_columns(['index'])
    
    # *****train and evaluate*****
    checkpoint_path = train(train_dataset, val_dataset)
    
    
    checkpoint_path = f"./checkpoints/{dataset_name}/lora_model"
    # checkpoint_path = f"./checkpoints/{dataset_name}/lora_model_{model_name}"
    evaluate(test_dataset, checkpoint_path)

if __name__ == "__main__":
    dataset_name_list = ['mag', 'amazon', 'prime']
    model_names = ["4o"] 
    for dataset_name in dataset_name_list:
        for model_name in model_names:        
            main(dataset_name, model_name)