import time

# Track the total time to run all of the package installations.
t0 = time.time()

# Required for using 4-bit quantization (to make our model fit in the GPU
# memory!) and the 8-bit optimizer (which helps with running training on the
# limited memory of the free T4 GPU).
!pip install bitsandbytes

!pip install peft

from google.colab import userdata

# I used the "secrets" panel in Colab and defined the variable
# "hf_hub_all_notebooks" and set it to my personal huggingface key.
# You could just paste in your key directly here if this is a private copy.
hf_api_key = userdata.get('hf_hub_all_notebooks')

!huggingface-cli login --token $hf_api_key

import torch

# There's a nice function now for retrieving the GPU info...
gpu_stats = torch.cuda.get_device_properties(0)

print(f"GPU = {gpu_stats.name}.")

# A key attribute is the amount of memory, in GB, the GPU has, so let's
# calculate that by dividing the total bytes by 2^30.
gpu_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)

print(f"GPU Memory = {gpu_memory} GB")

# Finally, we'll store a shorthand version of the name and use this as a switch
# for other settings (e.g., if gpu = "T4", then ...)
if "T4" in gpu_stats.name:
    gpu = "T4"

elif "L4" in gpu_stats.name:
    gpu = "L4"

elif "A100" in gpu_stats.name:
    gpu = "A100"

else:
    raise ValueError(f"Unsupported GPU: {gpu_stats.name}")

print("Shorthand Name:", gpu_stats.name, ' -->', gpu )

import time
import datetime

def format_time(elapsed):
    '''
    Takes a time in seconds and returns a string hh:mm:ss
    '''
    # Round to the nearest second.
    elapsed_rounded = int(round((elapsed)))

    # Format as hh:mm:ss
    return str(datetime.timedelta(seconds=elapsed_rounded))

print("All of that package installation stuff took:", format_time(time.time() - t0))

def format_size(num):
    """
    This function iterates through a list of suffixes ('K', 'M', 'B') and
    divides the input number by 1024 until the absolute value of the number is
    less than 1024. Then, it formats the number with the appropriate suffix and
    returns the result. If the number is larger than "B", it uses 'T'.
    """
    suffixes = ['', 'K', 'M', 'B']
    base = 1024

    for suffix in suffixes:
        if abs(num) < base:
            if num % 1 != 0:
                return f"{num:.2f}{suffix}"

            else:
                return f"{num:.0f}{suffix}"

        num /= base

    # Use "T" for anything larger.
    if num % 1 != 0:
        return f"{num:.2f}T"

    else:
        return f"{num:.0f}T"

def format_lr_as_multiple(lr):
    # Convert the learning rate into a multiple of 1e-5.
    multiple = lr / 1e-5

    # Return the formatted string.
    return "{:.1f} x 1e-5".format(multiple)

import os
import torch

def gpu_mem_used():
    """
    Returns the current GPU memory usage as a string, e.g., "5.02 GB"
    """

    # This approach doesn't work, because PyTorch only tracks its own memory
    # usage, not the total memory consumption of the GPU.
    #gpu_bytes_used = torch.cuda.memory_allocated()

    # Run the nvidia-smi command line tool to get memory used in megabytes.
    buf = os.popen('nvidia-smi --query-gpu=memory.used, --format=csv,noheader,nounits')

    # It returns an unformated integer number of "MiB" (2^20 bytes).
    gpu_mb_used = float(buf.read())

    # Divide that by 1024 to get GB.
    mem_used = gpu_mb_used / float(1024)

    return ("{0:.2f} GB".format(mem_used))

print("GPU memory used: {:}".format(gpu_mem_used()))

# (Not installed by default)
!pip install wget

import wget
import os

print('Downloading dataset...')

# The URL for the dataset zip file.
url = 'https://nyu-mll.github.io/CoLA/cola_public_1.1.zip'

# Download the file (if we haven't already)
if not os.path.exists('./cola_public_1.1.zip'):
    wget.download(url, './cola_public_1.1.zip')

# Unzip the dataset (if we haven't already)
if not os.path.exists('./cola_public/'):
    !unzip cola_public_1.1.zip

import pandas as pd

# Load the dataset into a pandas dataframe.
df = pd.read_csv(
    "./cola_public/raw/in_domain_train.tsv",
    delimiter = '\t',
    header = None,
    names = ['sentence_source', 'label', 'label_notes', 'sentence']
)

# Report the number of sentences.
print('Number of training sentences: {:,}\n'.format(df.shape[0]))

# Display 10 random rows from the data.
df.sample(10)

df.loc[df.label == 0].sample(5)[['sentence', 'label']]

# Since the positive samples have the label value "1", we can sum them to count
# them.
num_positive_samples = df.label.sum()

num_samples = len(df.label)

prcnt_positive = float(num_positive_samples / num_samples) * 100.0

print(
    "Number of positive samples: {:,} of {:,}. {:.3}%".format( \
    num_positive_samples, num_samples, prcnt_positive)
)

# Get the lists of sentences and their labels.
sentences = df.sentence.values
labels = df.label.values

from transformers import LlamaTokenizer
from transformers import AutoTokenizer
from transformers import LlamaTokenizerFast

# Load the tokenizer.
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")

# IMPORTANT: Add the leading space to the sentence!
example_sen = ' ' + sentences[0]

# Print the original sentence.
print(' Original: ', example_sen)

tokens = tokenizer.tokenize(example_sen)

# Print the sentence split into tokens.
print('Tokenized: ', tokens)

token_ids = tokenizer.convert_tokens_to_ids(tokens)

# Print the sentence mapped to token ids.
print('Token IDs: ', token_ids)

# Assume tokenizer is your LlamaTokenizer instance
tokenizer_properties = {
    'name_or_path': tokenizer.name_or_path,
    'vocab_size': tokenizer.vocab_size,
    'model_max_length': tokenizer.model_max_length,
    'is_fast': tokenizer.is_fast,
    'padding_side': tokenizer.padding_side,
    'truncation_side': tokenizer.truncation_side,
    'clean_up_tokenization_spaces': tokenizer.clean_up_tokenization_spaces,
    'bos_token': tokenizer.bos_token,
    'eos_token': tokenizer.eos_token,
    'unk_token': tokenizer.unk_token,
    'sep_token': tokenizer.sep_token,
    'pad_token': tokenizer.pad_token,
    'mask_token': tokenizer.mask_token,
}

# Print each property in a formatted way
for key, value in tokenizer_properties.items():
    print(f"{key.rjust(30)}:  {value}")

tokenizer.pad_token = tokenizer.eos_token

print("Token representations of our label words:")

print("    ' acceptable' = ", tokenizer.tokenize(' acceptable'))
print("  ' unacceptable' = ", tokenizer.tokenize(' unacceptable'))

print("Token representations of our label words:")

print("    ' correct' = ", tokenizer.tokenize(' correct'))
print("  ' incorrect' = ", tokenizer.tokenize(' incorrect'))

print("Token representations if we forget the leading space:")

print("    'acceptable' = ", tokenizer.tokenize("acceptable"))
print("  'unacceptable' = ", tokenizer.tokenize("unacceptable"))

label_val_to_word = {}
label_val_to_token_id = {}

pos_token_id = -1
neg_token_id = -1

# Select our word for the positive label (1)
label_val_to_word[1] = ' acceptable'

# Show it as a token.
pos_token = tokenizer.tokenize(label_val_to_word[1])[0]
print(pos_token)

# Lookup and store its token id.
pos_token_id = tokenizer.convert_tokens_to_ids(pos_token)
label_val_to_token_id[1] = pos_token_id
print(str(pos_token_id) + '\n')

# Select our word for the negative label (0)
label_val_to_word[0] = ' unacceptable'

# Show it as a token.
neg_token = tokenizer.tokenize(label_val_to_word[0])[0]
print(neg_token)

# Look up and store its token id.
neg_token_id = tokenizer.convert_tokens_to_ids(neg_token)
label_val_to_token_id[0] = neg_token_id
print(str(neg_token_id) + '\n')

""" Classify this sentence as grammatically acceptable or unacceptable: Him and me are going to the store.
A: unacceptable

Classify this sentence as grammatically acceptable or unacceptable: Him and I are going to the store.
A: acceptable

Classify this sentence as grammatically acceptable or unacceptable: {sentence}
A:{label_word}"""

""" Classify the grammar of this sentence as ' acceptable' or ' unacceptable': Him and me are going to the store.
 A: unacceptable

 Classify the grammar of this sentence as ' acceptable' or ' unacceptable': Him and I are going to the store.
 A: acceptable

 Classify the grammar of this sentence as ' acceptable' or ' unacceptable': {sentence}
 A:{label_word}"""

""" Examples of sentences that are grammatically ' acceptable' or ' unacceptable':
 Him and me are going to the store. - unacceptable
 Him and I are going to the store. - acceptable
 {sentence} -{label_word}"""

prompt_template = \
""" Examples of sentences that are grammatically ' acceptable' or ' unacceptable':
 Him and me are going to the store. - unacceptable
 Him and I are going to the store. - acceptable
 {sentence} -{label_word}"""

labeled_sentences = []
labels_as_ids = []

# For each sentence in the dataset...
for i in range(len(sentences)):

    sentence = sentences[i]
    label_val = labels[i]

    # Map the numerical label (0, 1) to the word we chose.
    label_word = label_val_to_word[label_val]

    # Look up the token id for the label.
    label_id = label_val_to_token_id[label_val]

    # Insert the sample and its label into the template.
    labeled_sentence = prompt_template.format(
        sentence = sentence,
        label_word = label_word
    )

    # Add to our new lists.
    labeled_sentences.append(labeled_sentence)
    labels_as_ids.append(label_id)

print("\"{:}\"".format(labeled_sentences[0]))

print("\"{:}\"".format(labeled_sentences[1]))

# If you want to quickly test out your code, you can slice this down to just a
# small subset.

#labeled_sentences = labeled_sentences[:500]
#labels_as_ids = labels_as_ids[:500]

# The tokenizer is a "callable object"--this invokes its __call__ function,
# which will tokenize and encode all of the input strings.
encodings = tokenizer(

    labeled_sentences, # List of strings.

    padding = 'longest',  # Pad out all of the samples to match the longest one
                          # in the data.

    #max_length = 64,      # An alternative strategy is to specify a maximum
    #padding='max_length', # length, but it makes sense to let the tokenizer
                           # figure that out.

    truncation = False, # These samples are too short to need truncating. If we
                        # did need to truncate, we'd also need to specify the
                        # max_length above.

    add_special_tokens = True, # Add the bos and eos tokens.

    return_attention_mask = True, # To distinguish the padding tokens.

    return_tensors = "pt" # Return the results as pytorch tensors.
)

max_seq_len = len(encodings['input_ids'][0])

print("Longest training sequence length:", max_seq_len)

# I'll add the prefix 'all' to these variables, since they still contain both
# the training and validation data.
all_input_ids = []
all_attention_masks = []
all_target_words = []
all_label_ids = []
all_label_pos = []

# For each of the encoded samples...
for i in range(len(labels_as_ids)):

    # Extract input_ids and attention_mask
    input_ids = encodings['input_ids'][i]
    attention_mask = encodings['attention_mask'][i]

    # Find the position of the last non-padding token using the attention mask
    # Because we appended the label to the end of the input, this is the
    # position of our label word.
    label_position = attention_mask.nonzero()[-1].item()

    # This will tell the model what to token to predict at each position.
    # (i.e., at position 12, the model should predict target_words[12])
    # You can set the value to -100 for any tokens you don't want to train on,
    # and in our case, we only want to train on the label.
    # Start by filling it all out with -100s
    target_words = torch.full_like(input_ids, -100)  # Initialize all labels to -100

    # Get the token id for the label
    label_id = labels_as_ids[i]

    # We want all of the words / tokens masked out, except for the label.
    target_words[label_position] = label_id

    # Store everything.
    all_input_ids.append(input_ids)
    all_attention_masks.append(attention_mask)
    all_target_words.append(target_words)
    all_label_pos.append(label_position)
    all_label_ids.append(label_id)

# Calculate the number of samples to include in each set.
train_size = int(0.9 * len(all_input_ids))
val_size = len(all_input_ids) - train_size

print('{:>5,} training samples'.format(train_size))
print('{:>5,} validation samples'.format(val_size))

import torch
from torch.utils.data import TensorDataset

# Convert the lists into PyTorch tensors and put into a TensorDataset object

dataset = TensorDataset(
    # These variables are currently lists of vectors. `stack` will combine
    # the vectors into a matrix with shape
    #   8551 x 123 ([number of samples] x [sequence length])
    torch.stack(all_input_ids),
    torch.stack(all_attention_masks),
    torch.stack(all_target_words),

    # These are lists, and need to be vectors.
    torch.tensor(all_label_ids),
    torch.tensor(all_label_pos)
)

# For reference, this is how we'll unpack a batch:
#     input_ids = batch[0]
#     attn_masks = batch[1]
#     targets = batch[2]
#     label_ids = batch[3]
#     label_pos = batch[4]

import numpy as np
import random

# Set the seed value all over the place to make this reproducible.
seed_val = 42

random.seed(seed_val)
np.random.seed(seed_val)
torch.manual_seed(seed_val)
torch.cuda.manual_seed_all(seed_val)

from torch.utils.data import random_split

train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

# The DataLoader needs to know our batch size for training, so we specify it
# here.

# First, specify the mathematical batch size that we want.
train_batch_size = 8

# Second, specify the number of samples we want to give to the GPU at a time.
# (We set this to a lower number if we're running out of memory)

# Because our sequence length is so short, we can actually manage to use our
# full batch size of 8 even on the 15GB T4--just barely!
# If you try out a longer prompt, though, you may go over.
if gpu == "T4" and max_seq_len > 90:
    gpu_batch_size = 4
# If you run this notebook as-is, then this size will work on all devices:
else:
    gpu_batch_size = 8

# These must be evenly divisible.
assert(train_batch_size % gpu_batch_size == 0)

# Calculate how many batches to accumulate the gradients over.
accumulate_passes = int(train_batch_size / gpu_batch_size)

print("For the math, we are using an (effective) batch size "\
      "of {:}.".format(train_batch_size))
print("For memory constraints, we will only give the GPU {:} samples at " \
      "a time.".format(gpu_batch_size))
print("We'll do this by accumulating the gradients over {:} GPU " \
      "passes before updating.".format(accumulate_passes))

# When running inference, there's no need for accumulation.
val_batch_size = gpu_batch_size

# The A100 can handle a bigger batch size to make the validation go quicker.
if gpu == "A100":
    val_batch_size = 16

from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from torch.utils.data import TensorDataset

# Create the DataLoaders for our training and validation sets.
# These will handle the creation of batches for us.
# Note that we want to use the GPU batch size here, not the effective one.

# We'll take training samples in random order.
train_dataloader = DataLoader(
            train_dataset,  # The training samples.
            sampler = RandomSampler(train_dataset), # Select batches randomly
            batch_size = gpu_batch_size
        )

# For validation the order doesn't matter, so we'll just read them sequentially.
validation_dataloader = DataLoader(
            val_dataset, # The validation samples.
            sampler = SequentialSampler(val_dataset), # Pull out batches sequentially.
            batch_size = val_batch_size # Evaluate with this batch size.
        )

if gpu == "T4":
    torch_dtype = torch.float16
elif gpu == "L4" or gpu == "A100":
    torch_dtype = torch.bfloat16

import torch
from transformers import BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    # Enable 4-bit quantization.
    load_in_4bit = True,

    # With "double quantization" we not only compress the weights, we compress
    # the scaling factors as well!
    bnb_4bit_use_double_quant = False,

    # The authors did some analysis of the distribution of weight values in
    # popular transformer models, and chose some hardcoded values to use as the
    # 16 "base values". They are normally distributed around 0.
    # They refer to this as the "nf4" data type.
    #
    # The alternative choice is "float4", which only has 15 unique values, but
    # also seems to be normally distributed.
    bnb_4bit_quant_type = "nf4",

    # The 16-bit data type for the scaling factors and math.
    bnb_4bit_compute_dtype = torch_dtype

    # Note that the "block size" of 64, which determines the compression ratio,
    # doesn't appear to be configurable!
)

# Tell PyTorch to use the GPU.
device = torch.device("cuda")

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

t0 = time.time()

# Load the model.
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",

    # Our 4-bit quantization setup defined in the previous section.
    quantization_config = bnb_config,

    # Squared Dot-Product Attention is the generic name for the key innovation
    # in FlashAttention. It's implemented in pytorch now. It's the default
    # value for this parameter.
    # You can change it to "eager" for the straightforward implementation if
    # you want to compare the two!
    attn_implementation = "sdpa",

    # I assume it's critical that this datatype match the one used in
    # the quantization configuration, and in LoRA!
    torch_dtype = torch_dtype,

    # I was getting this output message:
    #    "`low_cpu_mem_usage` was None, now set to True since model is
    #      quantized."
    # So I'm setting it to "True" as the message suggests. :)
    #
    # I tried setting it to False out of curiousity, and it crashed with:
    #    "Your session crashed after using all available RAM."
    low_cpu_mem_usage = True,

    # The model needs to know that we've set a pad token.
    pad_token_id = tokenizer.pad_token_id,
)

# Typically I would tell the model to run on the GPU at this point, but this
# actually throws an error:
#   "ValueError: Calling `cuda()` is not supported for `4-bit` or `8-bit`
#    quantized models. Please use the model as it is, since the model has
#    already been set to the correct devices and casted to the correct `dtype`."
#model.cuda()

print("\nDownloading, compressing, and loading took", format_time(time.time() - t0))

gpu_mem_model = gpu_mem_used()

print("\nGPU memory used to store model: {:}".format(gpu_mem_model))

# Get all of the model's parameters as a list of tuples.
params = list(model.named_parameters())

print('The model has {:} different named parameters.\n'.format(len(params)))

print("Parameter Name                                              Dimensions       Total Values    Trainable\n")

# For the first xx named parameters...
for i in range(len(params)):

    # First param is the embeddings
    if i == 0:
        print('==== Embedding Layer ====\n')

    # Next 10 parameters are the first decoder layer.
    elif i == 1:
        print('\n==== First Decoder ====\n')

    # Next 10 parameters are the second decoder layer.
    elif i == 10:
        print('\n==== Second Decoder ====\n')

    # Skip the other layers
    elif i >= 19 and i < len(params)-2:
        continue

    # Final 2 are the output layer.
    elif i == len(params)-2:
        print('\n==== Output Layer ====\n')

    # Get the name and the parameter.
    p_name, p = params[i]

    # Strip some unnecessary prefixes.
    if 'base_model.model.model.' in p_name:
        p_name = p_name[len('base_model.model.'):]


    # For 1D parameters, put '-' as the second dimension.
    if len(p.size()) == 1:
        p_dims = "{:>10,} x {:<10}".format(p.size()[0], "-")

    # For 2D parameters...
    if len(p.size()) == 2:
        p_dims = "{:>10,} x {:<10,}".format(p.size()[0], p.size()[1])

    # Print the parameter name, shape, number of elements, and whether it's been
    # 'frozen' or not.
    print("{:<55} {:}    {:>6}    {:}".format(p_name, p_dims, format_size(p.numel()), p.requires_grad))

def evaluate(model, validation_dataloader):
    """
    Run the model against a validation or test set and return a list of true
    labels and a list of predicted labels.
    """

    #t0 = time.time()

    batch_num = 0

    total_loss = 0

    # Put the model in evaluation mode--the dropout layers behave differently
    # during evaluation.
    model.eval()

    # Gather the true labels and predictions
    true_labels = []
    predictions = []

    # For each batch...
    for batch in validation_dataloader:

        # Print progress
        if batch_num % 40 == 0:
            print("  Batch {:>5,}  of  {:>5,}.".format(batch_num, len(validation_dataloader)))

        # Unpack this training batch from our dataloader.
        #
        # As we unpack the batch, we'll also copy each tensor to the GPU using
        # the `to` method.

        # These three all have the same shape of [batch_size x sequence_length]
        # e.g., [8 x 128]
        input_ids = batch[0].to(device)
        attention_mask = batch[1].to(device)
        targets = batch[2].to(device)
        true_label_ids = batch[3].to(device) # Shape: [8] (batch_size)
        label_positions = batch[4].to(device) # Shape: [8] (batch_size)

        # Tell pytorch not to bother with constructing the compute graph during
        # the forward pass, since this is only needed for backprop (training).
        with torch.no_grad():

            # Get the model's predictions
            outputs = model(
                input_ids = input_ids,
                attention_mask = attention_mask,
                labels = targets
            )

        # The output predictions for every token in the input.
        # Shape: [8, 128, vocab_size] (batch_size, sequence_length, vocab_size)
        logits = outputs.logits

        loss = outputs.loss
        total_loss += loss.item()

        # Extract the predicted token for each sample in the batch, one at a time.
        # For each sample in the batch...
        for i in range(input_ids.shape[0]):

            # The index of the final (non-padding) token.
            label_position = label_positions[i].item()

            # Extract logits for the prediction.
            # The logits for the position *just before* the label are the
            # predictions we want.
            # Shape: [vocab_size]
            label_logits = logits[i, label_position - 1]

            # Make our prediction by comparing the confidence for the two
            # label words.
            if label_logits[pos_token_id] > label_logits[neg_token_id]:
                predicted_token_id = pos_token_id
            else:
                predicted_token_id = neg_token_id

            # Append the true and predicted token IDs for comparison
            true_labels.append(true_label_ids[i].item())
            predictions.append(predicted_token_id)

        # Increment the batch counter
        batch_num += 1

    # Report the average validation loss, which we can use to compare to our
    # training loss and detect overfitting.
    avg_loss = total_loss / batch_num

    #print("\n  Validation took: {:}".format(format_time(time.time() - t0)))

    return(true_labels, predictions, avg_loss)

print('Running validation...')

true_labels, predictions, avg_loss = evaluate(model, validation_dataloader)

import numpy as np

from sklearn.metrics import matthews_corrcoef
from sklearn.metrics import confusion_matrix

true_labels = np.asarray(true_labels)
predictions = np.asarray(predictions)

# Report the final accuracy for this validation run.
val_accuracy = float(np.sum(true_labels == predictions)) / len(true_labels) * 100.0
print("Flat accuracy: {0:.2f}%".format(val_accuracy))

# Generate confusion matrix
conf_matrix = confusion_matrix(true_labels, predictions)

# Print or log the confusion matrix
print("\nConfusion Matrix:\n")
print(conf_matrix)
print("")
print("  TP    FP")
print("  FN    TN")

# Measure the MCC on the validation set.
few_shot_mcc = matthews_corrcoef(true_labels, predictions)

print("\nMCC on the validation set without fine-tuning: {0:.3f}".format(few_shot_mcc))

# Get the current GPU memory, which will reflect how much memory is required
# to store the activation values.
gpu_mem_forward_pass = gpu_mem_used()

print("\nGPU memory used for forward passes: {:}".format(gpu_mem_forward_pass))
print(
    "(With GPU batch size {:} and sequence length {:})".format(val_batch_size,
                                                               len(encodings['input_ids'][0]))

)

from peft import LoraConfig

# This is the key parameter. Think of it like the number of neurons you're
# using in each hidden layer of a neural network. More will allow for
# bigger changes to the model behavior but require more training data.
lora_rank = 8

# LoRA multiplies the gradients by a scaling factor, s = alpha / rank.
# When you double the rank, the size gradients of the gradients tend to double
# as well, so `s` is used to scale them back down.
# If you want to play with different rank values, you can keep alpha constant in
# order to maintain the same (effective) scaling factor.
#
# An oddity here is that `s` is redundant with the learning rate... both are
# just multipliers applied to the gradients.
# They could have addressed this instead by saying, e.g., "when doubling the
# rank, it's suggested that you cut the learning rate in half"
lora_alpha = 16

lora_config = LoraConfig(
    r = lora_rank,

    lora_alpha = lora_alpha,

    lora_dropout = 0.05,

    # The original convention with LoRA was to apply it to just the Query and
    # Value matrices. An important contribution of the QLoRA paper was that it's
    # important (and not very costly) to apply it to "everything" (the attention
    # matrices and the FFN)
    #target_modules=["q_proj", "v_proj"],
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    bias="none",

    # This is a crucial parameter that I missed when initially trying to use the
    # LlamaForSequenceClassification class. I couldn't figure out why my model
    # wasn't learning, until I discovered that LoRA was freezing the parameters
    # of the linear classifier head!
    task_type="CAUSAL_LM"
)

from peft import get_peft_model

# Wrap with PEFT Model applying LoRA
model = get_peft_model(model, lora_config).to(device)

# Tell pytorch to run this model on the GPU.
model.cuda()

# This is a bad idea / doesn't work when mixed with quantization.
# Also misleading about what LoRA is for.
#model.print_trainable_parameters()

gpu_mem_lora = gpu_mem_used()
print("\nGPU memory used after adding LoRA: {:}".format(gpu_mem_lora))

# Get all of the model's parameters as a list of tuples.
params = list(model.named_parameters())

print('The model has {:} different named parameters.\n'.format(len(params)))

print("Parameter Name                                              Dimensions       Total Values    Trainable\n")

# For the first xx named parameters...
for i in range(len(params)):

    # First param is the embeddings
    if i == 0:
        print('==== Embedding Layer ====\n')

    # Next 24 parameters are the first decoder layer.
    elif i == 1:
        print('\n==== First Decoder ====\n')

    # Next 24 parameters are the second decoder layer.
    elif i == 24:
        print('\n==== Second Decoder ====\n')

    # Skip the other layers
    elif i > 46 and i < len(params)-2:
        continue

    # Final 2 are the output layer.
    elif i == len(params)-2:
        print('\n==== Output Layer ====\n')

    # Get the name and the parameter.
    p_name, p = params[i]

    # Strip some unnecessary prefixes.
    if 'base_model.model.model.' in p_name:
        p_name = p_name[len('base_model.model.'):]

    # For 1D parameters, put '-' as the second dimension.
    if len(p.size()) == 1:
        p_dims = "{:>10,} x {:<10}".format(p.size()[0], "-")

    # For 2D parameters...
    if len(p.size()) == 2:
        p_dims = "{:>10,} x {:<10,}".format(p.size()[0], p.size()[1])

    # Print the parameter name, shape, number of elements, and whether it's been
    # 'frozen' or not.
    print("{:<55} {:}    {:>6}    {:}".format(p_name, p_dims, format_size(p.numel()), p.requires_grad))

import bitsandbytes as bnb
from transformers import AdamW

# Note that the LoRA ratio has a linear impact on the lr, so with 16/8,
# 1e-4 is effectively 2e-4.
lr = 1e-4

# At these lower sequence lengths, the 8-bit optimizer makes a significant
# difference.
if gpu == 'T4':
    #optimizer_class = bnb.optim.PagedAdamW8bit
    optimizer_class = bnb.optim.AdamW8bit
else:
    optimizer_class = AdamW

# Create the optimizer
optimizer = optimizer_class(

    model.parameters(), # The optimizer is responsible for making the updates to
                        # to our model parameters.
    lr = lr,

    weight_decay = 0.05, # Weight decay (related to regularization)
)

from transformers import get_linear_schedule_with_warmup
import math

# len(train_dataloader) gives us the number of GPU batches the data loader will
# divide our dataset into.
num_gpu_batches = len(train_dataloader)

# We need to divide the number of GPU batches by the number of accumulation
# passes to get the true number of update steps.
# Round up since there may be one partial batch at the end.
num_update_steps = math.ceil(num_gpu_batches / accumulate_passes)

# Our LoRA weights are completely random at first, so the initial gradients will
# be large. To counter this, we do some warmup steps with a tiny learning rate,
# gradually building up to the actual lr. Then the scheduler takes over.
num_warmup_steps = 100

# Create the learning rate scheduler.
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps = num_warmup_steps,
    num_training_steps = num_update_steps
)

# Without batch accumulation...

# Total number of training steps is [number of batches] x [number of epochs].
# (Note that this is not the same as the number of training samples).
#total_steps = len(train_dataloader) * epochs

# Create the learning rate scheduler.
#scheduler = get_linear_schedule_with_warmup(
#    optimizer,
#    num_warmup_steps = 100,
#    num_training_steps = total_steps
#)

import random
import numpy as np
from sklearn.metrics import confusion_matrix

# Set the seed value again to make sure this is reproducible--the data
# loader will select the same random batches every time we run this.
seed_val = 42

# Set it all over the place for good measure.
random.seed(seed_val)
np.random.seed(seed_val)
torch.manual_seed(seed_val)
torch.cuda.manual_seed_all(seed_val)

# We'll store a number of quantities such as training and validation loss,
# validation accuracy, and timings.
training_stats = []

# Measure the total training time for the whole run.
t0 = time.time()

total_train_loss = 0

# Track the current step (because gradient accumulation means that this is
# different than the number of forward passes on the GPU).
optim_step = 0

# Put the model into training mode. `dropout` and `batchnorm` layers behave
# differently during training vs. test.
model.train()

# ========================================
#               Training
# ========================================

# Perform one full pass over the training set.
print('Training...\n')

print(
"'step' refers to one optimizer step (an important distinction if you're \
using gradient accumulation).\n"
)

# Print the header row for status updates.
print("{:>10}  {:>10}  {:>10}  {:>12}  {:>10}  {:>12}".format(
    "Step", "of Total", "Elapsed", "Loss", "GPU Mem", "lr"))

# For each GPU batch of training data...
for gpu_batch_i, batch in enumerate(train_dataloader):

    # ===================
    #    Forward Pass
    # ===================

    # Unpack this training batch from our dataloader.
    #
    # As we unpack the batch, we'll also copy each tensor to the GPU using
    # the `to` method.
    #
    # `batch` contains three pytorch tensors:
    #   [0]: input ids
    #   [1]: attention masks
    #   [2]: labels -- In this case, these are target words at each
    #                  position.
    b_input_ids = batch[0].to(device)
    b_input_mask = batch[1].to(device)
    b_labels = batch[2].to(device)

    # Perform a forward pass (evaluate the model on this training batch).
    # In PyTorch, calling `model` will in turn call the model's `forward`
    # function and pass down the arguments.
    # Specifically, we'll get the loss (because we provided labels) and the
    # "logits"--the model outputs prior to activation.
    result = model(
        b_input_ids,
        attention_mask = b_input_mask,
        labels = b_labels
    )

    loss = result.loss
    logits = result.logits

    # Accumulate the training loss over all of the batches.
    # Note that the loss is already averaged over the GPU batch size!
    total_train_loss += loss.item()

    # =======================
    #     Backward Pass
    # =======================

    # Further scale the loss by the number of accumulation passes
    loss = loss / accumulate_passes

    # Perform a backward pass to calculate the gradients.
    loss.backward()

    # Clip the norm of the gradients to 1.0.
    # This is to help prevent the "exploding gradients" problem.
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

    # Perform optimizer step after accumulating gradients.
    if ((gpu_batch_i + 1) % accumulate_passes == 0) or \
       ((gpu_batch_i + 1) == len(train_dataloader)):

        # Update the step count.
        optim_step += 1

        # Update parameters and take a step using the computed gradient.
        # The optimizer dictates the "update rule"--how the parameters are
        # modified based on their gradients, the learning rate, etc.
        optimizer.step()

        # Update the learning rate according to the schedule.
        scheduler.step()

        # Clear the gradients.
        optimizer.zero_grad()

        # ======================
        #       Progress
        # ======================

        # Report every xx steps
        if optim_step % 10 == 0:

            # Calculate and format elapsed time.
            elapsed = format_time(time.time() - t0)

            # Display the learning rate as a multiple of 1e-5.
            lr_multiple = format_lr_as_multiple(scheduler.get_last_lr()[0])

            # Print the current row with values.
            print("{:>10,}  {:>10,}  {:>10}  {:>12.4f}  {:>10}  {:>12}".format(
                optim_step,
                num_update_steps,
                elapsed,
                total_train_loss / optim_step,
                gpu_mem_used(),
                lr_multiple
            ))

        # ====================
        #      Validation
        # ====================
        # Periodically measure our performance on our validation set.

        if optim_step % 200 == 0:

            print("")
            print("Running Validation...")

            val_t0 = time.time()

            # Run the validation dataset through the model.
            true_labels, predictions, avg_val_loss = evaluate(model, validation_dataloader)

            # Gather metrics...
            true_labels = np.asarray(true_labels)
            predictions = np.asarray(predictions)

            val_mcc = matthews_corrcoef(true_labels, predictions)

            print("\n  Validation MCC: {0:.2f}".format(val_mcc))

            # Measure how long the validation run took.
            validation_time = format_time(time.time() - val_t0)

            print("  Validation took: {:}".format(validation_time))

            # Record statistics.
            training_stats.append(
                {
                    'Step': optim_step,
                    'Training Loss': total_train_loss / (optim_step + 1),
                    'Valid. Loss': avg_val_loss,
                    'Valid. MCC': val_mcc,
                    'Training Time': format_time(time.time() - t0),
                    'Validation Time': validation_time
                }
            )

# =======================
#   Training Complete
# =======================

print("")
print("Training complete!")
print("")

# Calculate the average loss over all of the batches.
avg_train_loss = total_train_loss / num_update_steps

print("  Average training loss: {0:.4f}".format(avg_train_loss))

training_time = format_time(time.time() - t0)

print("  Training took: {:}".format(training_time))

import pandas as pd

# Display floats with two decimal places.
pd.set_option('display.precision', 4)

# Create a DataFrame from our training statistics.
df_stats = pd.DataFrame(data=training_stats)

# Display the table.
df_stats

import matplotlib.pyplot as plt
#% matplotlib inline

import seaborn as sns

# Use plot styling from seaborn.
sns.set(style='darkgrid')

# Increase the plot size and font size.
sns.set(font_scale=1.5)
plt.rcParams["figure.figsize"] = (12,6)

# Plot the learning curve.
plt.plot(df_stats['Training Loss'], 'b-o', label="Training")
plt.plot(df_stats['Valid. Loss'], 'g-o', label="Validation")

# Label the plot.
plt.title("Training & Validation Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.xticks([1, 2, 3, 4])

plt.show()

import pandas as pd

# ======== Load Test Set ========

# Load the dataset into a pandas dataframe.
df = pd.read_csv("./cola_public/raw/out_of_domain_dev.tsv", delimiter='\t', header=None, names=['sentence_source', 'label', 'label_notes', 'sentence'])

# Report the number of sentences.
print('Number of test sentences: {:,}\n'.format(df.shape[0]))

# Create sentence and label lists
sentences = df.sentence.values
labels = df.label.values

print('Positive samples: %d of %d (%.2f%%)' % (df.label.sum(),
                                               len(df.label),
                                            (df.label.sum() / len(df.label) * 100.0)))

labeled_sentences = []
labels_as_ids = []

# For each sentence in the dataset...
for i in range(len(sentences)):

    sentence = sentences[i]
    label_val = labels[i]

    # Map the numerical label (0, 1) to the word we chose.
    label_word = label_val_to_word[label_val]

    # Look up the token id for the label.
    label_id = label_val_to_token_id[label_val]

    # Insert the sample and its label into the template.
    labeled_sentence = prompt_template.format(
        sentence = sentence,
        label_word = label_word
    )

    # Add to our new lists.
    labeled_sentences.append(labeled_sentence)
    labels_as_ids.append(label_id)

print("Here's what they look like now:\n")
print("'{:}'".format(labeled_sentences[0]))

# ======== Tokenize ========

# The tokenizer is a "callable object"--this invokes its __call__ function,
# which will tokenize and encode all of the input strings.
test_encodings = tokenizer(

    labeled_sentences, # List of strings.

    padding = 'longest',  # Pad out all of the samples to match the longest one
                          # in the data.

    #max_length = 64,      # An alternative strategy is to specify a maximum
    #padding='max_length', # length, but it makes sense to let the tokenizer
                           # figure that out.

    truncation = True, # Truncate any samples longer than the model's maximum
                       # sequence length.

    add_special_tokens = True, # Add the bos and eos tokens.
    return_token_type_ids = False, # These were used in BERT, but not in Mistral.
    return_attention_mask = True, # Mistral uses attention masks.

    return_tensors = "pt" # Return the results as pytorch tensors.
)

# I'll add the prefix 'all' to these variables, since they still contain both
# the training and validation data.
test_input_ids = []
test_attention_masks = []
test_target_words = []
test_label_ids = []
test_label_pos = []

# For each of the encoded samples...
for i in range(len(labels_as_ids)):

    # Extract input_ids and attention_mask
    input_ids = test_encodings['input_ids'][i]
    attention_mask = test_encodings['attention_mask'][i]

    # Find the position of the last non-padding token using the attention mask
    # Because we appended the label to the end of the input, this is the
    # position of our label word.
    label_position = attention_mask.nonzero()[-1].item()

    # This will tell the model what to token to predict at each position.
    # (i.e., at position 12, the model should predict target_words[12])
    # You can set the value to -100 for any tokens you don't want to train on,
    # and in our case, we only want to train on the label.
    # Start by filling it all out with -100s
    target_words = torch.full_like(input_ids, -100)  # Initialize all labels to -100

    # Get the token id for the label
    label_id = labels_as_ids[i]

    # We want all of the words / tokens masked out, except for the label.
    target_words[label_position] = label_id

    # Store everything.
    test_input_ids.append(input_ids)
    test_attention_masks.append(attention_mask)
    test_target_words.append(target_words)
    test_label_pos.append(label_position)
    test_label_ids.append(label_id)

test_dataset = TensorDataset(
    torch.stack(test_input_ids),
    torch.stack(test_attention_masks),
    torch.stack(test_target_words),
    torch.tensor(test_label_ids),
    torch.tensor(test_label_pos)
)

# Set the batch size.
test_batch_size = gpu_batch_size

# Create the DataLoader.
prediction_sampler = SequentialSampler(test_dataset)

prediction_dataloader = DataLoader(
    test_dataset,
    sampler = prediction_sampler,
    batch_size = test_batch_size
)

print('Predicting labels for {:,} test sentences...'.format(len(test_encodings['input_ids'])))

true_labels, predictions, val_loss = evaluate(model, prediction_dataloader)

print('    DONE.')

import numpy as np

from sklearn.metrics import matthews_corrcoef
from sklearn.metrics import confusion_matrix

true_labels = np.asarray(true_labels)
predictions = np.asarray(predictions)

# Report the final accuracy for this validation run.
val_accuracy = float(np.sum(true_labels == predictions)) / len(true_labels) * 100.0
print("Flat accuracy: {0:.2f}%".format(val_accuracy))

# Generate confusion matrix
conf_matrix = confusion_matrix(true_labels, predictions)

# Print or log the confusion matrix
print("\nConfusion Matrix:\n")
print(conf_matrix)
print("")
print("  TP    FP")
print("  FN    TN")

mcc = matthews_corrcoef(true_labels, predictions)

print("\nMCC: {0:.3f}".format(mcc))

# Let's see what you've got, Llama...
logic_test = """
 Llamas can output words. - Incorrect
 Llama 3.1 can output words. - Correct
 Llamas with version numbers output words. - Correct
 My pet llama, Tina, can talk. - Incorrect
 If I upgrade Tina to v3.1, she can talk. -
"""

# Tokenize the input
inputs = tokenizer(logic_test, return_tensors="pt")

# Generate the output from the model
outputs = model.generate(**inputs.to(device), max_new_tokens=1)

# Decode the generated tokens to text
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

# Print the result
print(generated_text)

gpu_mem_train =  gpu_mem_used()

gpu_memory = "14.748 GB"

print("GPU Memory:", gpu_memory)
print("Compressed model:", gpu_mem_model)
print("Forward pass:", gpu_mem_forward_pass)
print("Adding LoRA:", gpu_mem_lora)
print("Training:", gpu_mem_train)

import matplotlib.pyplot as plt
import numpy as np

# Convert the string values to floats
gm_model = float(gpu_mem_model.split()[0])
gm_forward_pass = float(gpu_mem_forward_pass.split()[0])
gm_lora = float(gpu_mem_lora.split()[0])
gm_train = float(gpu_mem_train.split()[0])
gm_total = float(gpu_memory.split()[0])

# Calculate the incremental memory added by each step
memory_additions = [
    gm_model,  # Compressed Model
    gm_forward_pass - gm_model,  # Forward Pass
    gm_lora - gm_forward_pass,   # Adding LoRA
    gm_train - gm_lora           # Training
]

# Update labels with memory added information
labels = [
    f'+{memory_additions[0]:.2f} GB - Compressed Model',
    f'+{memory_additions[1]:.2f} GB - Forward Pass',
    f'+{memory_additions[2]:.2f} GB - Adding LoRA',
    f'+{memory_additions[3]:.2f} GB - Training'
]

# Create a stacked bar plot with updated labels and proper y-limit
fig, ax = plt.subplots(figsize=(2, 5), dpi=150)  # Reduced figure width for a narrower bar, higher resolution

# Plot a horizontal line representing total used.
ax.axhline(y=gm_total, color='b', linestyle='--', label=f'Used: {gpu_mem_train}')

# Plot a horizontal line representing the total GPU memory available.
ax.axhline(y=gm_total, color='r', linestyle='--', label=f'{gpu:>4}: {gpu_memory}')

# Use a stacked bar chart where each bar adds to the total
ax.bar(".", memory_additions[0], label=labels[0], color='skyblue')
ax.bar(".", memory_additions[1], bottom=memory_additions[0], label=labels[1], color='lightgreen')
ax.bar(".", memory_additions[2], bottom=memory_additions[0] + memory_additions[1], label=labels[2], color='orange')
ax.bar(".", memory_additions[3], bottom=memory_additions[0] + memory_additions[1] + memory_additions[2], label=labels[3], color='salmon')


# Adjust the y-limit to make sure the total GPU memory line is visible
ax.set_ylim(0, gm_total + 1)  # Add some space above the total GPU memory

# Add labels and title
ax.set_ylabel('Memory Usage (GB)', fontsize=10)
ax.set_title('GPU Batch: {:}\n      Seq Len: {:}'.format(gpu_batch_size, len(encodings['input_ids'][0])), fontsize=10)

# Reverse the legend order to match the order of the bars, and move it outside the plot
handles, labels = ax.get_legend_handles_labels()

handles = list(reversed(handles))
reversed_labels = list(reversed(labels))

handles = handles[-2:] + handles[:-2]
reversed_labels = reversed_labels[-2:] + reversed_labels[:-2]

ax.legend(handles,
          reversed_labels,
          loc='center left', bbox_to_anchor=(1, 0.5),
          prop={'size': 10, 'family': 'monospace'})

# Adjust x-axis to make the single bar narrower visually
ax.set_xlim(-0.75, 0.75)

ax.tick_params(axis='y', labelsize=10)  # Adjust '8' to your desired font size

# Show the plot
plt.show()

from datetime import datetime

summary = {}

summary["Timestamp"] = datetime.now().strftime("%Y-%m-%d - %H:%M:%S")

summary["Model"] = "LLaMA 3 8B"

# ==== System ====
summary["GPU"] = gpu
summary["GPU Memory"] = gpu_memory
summary["GPU Memory Used"] = gpu_mem_used()
summary["Data Type"] = str(torch_dtype)
summary["Memory for Model"] = gpu_mem_model
summary["Memory after Forward Pass"] = gpu_mem_forward_pass
summary["Memory after Adding LoRA"] = gpu_mem_lora
summary["Memory after Training"] = gpu_mem_train

# ==== Training Parameters ====
summary["Maximum Sequence Length"] = len(encodings['input_ids'][0])
summary["Effective Batch Size"] = train_batch_size
summary["GPU Batch Size"] = gpu_batch_size
#summary["Accumulate Batches"] = accumulate_passes
summary["Gradient Checkpointing"] = False
summary["Optimizer"] = str(optimizer_class)
summary["Raw Learning Rate"] = lr
summary["Learning Rate"] = format_lr_as_multiple(lr)
summary["Learning Rate x LoRA"] = format_lr_as_multiple(
    lr * lora_config.lora_alpha / lora_config.r
)
# Include the ' to tell sheets to treat this as a literal string and include the
# surrounding quotes.
summary["Formatted Example"] = "'\"{:}\"".format(labeled_sentences[0])

#summary["Weight Decay"] = training_args.weight_decay
#summary["Scheduler"] = training_args.scheduler
#summary["Steps"] = training_args.max_steps

# ==== QLoRA Parameters ====
summary["LoRA r"] = lora_config.r
summary["LoRA alpha"] = lora_config.lora_alpha
summary["LoRA Dropout"] = lora_config.lora_dropout
#summary["LoRA Targets"] = lora_config.target_modules
summary["Quantization"] = "4-bit"

num_records = len(df_stats)

# ==== Results ====
summary["Training Loss"] = df_stats['Training Loss'][num_records - 1]
#summary["Valid. Loss"] = df_stats['Valid. Loss'][epochs]
summary["Training time"] = training_time
summary["MCC"] = mcc
summary["Few-Shot MCC"] = few_shot_mcc


# Convert the summary dictionary to a DataFrame
summary_df = pd.DataFrame(list(summary.items()), columns=['Metric', 'Value'])

summary_df

!pip install gspread

from google.colab import auth
auth.authenticate_user()

import gspread
from google.auth import default
creds, _ = default()

gc = gspread.authorize(creds)

# Open the Google Sheet by its id
# This is 'Fine-Tuning StackLLaMA - Run Summaries'
all_sheets = gc.open_by_key("1-EaX_HjYxZSU6BwJ1WT2eEC41nIYGeWoIjlYfD09_0s")

# Second sheet is for this CoLA notebook
sheet = all_sheets.worksheets()[1]

all_sheets.worksheets()

import pandas as pd

# Read the existing data in the sheet
existing_data = sheet.get_all_values()
existing_df = pd.DataFrame(existing_data)

# Determine the next available column
next_col = len(existing_df.columns) + 1

# Append the summary data to the sheet
for index, row in summary_df.iterrows():

    # Find the 0-indexed row number for this metric from the summary table.
    index_obj = existing_df[existing_df[0] == row['Metric']].index

    # If the index object is empty, the metric wasn't found.
    if len(index_obj) == 0:
        print("Metric not found:", row['Metric'])

    else:
        # The spreadsheet row number is 1-indexed.
        sheet_row = index_obj[0] + 1

        print(row['Metric'], '   ', row['Value'])

        # Update the cell in the next available column
        sheet.update_cell(sheet_row, next_col, row['Value'])