Upload run_lora_mistral.py with huggingface_hub

run_lora_mistral.py

@@ -0,0 +1,562 @@
+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")