Upload try_lora.py with huggingface_hub

try_lora.py

@@ -0,0 +1,299 @@
+import json
+import os
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AdamW, default_data_collator, get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup,set_seed
+from peft import LoraConfig, PeftModel, get_peft_model
+import torch.nn.functional as F
+import json
+from datasets import load_dataset
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+import random
+import argparse
+import sys
+# Ensure the src directory is in the Python path
+current_dir = os.path.dirname(__file__)
+src_dir = os.path.abspath(os.path.join(current_dir, '..', '..'))
+if src_dir not in sys.path:
+    sys.path.append(src_dir)
+#from src.prompts.prompts import instruction_prompts
+
+# import pydevd_pycharm
+# pydevd_pycharm.settrace('9.195.44.45', port=12345, stdoutToServer=True, stderrToServer=True)
+
+def parse_arguments(arg_list: list = None):
+    parser = argparse.ArgumentParser(description='Lora fine-tuning of base model.')
+    parser.add_argument('--seed', type=int, default=42, help='random seed for reproducibility')
+    parser.add_argument('--initial_checkpoint', type=str, default="ibm-granite/granite-20b-code-instruct",
+                        help='Path to the intial checkpoint.')
+    parser.add_argument('--training_data_file_path', type=str, default="../../data/generated_data/version_2_1/train_data_version_2_1_anonymized_new_conversation.jsonl",
+                        help='Path to input training data file.')
+    parser.add_argument('--checkpoint_store_dir_path', type=str, default="../../checkpoints_granite/version_2_1/",
+                        help='path to save the checkpoints.')
+
+    parser.add_argument('--batch_size_train', type=int, default=1, help='Batch size for training.')
+    parser.add_argument('--batch_size_val', type=int, default=1, help='Batch size for Validation.')
+    parser.add_argument('--max_length', type=int, default=4000, help='max length of input.')
+    parser.add_argument('--initial_lr', type=float, default=2e-5, help='Initial learning rate')
+    parser.add_argument('--num_epochs', type=int, default=20, help='number of epochs')
+    parser.add_argument('--lora_r', type=int, default=8, help='LoRA rank')
+    parser.add_argument('--lora_alpha', type=int, default=32, help='LoRA alpha parameter')
+    parser.add_argument('--lora_dropout', type=float, default=0.1, help='LoRA dropout rate')
+    parser.add_argument('--target_modules', nargs='+', default= ['c_attn'],
+                        help='List of target modules')
+    if arg_list is None:
+        args = parser.parse_args()
+    else:
+        args = parser.parse_args(arg_list)
+
+    return args
+
+def main():
+
+    args = parse_arguments()
+    set_seed(args.seed)
+
+
+    initial_checkpoint = args.initial_checkpoint
+    training_data_file_path = args.training_data_file_path
+    checkpoint_store_dir_path = args.checkpoint_store_dir_path
+    initial_lr = args.initial_lr
+
+
+    max_length = args.max_length
+    batch_size_train = args.batch_size_train
+    batch_size_test = args.batch_size_val
+    num_epochs = args.num_epochs
+    lora_r = args.lora_r
+    lora_alpha = args.lora_alpha
+    lora_dropout = args.lora_dropout
+    target_modules = args.target_modules
+
+    print("Intial Model: {}".format(initial_checkpoint))
+    print("Data path: {}".format(training_data_file_path))
+    print("Checkpoint base dir path: {}".format(checkpoint_store_dir_path))
+    print("Initial learning rate: {}".format(initial_lr))
+    print("LoRA rank: {}".format(lora_r))
+    print("LoRA alpha: {}".format(lora_alpha))
+    print("LoRA dropout rate: {}".format(lora_dropout))
+    print("LoRA Target modules: {}".format(target_modules))
+
+    tokenizer = AutoTokenizer.from_pretrained(initial_checkpoint)
+    dataset = load_dataset('json', data_files=training_data_file_path)
+
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    past_context_code = 'past_context_code'
+    new_input = 'new_input'
+    label_column = 'output'
+
+
+    # def get_examples_all_context_granite(examples, only_code=False):
+
+    #     #inputs = [ f"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 only the provided APIs.\n\nAPIs:\n{apis}\n\nPrevious Conversation Context:\n{past_context_NL}\n\nPrevious Code Context:\n{past_context_code}\n\nInput:\n{new_input}\n"
+    #     # instruction = """Question:\nYou are an expert Python programmer.\nYou can update existing python code as per the provided instruction.\nYou first try to understand the existing code if provided and then decide where exactly to insert or update the code to make it fulfil the given instruction.\nA set of python functions are also provided that you can call. You cannot alter the functionality of any function.\nEach python function has a description explaining the task it performs.\nUse only simple python commands. You job is to only return the updated code.Do not include any explanation. Do not create your own function names.\nIf a function does not exist use the function called user_task with a parameter describing the task.\nIf Context is also provided for the instruction, then extract user preferences for functions from the context, and use them for answering the instruction.\n\nFunctions:"""
+    #     instruction = instruction_prompts['zeroshot_prompt']
+    #     inputs = [ instruction + f"\n\n{apis}\n\nPrevious Conversation Context:\n{past_context_NL}\n\nFollowing is the existing code for you to update:\n{past_context_code}\n\nInput:\n{new_input}\n\nAnswer:\n" for apis, past_context_NL, past_context_code, new_input in zip(examples['apis'], examples['past_context_NL'], examples['past_context_code'], examples['new_input'])]
+
+    #     if only_code == False:
+    #         targets = ["\nNew NL response:\n" + new_NL_response +  "\nUpdated Code:\n" + str(new_code) for new_NL_response, new_code in zip(examples['new_NL_response'],examples['new_code'])]
+    #     else:
+    #         targets = ["\nUpdated Code:\n" + str(new_code) for new_code in examples['new_code']]
+
+    #     return inputs, targets
+
+
+    def preprocess_function(examples):
+        batch_size = len(examples[past_context_code])
+        #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)
+        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"]
+        return model_inputs
+
+    processed_datasets = dataset.map(
+        preprocess_function,
+        batched=True,
+        num_proc=1,
+        remove_columns=dataset["train"].column_names,
+        load_from_cache_file=False,
+        desc="Running tokenizer on dataset",
+    )
+    train_dataset = processed_datasets["train"]
+    train_dataloader = DataLoader(
+        train_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[past_context_code])
+        #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(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
+
+
+    test_dataset = dataset["train"].map(
+        test_preprocess_function,
+        batched=True,
+        num_proc=1,
+        remove_columns=dataset["train"].column_names,
+        load_from_cache_file=False,
+        desc="Running tokenizer on dataset",
+    )
+
+    test_dataloader = DataLoader(test_dataset, collate_fn=default_data_collator, batch_size=batch_size_test, pin_memory=True)
+
+
+    model = AutoModelForCausalLM.from_pretrained(
+        initial_checkpoint,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2"
+    )
+
+    print(model)
+
+    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 = get_peft_model(model, peft_config)
+
+    print("Model loaded successfully!")
+
+    device = "cuda"
+
+    #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.eval()
+    # for step, batch_eval in enumerate(tqdm(test_dataloader)):
+    #     if step > 5:
+    #         break
+    #     batch_eval = {k: v.to(device) for k, v in batch_eval.items()}
+    #     outputs = model.generate(**batch_eval, max_new_tokens=1024)
+    #     out = tokenizer.batch_decode(outputs, skip_special_tokens=False)
+    #     for x in out:
+    #         print(x)
+    #     print("#" * 50)
+
+    for epoch in range(num_epochs):
+        model.train()
+        total_loss = 0
+        for step, batch in enumerate(tqdm(train_dataloader)):
+            batch = {k: v.to(device) for k, v in batch.items()}
+            #         print(batch)
+            #print(batch["input_ids"].shape)
+            # if step > 5:
+            #     break
+            outputs = model(**batch)
+            loss = outputs.loss
+            total_loss += loss.detach().float()
+            loss.backward()
+            #print("Loss {}".format(loss.item()))
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+
+        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 step, batch_eval in enumerate(tqdm(test_dataloader)):
+                if step > 5:
+                    break
+                batch_eval = {k: v.to(device) for k, v in batch_eval.items()}
+                outputs = model.generate(**batch_eval, max_new_tokens=1024)
+                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=}")
+
+    args_dict = vars(args)
+    # Path to the JSON file where the configuration will be saved
+    config_file_name = "training_config.json"
+    config_file_path = os.path.join(checkpoint_store_dir_path, config_file_name)
+
+    # Write the configuration dictionary to the JSON file
+    with open(config_file_path, "w") as json_file:
+        json.dump(args_dict, json_file, indent=4)
+
+if __name__ == "_main_":
+    main()