model_name = "mistralai/Mistral-7B-Instruct-v0.3"
#model_name = "bigcode/starcoder2-7b"
#model_name =  "dorkai/codeX-1.0" #"Alibaba-NLP/gte-Qwen1.5-7B-instruct"  #"google/flan-t5-small" #"microsoft/Phi-3-medium-128k-instruct" #"google/gemma-2-9b-it" # "meta-llama/CodeLlama-7b-hf" #"deepseek-ai/DeepSeek-Coder-V2-Instruct" #"deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"
#out_name = "HPC_2_mistral_iffp_20k_5_lora" #"meta-llama/Meta-Llama-3-8B" #"tiiuae/falcon-40b" #"Phind/Phind-CodeLlama-34B-v2" # "deepseek-ai/DeepSeek-Coder-V2-Instruct" # 
from datasets import load_dataset, Dataset
import pandas as pd
import json
import traceback
import peft
import os
from tqdm import tqdm
import sys
import math
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AdamW, default_data_collator, get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup,set_seed
from torch.utils.data import DataLoader
import numpy as np
import os
import argparse
import torch
import datetime
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import LoraConfig, PeftModel
from trl import SFTTrainer
import os
os.environ['WANDB_MODE'] = 'offline'
import wandb
import socket
import random

def set_seed(seed: int = 42):
    random.seed(seed)  # Python’s built-in random module
    np.random.seed(seed)  # NumPy
    torch.manual_seed(seed)  # PyTorch CPU
    torch.cuda.manual_seed(seed)  # PyTorch GPU
    torch.cuda.manual_seed_all(seed)  # If using multi-GPU
    torch.backends.cudnn.deterministic = True  # Ensures deterministic behavior in CuDNN
    torch.backends.cudnn.benchmark = False  # Disables benchmarking to maintain consistency

# Example usage
set_seed(42)
# import json
# filepath = "/kaggle/input/code-sim-try1/mutated_graph_all_lang_eq.json"
# examples = []
# with open(filepath, 'r') as file:
    # for l in file:
        # examples.append(json.loads(l))

# len(examples)
# examples[0]

# instruct_tune_dataset = load_dataset("mosaicml/instruct-v3",cache_dir = "/scratch/scai/mtech/aib222688/HF")
# instruct_tune_dataset = instruct_tune_dataset.filter(lambda x: x["source"] == "dolly_hhrlhf")


# traindataset_file = "./dataset_lfs/allpairs_data_large_900_loop2.json"
# valdataset_file = "./dataset_lfs/allpairs_data_val_900_loop2.json"
# testdataset_file = "./llm_for_code/datasets/codecontests/verified_iffp_900_loop2.json"

# initial_lr = 5e-6
# checkpoint_store_dir_path = "./HPC_2_mistral_iffp_20k_3_lora_900_loop2"

# num_epochs = 5
# batch_size_train = 1
# max_length = 2000
# ckpnt_NUM = 2000 
# SAVEALL = False #True


parser = argparse.ArgumentParser(description='Run lora finetuning..., NOTE: UPDATE PEFT CONFIG if needed')
parser.add_argument('--model_name', default="mistralai/Mistral-7B-Instruct-v0.3",type=str)
parser.add_argument('--traindataset_files',  nargs='+',  type=str, default="./dataset_lfs/allpairs_data_large_900_loop2.json")
parser.add_argument('--valdataset_file', type=str, default="./dataset_lfs/allpairs_data_val_900_loop2.json")
parser.add_argument('--testdataset_file', type=str, default="./llm_for_code/datasets/codecontests/verified_iffp_900_loop2.json")
parser.add_argument('--checkpoint_store_dir_path', type=str, default="./HPC_3_mistral_iffp_20k_3_lora_900_loop2")
parser.add_argument('--initial_lr', type=float, default=5e-6)
parser.add_argument('--num_epochs', type=int, default=5)
parser.add_argument('--batch_size_train', type=int, default=1)
parser.add_argument('--ckpnt_num', type=int, default=2000)
parser.add_argument('--saveall', type=int, default=0)
parser.add_argument('--prompt_file_path', type=str, default='./loop_prompt.txt')
parser.add_argument('--max_length', type=int, default=2000)
parser.add_argument('--max_new_tok', type=int, default=50)

args = parser.parse_args()
print(f"{len(vars(args))=}")





model_name = args.model_name.replace('\r', '')
traindataset_files = args.traindataset_files
for i in range(len(traindataset_files)):
    traindataset_files[i] = traindataset_files[i].replace('\r', '')
valdataset_file = args.valdataset_file.replace('\r', '')
testdataset_file = args.testdataset_file.replace('\r', '')
checkpoint_store_dir_path = args.checkpoint_store_dir_path.replace('\r', '')
initial_lr = args.initial_lr
num_epochs = args.num_epochs
batch_size_train = args.batch_size_train
ckpnt_NUM = args.ckpnt_num
SAVEALL = args.saveall
prompt_file_path = args.prompt_file_path.replace('\r', '')
max_length = args.max_length
max_new_tok = args.max_new_tok



hostname = socket.gethostname()
ip_address = socket.gethostbyname(hostname)
node_name = os.uname().nodename
system_info = os.uname()
machine_info = {
    "hostname": hostname,
    "ip_address": ip_address,
    "node_name": node_name,
    "system_info": {
        "sysname": system_info.sysname,
        "nodename": system_info.nodename,
        "release": system_info.release,
        "version": system_info.version,
        "machine": system_info.machine,
    },
}

os.makedirs(checkpoint_store_dir_path, exist_ok=True)
current_time = datetime.datetime.now()

with open(checkpoint_store_dir_path+'/lora_logs.txt', 'a') as log_file:
    log_file.write(f"{current_time}: running lora\n {vars(args)}\n")
    log_file.write(f"{machine_info}-----\n")

traindata = []
nf4_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_quant_type="nf4",
   bnb_4bit_use_double_quant=True,
   bnb_4bit_compute_dtype=torch.bfloat16
)

#mpath  = './codellama'
# model = AutoModelForCausalLM.from_pretrained(
    # model_name,
    # #device_map='auto',
    # #quantization_config=nf4_config,
    # use_cache=True,
    # #cache_dir = "../aib222688.scratch/HF/",
    # attn_implementation="sdpa", #"flash_attention_2",
    # torch_dtype=torch.float16,
    # #trust_remote_code=True,
# )
model = AutoModelForCausalLM.from_pretrained(
        #mpath,
        model_name,
        #use_cache=True,
        #cache_dir = "../aib222688.scratch/HF/",
        #attn_implementation="flash_attention_2",
        torch_dtype=torch.float16,
        #device_map='auto',
        #quantization_config=nf4_config,
        #use_cache=False
    )

print(f"Shards loaded for {model_name}")
# for name, module in model.named_modules():
    # print(f"{name}: {module}")

# model = AutoModelForCausalLM.from_pretrained(
    # "./HPC_2_mistral_iffp_20k_2_lora_1200/checkpoint_0_18000/"
# )

tokenizer = AutoTokenizer.from_pretrained(model_name)
#tokenizer = AutoTokenizer.from_pretrained(mpath)

tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"


#wandb.login(key='e7fdeef2a423ceed55ae12d2c9f1bc530a9e9331')
wandb.init(project=checkpoint_store_dir_path[3:]+'_wandb', config={
    'args' : vars(args),
    'machine' : machine_info,
})


for traindataset_file in traindataset_files:
    with open(traindataset_file, 'r') as file:
        for l in file:
            traindata.append(json.loads(l))
        
valdata = []

with open(valdataset_file, 'r') as file:
    for l in file:
        valdata.append(json.loads(l))
        
testdata = []


with open(testdataset_file, 'r') as file:
    for l in file:
        testdata.append(json.loads(l))
print(len(traindata), len(valdata), len(testdata))


  


file_content = ""
with open(prompt_file_path, 'r') as file:
    file_content = file.read()


def create_prompt(pair):
    bos_token = "<s>"
    eos_token = "</s>"
    if pair['prog1']['probid'] == pair['prog2']['probid']:
        response = "Yes"
    else:
        response = "No"
        
    full_prompt = ""
    full_prompt += bos_token
    #print(f"{pair['prog1']['scode']=}")
    full_prompt += file_content + pair['prog1']['scode'] + "\nProgram 2:" 
    full_prompt += pair['prog2']['scode']+ "\n### Response:"
    full_prompt += "\n" #+ response
    full_prompt += eos_token

    return full_prompt, response
#print(create_prompt(instruct_tune_dataset["train"][1]))
  
traindata1 = list(traindata) #[0:]
valdata1 = list( valdata)
testdata1 = list(testdata) #[0:]  

traindata = []
pos_cnt = 0
for tdata in traindata1:
    inp, trg = create_prompt(tdata)
    if tdata['prog1']['probid'] == tdata['prog2']['probid']: # if pair['label'] == 1: # 
        pos_cnt += 1
    traindata.append({
    'inputs' : inp,
    'targets' : trg
    })
print("traindata[0] ", traindata[0], pos_cnt)     
    
valdata = []

for tdata in valdata1:
    inp, trg = create_prompt(tdata)
    valdata.append({
    'inputs' : inp,
    'targets' : trg
    })
    
testdata = []

for tdata in testdata1:
    inp, trg = create_prompt(tdata)
    testdata.append({
    'inputs' : inp,
    'targets' : trg
    })

traindataset = Dataset.from_pandas(pd.DataFrame(traindata))
valdataset = Dataset.from_pandas(pd.DataFrame(valdata))
testdataset = Dataset.from_pandas(pd.DataFrame(testdata))   

instruct_tune_dataset = {"train": traindataset,
                        "val" : valdataset,
                        "test" : testdataset}
 
   
   






def preprocess_function(examples):
    batch_size = len(examples['inputs'])
    #inputs = [f"<s>[INST] Question : {x} [/INST] \\n Answer : " for x in examples[past_context_code]]
    #inputs = [f"\n<|user|>\n You are given a set of APIs and previously generated Code as context. The task is given a new requirement from Bob modify or expand the given code using the provided APIs.\n\nAPIs:\n{apis}\n\nContext:\n{past_context}\n\nInput:\n{new_input} \n<|assistant|>\n " for apis, past_context, new_input in zip(examples['apis'], examples['past_context_code'], examples['new_input'])]
    #targets = [str(x) for x in examples[label_column]]
    #inputs, targets = get_examples_all_context(examples)

    #inputs, targets = get_examples_all_context(examples)
    #inputs, targets = get_examples_all_context_granite(examples, only_code=False)
    inputs = []
    targets = []
    # for eg in examples:
        # print(eg)
        # #inp, trg = create_prompt(eg)
        # #inputs.append(inp)
        # #targets.append(trg)
    inputs = examples['inputs']
    targets = examples['targets']
    model_inputs = tokenizer(inputs)

    #print("Input example:\n{}".format(inputs[0]))
    #print("Output example:\n{}".format(targets[0]))

    input_sizes = [len(tokens) for tokens in model_inputs['input_ids']]
    #print("Input sizes {}".format(input_sizes))
    labels = tokenizer(targets, add_special_tokens=False)  # don't add bos token because we concatenate with inputs
    label_sizes = [len(tokens) for tokens in labels['input_ids']]
    #print("Label sizes {}".format(label_sizes))

    for i in range(batch_size):
        sample_input_ids = model_inputs["input_ids"][i]
        label_input_ids = labels["input_ids"][i] + [tokenizer.eos_token_id]
        # print(i, sample_input_ids, label_input_ids)
        model_inputs["input_ids"][i] = sample_input_ids + label_input_ids
        labels["input_ids"][i] = [-100] * len(sample_input_ids) + label_input_ids
        model_inputs["attention_mask"][i] = [1] * len(model_inputs["input_ids"][i])
    # print(model_inputs)
    for i in range(batch_size):
        sample_input_ids = model_inputs["input_ids"][i]
        label_input_ids = labels["input_ids"][i]
        model_inputs["input_ids"][i] = [tokenizer.pad_token_id] * (
            max_length - len(sample_input_ids)
        ) + sample_input_ids
        model_inputs["attention_mask"][i] = [0] * (max_length - len(sample_input_ids)) + model_inputs["attention_mask"][i]
        labels["input_ids"][i] = [-100] * (max_length - len(sample_input_ids)) + label_input_ids
        model_inputs["input_ids"][i] = torch.tensor(model_inputs["input_ids"][i][:max_length])
        model_inputs["attention_mask"][i] = torch.tensor(model_inputs["attention_mask"][i][:max_length])
        labels["input_ids"][i] = torch.tensor(labels["input_ids"][i][:max_length])
    model_inputs["labels"] = labels["input_ids"]
    input_sizes = [len(tokens) for tokens in model_inputs['input_ids']]
    #print("Input sizes {}".format(input_sizes))
    return model_inputs

processed_datasets = traindataset.map(
    preprocess_function,
    batched=True,
    num_proc=1,
    remove_columns=traindataset.column_names,
    load_from_cache_file=False,
    desc="Running tokenizer on dataset",
)
train_dataset = processed_datasets
train_dataloader = DataLoader(
    train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=batch_size_train, pin_memory=True
)


processed_datasets = valdataset.map(
    preprocess_function,
    batched=True,
    num_proc=1,
    remove_columns=valdataset.column_names,
    load_from_cache_file=False,
    desc="Running tokenizer on dataset",
)
val_dataset = processed_datasets
val_dataloader = DataLoader(
    val_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=batch_size_train, pin_memory=True
)




def test_preprocess_function(examples):
    batch_size = len(examples['inputs'])
    #inputs, targets = get_examples_all_context(examples)
    #inputs, targets = get_examples_all_context_granite(examples)
    #inputs = [f"\n<|user|>\n You are given a set of APIs and previously generated Code as context. The task is given a new requirement from Bob modify or expand the given code using the provided APIs.\n\nAPIs:\n{apis}\n\nContext:\n{past_context}\n\nInput:\n{new_input} \n<|assistant|>\n " for apis, past_context, new_input in zip(examples['apis'], examples['past_context_code'], examples['new_input'])]
    model_inputs = tokenizer(examples['inputs'])
    # print(model_inputs)
    for i in range(batch_size):
        sample_input_ids = model_inputs["input_ids"][i]
        model_inputs["input_ids"][i] = [tokenizer.pad_token_id] * (
            max_length - len(sample_input_ids)
        ) + sample_input_ids
        model_inputs["attention_mask"][i] = [0] * (max_length - len(sample_input_ids)) + model_inputs["attention_mask"][i]
        model_inputs["input_ids"][i] = torch.tensor(model_inputs["input_ids"][i][:max_length])
        model_inputs["attention_mask"][i] = torch.tensor(model_inputs["attention_mask"][i][:max_length])
    return model_inputs
    

processed_datasets = testdataset.map(
    preprocess_function,
    batched=True,
    num_proc=1,
    remove_columns=testdataset.column_names,
    load_from_cache_file=False,
    desc="Running tokenizer on dataset",
)
test_dataset = processed_datasets
test_dataloader = DataLoader(
    test_dataset, shuffle=False, collate_fn=default_data_collator, batch_size=batch_size_train, pin_memory=True
)




peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    #target_modules = ['c_attn'],
    target_modules = ['q_proj', 'k_proj', 'v_proj', 'o_proj'],   #qwen
    #target_modules = ['q_proj',  'v_proj'],   #mistral
    r=64,
    bias="none",
    task_type="CAUSAL_LM"
)

# peft_config = LoraConfig(
        # r=lora_r, 
        # lora_alpha=lora_alpha, 
        # lora_dropout=lora_dropout, 
        # target_modules= target_modules,  
        # bias="none", 
        # task_type="CAUSAL_LM"
    # )

model = peft.get_peft_model(model, peft_config)

wandb.watch(model, log='all')
print("Model loaded successfully!")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#optimizer = AdamW(model.parameters(), lr=3e-4)
optimizer = AdamW(model.parameters(), lr=initial_lr)

# Instantiate scheduler
lr_scheduler = get_cosine_schedule_with_warmup(
    optimizer=optimizer,
    num_warmup_steps=0.06 * (len(train_dataloader) * num_epochs),
    num_training_steps=(len(train_dataloader) * num_epochs),
)

model.to(device)
model.to('cuda')

#model = torch.nn.DataParallel(model)
#model = model.cuda()
the_best_eval_loss = 10000
for epoch in range(num_epochs):
    try:
        model.train()
        total_loss = 0
        best_eval_loss = 10000 #np.inf

        for step, batch in enumerate(tqdm(train_dataloader)):
            batch = {k: v.to(device) for k, v in batch.items()}
            #batch = {k: v.cuda() for k, v in batch.items()}
            #         print(batch)
            #print(batch["input_ids"].shape)
            # if step > 5:
            #     break
            #batch.to(device)
            outputs = model(**batch)
            loss = outputs.loss
            total_loss += loss.detach().float()
            wandb.log({'train_loss': loss})
            wandb.log({'lr': lr_scheduler.get_last_lr()[0], 'step': step})
            if step % 100 == 0:
                wandb.log({'train_step_loss': loss})
                print(loss)
            loss.backward()
            #print("Loss {}".format(loss.item()))
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            # if step % ckpnt_NUM == 0:
                # checkpoint_dir = os.path.join(checkpoint_store_dir_path, f"checkpoint_{epoch}_{step}/")
                # os.makedirs(checkpoint_dir, exist_ok=True)
                # model.save_pretrained(checkpoint_dir)
                
            if step % ckpnt_NUM == 0:
                
                model.eval()
                eval_loss = 0
                eval_preds = []
                eval_cnt = 1
                for step1, batch_eval in enumerate(tqdm(val_dataloader)):
                    
                    batch_eval = {k: v.to(device) for k, v in batch_eval.items()}
                    #batch_eval = {k: v.cuda() for k, v in batch_eval.items()}
                    
                    #outputs = model.generate(**batch_eval, max_new_tokens=48)
                    #out = tokenizer.batch_decode(outputs, skip_special_tokens=True)
                    # for x in out:
                        # print(x)
                    # print("#" * 50)
                    with torch.no_grad():
                        outputs = model(**batch_eval)
                    loss = outputs.loss
                    if not math.isnan(loss.detach().float()) : 
                        eval_loss += loss.detach().float()
                        eval_cnt += 1
                    # eval_preds.extend(
                        # tokenizer.batch_decode(torch.argmax(outputs.logits, -1).detach().cpu().numpy(),
                                               # skip_special_tokens=True)
                    # )
                print(eval_loss)
                wandb.log({'eval_loss': eval_loss})
                eval_epoch_loss = eval_loss / len(val_dataloader)
                if ((eval_loss/eval_cnt) < best_eval_loss) or SAVEALL==1:
                    
                    best_eval_loss = eval_loss/eval_cnt
                    print(f"saving...{best_eval_loss} to checkpoint_{epoch}_{step}\n")
                    checkpoint_dir = os.path.join(checkpoint_store_dir_path, f"checkpoint_{epoch}_{step}/")
                    os.makedirs(checkpoint_dir, exist_ok=True)
                    model.save_pretrained(checkpoint_dir)
                    
                if ((eval_loss/eval_cnt) < the_best_eval_loss):
                    the_best_eval_loss = eval_loss/eval_cnt
                    print(f"saving...{best_eval_loss} to checkpoint_{epoch}_{step} is best so far\n")
                    wandb.log({'thebest_loss': epoch, 'thebest_step' : step})
                    
                
                eval_ppl = torch.exp(eval_epoch_loss)
                train_epoch_loss = total_loss #/ len(train_dataloader)
                train_ppl = torch.exp(train_epoch_loss)
                print(f"{epoch=}: {train_ppl=} {train_epoch_loss=} {eval_ppl=} {eval_epoch_loss=} {eval_loss/eval_cnt=} {eval_cnt=}")
                #print(f"{epoch=}: {train_ppl=} {train_epoch_loss=}")
                
        print("Total Loss {}".format(total_loss.item()))
        if ((epoch+1) % 1) == 0:
            checkpoint_dir = os.path.join(checkpoint_store_dir_path, f"checkpoint_{epoch}/")
            os.makedirs(checkpoint_dir, exist_ok=True)
            model.save_pretrained(checkpoint_dir)
            model.eval()
            eval_loss = 0
            eval_preds = []
            for step1, batch_eval in enumerate(tqdm(test_dataloader)):
                if step1 > 5:
                    break
                batch_eval = {k: v.to(device) for k, v in batch_eval.items()}
                #batch_eval = {k: v.cuda() for k, v in batch_eval.items()}
                
                outputs = model.generate(**batch_eval, max_new_tokens=max_new_tok)
                out = tokenizer.batch_decode(outputs, skip_special_tokens=True)
                for x in out:
                    print(x)
                print("#" * 50)
                # with torch.no_grad():
                #     outputs = model(**batch)
                # loss = outputs.loss
                # eval_loss += loss.detach().float()
                # eval_preds.extend(
                #     tokenizer.batch_decode(torch.argmax(outputs.logits, -1).detach().cpu().numpy(),
                #                            skip_special_tokens=True)
                # )

            # eval_epoch_loss = eval_loss / len(test_dataloader)
            # eval_ppl = torch.exp(eval_epoch_loss)
            train_epoch_loss = total_loss / len(train_dataloader)
            train_ppl = torch.exp(train_epoch_loss)
            #print(f"{epoch=}: {train_ppl=} {train_epoch_loss=} {eval_ppl=} {eval_epoch_loss=}")
            print(f"{epoch=}: {train_ppl=} {train_epoch_loss=}")




    except KeyboardInterrupt:
            
        checkpoint_dir = os.path.join(checkpoint_store_dir_path, f"checkpoint_{epoch}_interrupt/")
        os.makedirs(checkpoint_dir, exist_ok=True)
        model.save_pretrained(checkpoint_dir)