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 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.json"
valdataset_file = "./dataset_lfs/allpairs_data_val.json"
testdataset_file = "./llm_for_code/datasets/codecontests/verified_iffp_1200.json"
traindata = []
initial_lr = 5e-6
checkpoint_store_dir_path = "./HPC_2_codellama_iffp_20k_4_lora_1200"
os.makedirs(checkpoint_store_dir_path, exist_ok=True)
num_epochs = 5
batch_size_train = 1
max_length = 2000
ckpnt_NUM = 2000 
#ckpnt_NUM2 = 5000 
SAVEALL = False #True





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

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

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    #device_map='auto',
    #quantization_config=nf4_config,
    use_cache=True,
    cache_dir = "/scratch/scai/mtech/aib222688/HF",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.float16,
    trust_remote_code=True,
)
print(model)
# model = AutoModelForCausalLM.from_pretrained(
    # "./HPC_2_codellama_iffp_20k_3_lora_900/checkpoint_0_99/",
    # use_cache=True,
    # cache_dir = "/scratch/scai/mtech/aib222688/HF",
# )

tokenizer = AutoTokenizer.from_pretrained(model_name)

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





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_path = './prompt_file.txt'
file_content = ""
with open(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
    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 = []

for tdata in traindata1:
    inp, trg = create_prompt(tdata)
    traindata.append({
    'inputs' : inp,
    'targets' : trg
    })
    
    
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', 'kv_a_proj_with_mqa', 'o_proj'], #['q_proj', 'kv_a_proj_with_mqa', 'kv_b_proj', 'o_proj'],
    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)


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 = torch.nn.DataParallel(model)
#model = model.cuda()

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()
            if step % 100 == 0:
                print(loss)
            loss.backward()
            #print("Loss {}".format(loss.item()))
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            # if step % ckpnt_NUM2 == 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)
                eval_epoch_loss = eval_loss / len(val_dataloader)
                if ((eval_loss/eval_cnt) < best_eval_loss) or SAVEALL:
                    
                    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)
                
                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=}")


    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)
        


        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=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)
                # 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=}")




sys.exit(0)
#model = peft.prepare_model_for_kbit_training(model)
#model = peft.get_peft_model(model, peft_config)

args = TrainingArguments(
  output_dir = out_name,
  num_train_epochs=5,
  #max_steps = 100, 
  per_device_train_batch_size = 4,
  warmup_steps = 0.03,
  logging_steps=10,
  save_strategy="epoch",
  #evaluation_strategy="epoch",
  evaluation_strategy="steps",
  eval_steps=100,
  learning_rate=2e-4,
  bf16=True,
  lr_scheduler_type='constant',
)

max_seq_length = 2048

trainer = SFTTrainer(
  model=model,
  peft_config=peft_config,
  max_seq_length=max_seq_length,
  tokenizer=tokenizer,
  packing=True,
  formatting_func=create_prompt,
  args=args,
  train_dataset=instruct_tune_dataset["train"],
  eval_dataset=instruct_tune_dataset["val"]
)

# predictions, labels, _ = trainer.predict(instruct_tune_dataset["train"][1])
    
    # # Convert predictions and labels to text
# predicted_texts = tokenizer.batch_decode(predictions.argmax(axis=-1), skip_special_tokens=True)
# true_texts = tokenizer.batch_decode(labels, skip_special_tokens=True)
# print(f" Pred {predicted_texts}, True {true_texts}")

import time
start = time.time()
try:
    trainer.train()
except Exception as e:
    print(f"ERRORR\n {e}")
    traceback.print_exc()

print(time.time()- start)


def generate_response(prompt, model):
    encoded_input = tokenizer(prompt,  return_tensors="pt", add_special_tokens=True)
    model_inputs = encoded_input.to('cuda')

    generated_ids = model.generate(**model_inputs, max_new_tokens=1000, do_sample=True, pad_token_id=tokenizer.eos_token_id)

    decoded_output = tokenizer.batch_decode(generated_ids)

    return decoded_output[0].replace(prompt, "")
    


try:
    print(generate_response("### Instruction:\nUse the provided input to create an instruction that could have been used to generate the response with an LLM.### Input:\nThere are more than 12,000 species of grass. The most common is Kentucky Bluegrass, because it grows quickly, easily, and is soft to the touch. Rygrass is shiny and bright green colored. Fescues are dark green and shiny. Bermuda grass is harder but can grow in drier soil.\n\n### Response:", trainer.model))
except Exception as e:
    print(f"ERRORR\n {e}")
    traceback.print_exc()


try:
    trainer.save_model(out_name)
except Exception as e:
    print(f"ERRORR\n {e}")
    traceback.print_exc()

    
try:
    tester = SFTTrainer(
      model=trainer.model,
      peft_config=peft_config,
      max_seq_length=max_seq_length,
      tokenizer=tokenizer,
      packing=True,
      formatting_func=create_prompt,
      args=args,
      train_dataset=instruct_tune_dataset["train"],
      eval_dataset=instruct_tune_dataset["test"]
    )
    print(tester.evaluate())
    
except Exception as e:
    print(f"ERRORR\n {e}")
    traceback.print_exc()

trainer.push_to_hub("FT/mistral-instruct-generation")