day2/llm-day2-notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "41b1710c-d5d9-4dcc-954e-2938ad1f32cf",
   "metadata": {},
   "source": [
    "# BUILD LLM Bootcamp Day 2\n",
    "\n",
    "## Setup\n",
    "\n",
    "### Install Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec84f4d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install snowflake-ml-python==1.0.12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "876bbda5-b330-4d71-a335-ac178274b6ff",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install peft transformers==4.34.0 tokenizers vllm==0.2.1.post1 bitsandbytes datasets absl-py==1.3.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d875a77-e12b-4fb7-bbef-e54e8b2113b8",
   "metadata": {},
   "source": [
    "### Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29fb24a1-da33-43e2-9924-c9ccae223717",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import Dataset, ClassLabel\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "import sys\n",
    "from utils import Concatenator\n",
    "import pandas as pd\n",
    "pd.set_option('display.max_colwidth', None)\n",
    "\n",
    "import torch\n",
    "from transformers import LlamaForCausalLM, LlamaTokenizer\n",
    "import os\n",
    "import json\n",
    "from transformers import TrainerCallback\n",
    "from contextlib import nullcontext\n",
    "from transformers import default_data_collator, Trainer, TrainingArguments\n",
    "\n",
    "from snowflake.snowpark.session import Session\n",
    "from snowflake.snowpark import VERSION\n",
    "import snowflake.snowpark.functions as F\n",
    "from snowflake.ml.registry import model_registry\n",
    "from snowflake.ml.model import deploy_platforms\n",
    "from snowflake.ml.model.models import llm\n",
    "\n",
    "import logging \n",
    "logger = logging.getLogger(\"snowflake.snowpark.session\")\n",
    "logger.setLevel(logging.ERROR)\n",
    "logger = logging.getLogger(\"snowflake.ml\")\n",
    "logger.setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0306b93f-c225-40ea-8480-10739c887b4d",
   "metadata": {},
   "source": [
    "### Load Base Model and Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a599bfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Set your Hugging Face token\n",
    "!huggingface-cli login --token YOUR_HG_TOKEN_GOES_HERE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21622331-74b3-48b7-9b8f-5f1b7fba3e74",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_id=\"meta-llama/Llama-2-7b-chat-hf\"\n",
    "print('loading tokenizer')\n",
    "tokenizer = LlamaTokenizer.from_pretrained(model_id)\n",
    "print('loading model')\n",
    "model =LlamaForCausalLM.from_pretrained(model_id, load_in_8bit=True, device_map='auto', torch_dtype=torch.float16)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "419a142a-182d-4a10-a92d-85410828226a",
   "metadata": {},
   "source": [
    "## Prepare Datasets\n",
    "### Load Datasets\n",
    "\n",
    "*IMP NOTE: Please do NOT not run the following cell during the bootcamp session. Fine-tuning on the entire dataset could take over an hour.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c1c2df6-84e2-4f5d-a019-4a31625e65bb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# def prepare_stratified_dataset(path, seed = 42):\n",
    "#     raw_df = pd.read_json(path, lines=True)\n",
    "#     raw_df['id'] = raw_df.index\n",
    "#     ds = Dataset.from_pandas(raw_df, split='train')\n",
    "#     cl = ClassLabel(num_classes=4, names=[\"EN\", \"FR\", \"DE\", \"ES\"])\n",
    "#     new_features = ds.features.copy()\n",
    "#     new_features['lang_label'] = cl\n",
    "#     cl_d = {l : cl.str2int(l) for l in [\"EN\", \"FR\", \"DE\", \"ES\"]}\n",
    "#     def convert_lang(sample):\n",
    "#         sample['lang_label'] = cl_d[sample['language']]\n",
    "#         return sample\n",
    "#     ds = ds.map(convert_lang, features=new_features)\n",
    "#     ds_split = ds.train_test_split(test_size=0.15, stratify_by_column='lang_label', seed=42)\n",
    "#     test_ds_split = ds_split['test'].train_test_split(test_size=2/3, stratify_by_column='lang_label', seed=42)\n",
    "#     return ds_split['train'].to_pandas(), test_ds_split['train'].to_pandas(), test_ds_split['test'].to_pandas()\n",
    "\n",
    "# train_df, eval_df, test_df = prepare_stratified_dataset(\n",
    "#     'transcripts.json'\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02ce0d6a-52f8-4a98-bb64-dffd9de39f1d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = pd.read_json(\"transcripts.json\", lines=True)\n",
    "print(f\"Total records: {df.shape[0]}\")\n",
    "train_df = df.head(100)\n",
    "print(f\"Train        : {train_df.shape[0]}\")\n",
    "eval_df = df[200:300]\n",
    "print(f\"Eval         : {eval_df.shape[0]}\")\n",
    "test_df = df.tail(100)\n",
    "print(f\"Test         : {test_df.shape[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48eb4986-85f7-4465-b56c-ef35f7edfa95",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79793829-d70f-4bed-93b2-1a4db2fbfaea",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "datasets = {\n",
    "    'train': Dataset.from_pandas(train_df),\n",
    "    'eval': Dataset.from_pandas(eval_df),\n",
    "    'test': Dataset.from_pandas(test_df)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cceb080d-a35e-4629-88d9-46c1d97c90dd",
   "metadata": {},
   "source": [
    "### Apply Prompt to Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e684e33f-d7db-4418-bab5-c3d490067fa3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_prompt = f\"\"\"\n",
    "[INST] <<SYS>>\n",
    "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n",
    "<</SYS>>\n",
    "### Instruction:\n",
    "{{instruction}}\n",
    "### Input:\n",
    "{{input_}}\n",
    "### Output:\n",
    "{{output}}\n",
    "{{eos_token}}\n",
    "\"\"\"\n",
    "\n",
    "eval_prompt = f\"\"\"\n",
    "[INST] <<SYS>>\n",
    "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n",
    "<</SYS>>\n",
    "### Instruction:\n",
    "{{instruction}}\n",
    "### Input:\n",
    "{{input_}}\n",
    "### Output:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01ff73e9-780a-44cd-b792-e351a9601f5b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def apply_train_template(sample):\n",
    "    return {\n",
    "        \"text\": train_prompt.format(\n",
    "            instruction=sample[\"instruction\"],\n",
    "            input_=sample[\"input\"].replace('\\\\n', '\\n'),\n",
    "            output=sample[\"output\"],\n",
    "            eos_token=tokenizer.eos_token,\n",
    "        )\n",
    "    }\n",
    "\n",
    "def apply_eval_template(sample):\n",
    "    return {\n",
    "        \"text\": eval_prompt.format(\n",
    "            instruction=sample[\"instruction\"],\n",
    "            input_=sample[\"input\"].replace('\\\\n', '\\n')\n",
    "        )\n",
    "    }\n",
    "\n",
    "\n",
    "\n",
    "#applying template\n",
    "\n",
    "datasets['train'] = datasets['train'].map(apply_train_template, remove_columns=list(datasets['train'].features))\n",
    "for k in ['eval', 'test']:\n",
    "    datasets[k] = datasets[k].map(apply_eval_template, remove_columns=list(datasets[k].features))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac551d54-0591-43bc-a96a-da5edd3aa565",
   "metadata": {},
   "source": [
    "### Tokenize Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2290f191-ba23-4b83-8653-3d9c97d512cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for k, v in datasets.items():\n",
    "    datasets[k] = v.map(\n",
    "        lambda sample: tokenizer(sample[\"text\"]),\n",
    "        batched=True,\n",
    "        remove_columns=list(v.features),\n",
    "    ).map(Concatenator(), batched=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88447074-21a8-433a-a937-a3e3514650c4",
   "metadata": {},
   "source": [
    "## Fine-tune Llama 2 Model\n",
    "\n",
    "\n",
    "Fine-tuning is one form of model training. We start training from a pre-trained model and adjust a set of model parameters to better solve for a concrete task based on task specific data. Today we are going to fine-tune 7B Llama 2 model using LoRA (Low-Rank Adaptation)--which is a parameter efficient way of fine-tuning LLM. \n",
    "\n",
    "Instead of adjusting all the ~7B parameters, LoRA allows us to adjust only a percent of model weights--which can save compute and memory resources dramatically. For this lab, we will fine-tune our model using LoRA on a single A10 GPU. This will demostrate how good the inference can be on fine-tuned models even with limited compute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba1f466-d668-4f68-9d8d-ea7367b0369d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#setting the model into training mode\n",
    "model.train()\n",
    "\n",
    "#setting up training\n",
    "def create_peft_config(model):\n",
    "    from peft import (\n",
    "        get_peft_model,\n",
    "        LoraConfig,\n",
    "        TaskType,\n",
    "        prepare_model_for_kbit_training,\n",
    "    )\n",
    "\n",
    "    peft_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False,\n",
    "        r=8,\n",
    "        lora_alpha=32,\n",
    "        lora_dropout=0.05,\n",
    "        target_modules = [\"q_proj\", \"v_proj\"]\n",
    "    )\n",
    "\n",
    "    # prepare int-8 model for training\n",
    "    model = prepare_model_for_kbit_training(model)\n",
    "    model = get_peft_model(model, peft_config)\n",
    "    model.print_trainable_parameters()\n",
    "    return model, peft_config\n",
    "\n",
    "# create peft config\n",
    "model, lora_config = create_peft_config(model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b15d38d4-c399-4daa-9771-df29edf8c3b9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# rm -r output_weights_dir # deletes prior fine tuning weights\n",
    "!mkdir output_weights_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9649d974-40bb-4679-b5e7-e9da3f9cc75e",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Configuration\n",
    "\n",
    "To achieve good performance for the task, you will need at least 1 `num_epochs`, feel free to explore this on your own and we've set `output_weights_dir` as the directory where the fine-tuned weights will be stored after fine-tuning job completes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "328c128c-5321-4132-a9b6-da362f62790c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "output_dir = \"output_weights_dir\"\n",
    "enable_profiler = False\n",
    "\n",
    "config = {\n",
    "    'lora_config': lora_config,\n",
    "    'learning_rate': 1e-4,\n",
    "    'num_train_epochs': 1,\n",
    "    'gradient_accumulation_steps': 2,\n",
    "    'per_device_train_batch_size': 2,\n",
    "    'gradient_checkpointing': False,\n",
    "}\n",
    "\n",
    "# Define training args\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    overwrite_output_dir=True,\n",
    "    bf16=True,  # Use BF16 if available\n",
    "    # logging strategies\n",
    "    logging_dir=f\"{output_dir}/logs\",\n",
    "    logging_strategy=\"steps\",\n",
    "    logging_steps=10,\n",
    "    save_strategy=\"no\",\n",
    "    optim=\"adamw_torch_fused\",\n",
    "    max_steps=total_steps if enable_profiler else -1,\n",
    "    **{k:v for k,v in config.items() if k != 'lora_config'}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "674205dc-ded7-45eb-9450-a767ae5cfe85",
   "metadata": {},
   "source": [
    "### Fine-tuning\n",
    "\n",
    "#### Configuration\n",
    "\n",
    "We're passing `train_dataset` and `eval_dataset` that are used to generate loss calculation during fine-tuning process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4188f41a-0513-4848-aab6-fb0226d557f4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "profiler = nullcontext() \n",
    "\n",
    "with profiler:\n",
    "    # Create Trainer instance\n",
    "    trainer = Trainer(\n",
    "        model=model,\n",
    "        args=training_args,\n",
    "        train_dataset=datasets['train'],\n",
    "        eval_dataset=datasets['eval'],\n",
    "        data_collator=default_data_collator,\n",
    "        callbacks=[profiler_callback] if enable_profiler else [],\n",
    "    )\n",
    "\n",
    "    # Start training\n",
    "    trainer.train()\n",
    "    \n",
    "model.save_pretrained(output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29e0f65d-5941-4b65-81e8-67dccf4c703b",
   "metadata": {},
   "source": [
    "## Log and Deploy Fine-tuned Llama 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d57af43-9b3c-45ad-9ce3-eefdf5b79d99",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Establish Secure Connection\n",
    "\n",
    "*NOTE: Update [connection.json](../connection.json) and set your password, Hugging Face token, and replace '####' with your user number.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "318991b7-c8f8-4c64-ba7f-d26be720e087",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create Snowflake Session object\n",
    "connection_parameters = json.load(open('connection.json'))\n",
    "session = Session.builder.configs(connection_parameters).create()\n",
    "session.sql_simplifier_enabled = True\n",
    "\n",
    "snowflake_environment = session.sql('select current_user(), current_version()').collect()\n",
    "snowpark_version = VERSION\n",
    "\n",
    "# Current Environment Details\n",
    "print('User                        : {}'.format(snowflake_environment[0][0]))\n",
    "print('Role                        : {}'.format(session.get_current_role()))\n",
    "print('Database                    : {}'.format(session.get_current_database()))\n",
    "print('Schema                      : {}'.format(session.get_current_schema()))\n",
    "print('Warehouse                   : {}'.format(session.get_current_warehouse()))\n",
    "print('Snowflake version           : {}'.format(snowflake_environment[0][1]))\n",
    "print('Snowpark for Python version : {}.{}.{}'.format(snowpark_version[0],snowpark_version[1],snowpark_version[2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e9cc712-2d3d-41c5-b417-d18d0f2712e1",
   "metadata": {},
   "source": [
    "### Model Registery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d42c0c2-39d7-4024-b4ce-a23558ef124e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "MODEL_NAME = \"LLAMA2_7b_CHAT\"\n",
    "MODEL_VERSION = \"FineTunedV1.1\"\n",
    "DEPLOYMENT_NAME = \"FINETUNED_LLAMA2\"\n",
    "MODEL_REGISTRY_DB = connection_parameters['database']\n",
    "MODEL_REGISTRY_SCHEMA = connection_parameters['schema']\n",
    "COMPUTE_POOL = connection_parameters['compute_pool']\n",
    "\n",
    "registry = model_registry.ModelRegistry(\n",
    "    session=session, \n",
    "    database_name=MODEL_REGISTRY_DB, \n",
    "    schema_name=MODEL_REGISTRY_SCHEMA, \n",
    "    create_if_not_exists=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3b75911-a7b7-4226-8b71-23a4a6dd9eda",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "registry.list_deployments(model_name=MODEL_NAME,model_version=MODEL_VERSION).to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b26134b3-9f3b-4834-91e6-d3115a605502",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# registry.delete_deployment(model_name=MODEL_NAME,model_version=MODEL_VERSION,deployment_name=DEPLOYMENT_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bc2f21a-cfcb-459b-afe8-c859acad9a2a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# registry.delete_model(model_name=MODEL_NAME,model_version=MODEL_VERSION,delete_artifact=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b9ef244-2df3-46e5-ac6d-b4c40bb56fc5",
   "metadata": {},
   "source": [
    "### Reference Llama 2 from Fine-tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ff039a5-36b9-4642-a242-b67e75a66e40",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "#referencing our fine tuned model\n",
    "#referencing huggingface token\n",
    "options = llm.LLMOptions(\n",
    "    token=connection_parameters['huggingface_token'],\n",
    "    max_batch_size=100,\n",
    ")\n",
    "#referencing our fine tuned weights and using the hugging face token to merge with base llama model\n",
    "llama_model = llm.LLM(\n",
    "    model_id_or_path='output_weights_dir',\n",
    "    options=options\n",
    ")\n",
    "\n",
    "# log model in registry\n",
    "llama_model_ref = registry.log_model(\n",
    "    model_name=MODEL_NAME,\n",
    "    model_version=MODEL_VERSION,\n",
    "    model=llama_model\n",
    ")\n",
    "\n",
    "# deploy model\n",
    "llama_model_ref.deploy(\n",
    "    deployment_name=DEPLOYMENT_NAME, \n",
    "    platform=deploy_platforms.TargetPlatform.SNOWPARK_CONTAINER_SERVICES,\n",
    "    permanent=True, \n",
    "   options={\"compute_pool\": COMPUTE_POOL, \"num_gpus\": 1})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae82d020-a7ce-4a88-858f-5467b88406db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# if its already deployed run the following:\n",
    "llama_model_ref = model_registry.ModelReference(\n",
    "    registry=registry, \n",
    "    model_name=MODEL_NAME, \n",
    "    model_version=MODEL_VERSION)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3654a572-97fb-48bb-a1b3-510a5d62690e",
   "metadata": {},
   "source": [
    "## Inference on Eval Dataset using fine-tuned Llama 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd360c91-5df0-44ac-bbb6-b07f13bac9a1",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "eval_df['input'] = eval_df.apply(\n",
    "    lambda x: eval_prompt.format(\n",
    "        instruction=x[\"instruction\"],\n",
    "        input_=x[\"input\"].replace('\\\\n', '\\n')\n",
    "    ), axis=1\n",
    ")\n",
    "eval_df.reset_index(drop=True, inplace=True)\n",
    "eval_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd02bca2-f2be-4d2e-b41c-9d0c6775cb35",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "eval_df['predicted'] = llama_model_ref.predict(deployment_name=DEPLOYMENT_NAME,data=eval_df)#.head()\n",
    "eval_df[['output', 'predicted']].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57ac4d8d-72c0-40ab-ae3d-54ac20f6874e",
   "metadata": {},
   "source": [
    "## Clean Up Resources\n",
    "\n",
    "Delete deployment and the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a53426a-93a8-4bb2-88d1-da6a885b9fd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_model_ref.delete_deployment(deployment_name=DEPLOYMENT_NAME)\n",
    "llama_model_ref.delete_model(delete_artifact=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}