{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "A100"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Project 3: Write code to trade stocks\n",
        "\n",
        "### An example code generator by fine-tuning StarCoder2 using QLoRA\n",
        "\n",
        "NOTE: This is a toy example to illustrate the technique – please don’t use\n",
        "any of this code to make trading decisions!\n",
        "\n",
        "Previously, we created a dataset and uploaded it to Hugging Face. Now we download the dataset and use it to fine-tune StarCoder2 using QLoRA.\n",
        "\n",
        "We'll see what kind of trade() functions our model can create before and after training."
      ],
      "metadata": {
        "id": "GHsssBgWM_l0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# pip installs\n",
        "\n",
        "!pip install -q requests==2.31.0 torch peft bitsandbytes transformers trl accelerate sentencepiece wandb"
      ],
      "metadata": {
        "id": "MDyR63OTNUJ6",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2515f54a-44c6-4ac2-b6a8-f41c57cf1ddb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m251.6/251.6 kB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m119.8/119.8 MB\u001b[0m \u001b[31m13.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m226.7/226.7 kB\u001b[0m \u001b[31m26.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m309.4/309.4 kB\u001b[0m \u001b[31m34.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.9/6.9 MB\u001b[0m \u001b[31m81.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.3/21.3 MB\u001b[0m \u001b[31m66.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m547.8/547.8 kB\u001b[0m \u001b[31m51.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m103.4/103.4 kB\u001b[0m \u001b[31m16.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m207.3/207.3 kB\u001b[0m \u001b[31m30.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m300.2/300.2 kB\u001b[0m \u001b[31m36.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m62.7/62.7 kB\u001b[0m \u001b[31m10.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m40.8/40.8 MB\u001b[0m \u001b[31m43.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m116.3/116.3 kB\u001b[0m \u001b[31m20.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m542.1/542.1 kB\u001b[0m \u001b[31m55.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m542.0/542.0 kB\u001b[0m \u001b[31m54.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m194.1/194.1 kB\u001b[0m \u001b[31m29.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m21.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# imports\n",
        "\n",
        "import os\n",
        "from google.colab import userdata\n",
        "from huggingface_hub import login\n",
        "import torch\n",
        "import transformers\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextStreamer, TrainingArguments\n",
        "from datasets import load_dataset, Dataset\n",
        "import wandb\n",
        "from peft import LoraConfig\n",
        "from trl import SFTTrainer, SFTConfig\n",
        "from datetime import datetime"
      ],
      "metadata": {
        "id": "-yikV8pRBer9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Constants\n",
        "\n",
        "BASE_MODEL = \"bigcode/starcoder2-3b\"   # choose 3b or 7b\n",
        "PROJECT_NAME = \"trading\"\n",
        "RUN_NAME =  f\"{datetime.now():%Y-%m-%d_%H.%M.%S}\"\n",
        "PROJECT_RUN_NAME = f\"{PROJECT_NAME}-{RUN_NAME}\"\n",
        "DATASET_NAME = \"ed-donner/trade_code_dataset\"\n",
        "\n",
        "# Hyperparameters for QLoRA Fine-Tuning\n",
        "# Details of QLoRA are out of scope for today, but there's\n",
        "# more information and links in the resources\n",
        "\n",
        "EPOCHS = 1\n",
        "LORA_ALPHA = 32\n",
        "LORA_R = 16\n",
        "LORA_DROPOUT = 0.1\n",
        "BATCH_SIZE = 1\n",
        "GRADIENT_ACCUMULATION_STEPS = 1\n",
        "LEARNING_RATE = 2e-4\n",
        "LR_SCHEDULER_TYPE = 'cosine'\n",
        "WEIGHT_DECAY = 0.001\n",
        "TARGET_MODULES = [\"q_proj\", \"v_proj\", \"k_proj\", \"o_proj\"]\n",
        "MAX_SEQUENCE_LENGTH = 320\n",
        "\n",
        "# Other config\n",
        "\n",
        "STEPS = 10\n",
        "SAVE_STEPS = 300"
      ],
      "metadata": {
        "id": "uuTX-xonNeOK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Log in to HuggingFace and Weights & Biases\n",
        "\n",
        "If you don't already have a HuggingFace account, visit https://huggingface.co to sign up and create a token.\n",
        "\n",
        "Then select the Secrets for this Notebook by clicking on the key icon in the left, and add a new secret called `HF_TOKEN` with the value as your token.\n",
        "\n",
        "Repeat this for weightsandbiases at https://wandb.ai and add a secret called `WANDB_API_KEY`"
      ],
      "metadata": {
        "id": "8JArT3QAQAjx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Log in to HuggingFace\n",
        "\n",
        "hf_token = userdata.get('HF_TOKEN')\n",
        "login(hf_token, add_to_git_credential=True)"
      ],
      "metadata": {
        "id": "WyFPZeMcM88v"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Log in to Weights & Biases\n",
        "wandb_api_key = userdata.get('WANDB_API_KEY')\n",
        "os.environ[\"WANDB_API_KEY\"] = wandb_api_key\n",
        "wandb.login()\n",
        "\n",
        "# Configure Weights & Biases to record against our project\n",
        "os.environ[\"WANDB_PROJECT\"] = PROJECT_NAME\n",
        "os.environ[\"WANDB_LOG_MODEL\"] = \"true\"\n",
        "os.environ[\"WANDB_WATCH\"] = \"false\""
      ],
      "metadata": {
        "id": "yJNOv3cVvJ68"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Now load the Tokenizer and Model"
      ],
      "metadata": {
        "id": "qJWQ0a3wZ0Bw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the Tokenizer and the Model\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "tokenizer.padding_side = \"right\"\n",
        "\n",
        "quant_config = BitsAndBytesConfig(load_in_8bit=True)\n",
        "\n",
        "base_model = AutoModelForCausalLM.from_pretrained(\n",
        "    BASE_MODEL,\n",
        "    quantization_config=quant_config,\n",
        "    device_map=\"auto\",\n",
        ")\n",
        "base_model.generation_config.pad_token_id = tokenizer.pad_token_id\n",
        "\n",
        "print(f\"Memory footprint: {base_model.get_memory_footprint() / 1e6:.1f} MB\")"
      ],
      "metadata": {
        "id": "R_O04fKxMMT-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Let's try out the model before we do fine-tuning"
      ],
      "metadata": {
        "id": "UObo1-RqaNnT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = \"\"\"\n",
        "# tickers is a list of stock tickers\n",
        "import tickers\n",
        "\n",
        "# prices is a dict; the key is a ticker and the value is a list of historic prices, today first\n",
        "import prices\n",
        "\n",
        "# Trade represents a decision to buy or sell a quantity of a ticker\n",
        "import Trade\n",
        "\n",
        "import random\n",
        "import numpy as np\n",
        "\n",
        "def trade():\n",
        "\"\"\""
      ],
      "metadata": {
        "id": "oaXOPhnySCcu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import TextStreamer\n",
        "streamer = TextStreamer(tokenizer)\n",
        "\n",
        "inputs = tokenizer.encode(prompt, return_tensors=\"pt\").to(\"cuda\")\n",
        "outputs = base_model.generate(inputs, max_new_tokens=100, streamer=streamer)"
      ],
      "metadata": {
        "id": "30lzJXBH7BcK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Load our dataset\n",
        "dataset = load_dataset(DATASET_NAME)['train']\n",
        "dataset"
      ],
      "metadata": {
        "id": "kVcmuZVgAAgr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# First, specify the configuration parameters for LoRA\n",
        "\n",
        "peft_parameters = LoraConfig(\n",
        "    lora_alpha=LORA_ALPHA,\n",
        "    lora_dropout=LORA_DROPOUT,\n",
        "    r=LORA_R,\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\",\n",
        "    target_modules=TARGET_MODULES,\n",
        ")\n",
        "\n",
        "# Next, specify the general configuration parameters for training\n",
        "\n",
        "train_params = SFTConfig(\n",
        "    output_dir=PROJECT_RUN_NAME,\n",
        "    num_train_epochs=EPOCHS,\n",
        "    per_device_train_batch_size=BATCH_SIZE,\n",
        "    per_device_eval_batch_size=1,\n",
        "    eval_strategy=\"no\",\n",
        "    gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,\n",
        "    optim=\"paged_adamw_32bit\",\n",
        "    save_steps=SAVE_STEPS,\n",
        "    save_total_limit=10,\n",
        "    logging_steps=STEPS,\n",
        "    learning_rate=LEARNING_RATE,\n",
        "    weight_decay=WEIGHT_DECAY,\n",
        "    fp16=False,\n",
        "    bf16=True,\n",
        "    max_grad_norm=0.3,\n",
        "    max_steps=-1,\n",
        "    warmup_ratio=0.03,\n",
        "    group_by_length=True,\n",
        "    lr_scheduler_type=LR_SCHEDULER_TYPE,\n",
        "    report_to=\"wandb\",\n",
        "    run_name=RUN_NAME,\n",
        "    max_seq_length=MAX_SEQUENCE_LENGTH,\n",
        "    dataset_text_field=\"text\",\n",
        ")\n",
        "\n",
        "# And now, the Supervised Fine Tuning Trainer will carry out the fine-tuning\n",
        "# Given these 2 sets of configuration parameters\n",
        "\n",
        "fine_tuning = SFTTrainer(\n",
        "    model=base_model,\n",
        "    train_dataset=dataset,\n",
        "    peft_config=peft_parameters,\n",
        "    tokenizer=tokenizer,\n",
        "    args=train_params\n",
        ")\n",
        "\n",
        "# Fine-tune!\n",
        "fine_tuning.train()\n",
        "\n",
        "# Push our fine-tuned model to Hugging Face\n",
        "fine_tuning.model.push_to_hub(PROJECT_RUN_NAME, private=True)"
      ],
      "metadata": {
        "id": "fCwmDmkSATvj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Code up a trade\n",
        "\n",
        "inputs = tokenizer.encode(prompt, return_tensors=\"pt\").to(\"cuda\")\n",
        "outputs = fine_tuning.model.generate(inputs, max_new_tokens=120, streamer=streamer)"
      ],
      "metadata": {
        "id": "3MGyNCSAFfy6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Another!\n",
        "\n",
        "outputs = fine_tuning.model.generate(inputs, max_new_tokens=120, streamer=streamer)"
      ],
      "metadata": {
        "id": "chiHKzbRtHed"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## That's the example of QLoRA Fine Tuning to write code to carry out a specific function (but don't actually use this for trading!)"
      ],
      "metadata": {
        "id": "QjktU3874KdY"
      }
    }
  ]
}