Upload run_lora_newchat.py with huggingface_hub

run_lora_newchat.py

@@ -0,0 +1,643 @@
+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'] = 'online'
+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))=}")
+
+np.random.seed(40)
+torch.manual_seed(40)
+torch.cuda.manual_seed_all(40)
+
+
+
+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', '')
+#traindataset_file = args.traindataset_file.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,
+        device_map='auto',
+        #use_cache=True,
+        #cache_dir = "../aib222688.scratch/HF/",
+        attn_implementation="flash_attention_2",
+        torch_dtype=torch.bfloat16,
+        
+        #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"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+#tokenizer = AutoTokenizer.from_pretrained(load_path)
+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 = "<|im_start|>system\n"+ "You are a helpful assistant.\n" + "<|im_end|>\n"+ "<|im_start|>user\n"
+    #eos_token = "<|im_end|>\n"+"<|im_start|>assistant\n"
+    bos_token = """<|im_start|>system
+You are a helpful assistant.<|im_end|>
+<|im_start|>user"""
+    eos_token = """<|im_end|>
+<|im_start|>assistant""" 
+    
+    if pair['label'] == 1: # 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
+    #response += "<|im_end|>"
+    
+    messages = [ {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": full_prompt}, ]
+
+    return full_prompt, response #messages, response
+#print(create_prompt(instruct_tune_dataset["train"][1]))
+
+
+if "Qwen" not in model_name:
+    print("DSC!!!")
+    def create_prompt(pair):
+        bos_token = "<｜begin▁of▁sentence｜>You are a helpful assistant.### Instruction:\n"
+        eos_token = "\n### Response:\n"
+        
+        """
+        if pair_label == 'true' or pair_label == True or pair_label==1:
+        #if category[2] == '1' and category[-1] == '1':
+        #if  prog0["probid"] ==  prog1["probid"]:
+            correct_answer = True
+        else:
+            correct_answer = False
+        if 'ori' in out_file and category[2] == '1' :
+            correct_answer = True
+        elif 'ori' in out_file and category[2] == '0':
+            correct_answer = False
+        """
+        
+        
+        if pair['prog1']['probid'] == pair['prog2']['probid']: # if pair['label'] == 1: # 
+            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
+        response += "\n"
+        
+       
+        return full_prompt, response 
+
+  
+traindata1 = list(traindata) #[0:]
+valdata1 = list( valdata)
+testdata1 = list(testdata) #[0:]  
+
+traindata = []
+pos_cnt = 0
+for tdata in traindata1:
+    if tdata['prog1']['probid'] == tdata['prog2']['probid']: # if pair['label'] == 1: # 
+        pos_cnt += 1
+    inp, trg = create_prompt(tdata)
+    traindata.append({
+    'inputs' : inp,
+    'targets' : trg
+    })
+ 
+print("traindata[0] ", traindata[0], pos_cnt) 
+#exit(0)
+    
+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']
+    # # tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
+    # model_input_ids = [tokenizer.apply_chat_template(convo, tokenize=True, add_generation_prompt=True) for convo in inputs]
+    # #model_inputs = [{'input_ids': lst} for lst in model_input_ids]
+    # model_inputs = {'input_ids': model_input_ids, 'attention_mask': [[1] * len(lst) for lst in model_input_ids]}
+    model_inputs = tokenizer(inputs, add_special_tokens=False)
+    #model_inputs = tokenizer(inputs)
+    
+    
+    #print("Input example:\n{}".format(inputs[0]))
+    #print("Output example:\n{}".format(targets[0]))
+    #print("model ip, tokeniszer of file ", "model_inputs", '\n', tokenizer([file_content, file_content]))
+    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}_tb is best so far\n")
+                    
+                
+                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)
+        
+
+
+