import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments, DataCollatorForLanguageModeling
from datasets import load_dataset
from trl import SFTTrainer
from peft import LoraConfig, prepare_model_for_kbit_training

dataset = load_dataset("csv", data_files="nowhere_training_input.csv", delimiter=";", split="train")

if torch.cuda.is_available():
    print("Cuda is available")

base_model_id = "abacaj/phi-2-super"
output_dir = "./results_phi-2-super"

tokenizer = AutoTokenizer.from_pretrained(base_model_id)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    print("pad_token was missing and has been set to eos_token")

# Configuration to load model in 4-bit quantized
bnb_config = BitsAndBytesConfig(load_in_4bit=True,
                                bnb_4bit_quant_type='nf4',
                                #bnb_4bit_compute_dtype='float16',
                                bnb_4bit_compute_dtype=torch.bfloat16,
                                bnb_4bit_use_double_quant=False)

model = AutoModelForCausalLM.from_pretrained(base_model_id, attn_implementation="flash_attention_2", quantization_config=bnb_config, torch_dtype="auto")
print(model)

# Gradient checkpointing to save memory
model.gradient_checkpointing_enable()

# Freeze base model layers and cast layernorm in fp32
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)

peft_config = LoraConfig(
    r=64,
    lora_alpha=64,
    target_modules= ["q_proj","k_proj","v_proj","dense","fc2","fc1"],
    bias="none",
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
)

training_args = TrainingArguments(
    output_dir=output_dir,  # Output directory for checkpoints and predictions
    overwrite_output_dir=True, # Overwrite the content of the output directory
    per_device_train_batch_size=2,  # Batch size for training
    per_device_eval_batch_size=2,  # Batch size for evaluation
    gradient_accumulation_steps=5, # number of steps before optimizing
    gradient_checkpointing=True,   # Enable gradient checkpointing
    gradient_checkpointing_kwargs={"use_reentrant": False},
    warmup_steps=10,  # Number of warmup steps
    #max_steps=1000,  # Total number of training steps
    num_train_epochs=100,  # Number of training epochs
    learning_rate=5e-5,  # Learning rate
    weight_decay=0.01,  # Weight decay
    optim="paged_adamw_8bit", #Keep the optimizer state and quantize it
    bf16=True, #Use mixed precision training
    #For logging and saving
    logging_dir='./logs',
    logging_strategy="epoch",
    logging_steps=10,
    save_strategy="epoch",
    save_steps=10,
    save_total_limit=2,  # Limit the total number of checkpoints
    evaluation_strategy="epoch",
    eval_steps=10,
    load_best_model_at_end=True, # Load the best model at the end of training
    lr_scheduler_type="linear",
)

def formatting_func(data):
    # text = f"[INST] {data['prompt']} [/INST]{data['completion']}{tokenizer.eos_token} "
    chat = [
        { "role": "user", "content": data['prompt'] },
        { "role": "assistant", "content": data['completion'] },
    ]

    text = tokenizer.apply_chat_template(chat, tokenize=False)
    print(text)
    data['text'] = text
    
    return data

dataset = dataset.map(formatting_func)

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    eval_dataset=dataset,
    peft_config=peft_config,
    args=training_args,
    max_seq_length=1024,
    packing=True,
    dataset_text_field="text"
)

model.config.use_cache = False  # silence the warnings. Please re-enable for inference!

start_time = time.time()  # Record the start time
trainer.train()
end_time = time.time()  # Record the end time

training_time = end_time - start_time  # Calculate total training time

trainer.save_model(output_dir)
print(f"Training completed in {training_time} seconds.")