Upload run_mix_lora.py with huggingface_hub

run_mix_lora.py

@@ -0,0 +1,601 @@
+
+from run_mix_lora_config import *
+#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 random
+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 wandb
+import socket
+# 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_file', 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_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("./run_mix_lora_config.py", 'r') as config_file:
+    args = config_file.read()
+
+with open(checkpoint_store_dir_path+'/lora_logs.txt', 'a') as log_file:
+    log_file.write(f"{current_time}: running lora\n {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
+)
+
+random.seed(40)
+per_file = 10000
+for traindataset_file in traindataset_files:
+    cnt = 0
+    with open(traindataset_file, 'r') as file:
+        for l in file:
+            sample = json.loads(l)
+            if sample['category'][-1] == '1':
+                cnt += 1
+                traindata.append(sample)
+            if cnt > per_file:
+                break
+
+for traindataset_file in traindataset_files:
+    #cnt = 0
+    with open(traindataset_file, 'r') as file:
+        for l in file:
+            sample = json.loads(l)
+            if sample['category'][-1] == '0':
+                cnt -= 1
+                traindata.append(sample)
+            if cnt <= 0:
+                break
+                
+            
+random.shuffle(traindata)
+print(traindata[0])
+
+            
+valdata = []
+per_file = 100
+for valdataset_file in valdataset_files:
+    cnt = 0
+    with open(valdataset_file, 'r') as file:
+        for l in file:
+            cnt += 1
+            if cnt > per_file:
+                break
+            valdata.append(json.loads(l))
+random.shuffle(valdata)      
+testdata = []
+
+testdataset_file = testdataset_files[0]
+with open(testdataset_file, 'r') as file:
+    for l in file:
+        testdata.append(json.loads(l))
+
+
+
+
+def filter_ds(data):
+    filtered_data = [item for item in data if item['category'] in category_included]
+    print(len(filtered_data), len(data))
+    pos_data = [item for item in filtered_data if item['category'][-1] == '1']
+    random.shuffle(pos_data)
+    max_len = len(pos_data)
+    print("Max pos: ", max_len)
+    neg_data = random.sample([item for item in filtered_data if item['category'][-1] == '0'], max_len)
+    train_size = int(0.999* max_len)
+    new_train = pos_data[:train_size] + neg_data[:train_size]
+    new_val = pos_data[train_size:] + neg_data[train_size:]
+    random.shuffle(new_train)
+    random.shuffle(new_val)
+    return new_train, new_val
+    
+traindata, valdata = filter_ds(traindata)
+    
+print(len(traindata), len(valdata), len(testdata))
+
+
+
+
+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,
+)
+if not (load_path is None):
+    model = AutoModelForCausalLM.from_pretrained(
+            load_path,
+            #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
+        )
+
+# model = AutoModelForCausalLM.from_pretrained(
+    # "./HPC_2_mistral_iffp_20k_2_lora_1200/checkpoint_0_18000/"
+# )
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
+
+
+
+wandb.init(project=checkpoint_store_dir_path[3:]+'_wandb', config={
+    'args' : args,
+    'machine' : machine_info,
+})  
+
+
+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)[:250]
+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'],
+    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()
+
+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})
+            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)
+                
+                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)
+        
+
+
+