import os
import torch
import pandas as pd
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import LoraConfig, PeftModel
from trl import SFTTrainer

def custom_format(row):
    custom_dict = {
        'Text': row['text'],
        'Question': row['question'],
    }

    options_columns = ['key_1', 'key_2', 'key_3']

    for column in options_columns:
        right_key = f"{column}(right)"
        wrong_key = f"{column}(wrong)"
                
        right = f"({row[column]}(right))"
        wrong = f"({row[column]}(wrong))"
                
        custom_dict.update({
            right: row[right_key],
            wrong: row[wrong_key],
        })

    return custom_dict

if __name__ == "__main__":
    # Downloaded model path from local 
    path = "......."
    base_model = "path"

    # Load the dataset
    dataset = load_dataset('csv', data_files='/home/ubuntu/item_multiple_group/data/main_data_.csv')
    train_df = pd.DataFrame(dataset['train']) 
    train_df = train_df.astype(str)
    formatted_dataset = train_df.apply(custom_format, axis=1) # Apply custom formatting to each row
    
    

    compute_dtype = getattr(torch, "float16")

    quant_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=compute_dtype,
        bnb_4bit_use_double_quant=False,
    )


    # Load base model
    model = AutoModelForCausalLM.from_pretrained(
        base_model,
        quantization_config=quant_config,
        device_map={"": 0}
    )
    model.config.use_cache = False
    model.config.pretraining_tp = 1

    # Load LLaMA tokenizer
    tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    # Load LoRA config
    peft_args = LoraConfig(
        lora_alpha=16,
        lora_dropout=0.1,
        r=64,
        bias="none",
        task_type="CAUSAL_LM",
    )

    # Set training parameters
    training_params = TrainingArguments(
        output_dir="./results",
        num_train_epochs=1,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=1,
        optim="paged_adamw_32bit",
        save_steps=25,
        logging_steps=25,
        learning_rate=2e-4,
        weight_decay=0.001,
        fp16=False,
        bf16=False,
        max_grad_norm=0.3,
        max_steps=-1,
        warmup_ratio=0.03,
        group_by_length=True,
        lr_scheduler_type="constant",
        report_to="tensorboard"
    )

    # Set supervised fine-tuning params
    trainer = SFTTrainer(
        model=model,
        train_dataset=dataset,
        peft_config=peft_args,
        dataset_text_field="text",
        max_seq_length=None,
        tokenizer=tokenizer,
        args=training_params,
        packing=False,
    )
   
    trainer.train()

    # Save model
    trainer.model.save_pretrained(new_model)