src/cms_main.py

import yaml
import draccus

from typing import List, Tuple
from dataclasses import field, dataclass, asdict

from .config import MainConfig, convert_to_trainer_args

import random
import numpy as np
import torch
import transformers

import wandb
from datasets import load_dataset
import os
import json
from datetime import datetime

import torch
import torch.optim as optim
from typing import Sequence, Literal, Dict
from torch.nn.utils.rnn import pad_sequence
torch._dynamo.config.cache_size_limit = 128

from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    Trainer,
    set_seed,
    get_linear_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    EarlyStoppingCallback,
    
)

import copy 

from smpeft.sama import SamaConfig #RotationTuner
from smpeft import get_peft_model, PeftModel

from .utils import trainable_parameters_to_file, set_seed_all
IGNORE_INDEX=-100

# Define a standard prompt template for instruction tuning
PROMPT_TEMPLATE = (
    "Below is an instruction that describes a task, paired with an input that provides further context. "
    "Write a response that appropriately completes the request.\n\n"
    "### Instruction:\n{instruction}\n\n{input_section}"
    "### Response:\n"
)
def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
    """
    Tokenize a list of strings.
    Returns python lists (not tensors) for better compatibility with HF dataset.map().
    """
    tokenized_list = [
        tokenizer(
            text,
            return_tensors=None, # Return python lists
            padding=False,       # Do not pad here, pad in DataCollator
            max_length=tokenizer.model_max_length,
            truncation=True,
        )
        for text in strings
    ]
    
    input_ids = [tokenized['input_ids'] for tokenized in tokenized_list]
    
    # Calculate length of valid tokens
    input_ids_lens = [len(x) for x in input_ids]
    
    return dict(
        input_ids=input_ids,
        labels=input_ids, # Placeholder, will be modified later
        input_ids_lens=input_ids_lens,
    )

def preprocess(
    sources: Sequence[str],
    targets: Sequence[str],
    tokenizer: transformers.PreTrainedTokenizer,
) -> Dict:
    """
    Preprocess the data by tokenizing and masking the source (instruction).
    Logic:
    1. Concatenate Source + Target.
    2. Tokenize Source separatey to find the boundary.
    3. Mask the Source part in the labels using IGNORE_INDEX.
    """
    # 1. Concatenate source (instruction) and target (answer)
    examples = [s + t for s, t in zip(sources, targets)]
    
    # 2. Tokenize the full examples (Prompt + Answer)
    examples_tokenized = _tokenize_fn(examples, tokenizer)
    
    # 3. Tokenize just the sources (Prompt only) to find masking boundary
    sources_tokenized = _tokenize_fn(sources, tokenizer)

    input_ids = examples_tokenized["input_ids"]
    labels = copy.deepcopy(input_ids)


    # # 3. Masking: Set the labels corresponding to the source text to IGNORE_INDEX
    # for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
    #     # Safety check: ensure we don't mask everything if truncation happened badly
    #     if source_len < len(label):
    #         label[:source_len] = [IGNORE_INDEX] * source_len
    #     else:
    #         # If source is longer than max_seq_length (truncated), we might mask everything.
    #         # In practice, you might want to log a warning here.
    #         label[:] = [IGNORE_INDEX] * len(label)
    #         raise UserWarning(f"Truncated prompt: source_len = {source_len}, label_len = {len(label)}")
    
    # 4. Masking Loop
    for i, source_len in enumerate(sources_tokenized["input_ids_lens"]):
        # The full sequence length
        full_len = len(input_ids[i])
        
        # Determine how many tokens to mask. 
        # Ideally, we mask strictly the source_len.
        mask_len = source_len

        # Safety: If truncation cut off the answer entirely, mask everything
        if mask_len >= full_len:
            mask_len = full_len
            
        # Apply IGNORE_INDEX to the prompt part
        labels[i][:mask_len] = [IGNORE_INDEX] * mask_len
        
    return dict(input_ids=input_ids, labels=labels)

def train_tokenize_function(examples, tokenizer):
    """
    Adaptation for COMMON SENSE dataset structure.
    JSON fields: 'instruction', 'input', 'output', 'answer'
    """
    sources = []
    targets = []
    
    # Iterate through the batch
    for instruction, inp, output in zip(examples['instruction'], examples['input'], examples['output']):
        
        # 1. Format the Input (Source)
        # Check if optional context 'input' exists
        if inp and str(inp).strip():
            input_section = f"### Input:\n{inp}\n\n"
        else:
            input_section = ""
            
        # Format the source string
        source_text = PROMPT_TEMPLATE.format_map({
            "instruction": instruction,
            "input_section": input_section
        })
        sources.append(source_text)

        # 2. Format the Output (Target)
        # We use 'output' field as it usually contains the full sentence explanation
        # Ensure we add the EOS token so the model learns to stop
        target_text = f"{output}{tokenizer.eos_token}"
        targets.append(target_text)

    # 3. Run the preprocessing logic with masking
    data_dict = preprocess(sources, targets, tokenizer)
    return data_dict

    
####

@dataclass
class DataCollatorForSupervisedDataset():
    tokenizer: transformers.PreTrainedTokenizer
    max_length: int = field(default=512)
    mode: str = field(default="fixed")  # "dynamic" or "fixed"

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        # Extract inputs and labels
        # Assuming instances is a list of dicts like {'input_ids': [...], 'labels': [...]}
        input_ids_list = [torch.tensor(x["input_ids"], dtype=torch.long) for x in instances]
        labels_list = [torch.tensor(x["labels"], dtype=torch.long) for x in instances]

        # 1. Determine padding logic
        if self.mode == "dynamic":
            # Dynamic padding: pad to the longest sequence in the batch
            # But cap it at self.max_length to prevent OOM
            batch_max_len = max([len(x) for x in input_ids_list])
            target_len = min(batch_max_len, self.max_length)
        else:
            # Fixed padding: always pad to max_length
            target_len = self.max_length

        # 2. Helper to pad and truncate
        def pad_and_truncate(tensors, padding_value):
            # First, pad everything using PyTorch's optimized utility (batch_first=True)
            padded = pad_sequence(tensors, batch_first=True, padding_value=padding_value)
            
            # Handle truncation/extending to exact target_len
            curr_len = padded.shape[1]
            if curr_len > target_len:
                # Truncate if too long (rare if filtered beforehand)
                return padded[:, :target_len]
            elif curr_len < target_len:
                # Pad more if shorter than target_len (happens in fixed mode)
                diff = target_len - curr_len
                padding = torch.full((padded.shape[0], diff), padding_value, dtype=padded.dtype)
                return torch.cat([padded, padding], dim=1)
            else:
                return padded

        # 3. Apply padding
        # Critical: tokenizer.pad_token_id must NOT be None here
        if self.tokenizer.pad_token_id is None:
            raise ValueError("Tokenizer.pad_token_id is None. Please set it to eos_token_id or unk_token_id.")
            
        input_ids = pad_and_truncate(input_ids_list, self.tokenizer.pad_token_id)
        labels = pad_and_truncate(labels_list, IGNORE_INDEX)

        # 4. Create Attention Mask explicitly
        # .ne() creates Bools, .long() casts to 0s and 1s for compatibility
        attention_mask = input_ids.ne(self.tokenizer.pad_token_id).long()

        return {
            "input_ids": input_ids,
            "labels": labels,
            "attention_mask": attention_mask
        }
        
        
@draccus.wrap()
def main(mainCfg: MainConfig):
    print('='*120)
    # print(mainCfg)
    # print(draccus.dump(mainCfg, default_flow_style=False))
    set_seed_all(mainCfg.seed)
    training_args = convert_to_trainer_args(mainCfg)
    # training_args.project = f'Rotation-Llama2-{mainCfg.data.dataset_name}'

    # print(training_args.to_json_string())

    task_name = mainCfg.data.dataset_name # CMS
    # wandb 
    ENTITY = "nvan-13-korea-university" 
    PROJECT = os.environ.get("WANDB_PROJECT")
    api = wandb.Api()
    try:
        runs_list = api.runs(f"{ENTITY}/{PROJECT}")
        next_run_num = len(runs_list) + 1
    except Exception as e:
        next_run_num = 1

    # training_args.run_name = f'[{next_run_num}-{task_name}]mlr={training_args.learning_rate:.1e},'\
    #                         f'b={mainCfg.trainer_args.per_device_train_batch_size},'\
    #                         f'n={mainCfg.sama_adapter.num_unique_blocks_L, mainCfg.sama_adapter.num_unique_blocks_R},'\
    #                         f'cL={mainCfg.sama_adapter.col_L},'\
    #                         f'rR={mainCfg.sama_adapter.row_R}, s={mainCfg.sama_adapter.scaling}', \
    #                         f'init={mainCfg.run_text},dr{mainCfg.sama_adapter.drop_out}'
    
    model = AutoModelForCausalLM.from_pretrained(mainCfg.model.model_name,
                                                 device_map="auto", low_cpu_mem_usage=True,
                                                 dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
                                                 # attn_implementation="sdpa",
                                                )
    total_params_now = sum(p.numel() for p in model.parameters())
    print(f'#params of the pretrained model, {total_params_now:,}')



    if mainCfg.model.adapter_path is not None:
        print('___ Loading from:  ', mainCfg.model.adapter_path)
        model = PeftModel.from_pretrained(model, mainCfg.model.adapter_path, is_trainable = True)
    elif mainCfg.sama_adapter.col_L is not None:
        sama_adapter_config = asdict(mainCfg.sama_adapter)
        # rotation_adapter_config[peft_type]
        
        for adapter_name in mainCfg.data.adapter_names:
            print("Init from Sama Config:", json.dumps(sama_adapter_config, indent=4, sort_keys=True))
            sama_config = SamaConfig(**sama_adapter_config)
            model = get_peft_model(model, sama_config, adapter_name=adapter_name)
            # model.set_adapter(adapter_name)
            
        # from peft import LoraConfig, get_peft_model
        # peft_config = LoraConfig(
        #     r=4,
        #     lora_alpha= 11.313,
        #     lora_dropout=0.05,
        #     use_rslora = False,
        #     # init_lora_weights = 'pissa',
        #     bias="none",
        #     task_type="CAUSAL_LM",
        #     target_modules=["q_proj", "v_proj",]
        #     )
        # model = get_peft_model(model, peft_config)
    else:
        print("Full Parameter Fine-Tuning")

    model.print_trainable_parameters()
    # for name, param in model.named_parameters():
    #     if param.requires_grad == True:
    #         print(f'name {name}, shape {param.shape}')

    # Downloading and loading a dataset from the hub.
    raw_datasets = load_dataset(
        "drop",
        # use_auth_token=True if model_args.use_auth_token else None,
    )

    sama_trainable_layers = filter(
            lambda p: p.requires_grad, model.parameters()
        )    

    tokenizer = AutoTokenizer.from_pretrained(
        mainCfg.model.model_name,
        model_max_length=mainCfg.model.model_max_seq_length,
        padding_side="right",
        use_fast=True,
    )

    if tokenizer.pad_token is None:
        if tokenizer.unk_token_id is not None:
            tokenizer.pad_token_id = tokenizer.unk_token_id
            tokenizer.pad_token = tokenizer.unk_token
            print("Set PAD token to UNK token.")
        elif tokenizer.eos_token_id is not None:
            tokenizer.pad_token_id = tokenizer.eos_token_id
            tokenizer.pad_token = tokenizer.eos_token
            print("Set PAD token to EOS token.")

        if model is not None:
            model.config.pad_token_id = tokenizer.pad_token_id
            if model.config.pad_token_id != tokenizer.pad_token_id:
                raise ValueError("Failed to sync pad_token_id between tokenizer and model config")

    if mainCfg.data.path.endswith(".json"):  # todo: support jsonl
        raw_data = load_dataset("json", data_files=mainCfg.data.path, split='train')
    else:
        raw_data = load_dataset(mainCfg.data.path)
        
    if mainCfg.data.split_ratio > 0:
        # raw_data['train'] json file or above
        split_data = raw_data.train_test_split(test_size=mainCfg.data.split_ratio, seed=42)
        
    train_dataset_raw = split_data['train']
    dev_dataset_raw = split_data['test']


    ### Process datasets
    train_dataset = train_dataset_raw.map(
        train_tokenize_function,
        batched=True,
        batch_size=20000,
        num_proc=32, # Adjust based on your CPU
        remove_columns=train_dataset_raw.column_names,
        load_from_cache_file=True, # Set False for debugging new logic
        desc="Running tokenizer on train dataset",
        fn_kwargs={"tokenizer": tokenizer}
    )
    dev_dataset = dev_dataset_raw.map(
        train_tokenize_function,
        batched=True,
        batch_size=20000,
        num_proc=32,
        load_from_cache_file=True,
        remove_columns=dev_dataset_raw.column_names,
        fn_kwargs={"tokenizer": tokenizer}
    )

    print('- Train dataset size: ', len(train_dataset))
    print('- Dev dataset size: ', len(dev_dataset))
    
    # debug_masking_visualizer(dev_dataset, tokenizer)
    # exit()
    
    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer, max_length=mainCfg.model.model_max_seq_length, 
                                                    #mode=mainCfg.model.data_collator_mode,
                                                    )
    data_module = dict(train_dataset=train_dataset, data_collator=data_collator, eval_dataset=dev_dataset)

    optimizer = optim.AdamW(
        sama_trainable_layers, 
        lr=mainCfg.trainer_args.learning_rate, #
        eps=1e-8,
        # betas=(0.9, 0.9995)
    )

    num_devices = training_args.n_gpu if training_args.n_gpu > 0 else 1
    per_device_train_batch_size = training_args.per_device_train_batch_size
    gradient_accumulation_steps = training_args.gradient_accumulation_steps

    # Effective batch size used for updates
    total_train_batch_size = per_device_train_batch_size * num_devices * gradient_accumulation_steps
    
    # Calculate steps
    num_update_steps_per_epoch = len(train_dataset) // total_train_batch_size
    max_steps = int(training_args.num_train_epochs * num_update_steps_per_epoch)

    print(f"___ Estimated Total Training Steps: {max_steps}")
    
    if training_args.lr_scheduler_type == "cosine": # Change to "linear" if preferred
        lr_scheduler = get_cosine_schedule_with_warmup(
            optimizer=optimizer,
            num_warmup_steps=training_args.warmup_steps,
            num_training_steps=max_steps,
        )
    else:
        # Default to Linear Decay
        lr_scheduler = get_linear_schedule_with_warmup(
            optimizer=optimizer,
            num_warmup_steps=training_args.warmup_steps,
            num_training_steps=max_steps,
        )
    
    start_time = datetime.now()
    date_str = start_time.strftime("%y%m%dd%Hh%Mm%S")[1:]
    output_dir = f'{training_args.output_dir}/{task_name}/'\
                f't{date_str},ep={training_args.num_train_epochs},' \
                f'mlr{training_args.learning_rate:.1e},'\
                f'b{mainCfg.trainer_args.per_device_train_batch_size},{mainCfg.trainer_args.gradient_accumulation_steps},'\
                f'nb{mainCfg.sama_adapter.num_unique_blocks_L},{mainCfg.sama_adapter.num_unique_blocks_R},'\
                f'cL{mainCfg.sama_adapter.col_L},'\
                f'rR{mainCfg.sama_adapter.row_R},s{mainCfg.sama_adapter.scaling},'\
                f'init{mainCfg.run_text},dr{mainCfg.sama_adapter.drop_out},'\
                f'size{len(train_dataset)},{len(mainCfg.sama_adapter.target_modules)}'

    training_args.output_dir=output_dir
    print(f'Current output_dir: {output_dir}')
    # Save infor to a file
    trainable_parameters_to_file(model, output_dir)

    training_args.run_name = f'[{next_run_num}]mlr{training_args.learning_rate:.1e},'\
                            f'b{mainCfg.trainer_args.per_device_train_batch_size},{mainCfg.trainer_args.gradient_accumulation_steps},'\
                            f'nb{mainCfg.sama_adapter.num_unique_blocks_L},{mainCfg.sama_adapter.num_unique_blocks_R},'\
                            f'cL{mainCfg.sama_adapter.col_L},'\
                            f'rR={mainCfg.sama_adapter.row_R},s{mainCfg.sama_adapter.scaling},'\
                            f'init{mainCfg.run_text},dr{mainCfg.sama_adapter.drop_out}' \
                            f't{date_str},ep{training_args.num_train_epochs},{task_name},' \
                            f'size{len(train_dataset)},{len(mainCfg.sama_adapter.target_modules)}'
    print('Run nume: ', training_args.run_name)
    print(f'data: {task_name}, train: {len(train_dataset)}, valid: {len(dev_dataset)}')
    
    # from .utils import ExperimentMonitorCallback
    # monitor = ExperimentMonitorCallback(
    #         log_file_path="./training_metrics_bs8.json",
    #         run_name="Experiment_BatchSize_8",
    #         log_interval=10  # Will calculate average over every 100 steps
    #     )

    # early_stopping_callback = EarlyStoppingCallback(
    #     early_stopping_patience=3,       # Number of evaluations to wait without improvement
    #     early_stopping_threshold=0.0     # Minimum change to qualify as an improvement
    # )
    trainer = Trainer(
        model=model,
        args=training_args,
        # compute_metrics=compute_metrics,
        processing_class=tokenizer,
        optimizers=(optimizer, lr_scheduler),
        **data_module,
        # callbacks=[monitor],
        # callbacks=[early_stopping_callback],
        # callbacks=[EarlyStoppingCallback(early_stopping_patience=19)]
    )
    model.config.use_cache = False
    # trainer.train(resume_from_checkpoint='')
    trainer.train()
    end_time = datetime.now()
    print('end time: ', end_time.strftime("%Y-%m-%d %H:%M:%S"), '| duration: ', end_time - start_time)

    tokenizer.save_pretrained(os.path.join(training_args.output_dir, 'ft'))
    trainer.save_state()
    model.peft_config.save_pretrained(os.path.join(training_args.output_dir, 'ft'))
    model.save_pretrained(os.path.join(training_args.output_dir, 'ft2'))
    
    is_valid_test = False
    is_infernce = False
    if is_infernce:
        model = model.merge_and_unload()
        model.eval()
        
        if is_valid_test:
            eval_trainer = Trainer(
                model=model,
                args=training_args,
                # compute_metrics=compute_metrics,
                processing_class=tokenizer,
                optimizers=(optimizer, lr_scheduler),
                **data_module
            )
            metrics = eval_trainer.evaluate()
            print(">>> Evaluation Metrics (Loss):", metrics)
    
            



if __name__ == "__main__":
    main()