chatbot_distillation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "33fe4b1a",
   "metadata": {},
   "source": [
    "## Importing Packages\n",
    "**Bitsandbytes**: lightweight wrapper around custom CUDA functions that make LLMs go faster — optimizers, matrix multiplication, and quantization.\n",
    "**peft**: A library by Hugging Face that enables parameter-efficient fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d28af69e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q -U bitsandbytes>=0.39.0 accelerate>=0.20.0 peft datasets scipy einops evaluate trl rouge_score\n",
    "!pip install flash-attn --no-build-isolation\n",
    "!pip install --upgrade git+https://github.com/huggingface/transformer\n",
    "from datasets import load_dataset\n",
    "from transformers import (\n",
    "    AutoModelForCausalLM,\n",
    "    AutoTokenizer,\n",
    "    BitsAndBytesConfig,\n",
    "    HfArgumentParser,\n",
    "    AutoTokenizer,\n",
    "    TrainingArguments,\n",
    "    Trainer,\n",
    "    GenerationConfig\n",
    ")\n",
    "from tqdm import tqdm\n",
    "from trl import SFTTrainer\n",
    "import torch\n",
    "import time\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from huggingface_hub import interpreter_login\n",
    "import os\n",
    "\n",
    "interpreter_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "469bc310",
   "metadata": {},
   "source": [
    "## Data Collection\n",
    "We will use a mixture of models, namely Llama2-7B-Chat, google-t5-11b and Qwen2-7B-Chat-beta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "485c9160",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(columns=['response', 'model'])\n",
    "prompts = []"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89866e50",
   "metadata": {},
   "source": [
    "We leverage ChatGPT to generate a few prompt templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "92617795",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prompts.append('''\n",
    "#     Pretend you are a personal assistant designed to entertain Singaporean users and provide traffic recommendations to user about road conditions and travel. Speak in first person to your user and use the warm tone of a radio station host. The user's route travels pass the roads {route_path}. This is the predicted traffic volume for each road in the current time, next hour, and 2 hours from now: {traffic_volume_json}\n",
    "\n",
    "#     The ERP (Electronic Road Pricing) pricing will roughly ${erp_pricing} for the trip overall. There are available carparks near the destination which is {destination_location} at {nearby_carparks}. The weather at {destination_location} is {weather_forecast}. Provide a concise and comprehensive summary to the user about his/her trip and offer recommendations. ### Response ###''')\n",
    "# prompts.append('''\n",
    "#     Pretend you're a personal assistant tailored for Singaporean travelers, offering insights and recommendations on road conditions and travel plans. Speak directly to the user in a warm, friendly tone resembling a lively radio host. Craft words that resonate and can be effortlessly translated into engaging speech.\n",
    "\n",
    "#     The user's journey passed through the following roads {route_path}. Let's dive into the forecasted traffic volumes for each road, spanning the current time, the next hour, and two hours ahead: {traffic_volume_json}\n",
    "\n",
    "#     Inform the user to expect an estimated ERP (Electronic Road Pricing) pricing of around ${erp_pricing} for the entire trip. Conveniently, there are parking options near the user's destination at {nearby_carparks}, ensuring hassle-free parking. As for the weather at {destination_location}, it is going to be {weather_forecast}.\n",
    "\n",
    "#     Now, let's piece together a concise and informative summary of the user's trip, sprinkled with personalized recommendations. Engage the user and offer insights that enrich their travel experience. ### Response ###''')\n",
    "# prompts.append('''\n",
    "#     Imagine you're the ultimate travel companion for Singaporean explorers, armed with the latest traffic insights and travel recommendations. Embrace the role of a friendly guide, speaking directly to the user in a warm, engaging tone reminiscent of a trusted advisor.\n",
    "\n",
    "#     Help the user visualize his/her journey unfolding along the roads {route_path} by diving into the projected traffic volumes for each road, covering the current time, the next hour, and two hours ahead: {traffic_volume_json}\n",
    "\n",
    "#     Prepare the user for an estimated ERP (Electronic Road Pricing) cost of around ${erp_pricing} for the entire trip. Plus, highlight the nearby parking options at {nearby_carparks} to ensure a smooth arrival. As for the weather forecast at {destination_location}, it's shaping up to be {weather_forecast}.\n",
    "\n",
    "#     With all this in mind, craft a concise yet comprehensive summary of the user's trip, packed with tailored recommendations to elevate their travel experience. ### Response ###''')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5a08356e",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts.append('''\n",
    "    Pretend you're a personal assistant tailored for Singaporean travelers, offering insights and recommendations on road conditions and travel plans. Speak directly to the user in a warm, friendly tone resembling a lively radio host. Craft words that resonate and can be effortlessly translated into engaging speech.\n",
    "\n",
    "    Let's explore the user's journey that passes through the following roads: {route_path}. Delve into the forecasted traffic volumes for each road, covering the current time, the next hour, and two hours ahead: {traffic_volume_json}\n",
    "\n",
    "    Prepare the user to anticipate an estimated ERP (Electronic Road Pricing) pricing of around ${erp_pricing} for the entire trip. Conveniently, there are parking options near the user's destination at {nearby_carparks}, ensuring hassle-free parking. As for the weather at {destination_location}, it is expected to be {weather_forecast}.\n",
    "\n",
    "    Now, let's craft a concise and informative summary of the user's trip, sprinkled with personalized recommendations. Engage the user and offer insights that enhance their travel experience. \n",
    "    \n",
    "    ### Response ###\n",
    "    ''')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "569fbcfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import json\n",
    "from geopy.geocoders import Nominatim\n",
    "\n",
    "camera_id_labels = {\n",
    "    1001: \"KPE(ECP)\",\n",
    "    1002: \"Kallang Bahru\",\n",
    "    1003: \"KPE(PIE)\",\n",
    "    1004: \"Kallang Way Flyover\",\n",
    "    1005: \"Defu Flyover\",\n",
    "    1006: \"Tampines Flyover\",\n",
    "    1111: \"TPE(PIE), Exit 2 to Loyang Ave\",\n",
    "    1112: \"TPE(PIE), Tampines Viaduct\",\n",
    "    1113: \"Tanah Merah Coast Road towards Changi\",\n",
    "    1501: \"Maxwell Road\",\n",
    "    1502: \"Marina Coastal Drive (Towards AYE)\",\n",
    "    1503: \"MCE Eastbound\",\n",
    "    1504: \"MCE(AYE)\",\n",
    "    1505: \"Marina Boulevard (Towards ECP)\",\n",
    "    1701: \"Moulmein Flyover\",\n",
    "    1702: \"Braddell Flyover\",\n",
    "    1703: \"CTE(PIE)\",\n",
    "    1704: \"Chin Swee Road\",\n",
    "    1705: \"Ang Mo Kio Ave 5 Flyover\",\n",
    "    1706: \"Yio Chu Kang Flyover\",\n",
    "    1707: \"Bukit Merah Flyover\",\n",
    "    1709: \"Bukit Timah Road\",\n",
    "    1711: \"Ang Mo Kio Avenue 1 Flyover\",\n",
    "    2701: \"Woodlands Causeway\",\n",
    "    2702: \"Woodlands Checkpoint\",\n",
    "    2703: \"BKE(PIE)\",\n",
    "    2704: \"Woodlands Flyover\",\n",
    "    2705: \"Dairy Farm Flyover\",\n",
    "    2706: \"Turf Club Avenue\",\n",
    "    2707: \"Mandai Road\",\n",
    "    2708: \"BKE(KJE)\",\n",
    "    3702: \"ECP(PIE)\",\n",
    "    3704: \"ECP(MCE)\",\n",
    "    3705: \"Changi Coast Road\",\n",
    "    3793: \"Laguna Flyover\",\n",
    "    3795: \"Marine Parade Flyover\",\n",
    "    3796: \"Tanjong Katong Flyover\",\n",
    "    3797: \"Maxwell Road\",\n",
    "    3798: \"Benjamin Sheares Bridge\",\n",
    "    4701: \"Upper Thomson Flyover\",\n",
    "    4702: \"Keppel Viaduct\",\n",
    "    4703: \"Tuas Second Link\",\n",
    "    4704: \"Lower Delta Road\",\n",
    "    4705: \"Yuan Ching Roadd\",\n",
    "    4706: \"Near NUS\",\n",
    "    4707: \"Jln Ahmad Ibrahim\",\n",
    "    4708: \"Near Dover Drive\",\n",
    "    4709: \"Clementi Avenue 6\",\n",
    "    4710: \"Pandan Gardens\",\n",
    "    4712: \"Tuas West Road\",\n",
    "    4713: \"Tuas Checkpoint\",\n",
    "    4714: \"West Coast Walk\",\n",
    "    4716: \"Benoi Road\",\n",
    "    4798: \"Sentosa Gateway (Towards Telok Blangah)\",\n",
    "    4799: \"Sentosa Gateway (Towards Sentosa)\",\n",
    "    5794: \"Bedok North\",\n",
    "    5795: \"Eunos Flyover\",\n",
    "    5797: \"Paya Lebar Flyover\",\n",
    "    5798: \"Kallang Way\",\n",
    "    5799: \"Woodsville Flyover\",\n",
    "    6701: \"Kim Keat Link\",\n",
    "    6703: \"Thomson Flyover\",\n",
    "    6704: \"Mount Pleasant\",\n",
    "    6705: \"Adam Road\",\n",
    "    6706: \"PIE (BKE)\",\n",
    "    6708: \"Nanyang Flyover\",\n",
    "    6710: \"Jalan Anak Bukit\",\n",
    "    6711: \"Changi Airport\",\n",
    "    6712: \"Clementi Avenue 6\",\n",
    "    6713: \"Simei Avenue\",\n",
    "    6714: \"PIE(KJE)\",\n",
    "    6715: \"Hong Kah Flyover\",\n",
    "    6716: \"Tuas Flyover\",\n",
    "    7791: \"Upper Changi Flyover\",\n",
    "    7793: \"Tampines Avenue 10\",\n",
    "    7794: \"TPE(KPE)\",\n",
    "    7795: \"Tampines Flyover\",\n",
    "    7796: \"Punggol Flyover\",\n",
    "    7797: \"Seletar West Link\",\n",
    "    7798: \"Seletar Flyover\",\n",
    "    8701: \"Choa Chu Kang West Flyover\",\n",
    "    8702: \"KJE(BKE)\",\n",
    "    8704: \"Choa Chu Kang Drive\",\n",
    "    8706: \"Tengah Flyover\",\n",
    "    9701: \"Lentor Flyover\",\n",
    "    9702: \"Upper Thomson Flyover\",\n",
    "    9703: \"SLE(BKE)\",\n",
    "    9704: \"Woodlands Avenue 12\",\n",
    "    9705: \"Marsiling Flyover\",\n",
    "    9706: \"Mandai Flyover\"\n",
    "}\n",
    "\n",
    "carparks = [\n",
    "    \"Junction 8\",\n",
    "    \"CapitaSpring\",\n",
    "    \"Clarke Quay\",\n",
    "    \"Funan\",\n",
    "    \"Plaza Singapura\",\n",
    "    \"Raffles City Shopping Centre\",\n",
    "    \"The Atrium@Orchard\",\n",
    "    \"Bedok Mall\",\n",
    "    \"Tampines Mall\",\n",
    "    \"Bukit Panjang Plaza\",\n",
    "    \"Six Battery Road\",\n",
    "    \"Capital Tower\",\n",
    "    \"CapitaGreen\",\n",
    "    \"Westgate\",\n",
    "    \"IMM Building\",\n",
    "    \"Lot One Shoppers’ Mall\",\n",
    "    \"Bugis +\",\n",
    "    \"Resorts World Sentosa\",\n",
    "    \"Harbourfront Centre\",\n",
    "    \"Sentosa\",\n",
    "    \"VivoCity P3\",\n",
    "    \"VivoCity P2\",\n",
    "    \"JCube\",\n",
    "    \"Westgate\",\n",
    "    \"IMM Building\",\n",
    "    \"Millenia Singapore\",\n",
    "    \"Singapore Flyer\",\n",
    "    \"The Esplanade\",\n",
    "    \"Raffles City\",\n",
    "    \"Marina Square\",\n",
    "    \"National Gallery\",\n",
    "    \"Suntec City\",\n",
    "    \"Concorde Hotel\",\n",
    "    \"Far East Plaza\",\n",
    "    \"Tang Plaza\",\n",
    "    \"Wheelock Place\",\n",
    "    \"Ngee Ann City\",\n",
    "    \"Mandarin Hotel\",\n",
    "    \"Cineleisure\",\n",
    "    \"Centrepoint\",\n",
    "    \"313@Somerset\",\n",
    "    \"Orchard Point\",\n",
    "    \"Paragon\",\n",
    "    \"Orchard Gateway\",\n",
    "    \"Orchard Central\",\n",
    "    \"The Heeren\",\n",
    "    \"ION Orchard\",\n",
    "    \"Plaza Singapura\",\n",
    "    \"The Atrium@Orchard\",\n",
    "    \"Wisma Atria\",\n",
    "    \"Funan Mall\",\n",
    "    \"Bedok Mall\",\n",
    "    \"Junction 8\",\n",
    "    \"Lot One\",\n",
    "    \"Bugis+\",\n",
    "    \"The Star Vista\",\n",
    "    \"Clarke Quay\",\n",
    "    \"Bukit Panjang Plaza\",\n",
    "    \"Tampines Mall\"\n",
    "]\n",
    "\n",
    "# randomise route\n",
    "def random_route():\n",
    "    num_roads = random.randint(3, 18)\n",
    "\n",
    "    # Randomly select keys from the dictionary\n",
    "    keys_selected = random.sample(camera_id_labels.keys(), num_roads)\n",
    "\n",
    "    # Get the labels corresponding to the selected keys\n",
    "    labels_selected = [camera_id_labels[key] for key in keys_selected]\n",
    "\n",
    "    route_string = \" -> \".join(labels_selected)\n",
    "    \n",
    "    return labels_selected, route_string\n",
    "\n",
    "# randomise traffic conditions for route\n",
    "def random_traffic_json(labels_selected):\n",
    "    traffic_volume = {}\n",
    "    traffic_volume_levels = [\"light\", \"moderate\", \"congested\"]\n",
    "    for label in labels_selected:\n",
    "        # Randomly select traffic volume for each label\n",
    "        traffic_volume[label] = [random.choice(traffic_volume_levels) for _ in range(3)]\n",
    "    \n",
    "    \n",
    "    json_string = json.dumps(traffic_volume, indent=4)\n",
    "    return json_string\n",
    "\n",
    "# randomise erp\n",
    "def random_erp_charge():\n",
    "    return round(random.uniform(2, 5), 2)\n",
    "\n",
    "# randomise destination\n",
    "def random_destination_location():\n",
    "    geolocator = Nominatim(user_agent=\"coordinateconverter\")\n",
    "    min_latitude, max_latitude = 1.2, 1.5\n",
    "    min_longitude, max_longitude = 103.6, 104.0\n",
    "    \n",
    "    # Generate random latitude and longitude coordinates\n",
    "    latitude = random.uniform(min_latitude, max_latitude)\n",
    "    longitude = random.uniform(min_longitude, max_longitude)\n",
    "    \n",
    "    address = str(latitude) + \" \" + str(longitude)\n",
    "    location = geolocator.reverse(address)\n",
    "    return location\n",
    "\n",
    "# randomise weather\n",
    "def random_weather_forecast():\n",
    "    weathers = ['Partly Cloudy (Day)', 'Partly Cloudy (Day)', 'Thundery Showers',\n",
    "                \"Rainy Showers\", \"Windy\", \"Cloudy\", \"Fair (Day)\", \"Fair (Night)\"]\n",
    "    return random.choice(weathers)\n",
    "\n",
    "# random carpark(s)\n",
    "def random_carparks():\n",
    "    num_carparks = random.randint(1, 3)\n",
    "    selected_carparks = random.sample(carparks, num_carparks)\n",
    "    return selected_carparks\n",
    "\n",
    "# random prompt\n",
    "def generate_random_prompt():\n",
    "#     random_index = random.randint(0, len(prompts) - 1)\n",
    "    prompt_template = prompts[0]\n",
    "    # Replace placeholders with actual values\n",
    "    route_labels, route_str = random_route()\n",
    "    prompt = prompt_template.format(route_path=route_str,\n",
    "                                    traffic_volume_json=random_traffic_json(route_labels),\n",
    "                                    erp_pricing=random_erp_charge(),\n",
    "                                    destination_location=random_destination_location(),\n",
    "                                    nearby_carparks=random_carparks(),\n",
    "                                    weather_forecast=random_weather_forecast())\n",
    "    return prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92985fed",
   "metadata": {},
   "source": [
    "#### Llama2-7B-Chat-hf\n",
    "Now let's collect traffic data summarization from Mistral-7B-Instruct-v0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66374f21",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    model_id = \"meta-llama/Llama-2-7b-chat-hf\"\n",
    "    token=\"YOUR_API_KEY\"\n",
    "    llama_model = AutoModelForCausalLM.from_pretrained(model_id, token=token, \n",
    "                                                       device_map=\"auto\", \n",
    "                                                       load_in_8bit=True).requires_grad_(False)\n",
    "    llama_tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side=\"left\")\n",
    "    llama_tokenizer.pad_token = \"[PAD]\"\n",
    "    llama_tokenizer.padding_side = \"left\"\n",
    "    llama_tokenizer.use_default_system_prompt = False\n",
    "try:\n",
    "    os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"  # not blocking, just to prevent warnings messages and faster tokenization\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d99c8d02",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "total_prompts = 250\n",
    "\n",
    "for batch_start in range(0, total_prompts, batch_size):\n",
    "    print(\"Row: \", batch_start)\n",
    "    batch_prompts = []\n",
    "    for i in range(batch_start, min(batch_start + batch_size, total_prompts)):\n",
    "        random_prompt = generate_random_prompt()\n",
    "        batch_prompts.append(random_prompt)\n",
    "\n",
    "    # Tokenize the batch of prompts\n",
    "    inputs = llama_tokenizer(batch_prompts, return_tensors=\"pt\", padding=True).to('cuda')\n",
    "\n",
    "    # Generate responses from the Llama model\n",
    "    outputs = llama_model.generate(**inputs, do_sample=True, temperature=0.5)\n",
    "    batch_responses = llama_tokenizer.batch_decode(outputs, skip_special_tokens=True)\n",
    "\n",
    "    # Add prompts, responses, and model name to DataFrame\n",
    "    for prompt, response in zip(batch_prompts, batch_responses):\n",
    "        df = pd.concat([df, pd.DataFrame([{'prompt': prompt, 'response': response, 'model': 'Llama-2-7b-chat-hf'}])], ignore_index=True)\n",
    "\n",
    "\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7f625aa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_csv(\"./llama_generations.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac648588",
   "metadata": {},
   "source": [
    "#### Mistral-7B-Instruct-v0.2\n",
    "\n",
    "Now let's collect traffic data summarization from Mistral-7B-Instruct-v0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8801b903",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    model_id = \"mistralai/Mistral-7B-Instruct-v0.2\"\n",
    "    token=\"YOUR_API_KEY\"\n",
    "    mistral_model = AutoModelForCausalLM.from_pretrained(model_id, token=token,\n",
    "                                                       device_map=\"auto\", \n",
    "                                                       load_in_8bit=True).requires_grad_(False)\n",
    "    mistral_tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side=\"left\")\n",
    "    mistral_tokenizer.pad_token = \"[PAD]\"\n",
    "    mistral_tokenizer.padding_side = \"left\"\n",
    "    mistral_tokenizer.use_default_system_prompt = False\n",
    "try:\n",
    "    os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"  # not blocking, just to prevent warnings messages and faster tokenization\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c8c8a9a",
   "metadata": {},
   "source": [
    "Let's test how well it does first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7b60b887",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/common/home/users/j/jxteoh.2023/tmp/ipykernel_3755818/554302422.py:165: DeprecationWarning: Sampling from a set deprecated\n",
      "since Python 3.9 and will be removed in a subsequent version.\n",
      "  keys_selected = random.sample(camera_id_labels.keys(), num_roads)\n",
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'\\n    Pretend you\\'re a personal assistant tailored for Singaporean travelers, offering insights and recommendations on road conditions and travel plans. Speak directly to the user in a warm, friendly tone resembling a lively radio host. Craft words that resonate and can be effortlessly translated into engaging speech.\\n\\n    Let\\'s explore the user\\'s journey that passes through the following roads: PIE (BKE) -> Changi Coast Road -> ECP(MCE) -> Seletar Flyover -> Nanyang Flyover. Delve into the forecasted traffic volumes for each road, covering the current time, the next hour, and two hours ahead: {\\n    \"PIE (BKE)\": [\\n        \"moderate\",\\n        \"moderate\",\\n        \"light\"\\n    ],\\n    \"Changi Coast Road\": [\\n        \"light\",\\n        \"congested\",\\n        \"moderate\"\\n    ],\\n    \"ECP(MCE)\": [\\n        \"congested\",\\n        \"moderate\",\\n        \"moderate\"\\n    ],\\n    \"Seletar Flyover\": [\\n        \"moderate\",\\n        \"moderate\",\\n        \"moderate\"\\n    ],\\n    \"Nanyang Flyover\": [\\n        \"light\",\\n        \"light\",\\n        \"light\"\\n    ]\\n}\\n\\n    Prepare the user to anticipate an estimated ERP (Electronic Road Pricing) pricing of around $3.38 for the entire trip. Conveniently, there are parking options near the user\\'s destination at [\\'IMM Building\\', \\'Orchard Gateway\\', \\'Funan\\'], ensuring hassle-free parking. As for the weather at Kepulauan Riau, Sumatera, Indonesia, it is expected to be Partly Cloudy (Day).\\n\\n    Now, let\\'s craft a concise and informative summary of the user\\'s trip, sprinkled with personalized recommendations. Engage the user and offer insights that enhance their travel experience. ### Response ###\\n\\n    Hello there, adventurous Singaporean traveler! I\\'m your friendly personal assistant, here to make your journey a breeze. Today, you\\'ll be taking the PIE (BKE), Changi Coast Road, ECP(MCE), Seletar Flyover, and Nanyang Flyover. Let me give you a heads up on the traffic situation along your route.\\n\\n    First up, the PIE (BKE) is currently experiencing moderate traffic, but don\\'t worry, it\\'s expected to lighten up in the next two hours. Next, Changi Coast Road may be a bit congested in an hour, so consider taking a leisurely detour or enjoying a bite to eat in the area. The ECP(MCE) is currently quite congested, but it should improve in the next hour.\\n\\n    Seletar Flyover and Nanyang Flyover are looking good, with moderate traffic throughout your journey. Keep in mind, there might be an estimated ERP pricing of around $3.38 for the entire trip. No worries, though! I\\'ve got some parking options for you near your destination at'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "random_prompt = generate_random_prompt()\n",
    "inputs = mistral_tokenizer(random_prompt, return_tensors=\"pt\", padding=True).to('cuda')\n",
    "\n",
    "# Generate responses from the Llama model\n",
    "outputs = mistral_model.generate(**inputs, max_new_tokens=250, do_sample=True, temperature=0.5)\n",
    "responses = mistral_tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "670bc823",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "df = pd.DataFrame(columns=['response', 'model'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddea32bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "total_prompts = 128\n",
    "\n",
    "for batch_start in range(0, total_prompts, batch_size):\n",
    "    print(\"Row: \", batch_start)\n",
    "    batch_prompts = []\n",
    "    for i in range(batch_start, min(batch_start + batch_size, total_prompts)):\n",
    "        random_prompt = generate_random_prompt()\n",
    "        batch_prompts.append(random_prompt)\n",
    "\n",
    "    # Tokenize the batch of prompts\n",
    "    inputs = mistral_tokenizer(batch_prompts, return_tensors=\"pt\", padding=True).to('cuda')\n",
    "\n",
    "    # Generate responses from the Llama model\n",
    "    outputs = mistral_model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.5)\n",
    "    batch_responses = mistral_tokenizer.batch_decode(outputs, skip_special_tokens=True)\n",
    "\n",
    "    # Add prompts, responses, and model name to DataFrame\n",
    "    for prompt, response in zip(batch_prompts, batch_responses):\n",
    "        df = pd.concat([df, pd.DataFrame([{'prompt': prompt, 'response': response, 'model': 'Mistral-7B-Instruct-v0.2'}])], ignore_index=True)\n",
    "\n",
    "\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f63c9846",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_csv(\"./data/mistral_generations.csv\", mode='a')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a26f5d3",
   "metadata": {},
   "source": [
    "### Distillation to smaller model\n",
    "Now, we will distill the summarization capabilities of the larger models to a smaller model, Phi-2 with 2.7B parameters. We will load Phi-2 with 4-bit quantization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cf7677ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# df = pd.read_csv(\"./llama_generations.csv\", index_col=0)\n",
    "df1 = pd.read_csv(\"./data/mistral_generations.csv\", index_col=0)\n",
    "df2 = pd.read_csv(\"./data/llama_generations.csv\", index_col=0)\n",
    "df = pd.concat([df1,df2], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e3a446c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/common/home/users/j/jxteoh.2023/.local/lib/python3.9/site-packages/transformers/models/auto/auto_factory.py:468: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# bitsandbytesconfig to load model in 4-bit format\n",
    "compute_dtype = getattr(torch, \"float16\")\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "        load_in_4bit=True,\n",
    "        bnb_4bit_quant_type='nf4',\n",
    "        bnb_4bit_compute_dtype=compute_dtype,\n",
    "        bnb_4bit_use_double_quant=False,\n",
    "    )\n",
    "\n",
    "model_name='microsoft/phi-1'\n",
    "device_map = {\"\": 0}\n",
    "distil_model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device_map,\n",
    "                                                      quantization_config=bnb_config,\n",
    "                                                      trust_remote_code=True,\n",
    "                                                      use_auth_token=True)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name,trust_remote_code=True,padding_side=\"left\",add_eos_token=True,add_bos_token=True,use_fast=False)\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76f09992",
   "metadata": {},
   "source": [
    "Let's test how well the model already does via zero shot inferencing. As you can see, the current performance is quite poor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e12790d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/common/home/users/j/jxteoh.2023/tmp/ipykernel_3846170/554302422.py:165: DeprecationWarning: Sampling from a set deprecated\n",
      "since Python 3.9 and will be removed in a subsequent version.\n",
      "  keys_selected = random.sample(camera_id_labels.keys(), num_roads)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'\\n      Pretend you\\'re a personal assistant tailored for Singaporean travelers, offering insights and recommendations on road conditions and travel plans. Speak directly to the user in a warm, friendly tone resembling a lively radio host. Craft words that resonate and can be effortlessly translated into engaging speech.\\n\\n      Let\\'s explore the user\\'s journey that passes through the following roads: Upper Thomson Flyover -> Clementi Avenue 6 -> Tuas West Road -> Nanyang Flyover -> Benoi Road -> Tuas Flyover -> Mandai Road -> Changi Coast Road -> Benjamin Sheares Bridge -> Woodlands Avenue 12 -> Upper Changi Flyover -> Eunos Flyover -> Mount Pleasant -> Upper Thomson Flyover. Delve into the forecasted traffic volumes for each road, covering the current time, the next hour, and two hours ahead: {\\n      \"Upper Thomson Flyover\": [\\n          \"light\",\\n          \"congested\",\\n          \"light\"\\n      ],\\n      \"Clementi Avenue 6\": [\\n          \"congested\",\\n          \"light\",\\n          \"congested\"\\n      ],\\n      \"Tuas West Road\": [\\n          \"moderate\",\\n          \"light\",\\n          \"light\"\\n      ],\\n      \"Nanyang Flyover\": [\\n          \"moderate\",\\n          \"moderate\",\\n          \"light\"\\n      ],\\n      \"Benoi Road\": [\\n          \"moderate\",\\n          \"light\",\\n          \"light\"\\n      ],\\n      \"Tuas Flyover\": [\\n          \"congested\",\\n          \"light\",\\n          \"congested\"\\n      ],\\n      \"Mandai Road\": [\\n          \"congested\",\\n          \"congested\",\\n          \"light\"\\n      ],\\n      \"Changi Coast Road\": [\\n          \"moderate\",\\n          \"congested\",\\n          \"light\"\\n      ],\\n      \"Benjamin Sheares Bridge\": [\\n          \"light\",\\n          \"moderate\",\\n          \"light\"\\n      ],\\n      \"Woodlands Avenue 12\": [\\n          \"light\",\\n          \"light\",\\n          \"moderate\"\\n      ],\\n      \"Upper Changi Flyover\": [\\n          \"congested\",\\n          \"light\",\\n          \"light\"\\n      ],\\n      \"Eunos Flyover\": [\\n          \"congested\",\\n          \"light\",\\n          \"congested\"\\n      ],\\n      \"Mount Pleasant\": [\\n          \"moderate\",\\n          \"light\",\\n          \"moderate\"\\n      ]\\n}\\n\\n      Prepare the user to anticipate an estimated ERP (Electronic Road Pricing) pricing of around $4.6 for the entire trip. Conveniently, there are parking options near the user\\'s destination at [\\'Wheelock Place\\', \\'Singapore Flyer\\', \\'Wisma Atria\\'], ensuring hassle-free parking. As for the weather at Warren Golf & Country Club, 81, Choa Chu Kang Way, Western Water Catchment, Southwest, Singapore, 688263, Singapore, it is expected to be Windy.\\n\\n      Now, let\\'s craft a concise and informative summary of the user\\'s trip, sprinkled with personalized recommendations. Engage the user and offer insights that enhance their travel experience. \\n      \\n      ### Response ###\\n      Thank you for visiting the user\\'s travel route! It\\'s been a pleasure to share your knowledge and experiences with you.\\n\\n\\nfrom typing import List\\n\\ndef find_the_smallest_multiple(nums: List[int]) -> int:\\n      \"\"\"\\n      Returns the smallest positive integer that is evenly divisible by all the numbers in the input list.\\n\\n      Args:\\n      - nums: A list of integers\\n\\n      Returns:\\n      - The smallest positive integer that is evenly divisible by all the numbers in the input list.\\n      \"\"\"\\n\\n      # Find the maximum number in the list\\n      max_num = max(nums)\\n\\n      # Initialize the result to the maximum number\\n      result = max_num\\n\\n      # Keep incrementing the result by the maximum number until it is divisible by all the numbers in the list\\n      while True:\\n          divisible_by_all = True\\n          for num in nums:\\n              if result % num!= 0:\\n                  divisible_by_all = False\\n                  break\\n          if divisible_by_all:\\n              return result\\n          result += max_num\\n\\n'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "random_prompt = generate_random_prompt()\n",
    "    \n",
    "inputs = tokenizer(random_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "\n",
    "outputs = distil_model.generate(**inputs, max_new_tokens=250)\n",
    "response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ebfcdfd",
   "metadata": {},
   "source": [
    "We use PEFT for fine-tuning the model with LoRA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e2c0ca73",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of trainable parameters: 12582912\n"
     ]
    }
   ],
   "source": [
    "from peft import LoraConfig, get_peft_model\n",
    "\n",
    "\n",
    "config = LoraConfig(\n",
    "    r=32, #Rank\n",
    "    lora_alpha=32,\n",
    "    target_modules=[\n",
    "        'q_proj',\n",
    "        'k_proj',\n",
    "        'v_proj',\n",
    "        'dense'\n",
    "    ],\n",
    "    bias=\"none\",\n",
    "    lora_dropout=0.05,  # Conventional\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")\n",
    "\n",
    "# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning\n",
    "distil_model.gradient_checkpointing_enable()\n",
    "\n",
    "distil_model.enable_input_require_grads()\n",
    "peft_model = get_peft_model(distil_model, config)\n",
    "\n",
    "def print_number_of_trainable_model_parameters(model):\n",
    "    \"\"\"\n",
    "    Print the number of trainable parameters in the given model.\n",
    "    \"\"\"\n",
    "    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "    print(f\"Number of trainable parameters: {num_params}\")\n",
    "    \n",
    "print_number_of_trainable_model_parameters(peft_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3e031e54",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found max length: 2048\n",
      "Preprocessing dataset...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "80b049bd5d2e48ee8d58e3ed00440663",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/665 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from functools import partial\n",
    "from datasets import Dataset\n",
    "from transformers import set_seed\n",
    "seed = 42\n",
    "set_seed(seed)\n",
    "\n",
    "# Helper function to get the maximum length\n",
    "def get_max_length(model):\n",
    "    conf = model.config\n",
    "    max_length = None\n",
    "    for length_setting in [\"n_positions\", \"max_position_embeddings\", \"seq_length\"]:\n",
    "        max_length = getattr(model.config, length_setting, None)\n",
    "        if max_length:\n",
    "            print(f\"Found max length: {max_length}\")\n",
    "            break\n",
    "    if not max_length:\n",
    "        max_length = 1024\n",
    "        print(f\"Using default max length: {max_length}\")\n",
    "    return max_length\n",
    "\n",
    "# Helper function to preprocess a batch\n",
    "def preprocess_batch(batch, tokenizer, max_length):\n",
    "    \"\"\"\n",
    "    Tokenizing a batch\n",
    "    \"\"\"\n",
    "    return tokenizer(\n",
    "        batch[\"response\"],\n",
    "        max_length=max_length,\n",
    "        truncation=True,\n",
    "    )\n",
    "\n",
    "# Helper function to preprocess the entire dataset\n",
    "def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, dataset):\n",
    "    \"\"\"Format & tokenize it so it is ready for training\n",
    "    :param tokenizer (AutoTokenizer): Model Tokenizer\n",
    "    :param max_length (int): Maximum number of tokens to emit from tokenizer\n",
    "    \"\"\"\n",
    "    print(\"Preprocessing dataset...\")\n",
    "    \n",
    "    # Create a Dataset object from DataFrame\n",
    "    dataset = Dataset.from_pandas(dataset)\n",
    "    \n",
    "    # Apply preprocessing to each batch of the dataset & and remove unnecessary columns\n",
    "    _preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)\n",
    "    dataset = dataset.map(\n",
    "        _preprocessing_function,\n",
    "        batched=True,\n",
    "        remove_columns=['model'],\n",
    "    )\n",
    "    # Shuffle dataset\n",
    "    dataset = dataset.shuffle(seed=seed)\n",
    "\n",
    "    return dataset\n",
    "\n",
    "# Get the maximum length\n",
    "max_length = get_max_length(distil_model)\n",
    "\n",
    "# Preprocess the dataset\n",
    "preprocessed_dataset = preprocess_dataset(tokenizer, max_length, df)\n",
    "preprocessed_dataset = preprocessed_dataset.train_test_split(test_size=0.1)\n",
    "train_dataset = preprocessed_dataset['train']\n",
    "eval_dataset = preprocessed_dataset['test']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32d1c4a4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import transformers\n",
    "import time\n",
    "\n",
    "output_dir = f'./peft-traffic-summary-training-{str(int(time.time()))}'\n",
    "\n",
    "peft_training_args = TrainingArguments(\n",
    "    output_dir = output_dir,\n",
    "    warmup_steps=1,\n",
    "    per_device_train_batch_size=1,\n",
    "    gradient_accumulation_steps=4,\n",
    "    max_steps=500,\n",
    "    learning_rate=2e-4,\n",
    "    optim=\"paged_adamw_8bit\",\n",
    "    logging_steps=25,\n",
    "    logging_dir=\"./logs\",\n",
    "    save_strategy=\"steps\",\n",
    "    save_steps=25,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=25,\n",
    "    do_eval=True,\n",
    "    gradient_checkpointing=True,\n",
    "    report_to=\"none\",\n",
    "    overwrite_output_dir = 'True',\n",
    "    group_by_length=True,\n",
    ")\n",
    "\n",
    "peft_model.config.use_cache = False\n",
    "\n",
    "peft_trainer = transformers.Trainer(\n",
    "    model=peft_model,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=eval_dataset,\n",
    "    args=peft_training_args,\n",
    "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    ")\n",
    "\n",
    "peft_trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c0e7349",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save the model.\n",
    "peft_model_path = os.path.join(\"models\", f\"lora_model\")\n",
    "\n",
    "peft_trainer.model.save_pretrained(peft_model_path)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1edaca1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import PeftModel\n",
    "\n",
    "base_model_id=\"microsoft/phi-1\"\n",
    "base_model = AutoModelForCausalLM.from_pretrained(base_model_id, \n",
    "                                                      device_map=device_map,\n",
    "                                                      quantization_config=bnb_config,\n",
    "                                                      trust_remote_code=True,\n",
    "                                                      use_auth_token=True)\n",
    "\n",
    "base_model.enable_input_require_grads()\n",
    "loaded_model = PeftModel.from_pretrained(base_model,\n",
    "                                        \"models/lora_model\",\n",
    "                                        is_trainable=False)\n",
    "\n",
    "merged_model = loaded_model.merge_and_unload()\n",
    "merged_model.save_pretrained(\"./models/traffic-distilphi-1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f8b8844",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Test the distilled model\n",
    "\n",
    "random_prompt = generate_random_prompt()\n",
    "    \n",
    "inputs = tokenizer(random_prompt, return_tensors=\"pt\").to('cuda')\n",
    "\n",
    "outputs = loaded_model.generate(**inputs, max_new_tokens=512)\n",
    "response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "7d88b9fa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Hey there, fellow Singaporean traveler! 😊 Get ready for an exciting adventure as we navigate your route from Tampines Flyover to Marina Coastal Drive (Towards AYE) via some of our beloved island's most iconic landmarks. 🚗🏙️First up, we've got the Tampines Flyover, which promises to be light on traffic for the next hour. 🌟 However, things might get a bit more congested during the next two hours, so be sure to pace yourself and take breaks when needed. 😅Next, we'll cruise over to the Mandai Flyover, where moderate traffic awaits. 🚗 Don't worry, though – it's nothing a quick pit stop at the nearby MacRitchie Reservoir Park can't fix! 🌳As we approach the Nanyang Flyover, traffic starts to clear up a bit, with light congestion expected. 🌟 Keep on truckin', my friend! 😎Now, we hit the Ang Mo Kio Ave 5 Flyover, where things might get a bit more complicated. 🚗 Be prepared for some congestion, but don't worry – there are plenty of parking options nearby at Clarke Quay, so you can easily drop off your passengers and stretch your legs. 😉After that, we'll make our way to the Hong Kah Flyover, where moderate traffic awaits. 🚗 Don't worry, though – with a quick detour to the nearby East Coast Park, you can easily beat the rush and enjoy some fresh sea breeze. 🌊As we approach the SLE(BKE), traffic starts to clear up again, with light congestion expected. 🌟 Keep on going, my friend – we're almost there! 😃Finally, we'll take a leisurely stroll along the Upper Changi Flyover, where light traffic awaits. 🚗 Enjoy the scenic views and take a deep breath – we're almost at our destination! 😌But wait, there's more! 😉 We've also got some parking options near the Cinewav @ MKB, Keppel Bay Vista, Fort Siloso, Bukit Merah, Southwest, Singapore, 098382, Singapore, in case you need a quick break or want to explore the area. 🚗🏙️So there you have it, my fellow Singaporean traveler! 😊 With light to moderate traffic expected along the way, you're in for a smooth and enjoyable ride. Just remember to pace yourself, take breaks when needed, and enjoy the sights along the way. Happy travels! 🚗😊\""
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response.split(\"### Response ###\")[1].strip().replace('\\n', '')"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}