Fine-Tuning LLaMa 2 for Text Summarization

from warnings import filterwarnings
filterwarnings('ignore')
import json
import re
import wandb
import os
from pprint import pprint
import pandas as pd
import torch
from kaggle_secrets import UserSecretsClient
from datasets import Dataset, load_dataset
from huggingface_hub import notebook_login, login
from peft import LoraConfig, PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from trl import SFTTrainer

import wandb

# Start a W&B run
wandb.init(project="dialog-summarization", entity="enter your w&b name here")

DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
MODEL_NAME = "/kaggle/input/llama-2/pytorch/7b-hf/1"

df = pd.read_csv("/kaggle/input/dialogsum/CSV/train.csv", nrows=500)
df.columns = [str(q).strip() for q in df.columns]

dataset = Dataset.from_pandas(df)

DEFAULT_SYSTEM_PROMPT = """
Below is a conversation between a human and an AI agent. Write a summary of the conversation.
""".strip()


def generate_training_prompt(
    conversation: str, summary: str, system_prompt: str = DEFAULT_SYSTEM_PROMPT
) -> str:
    return f"""### Instruction: {system_prompt}

### Input:
{conversation.strip()}

### Response:
{summary}
""".strip()

def create_conversation_text(data_point):
    return data_point["dialogue"]

def generate_text(data_point):
    summary = data_point["summary"]
    conversation_text = create_conversation_text(data_point)
    return {
        "conversation": conversation_text,
        "summary": summary,
        "text": generate_training_prompt(conversation_text, summary),
    }

# Example usage with a new dataset format
example_data_point = {
    "id": "train_0",
    "dialogue": "#Person1#: Hi, Mr. Smith. I'm Doctor Hawkins. Why are you here today? #Person2#: I found it would...",
    "summary": "Mr. Smith's getting a check-up, and Doctor Hawkins advises him to have one every year. Hawkins'll gi...",
    "topic": "get a check-up"
}

example = generate_text(example_data_point)
print(example["text"])

# Example usage with a new dataset format
example_data_point = {
    "id": "train_0",
    "dialogue": "#Person1#: Hi, Mr. Smith. I'm Doctor Hawkins. Why are you here today? #Person2#: I found it would...",
    "summary": "Mr. Smith's getting a check-up, and Doctor Hawkins advises him to have one every year. Hawkins'll gi...",
    "topic": "get a check-up"
}

example = generate_text(example_data_point)
print(example["text"])

from datasets import Dataset

def process_dataset(data: Dataset) -> Dataset:
    """
    This function processes the dataset to include only the necessary columns.
    """
    # First, apply generate_text to each record in the dataset
    processed_data = data.map(generate_text)

    # Then, remove unnecessary columns
    columns_to_remove = [col for col in processed_data.column_names if col not in ["conversation", "summary", "text"]]
    return processed_data.remove_columns(columns_to_remove)

# Process the entire dataset
processed_dataset = process_dataset(dataset)

# Split the processed dataset into train, validation, and test sets
train_dataset = processed_dataset.shuffle(seed=42).select(range(0, int(0.8 * len(processed_dataset))))
validation_dataset = processed_dataset.shuffle(seed=42).select(range(int(0.8 * len(processed_dataset)), int(0.9 * len(processed_dataset))))
test_dataset = processed_dataset.shuffle(seed=42).select(range(int(0.9 * len(processed_dataset)), len(processed_dataset)))

dataset

def create_model_and_tokenizer():
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        use_safetensors=True,
        quantization_config=bnb_config,
        trust_remote_code=True,
        device_map="auto",
    )

    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    return model, tokenizer

model, tokenizer = create_model_and_tokenizer()
model.config.use_cache = False

model.config.quantization_config.to_dict()

lora_alpha = 32
lora_dropout = 0.05
lora_r = 16

peft_config = LoraConfig(
    lora_alpha=lora_alpha,
    lora_dropout=lora_dropout,
    r=lora_r,
    bias="none",
    task_type="CAUSAL_LM",
)

def generate_prompt(
    conversation: str, system_prompt: str = DEFAULT_SYSTEM_PROMPT
) -> str:
    return f"""### Instruction: {system_prompt}

### Input:
{conversation.strip()}

### Response:
""".strip()


def summarize(model, text: str):
    inputs = tokenizer(text, return_tensors="pt").to(DEVICE)
    inputs_length = len(inputs["input_ids"][0])
    with torch.inference_mode():
        outputs = model.generate(**inputs, max_new_tokens=256, temperature=0.0001)
    return tokenizer.decode(outputs[0][inputs_length:], skip_special_tokens=True)

def generate_summaries(model, dataset, tokenizer, num_samples=5):
    summaries = []
    for i, example in enumerate(dataset):
        if i >= num_samples:
            break
        print(i)
        prompt = generate_prompt(example['conversation'])
        summary = summarize(model, prompt)
        summaries.append({'conversation': example['conversation'], 'generated_summary': summary})
    return summaries

# Generate summaries before fine-tuning
original_summaries = generate_summaries(model, test_dataset, tokenizer, num_samples=5)

# Convert to DataFrame and log to W&B
df_original = pd.DataFrame(original_summaries)
wandb.log({"original_summaries": wandb.Table(dataframe=df_original)})

OUTPUT_DIR = "dialog-summarization-llama-2-finetuned"

training_arguments = TrainingArguments(
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    optim="paged_adamw_32bit",
    logging_steps=1,
    learning_rate=1e-4,
    fp16=True,
    max_grad_norm=0.3,
    num_train_epochs=20,
    evaluation_strategy="epoch",
    eval_steps=0.2,
    warmup_ratio=0.05,
    save_strategy="epoch",
    group_by_length=True,
    output_dir=OUTPUT_DIR,
    report_to="wandb",  # Set report_to here
    save_safetensors=True,
    lr_scheduler_type="cosine",
    seed=42,
    load_best_model_at_end=True,
    push_to_hub=True,
)

login(token="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")

trainer = SFTTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=validation_dataset,
    peft_config=peft_config,
    dataset_text_field="text",
    max_seq_length=4096,
    tokenizer=tokenizer,
    args=training_arguments,
)

# Fine-tune your model
trainer.train()

# Generate summaries after fine-tuning
fine_tuned_summaries = generate_summaries(trainer.model, test_dataset, tokenizer, num_samples=5)

# Convert to DataFrame and log to W&B
df_fine_tuned = pd.DataFrame(fine_tuned_summaries)
wandb.log({"fine_tuned_summaries": wandb.Table(dataframe=df_fine_tuned)})