mistral/mistral-small-24b-fsdp-qlora-remote-decorator_qa.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f50a1de6-c757-4392-a48d-e5840348537c",
   "metadata": {},
   "source": [
    "# Fine-tune Mistral-Small-24B-Instruct-2501 with PyTorch FSDP, QLoRA and SageMaker remote decorator\n",
    "\n",
    "## Instruction-based fine-tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a54122cf-17ff-4d98-bf4a-ac2330e8e63f",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "In this demo notebook, we demonstrate how to fine-tune the Mistral-Small-24B-Instruct-2501 model using PyTorch FSDP, QLoRA, Hugging Face PEFT, and bitsandbytes.\n",
    "\n",
    "We are using SageMaker remote decorator for runinng the fine-tuning job on Amazon SageMaker Training job\n",
    "---\n",
    "\n",
    "JupyterLab Instance Type: ml.t3.medium\n",
    "\n",
    "Python version: 3.11\n",
    "\n",
    "Fine-Tuning:\n",
    "* Instance Type: ml.g5.12xlarge\n",
    "* Instance Count: 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac5833a9-a34a-4c5c-a25f-7bcec69317f6",
   "metadata": {},
   "source": [
    "Install the required libriaries, including the Hugging Face libraries, and restart the kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "505031fd-34d1-4261-a3f6-88ede3fd16a7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-26T09:07:37.160628Z",
     "start_time": "2024-01-26T09:07:24.214358Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt --upgrade\n",
    "%pip install -q -U python-dotenv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95c59da5-3c66-48a1-8436-e59878e218eb",
   "metadata": {},
   "source": [
    "## Setup Configuration file path\n",
    "\n",
    "We are setting the directory in which the config.yaml file resides so that remote decorator can make use of the settings through [SageMaker Defaults](https://sagemaker.readthedocs.io/en/stable/overview.html#configuring-and-using-defaults-with-the-sagemaker-python-sdk).\n",
    "\n",
    "This notebook is using the Hugging Face container for the `us-east-1` region. Make sure you are using the right image for your AWS region, otherwise edit [config.yaml](./config.yaml). Container Images are available [here](https://github.com/aws/deep-learning-containers/blob/master/available_images.md)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95f6fd4d-c4e6-4439-bff5-0ffe05cc5260",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-15T09:24:51.291677Z",
     "start_time": "2023-11-15T09:24:51.282905Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "import os\n",
    "\n",
    "# Use .env in case of hidden variables\n",
    "load_dotenv()\n",
    "\n",
    "# Set path to config file\n",
    "os.environ[\"SAGEMAKER_USER_CONFIG_OVERRIDE\"] = os.getcwd()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82089d28-b97a-4956-83fb-d8c46d44fdb5",
   "metadata": {},
   "source": [
    "## Visualize and upload the dataset\n",
    "\n",
    "We are going to load [NousResearch/hermes-function-calling-v1](https://huggingface.co/datasets/NousResearch/hermes-function-calling-v1) dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63c1c1f7-eec0-441b-88be-ad6ecd19352e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ff9c879-adec-4f5d-a5f5-cce7f7597407",
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_session = sagemaker.Session()\n",
    "bucket_name = sagemaker_session.default_bucket()\n",
    "default_prefix = sagemaker_session.default_bucket_prefix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f6da4af-4041-4e63-b2e8-a085e4e62652",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "import pandas as pd\n",
    "\n",
    "dataset = load_dataset(\"NousResearch/hermes-function-calling-v1\", data_files={'train': ['json-mode-agentic.json']})\n",
    "df = pd.DataFrame(dataset['train'])\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "766e8efe-4dbc-4f32-a439-67f1f45b3fa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "train, test = train_test_split(df, test_size=0.1, random_state=42)\n",
    "\n",
    "print(\"Number of train elements: \", len(train))\n",
    "print(\"Number of test elements: \", len(test))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42a9132f",
   "metadata": {},
   "source": [
    "Create a prompt template based on the model's prompt style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf447f66",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_conversation(conversation):\n",
    "    transformed = []\n",
    "    for msg in conversation:\n",
    "        # Create a new dictionary with the renamed keys\n",
    "        new_msg = {\n",
    "            'role': 'user' if msg['from'] == 'human' \n",
    "                   else 'assistant' if msg['from'] == 'gpt' \n",
    "                   else 'system',\n",
    "            'content': msg['value']\n",
    "        }\n",
    "        transformed.append(new_msg)\n",
    "    return transformed\n",
    "\n",
    "train['conversations'] = train['conversations'].apply(transform_conversation)\n",
    "test['conversations'] = test['conversations'].apply(transform_conversation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c050eb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"mistralai/Mistral-Small-24B-Instruct-2501\")\n",
    "\n",
    "train['text'] = train['conversations'].apply(lambda x: tokenizer.apply_chat_template(x, tokenize=False))\n",
    "test['text'] = test['conversations'].apply(lambda x: tokenizer.apply_chat_template(x, tokenize=False))\n",
    "\n",
    "train = train[['text']]\n",
    "test = test[['text']]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e7f0a32-96c2-4285-b593-3de211f79933",
   "metadata": {},
   "source": [
    "Use the Hugging Face Trainer class to fine-tune the model. Define the hyperparameters we want to use. We also create a DataCollator that will take care of padding our inputs and labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e9cbedd-7403-467e-8cc6-1d2550d8b8e8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-03T00:02:10.364459Z",
     "start_time": "2023-09-03T00:02:09.672705Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import Dataset, DatasetDict\n",
    "from random import randint\n",
    "\n",
    "train_dataset = Dataset.from_pandas(train)\n",
    "test_dataset = Dataset.from_pandas(test)\n",
    "\n",
    "print(train_dataset[randint(0, len(dataset))][\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bd66949-f372-45c3-9a8e-77963207c8e5",
   "metadata": {},
   "source": [
    "Utility function for initializing the distribution across multiple GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3df38ffa-f093-4916-b516-dd8c9b54e490",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "def init_distributed():\n",
    "    # Initialize the process group\n",
    "    torch.distributed.init_process_group(\n",
    "        backend=\"nccl\", # Use \"gloo\" backend for CPU\n",
    "        timeout=datetime.timedelta(seconds=5400)\n",
    "    )\n",
    "    local_rank = int(os.environ[\"LOCAL_RANK\"])\n",
    "    torch.cuda.set_device(local_rank)\n",
    "\n",
    "    return local_rank"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee438299-2c6c-4ea7-9648-88852176672a",
   "metadata": {},
   "source": [
    "Utility function for model download"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6aeb2674-2932-4df8-bb0f-567a6c0b12d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import snapshot_download\n",
    "import os\n",
    "\n",
    "def download_model(model_name):\n",
    "    print(\"Downloading model \", model_name)\n",
    "\n",
    "    os.makedirs(\"/tmp/tmp_folder\", exist_ok=True)\n",
    "\n",
    "    snapshot_download(repo_id=model_name, local_dir=\"/tmp/tmp_folder\")\n",
    "\n",
    "    print(f\"Model {model_name} downloaded under /tmp/tmp_folder\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d7c6cf2-5b9b-4c34-8839-9d482d982948",
   "metadata": {},
   "source": [
    "To train our model, we need to convert our inputs (text) to token IDs. This is done by a Hugging Face Transformers Tokenizer. In addition to Lora, we will use bitsanbytes 4-bit precision to quantize out frozen LLM to 4-bit and attach LoRA adapters on it.\n",
    "\n",
    "Define the train function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6ee05aa-df07-4824-b13a-0d4395de460c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"mistralai/Mistral-Small-24B-Instruct-2501\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10cabb4d-b0b2-498c-95cb-41ed7d05ee65",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-03T00:02:21.382486Z",
     "start_time": "2023-09-03T00:02:20.962208Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from accelerate import Accelerator\n",
    "import datetime\n",
    "from huggingface_hub import snapshot_download\n",
    "import os\n",
    "from peft import AutoPeftModelForCausalLM, LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
    "from sagemaker.remote_function import remote\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, set_seed\n",
    "import transformers\n",
    "import traceback\n",
    "\n",
    "# Start training\n",
    "@remote(\n",
    "    keep_alive_period_in_seconds=0, #Warm-pool instance. Put 0 for avoiding additional costs\n",
    "    instance_type=\"ml.g5.12xlarge\",\n",
    "    instance_count=2,\n",
    "    volume_size=100,\n",
    "    job_name_prefix=f\"train-{model_id.split('/')[-1].replace('.', '-')}\",\n",
    "    use_torchrun=True,\n",
    "    disable_output_compression=True,\n",
    ")\n",
    "def train_fn(\n",
    "    model_name,             # Name or path of the base model to fine-tune\n",
    "    train_ds,               # Training dataset\n",
    "    test_ds=None,           # Optional test/validation dataset\n",
    "    torch_dtype=torch.bfloat16,  # Precision type for training\n",
    "    lora_r=8,               # LoRA rank - controls capacity of adaptations\n",
    "    lora_alpha=16,          # LoRA alpha - scales the adaptations\n",
    "    lora_dropout=0.1,       # Dropout probability for LoRA layers\n",
    "    per_device_train_batch_size=8,  # Batch size for training\n",
    "    per_device_eval_batch_size=8,   # Batch size for evaluation\n",
    "    gradient_accumulation_steps=1,  # Number of steps to accumulate gradients\n",
    "    learning_rate=2e-4,     # Learning rate for training\n",
    "    num_train_epochs=1,     # Number of training epochs\n",
    "    fsdp=\"\",                # Fully Sharded Data Parallel configuration\n",
    "    fsdp_config=None,       # Additional FSDP configurations\n",
    "    gradient_checkpointing=False,  # Whether to use gradient checkpointing\n",
    "    merge_weights=False,    # Whether to merge LoRA weights with base model\n",
    "    seed=42,                # Random seed for reproducibility\n",
    "    token=None              # HuggingFace token for model access\n",
    "):\n",
    "    # Initialize distributed training if multiple GPUs are available\n",
    "    if torch.cuda.is_available() and (torch.cuda.device_count() > 1 or int(os.environ.get(\"SM_HOST_COUNT\", 1)) > 1):\n",
    "        # Call this function at the beginning of your script\n",
    "        local_rank = init_distributed()\n",
    "\n",
    "        # Now you can use distributed functionalities\n",
    "        torch.distributed.barrier(device_ids=[local_rank])\n",
    "\n",
    "    # Enable HuggingFace transfer for model downloading\n",
    "    os.environ.update({\"HF_HUB_ENABLE_HF_TRANSFER\": \"1\"})\n",
    "\n",
    "    set_seed(seed)\n",
    "\n",
    "    accelerator = Accelerator()\n",
    "\n",
    "    # Set up HuggingFace token if provided\n",
    "    if token is not None:\n",
    "        os.environ.update({\"HF_TOKEN\": token})\n",
    "        accelerator.wait_for_everyone()\n",
    "\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "    # Set Tokenizer pad Token\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "    with accelerator.main_process_first():\n",
    "        # tokenize and chunk dataset\n",
    "        lm_train_dataset = train_ds.map(\n",
    "            lambda sample: tokenizer(sample[\"text\"]), remove_columns=list(train_ds.features)\n",
    "        )\n",
    "\n",
    "        print(f\"Total number of train samples: {len(lm_train_dataset)}\")\n",
    "\n",
    "        if test_ds is not None:\n",
    "            lm_test_dataset = test_ds.map(\n",
    "                lambda sample: tokenizer(sample[\"text\"]), remove_columns=list(test_ds.features)\n",
    "            )\n",
    "\n",
    "            print(f\"Total number of test samples: {len(lm_test_dataset)}\")\n",
    "        else:\n",
    "            lm_test_dataset = None\n",
    "\n",
    "    # Configure model settings for bfloat16 precision\n",
    "    # Setup flash_attention_2 for memory-efficient attention computation\n",
    "    if torch_dtype == torch.bfloat16:\n",
    "        print(\"flash_attention_2 init\")\n",
    "\n",
    "        model_configs = {\n",
    "            \"attn_implementation\": \"flash_attention_2\",\n",
    "            \"torch_dtype\": torch_dtype,\n",
    "        }\n",
    "    else:\n",
    "        model_configs = dict()\n",
    "\n",
    "    # Configure training settings based on FSDP usage\n",
    "    # Set up trainer configurations for FSDP or standard training\n",
    "    if fsdp != \"\" and fsdp_config is not None:\n",
    "        print(\"Configurations for FSDP\")\n",
    "\n",
    "        bnb_config_params = {\n",
    "            \"bnb_4bit_quant_storage\": torch_dtype\n",
    "        }\n",
    "\n",
    "        trainer_configs = {\n",
    "            \"fsdp\": fsdp,\n",
    "            \"fsdp_config\": fsdp_config,\n",
    "            \"gradient_checkpointing_kwargs\": {\n",
    "                \"use_reentrant\": False\n",
    "            }\n",
    "        }\n",
    "    else:\n",
    "        bnb_config_params = dict()\n",
    "        trainer_configs = {\n",
    "            \"gradient_checkpointing\": gradient_checkpointing, # Enable in case of DDP\n",
    "        }\n",
    "\n",
    "    # Enable Quantization\n",
    "    bnb_config = BitsAndBytesConfig(\n",
    "        load_in_4bit=True,\n",
    "        bnb_4bit_use_double_quant=True,\n",
    "        bnb_4bit_quant_type=\"nf4\",\n",
    "        bnb_4bit_compute_dtype=torch_dtype,\n",
    "        **bnb_config_params\n",
    "    )\n",
    "\n",
    "    model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_name,\n",
    "        trust_remote_code=True,\n",
    "        quantization_config=bnb_config,\n",
    "        use_cache=not gradient_checkpointing,\n",
    "        cache_dir=\"/tmp/.cache\",\n",
    "        **model_configs\n",
    "    )\n",
    "\n",
    "    # Configure gradient checkpointing based on FSDP usage\n",
    "    if fsdp == \"\" and fsdp_config is None:\n",
    "        print(\"Prepare model for quantization\")\n",
    "        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=gradient_checkpointing)\n",
    "\n",
    "        if gradient_checkpointing:\n",
    "            print(\"gradient_checkpointing enabled\")\n",
    "            model.gradient_checkpointing_enable()\n",
    "    else:\n",
    "        if gradient_checkpointing:\n",
    "            print(\"gradient_checkpointing enabled\")\n",
    "            model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={\"use_reentrant\": False})\n",
    "\n",
    "    config = LoraConfig(\n",
    "        r=lora_r,\n",
    "        lora_alpha=lora_alpha,\n",
    "        target_modules=\"all-linear\",\n",
    "        lora_dropout=lora_dropout,\n",
    "        bias=\"none\",\n",
    "        task_type=\"CAUSAL_LM\"\n",
    "    )\n",
    "\n",
    "    model = get_peft_model(model, config)\n",
    "\n",
    "    trainer = transformers.Trainer(\n",
    "        model=model,\n",
    "        train_dataset=lm_train_dataset,\n",
    "        eval_dataset=lm_test_dataset if lm_test_dataset is not None else None,\n",
    "        args=transformers.TrainingArguments(\n",
    "            per_device_train_batch_size=per_device_train_batch_size,\n",
    "            per_device_eval_batch_size=per_device_eval_batch_size,\n",
    "            gradient_accumulation_steps=gradient_accumulation_steps,\n",
    "            logging_strategy=\"steps\",\n",
    "            logging_steps=1,\n",
    "            log_on_each_node=False,\n",
    "            num_train_epochs=num_train_epochs,\n",
    "            learning_rate=learning_rate,\n",
    "            bf16=(\n",
    "                True if torch_dtype == torch.bfloat16 else False\n",
    "            ),  # Enable mixed-precision training\n",
    "            tf32=False,\n",
    "            ddp_find_unused_parameters=False,\n",
    "            save_strategy=\"no\",\n",
    "            output_dir=\"outputs\",\n",
    "            **trainer_configs\n",
    "        ),\n",
    "        data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    "    )\n",
    "\n",
    "    if trainer.accelerator.is_main_process:\n",
    "        trainer.model.print_trainable_parameters()\n",
    "\n",
    "    trainer.train()\n",
    "\n",
    "    if trainer.is_fsdp_enabled:\n",
    "        trainer.accelerator.state.fsdp_plugin.set_state_dict_type(\"FULL_STATE_DICT\")\n",
    "\n",
    "    if merge_weights:\n",
    "        output_dir = \"/tmp/model\"\n",
    "\n",
    "        # merge adapter weights with base model and save\n",
    "        # save int 4 model\n",
    "        trainer.model.save_pretrained(output_dir, safe_serialization=False)\n",
    "\n",
    "        if accelerator.is_main_process:\n",
    "            # clear memory\n",
    "            del model\n",
    "            del trainer\n",
    "\n",
    "            torch.cuda.empty_cache()\n",
    "\n",
    "            # load PEFT model\n",
    "            model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "                output_dir,\n",
    "                torch_dtype=torch.float16,\n",
    "                low_cpu_mem_usage=True,\n",
    "                trust_remote_code=True,\n",
    "                use_cache=True,\n",
    "                cache_dir=\"/tmp/.cache\",\n",
    "            )\n",
    "\n",
    "            # Merge LoRA and base model and save\n",
    "            model = model.merge_and_unload()\n",
    "            model.save_pretrained(\n",
    "                os.environ.get(\"SM_MODEL_DIR\", \"/opt/ml/model\"),\n",
    "                safe_serialization=True,\n",
    "                max_shard_size=\"2GB\"\n",
    "            )\n",
    "    else:\n",
    "        trainer.model.save_pretrained(\n",
    "            os.environ.get(\"SM_MODEL_DIR\", \"/opt/ml/model\"),\n",
    "            safe_serialization=True\n",
    "        )\n",
    "\n",
    "    if accelerator.is_main_process:\n",
    "        tokenizer.save_pretrained(os.environ.get(\"SM_MODEL_DIR\", \"/opt/ml/model\"))\n",
    "\n",
    "    accelerator.wait_for_everyone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a386bd9-172c-485c-af45-ebc1d126470b",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_fn(\n",
    "    model_id,\n",
    "    train_ds=train_dataset,\n",
    "    test_ds=test_dataset,\n",
    "    per_device_train_batch_size=2,\n",
    "    per_device_eval_batch_size=1,\n",
    "    gradient_accumulation_steps=2,\n",
    "    gradient_checkpointing=True,\n",
    "    num_train_epochs=1,\n",
    "    fsdp=\"full_shard auto_wrap offload\",\n",
    "    fsdp_config={\n",
    "        \"backward_prefetch\": \"backward_pre\",\n",
    "        \"cpu_ram_efficient_loading\": True,\n",
    "        \"offload_params\": True,\n",
    "        \"forward_prefetch\": False,\n",
    "        \"use_orig_params\": False,\n",
    "    },\n",
    "    merge_weights=True,\n",
    "    token=os.environ.get(\"HF_TOKEN\", None),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da81686e-d27f-4c7b-bec6-f596e7dbaa32",
   "metadata": {},
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeac4c44",
   "metadata": {},
   "outputs": [],
   "source": [
    "test.iloc()[0][\"text\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb892b35-ab33-4964-9947-f9487a1e50cb",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Load Fine-Tuned model\n",
    "\n",
    "Note: Run `train_fn` with `merge_weights=True` for merging the trained adapter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "288e2807-fe0e-46d8-8e79-8dd65f45757a",
   "metadata": {},
   "source": [
    "### Download model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b50ff8a-842c-45bc-aa55-4a5e87f2b190",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-20T18:41:08.515277Z",
     "start_time": "2023-11-20T18:41:08.503555Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57abfc30-b9f0-4cd2-8c2e-99bb77c80767",
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_session = sagemaker.Session()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35481eff-1142-46f3-8e38-50a1bdadba7b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-03T13:45:11.757861Z",
     "start_time": "2023-09-03T13:45:11.747993Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_id = \"mistralai/Mistral-Small-24B-Instruct-2501\"\n",
    "\n",
    "bucket_name = sagemaker_session.default_bucket()\n",
    "default_prefix = sagemaker_session.default_bucket_prefix\n",
    "job_prefix = f\"train-{model_id.split('/')[-1].replace('.', '-')}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ed118e7-1c80-4392-8ea5-147b63fc2f03",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_last_job_name(job_name_prefix):\n",
    "    sagemaker_client = boto3.client('sagemaker')\n",
    "\n",
    "    matching_jobs = []\n",
    "    next_token = None\n",
    "\n",
    "    while True:\n",
    "        # Prepare the search parameters\n",
    "        search_params = {\n",
    "            'Resource': 'TrainingJob',\n",
    "            'SearchExpression': {\n",
    "                'Filters': [\n",
    "                    {\n",
    "                        'Name': 'TrainingJobName',\n",
    "                        'Operator': 'Contains',\n",
    "                        'Value': job_name_prefix\n",
    "                    },\n",
    "                    {\n",
    "                        'Name': 'TrainingJobStatus',\n",
    "                        'Operator': 'Equals',\n",
    "                        'Value': \"Completed\"\n",
    "                    }\n",
    "                ]\n",
    "            },\n",
    "            'SortBy': 'CreationTime',\n",
    "            'SortOrder': 'Descending',\n",
    "            'MaxResults': 100\n",
    "        }\n",
    "\n",
    "        # Add NextToken if we have one\n",
    "        if next_token:\n",
    "            search_params['NextToken'] = next_token\n",
    "\n",
    "        # Make the search request\n",
    "        search_response = sagemaker_client.search(**search_params)\n",
    "\n",
    "        # Filter and add matching jobs\n",
    "        matching_jobs.extend([\n",
    "            job['TrainingJob']['TrainingJobName'] \n",
    "            for job in search_response['Results']\n",
    "            if job['TrainingJob']['TrainingJobName'].startswith(job_name_prefix)\n",
    "        ])\n",
    "\n",
    "        # Check if we have more results to fetch\n",
    "        next_token = search_response.get('NextToken')\n",
    "        if not next_token or matching_jobs:  # Stop if we found at least one match or no more results\n",
    "            break\n",
    "\n",
    "    if not matching_jobs:\n",
    "        raise ValueError(f\"No completed training jobs found starting with prefix '{job_name_prefix}'\")\n",
    "\n",
    "    return matching_jobs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18f4e9bd-de61-4806-b314-6bcf988a2c86",
   "metadata": {},
   "outputs": [],
   "source": [
    "job_name = get_last_job_name(job_prefix)\n",
    "\n",
    "job_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36eeba4b-cfc9-4aae-be86-8d2f3c4e5cb0",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "#### Inference configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b89ef485-b13b-43a7-a523-eacbdddbb81d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker import Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f29f2d8-b3b3-4b27-a0d2-e49812f56a21",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-20T18:41:10.492877Z",
     "start_time": "2023-11-20T18:41:10.488495Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "instance_count = 1\n",
    "instance_type = \"ml.g5.12xlarge\"\n",
    "health_check_timeout = 700"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f062e4-e829-40f4-a613-bf16ad503829",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_uri = sagemaker.image_uris.retrieve(\n",
    "    framework=\"djl-lmi\",\n",
    "    region=sagemaker_session.boto_session.region_name,\n",
    "    version=\"latest\"\n",
    ")\n",
    "\n",
    "image_uri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bdd6ef9-484a-4234-ba92-3c0ef95ccad0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-20T18:41:12.433399Z",
     "start_time": "2023-11-20T18:41:11.963091Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if default_prefix:\n",
    "    model_data = f\"s3://{bucket_name}/{default_prefix}/{job_name}/{job_name}/output/model/\"\n",
    "else:\n",
    "    model_data = f\"s3://{bucket_name}/{job_name}/{job_name}/output/model/\"\n",
    "\n",
    "model = Model(\n",
    "    image_uri=image_uri,\n",
    "    role=get_execution_role(),\n",
    "    env={\n",
    "        \"HF_MODEL_ID\": model_data,  # path to where sagemaker stores the model\n",
    "        \"OPTION_TRUST_REMOTE_CODE\": \"true\",\n",
    "        \"OPTION_ROLLING_BATCH\": \"vllm\",\n",
    "        \"OPTION_DTYPE\": \"bf16\",\n",
    "        \"OPTION_QUANTIZE\": \"fp8\",\n",
    "        \"OPTION_TENSOR_PARALLEL_DEGREE\": \"max\",\n",
    "        \"OPTION_MAX_ROLLING_BATCH_SIZE\": \"1\",\n",
    "        \"OPTION_MODEL_LOADING_TIMEOUT\": \"3600\",\n",
    "        \"OPTION_MAX_MODEL_LEN\": \"4096\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a85d8484-69c7-4a91-934b-7cc12613d563",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"mistralai/Mistral-Small-24B-Instruct-2501\"\n",
    "\n",
    "endpoint_name = f\"{model_id.split('/')[-1].replace('.', '-')}-djl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9035076-ac69-4859-9824-dcbf07c0f2b2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-20T18:45:49.265298Z",
     "start_time": "2023-11-20T18:41:14.621743Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "predictor = model.deploy(\n",
    "    endpoint_name=endpoint_name,\n",
    "    initial_instance_count=instance_count,\n",
    "    instance_type=instance_type,\n",
    "    container_startup_health_check_timeout=health_check_timeout,\n",
    "    model_data_download_timeout=3600\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "545d0cac-4af1-4034-9f88-35861396c228",
   "metadata": {},
   "source": [
    "#### Predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d6e8c1-5a3b-459f-8763-febab2b8f094",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "625af54c-198f-4a1a-8bf7-8e6d8dcfa460",
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_session = sagemaker.Session()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a61094a1-1fa9-495f-8343-aa27d9c4ba2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"mistralai/Mistral-Small-24B-Instruct-2501\"\n",
    "\n",
    "endpoint_name = f\"{model_id.split('/')[-1].replace('.', '-')}-djl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e7b040d-b8f4-4be0-9687-b28eff520236",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = sagemaker.Predictor(\n",
    "    endpoint_name=endpoint_name,\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    serializer=sagemaker.serializers.JSONSerializer(),\n",
    "    deserializer=sagemaker.deserializers.JSONDeserializer(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d5b03ad-a4aa-4107-be89-5d1160cd2d01",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_prompt = f\"\"\"\n",
    "<s>\n",
    "[SYSTEM_PROMPT]\n",
    "{{system}}\n",
    "[/SYSTEM_PROMPT]\n",
    "[INST]\n",
    "{{question}}\n",
    "[/INST]\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1586b889-6554-4d75-b295-e7dc99673cf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "system = \"\"\"\n",
    "You are a helpful assistant that answers in JSON. Here\\'s the json schema you must adhere to:\n",
    "<schema>\n",
    "{\n",
    "    \\'DateRange\\': {\\'type\\': \\'object\\', \\'properties\\': {\\'explain\\': {\\'type\\': \\'string\\', \\'description\\': \\'Explain the date range in the context of the text before generating the date range and the repeat pattern.\\'}, \n",
    "    \\'repeats\\': {\\'type\\': \\'string\\', \\'enum\\': [\\'daily\\', \\'weekly\\', \\'monthly\\', None], \\'description\\': \"If the date range repeats, and how often, this way we can generalize the date range to the future. If it\\'s special, then we can assume it is a one-time event.\"}, \n",
    "    \\'days_of_week\\': {\\'type\\': \\'array\\', \\'items\\': {\\'type\\': \\'string\\', \\'enum\\': [\\'monday\\', \\'tuesday\\', \\'wednesday\\', \\'thursday\\', \\'friday\\', \\'saturday\\', \\'sunday\\', None]}, \\'description\\': \\'If the date range repeats, which days of the week does it repeat on.\\'}, \n",
    "    \\'time_start\\': {\\'type\\': \\'string\\', \\'format\\': \\'date-time\\', \\'description\\': \\'The start of the first time range in the day.\\'}, \n",
    "    \\'time_end\\': {\\'type\\': \\'string\\', \\'format\\': \\'date-time\\', \\'description\\': \\'The end of the first time range in the day.\\'}}, \n",
    "    \\'required\\': [\\'explain\\', \\'days_of_week\\', \\'time_start\\', \\'time_end\\']}, \n",
    "    \\'AvailabilityResponse\\': {\\'type\\': \\'object\\', \\'properties\\': {\\'availability\\': {\\'type\\': \\'array\\', \\'items\\': {\\'$ref\\': \\'#/definitions/DateRange\\'}}}, \n",
    "    \\'required\\': [\\'availability\\']}\n",
    "}\n",
    "</schema>\n",
    "\"\"\"\n",
    "question = \"\"\"\n",
    "I\\'m coordinating the schedule for our community center and I need to parse the availability from the text provided. \n",
    "The text states: \n",
    "\\'The center is available for events from April 10th to April 20th, 9am to 2pm on weekdays. \n",
    "This period of availability repeats daily, specifically on Monday, Tuesday, Wednesday, Thursday, and Friday, starting from April 10th, 2023, at 9:00 AM until 2:00 PM on the same day. \n",
    "Additionally, there are special yoga classes every Wednesday evening at 6pm starting from April 5th and continuing for 8 weeks. These yoga classes are scheduled to repeat weekly on Wednesdays, with the first class beginning on April 5th, 2023, at 6:00 PM and ending at 7:00 PM. The center is also open for morning activities from dawn till 10am on weekends. \n",
    "Assuming \\'dawn\\' is approximately 6 am, this weekend availability repeats weekly on Saturday and Sunday, with the start time at 6:00 AM and end time at 10:00 AM, beginning from the first weekend within the provided date range.\\'\n",
    "\"\"\"\n",
    "\n",
    "prompt = base_prompt.format(\n",
    "    system=system,\n",
    "    question=question\n",
    ")\n",
    "\n",
    "response = predictor.predict(\n",
    "    {\n",
    "        \"inputs\": prompt,\n",
    "        \"parameters\": {\n",
    "            \"temperature\": 0.2,\n",
    "            \"top_p\": 0.9,\n",
    "            \"max_new_tokens\": 512,\n",
    "            \"return_full_text\": False,\n",
    "            \"stop\": [\"</s>\"],\n",
    "        },\n",
    "    }\n",
    ")\n",
    "\n",
    "response = response[\"generated_text\"].split(\"</s>\")[0]\n",
    "\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6ff5d47-62d5-4751-828c-c6f998c33b16",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "#### Delete Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06aebbc6-5761-44a7-949d-45b2fffd2522",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-20T18:48:55.153276Z",
     "start_time": "2023-11-20T18:48:54.165351Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "predictor.delete_model()\n",
    "predictor.delete_endpoint(delete_endpoint_config=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mlops-demo",
   "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}