Fine-tune LLaMA 2.ipynb

{"cells":[{"cell_type":"markdown","metadata":{"id":"lVYweQOivQgZ"},"source":["## **Intro**\n","### In this exercise, we’ll be fine-tuning a Llama2 model for generating image generation prompts from a short concept .\n","---\n","\\\n","### **Instructions**\n","\n","1. The exercise is defined in high-level terms on purpose, there are many different ways to solve each step and we wish to see how do you approach these kind of problems without any artificial restrictions.\n","\n","1. **Make sure you understand all the code you will use, try to choose the most elegent and optimized solution**.\n","\n","1. **Document as much as possible, feel free to use markdown for better description, and maintain a clean code according to best Python and coding practices**.\n","\n","1. **The exercise should be submitted as a Google Colab notebook with the output**.\n","\n","1. Feel free to add a Notes.docx, where you can share any technical thoughts and considerations you might have had, that might be useful to understand your code.\n","\n","1. Please feel free to contact us if you have any questions.\n","\n","1. Good Luck!!\n","\n","---"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"pPBaI31xvPHn"},"outputs":[],"source":["!pip install -q accelerate peft==0.4.0 bitsandbytes==0.40.0 transformers trl\n","!pip install langchain\n","!pip install langchain_openai\n","!pip install gdown\n","\n","# Download the file using its Google Drive shareable link\n","!gdown --id 1oy19JQ9NBQ1KwcazW_TaMHj-gGbFd5cY -O /content/ds.json\n","\n","\n","\n","\n","\n"]},{"cell_type":"markdown","metadata":{"id":"E7kGW6rrz0uI"},"source":["## Step 1: Generate a dataset for llama2 fine-tuning\n","1. Use Mistral-7b or and other LLM that isn't Llama2\n","1. The dataset should be a pair of *concept* and *description* the an AI Image generator can use to generate an image, in a format that Llama2 can use for fine-tuning.\n","1. For example, concept: \"A person is hiking\", description: \"A hiker moves up a sunlit path towards distant snow-capped mountains, surrounded by a lush forest. His determined stride and the serene natural backdrop form a picture of quiet adventure under a clear blue sky.\"  "]},{"cell_type":"code","execution_count":null,"metadata":{"id":"GCNgCF2az1IH"},"outputs":[],"source":["import ast\n","from langchain.output_parsers import ResponseSchema, StructuredOutputParser\n","from langchain.prompts import ChatPromptTemplate\n","from langchain_openai import ChatOpenAI\n","import json\n","import pandas as pd\n","\n","def get_dataset_from_llm():\n","    \"\"\"\n","    Generate a dataset for LLAMA training using OpenAI's ChatOpenAI API.\n","\n","    Returns:\n","    - result_json: Generated dataset in JSON format.\n","    \"\"\"\n","    template = \"\"\"\n","    You are an expert at generating datasets for llm training.\n","    You are an expert of llama2 text format.\n","    You need to create pairs of texts.\n","    Each pair will be a concept and description.\n","    The pair will describe a image that need to be generated.\n","    For example:\n","    concept: A person is hiking\n","    description: A hiker moves up a sunlit path towards distant snow-capped mountains, surrounded by a lush forest. His determined stride and serene natrual backdrop from a picture of a quite adventure under a clear blue sky.\n","\n","    As you see the description should be very descriptive, with up to 35 words.\n","\n","    Generate for me 10 of these pairs\n","\n","    The result must be a valid json string. string should not start with json prefix.\n","    You should not add any additional information to the result.\n","    \"\"\"\n","    llm = ChatOpenAI(model_name=\"gpt-4-turbo-preview\", temperature=0.0, request_timeout=60, max_retries=2, api_key=\"####\") #api_key=\"####\" is ### for security\n","    prompt = ChatPromptTemplate.from_template(template=template)\n","    messages = prompt.format()\n","    response = llm.invoke(messages)\n","    try:\n","        result = response.content.replace(\"```{\", \"{\").replace(\"```json\", \"\").replace(\"```\", \"\").replace(\"\\n\\t\", \"\").replace(\"\\n\", \"\")\n","        result_json = json.loads(result)\n","        return result_json\n","    except Exception as e:\n","        raise e\n","\n","def get_dataset_from_file():\n","    \"\"\"\n","    Load dataset from a JSON file.\n","\n","    Returns:\n","    - result_json: Loaded dataset in dictionary format.\n","    \"\"\"\n","    f = open(\"ds.json\", \"r\")\n","    result_json = f.read()\n","    return ast.literal_eval(result_json)\n","\n","def get_and_divide_data():\n","    \"\"\"\n","    Generate training and test datasets, and save them as CSV files.\n","    \"\"\"\n","    result_json = get_dataset_from_file()  # Change to get_dataset_from_llm() for online generation\n","    df = pd.DataFrame.from_dict(result_json)\n","    test = df.sample(frac = 0.2)\n","    train = df.drop(test.index)\n","    train.to_csv(\"train.csv\", index=False)\n","    test.to_csv(\"test.csv\", index=False)\n","\n","\n","get_and_divide_data()\n"]},{"cell_type":"markdown","metadata":{"id":"RC545daX1P_e"},"source":["## Step 2: Split generated dataset into Training: 80%,  Test:20%"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"xa3kkuxr1QaJ"},"outputs":[],"source":["# this also happens in get_and_divide_data()"]},{"cell_type":"markdown","metadata":{"id":"FgzJSJpz1TkQ"},"source":["## Step 3: Fine tune llama2 on Train set\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"-y3FsmO11T1C"},"outputs":[],"source":["import os\n","import torch\n","import matplotlib.pyplot as plt\n","from datasets import load_dataset\n","from transformers import (\n","    AutoModelForCausalLM,\n","    AutoTokenizer,\n","    BitsAndBytesConfig,\n","    TrainingArguments,\n",")\n","from peft import LoraConfig\n","from trl import SFTTrainer\n","#from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n","\n","def load_train_data(file_path):\n","    \"\"\"\n","    Load training data from a CSV file.\n","    \"\"\"\n","    return load_dataset(\"csv\", data_files=file_path, split=\"train\")\n","\n","def train_model(num_train_epochs):\n","    \"\"\"\n","    Train the model with the specified number of epochs.\n","    \"\"\"\n","    model_name = \"NousResearch/Llama-2-7b-chat-hf\"\n","    new_model = \"Llama-2-7b-chat-finetune\"\n","    lora_r = 64\n","    lora_alpha = 16\n","    lora_dropout = 0.1\n","    use_4bit = True\n","    bnb_4bit_compute_dtype = \"float16\"\n","    bnb_4bit_quant_type = \"nf4\"\n","    use_nested_quant = False\n","    output_dir = \"./results\"\n","    fp16 = False\n","    bf16 = False\n","    per_device_train_batch_size = 4\n","    per_device_eval_batch_size = 4\n","    gradient_accumulation_steps = 1\n","    max_grad_norm = 0.3\n","    learning_rate = 2e-4\n","    weight_decay = 0.001\n","    optim = \"paged_adamw_32bit\"\n","    path = \"/content/\"\n","\n","    # Load training data\n","    train_split = load_train_data(\"train.csv\")\n","\n","    # Load base model\n","    model = load_base_model(model_name, use_4bit, bnb_4bit_compute_dtype, bnb_4bit_quant_type, use_nested_quant)\n","\n","    # Fine-tune model\n","    trainer ,tokenizer = fine_tune_model(model, train_split, lora_r, lora_alpha, lora_dropout, output_dir, num_train_epochs,\n","                                       fp16, bf16, per_device_train_batch_size, per_device_eval_batch_size,\n","                                       gradient_accumulation_steps, max_grad_norm, learning_rate, weight_decay,\n","                                       optim)\n","\n","    # Save trained model\n","    trainer.model.save_pretrained(path)\n","    return trainer,tokenizer  # Return the trainer object\n","\n","\n","def load_base_model(model_name, use_4bit, bnb_4bit_compute_dtype, bnb_4bit_quant_type, use_nested_quant):\n","    \"\"\"\n","    Load the base model with quantization configuration.\n","    \"\"\"\n","    compute_dtype = getattr(torch, bnb_4bit_compute_dtype)\n","    bnb_config = BitsAndBytesConfig(\n","        load_in_4bit=use_4bit,\n","        bnb_4bit_quant_type=bnb_4bit_quant_type,\n","        bnb_4bit_compute_dtype=compute_dtype,\n","        bnb_4bit_use_double_quant=use_nested_quant,\n","    )\n","    model = AutoModelForCausalLM.from_pretrained(\n","        model_name,\n","        quantization_config=bnb_config,\n","        device_map={\"\": 0},output_hidden_states=True\n","    )\n","    model.config.use_cache = False\n","    model.config.pretraining_tp = 1\n","    return model\n","\n","\n","\n","\n","def fine_tune_model(model, train_dataset, lora_r, lora_alpha, lora_dropout, output_dir, num_train_epochs, fp16, bf16,\n","                    per_device_train_batch_size, per_device_eval_batch_size, gradient_accumulation_steps, max_grad_norm,\n","                    learning_rate, weight_decay, optim):\n","    \"\"\"\n","    Fine-tune the model using the specified parameters.\n","    \"\"\"\n","    model_name = \"NousResearch/Llama-2-7b-chat-hf\"\n","    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n","    tokenizer.pad_token = tokenizer.eos_token\n","    tokenizer.padding_side = \"right\"\n","\n","    peft_config = LoraConfig(\n","        lora_alpha=lora_alpha,\n","        lora_dropout=lora_dropout,\n","        r=lora_r,\n","        bias=\"none\",\n","        task_type=\"CAUSAL_LM\",\n","    )\n","\n","    training_arguments = TrainingArguments(\n","        output_dir=output_dir,\n","        num_train_epochs=num_train_epochs,\n","        per_device_train_batch_size=per_device_train_batch_size,\n","        gradient_accumulation_steps=gradient_accumulation_steps,\n","        optim=optim,\n","        logging_steps=25,\n","        learning_rate=learning_rate,\n","        weight_decay=weight_decay,\n","        fp16=fp16,\n","        bf16=bf16,\n","        max_grad_norm=max_grad_norm,\n","        warmup_ratio=0.03,\n","        group_by_length=True,\n","        lr_scheduler_type=\"cosine\",\n","        report_to=\"tensorboard\"\n","    )\n","\n","    trainer = SFTTrainer(\n","        model=model,\n","        train_dataset=train_dataset,\n","        peft_config=peft_config,\n","        dataset_text_field=\"concept\",\n","        tokenizer=tokenizer,\n","        args=training_arguments,\n","    )\n","    trainer.train()\n","\n","    return trainer,tokenizer\n","\n","\n","num_train_epochs = 1\n","train_split = load_train_data(\"train.csv\")\n","trainer ,tokenizer = train_model(num_train_epochs)  # Get the trainer object"]},{"cell_type":"markdown","metadata":{"id":"4XTrPsoW1Xn1"},"source":["\n","## Step 4: Evaluate the fine-tuned model on the Train set\n","1. For each entry in the Train set, compare the base model response and fine-tuned model response similarity to the expected results\n","1. Improvement ideas in *Notes.docx* would be welcome."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"FLEXPOb01XyA"},"outputs":[],"source":["import matplotlib.pyplot as plt\n","from sklearn.metrics.pairwise import cosine_similarity\n","\n","\n","def plot_similarity_distribution(train_split, trainer,tokenizer):\n","    similarities = []\n","    model_name = \"NousResearch/Llama-2-7b-chat-hf\"\n","\n","    for entry in train_split:\n","        concept = entry[\"concept\"]\n","        base_response = concept\n","        # Generate fine-tuned response\n","        inputs = tokenizer(concept, return_tensors=\"pt\")\n","        fine_tuned_response = trainer.model.generate(**inputs)\n","        fine_tuned_response_str = tokenizer.decode(fine_tuned_response[0], skip_special_tokens=True)\n","        similarity = compare_responses(base_response, fine_tuned_response_str,trainer,tokenizer)\n","        similarities.append(similarity)\n","    plt.hist(similarities, bins=20, color='blue', alpha=0.7)\n","    plt.xlabel('Cosine Similarity')\n","    plt.ylabel('Data example')\n","    plt.title('Distribution of Cosine Similarity between Base and Fine-tuned Responses')\n","    plt.grid(True)\n","    plt.show()\n","\n","def compare_responses(base_response, fine_tuned_response,trainer,tokenizer):\n","    base_embedding = generate_embedding(base_response,trainer,tokenizer)\n","    fine_tuned_embedding = generate_embedding(fine_tuned_response,trainer,tokenizer)\n","    return calculate_similarity(base_embedding, fine_tuned_embedding)\n","\n","def generate_embedding(text, trainer, tokenizer):\n","  # i tried to use the orignial embedding by useing HuggingFaceEmbedding(model_name=\"NousResearch/Llama-2-7b-chat-hf\")\n","  # but there werent enough avalibale RAM, so i used my model as the embedding function\n","  # in a real project i would have search for a better embedding function\n","    inputs = tokenizer(text, return_tensors=\"pt\")\n","    with torch.no_grad():\n","        outputs = trainer.model(**inputs)\n","\n","    # Retrieve the hidden states from the model output\n","    hidden_states = outputs.hidden_states\n","\n","    # Choose which hidden state to use for the embedding\n","    last_hidden_state = hidden_states[-1]\n","\n","    # Use the mean of the last hidden state as the embedding\n","    embedding = last_hidden_state.mean(dim=1).squeeze().tolist()\n","    return embedding\n","\n","\n","def calculate_similarity(base_embedding, fine_tuned_embedding):\n","    return cosine_similarity([base_embedding], [fine_tuned_embedding])[0][0] # shape (1,1)\n","\n","\n","plot_similarity_distribution(train_split, trainer,tokenizer)"]},{"cell_type":"markdown","metadata":{"id":"aYq6QJjM1a6Z"},"source":["\n","## Bonus Step: On \"Llama-2-7b-chat-hf-sharded\" or any other \"chat fine-tuned\" model, compare the first and last activation layers for 2 prompts\n","1. For each of the prompts: “How to build a snowman?” and “How to build a bomb?”\n","  1. Get the values vector of the first and the last activation layer on each forward-pass\n","  1. Calculate 2 average values vectors for the first and last activation layers (for all forward-passes)\n","1. Share your thoughts on the results for the 2 prompts and whether Alignment \\ Guardrails have anything to do with it.\n","1. The main objective of this exercise is to see if you're comfortable with interacting with the models underlying layers, the actual similarty function and other \"mathematical\" considerations, are less important for this step."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"9bHWrUEm6qNv"},"outputs":[],"source":["def extracting_layers(tokenizer):\n","    # Load the fine-tuned model and tokenizer\n","    model_name = \"NousResearch/Llama-2-7b-chat-hf\"\n","    use_4bit = True\n","    bnb_4bit_compute_dtype = \"float16\"\n","    bnb_4bit_quant_type = \"nf4\"\n","    use_nested_quant = False\n","\n","    # Load base model\n","    model = load_base_model(model_name, use_4bit, bnb_4bit_compute_dtype, bnb_4bit_quant_type, use_nested_quant)\n","\n","    # Define the prompts\n","    prompts = [\"How to build a snowman?\", \"How to build a bomb?\"]\n","\n","    # Tokenize the prompts and convert them to input IDs\n","    input_ids_list = [tokenizer.encode(prompt, return_tensors=\"pt\") for prompt in prompts]\n","\n","    # List to store activation vectors for the first and last layers\n","    first_layer_activations_list = [[], []]\n","    last_layer_activations_list = [[], []]\n","\n","    # Forward pass through the model for each prompt\n","    for i, input_ids in enumerate(input_ids_list):\n","        with torch.no_grad():\n","            outputs = model(input_ids)\n","\n","            # Extract activations for the first and last layers\n","            first_layer_activations = outputs.hidden_states[0]  # First layer activations\n","            last_layer_activations = outputs.hidden_states[-1]  # Last layer activations\n","\n","            # Select the first and last values vectors\n","            first_layer_activations_list[i].append(first_layer_activations)\n","            last_layer_activations_list[i].append(last_layer_activations)\n","\n","            # # Generate response from the model\n","            # logits_flat = [logit for sublist in outputs.logits[0] for logit in sublist]\n","            # generated_response = tokenizer.decode(logits_flat, skip_special_tokens=False)\n","            # generated_responses.append(generated_response)\n","\n","\n","    # Plot activations for each prompt\n","    for i, (prompt, first_activations, last_activations) in enumerate(zip(prompts, first_layer_activations_list, last_layer_activations_list), start=1):\n","        num_inputs = len(first_activations)\n","        plt.figure(figsize=(10, 5 * num_inputs))\n","\n","        for j, (first_layer_activation, last_layer_activation) in enumerate(zip(first_activations, last_activations), start=1):\n","            # Plot activations for the first layer\n","            plt.subplot(num_inputs, 2, 2 * j - 1)\n","            plt.plot(first_layer_activation.squeeze().tolist())\n","            plt.title(f'Prompt {i}, Input {j} - First Layer Activations')\n","            plt.xlabel('Token Position')\n","            plt.ylabel('Activation Value')\n","\n","            # Plot activations for the last layer\n","            plt.subplot(num_inputs, 2, 2 * j)\n","            plt.plot(last_layer_activation.squeeze().tolist())\n","            plt.title(f'Prompt {i}, Input {j} - Last Layer Activations')\n","            plt.xlabel('Token Position')\n","            plt.ylabel('Activation Value')\n","\n","        plt.tight_layout()\n","        plt.show()\n","\n","\n","extracting_layers(tokenizer)"]}],"metadata":{"accelerator":"GPU","colab":{"provenance":[{"file_id":"1PWw9zf5sqdFonFGE-D7_BgkvpaE-GCSl","timestamp":1712576532085}],"gpuType":"T4"},"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"name":"python"}},"nbformat":4,"nbformat_minor":0}