Upload run_lora_santa.py with huggingface_hub

run_lora_santa.py

@@ -0,0 +1,462 @@
+#model_name = "mistralai/Mistral-7B-Instruct-v0.3"
+#model_name = "bigcode/starcoder2-7b"
+model_name = "meta-llama/CodeLlama-7b-hf" # "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_codellama_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_900.json"
+traindata = []
+initial_lr = 1e-5
+checkpoint_store_dir_path = "./HPC_2_codellama_iffp_20k_4_lora_900"
+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=}")
+