FinGPT_Training_LoRA_with_Baichuan2_7B.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "4oLjc7bbv0tO"
      },
      "source": [
        "## Part 1: Preparing the Data\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-maUV8CH7JPB"
      },
      "outputs": [],
      "source": [
        "!pip install datasets torch torchvision torchaudio tqdm pandas huggingface_hub\n",
        "!pip install transformers==4.30\n",
        "!pip install sentencepiece\n",
        "!pip install protobuf cpm_kernels gradio mdtex2html sentencepiece accelerate\n",
        "!pip install loguru\n",
        "!pip install datasets\n",
        "!pip install peft\n",
        "!pip install bitsandbytes\n",
        "!pip install tensorboard\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "hJp_UOiB70o3"
      },
      "source": [
        "### 1.1 Initialize Directories:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "aBE7gRUJ3L8u"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import shutil\n",
        "\n",
        "jsonl_path = \"../data/dataset_new.jsonl\"\n",
        "save_path = '../data/dataset_new'\n",
        "\n",
        "\n",
        "if os.path.exists(jsonl_path):\n",
        "    os.remove(jsonl_path)\n",
        "\n",
        "if os.path.exists(save_path):\n",
        "    shutil.rmtree(save_path)\n",
        "\n",
        "directory = \"../data\"\n",
        "if not os.path.exists(directory):\n",
        "    os.makedirs(directory)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "OvIBHhS5pV8o"
      },
      "source": [
        "### 1.2 Load and Prepare Dataset:\n",
        "\n",
        "* Import necessary libraries from the datasets package: https://huggingface.co/docs/datasets/index\n",
        "* Load the Twitter Financial News Sentiment (TFNS) dataset and convert it to a Pandas dataframe. https://huggingface.co/datasets/zeroshot/twitter-financial-news-sentiment\n",
        "* Map numerical labels to their corresponding sentiments (negative, positive, neutral).\n",
        "* Add instruction for each data entry, which is crucial for Instruction Tuning.\n",
        "* Convert the Pandas dataframe back to a Hugging Face Dataset object.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wVVWAq54ohCT",
        "outputId": "e04ef3e5-fed9-4648-b988-61c918d592ba"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Dataset({\n",
              "    features: ['input', 'output', 'instruction'],\n",
              "    num_rows: 9543\n",
              "})"
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from datasets import load_dataset\n",
        "import datasets\n",
        "\n",
        "dic = {\n",
        "    0:\"negative\",\n",
        "    1:'positive',\n",
        "    2:'neutral',\n",
        "}\n",
        "\n",
        "tfns = load_dataset('zeroshot/twitter-financial-news-sentiment')\n",
        "tfns = tfns['train']\n",
        "tfns = tfns.to_pandas()\n",
        "tfns['label'] = tfns['label'].apply(lambda x:dic[x])\n",
        "tfns['instruction'] = 'What is the sentiment of this tweet? Please choose an answer from {negative/neutral/positive}.'\n",
        "tfns.columns = ['input', 'output', 'instruction']\n",
        "tfns = datasets.Dataset.from_pandas(tfns)\n",
        "tfns"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "BRXfrEc8-yWy"
      },
      "source": [
        "### 1.3 Concatenate and Shuffle Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nXND6llnpTQj",
        "outputId": "4b807ce3-1093-43f6-f43d-5fdb127efed7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "19086\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "(19086, 3)"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tmp_dataset = datasets.concatenate_datasets([tfns]*2)\n",
        "train_dataset = tmp_dataset\n",
        "print(tmp_dataset.num_rows)\n",
        "\n",
        "all_dataset = train_dataset.shuffle(seed = 42)\n",
        "all_dataset.shape"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "oOVsy4if8_et"
      },
      "source": [
        "Now that your training data is loaded and prepared."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Oob5cnSlpi7C"
      },
      "source": [
        "## Part 2: Dataset Formatting and Tokenization\n",
        "Once your data is prepared, the next steps involve formatting the dataset for model ingestion and tokenizing the input data. Below, we provide a step-by-step breakdown of the code snippets shared.\n",
        "\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "1vI4We-r9QXh"
      },
      "source": [
        "### 2.1 Dataset Formatting:\n",
        "You need to structure your data in a specific format that aligns with the training process.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "mnZFyhcvpTIv"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "from tqdm.notebook import tqdm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "DjmW94OTpjuU"
      },
      "outputs": [],
      "source": [
        "def format_example(example: dict) -> dict:\n",
        "    context = f\"Instruction: {example['instruction']}\\n\"\n",
        "    if example.get(\"input\"):\n",
        "        context += f\"Input: {example['input']}\\n\"\n",
        "    context += \"Answer: \"\n",
        "    target = example[\"output\"]\n",
        "    return {\"context\": context, \"target\": target}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "oomCwuggpjsF"
      },
      "outputs": [],
      "source": [
        "data_list = []\n",
        "for item in all_dataset.to_pandas().itertuples():\n",
        "    tmp = {}\n",
        "    tmp[\"instruction\"] = item.instruction\n",
        "    tmp[\"input\"] = item.input\n",
        "    tmp[\"output\"] = item.output\n",
        "    data_list.append(tmp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "8c47dc30172c4e24986893d9c624d1ce",
            "568b044661a74d0c8cdb7acd76247102",
            "239f6a26d4e044e284ad27dbfa0849c7",
            "631af510b008492294e00ad89d144818",
            "075ba55d210e400b9053b7b0d25d826e",
            "b3f7924f573b47e0ae1807ca479755fc",
            "fc81ffa52b344d7baea4504fda77e7e9",
            "38f6b3078f4f4364966b03409000ae57",
            "ab1460ae112f4b63970295aaf576393e",
            "703bdb4d9bc14241b52549b7539d6a73",
            "10877fd540614d6fa5d94cbbd9b6d032"
          ]
        },
        "id": "VuuRhlD1pjqH",
        "outputId": "f61069b9-a4f7-45c7-ec30-774628ef446d"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "8c47dc30172c4e24986893d9c624d1ce",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "formatting..:   0%|          | 0/19086 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# save to a jsonl file\n",
        "with open(\"../data/dataset_new.jsonl\", 'w') as f:\n",
        "    for example in tqdm(data_list, desc=\"formatting..\"):\n",
        "        f.write(json.dumps(format_example(example)) + '\\n')"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "YptAdKh0preG"
      },
      "source": [
        "### 2.2 Tokenization\n",
        "Tokenization is the process of converting input text into tokens that can be fed into the model.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "Sjh98Wr_pjmI"
      },
      "outputs": [],
      "source": [
        "import datasets\n",
        "from transformers import AutoTokenizer, AutoConfig\n",
        "\n",
        "model_name = \"baichuan-inc/Baichuan2-7B-Base\"\n",
        "jsonl_path = \"../data/dataset_new.jsonl\"  # updated path\n",
        "save_path = '../data/dataset_new'  # updated path\n",
        "max_seq_length = 512\n",
        "skip_overlength = True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "_WJtjr5MpjkA"
      },
      "outputs": [],
      "source": [
        "# The preprocess function tokenizes the prompt and target, combines them into input IDs,\n",
        "# and then trims or pads the sequence to the maximum sequence length.\n",
        "def preprocess(tokenizer, config, example, max_seq_length):\n",
        "    prompt = example[\"context\"]\n",
        "    target = example[\"target\"]\n",
        "    prompt_ids = tokenizer.encode(prompt, max_length=max_seq_length, truncation=True)\n",
        "    target_ids = tokenizer.encode(\n",
        "        target,\n",
        "        max_length=max_seq_length,\n",
        "        truncation=True,\n",
        "        add_special_tokens=False)\n",
        "    input_ids = prompt_ids + target_ids + [config.eos_token_id]\n",
        "    return {\"input_ids\": input_ids, \"seq_len\": len(prompt_ids)}\n",
        "\n",
        "# The read_jsonl function reads each line from the JSONL file, preprocesses it using the preprocess function,\n",
        "# and then yields each preprocessed example.\n",
        "def read_jsonl(path, max_seq_length, skip_overlength=False):\n",
        "    tokenizer = AutoTokenizer.from_pretrained(\n",
        "        model_name, trust_remote_code=True)\n",
        "    config = AutoConfig.from_pretrained(\n",
        "        model_name, trust_remote_code=True, device_map='auto')\n",
        "    with open(path, \"r\") as f:\n",
        "        for line in tqdm(f.readlines()):\n",
        "            example = json.loads(line)\n",
        "            feature = preprocess(tokenizer, config, example, max_seq_length)\n",
        "            if skip_overlength and len(feature[\"input_ids\"]) > max_seq_length:\n",
        "                continue\n",
        "            feature[\"input_ids\"] = feature[\"input_ids\"][:max_seq_length]\n",
        "            yield feature"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "j9M_NsM2-9Mk"
      },
      "source": [
        "### 2.3 Save the dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "6fc06d59622c4e578ccbea06294e8798",
            "7409ccca97a840f5b6b65cfe96ee8c17",
            "68bb67448d7e4d0c981d047f91e0a076",
            "7b074461c2af46439560e0d871ea821a",
            "84af93d2aa224b528a3905e4ad16d383",
            "a2d6e7aaef5a4dca9393f03c4578566a",
            "606533fc0262444ab838d61f3df8a071",
            "353bca75a83d4a2f97936f20c3978cb2",
            "8eb67374698449939a44d14dad9c02a2",
            "0cd02ac4e7dd4512b102fd82f36e40ba",
            "728c96bdf81f41e697a2c55bf0684e1a"
          ]
        },
        "id": "B5yRJ52HpjiD",
        "outputId": "5fa2a6c6-a39e-41f6-b6f5-0acfa64f6a3a"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "6fc06d59622c4e578ccbea06294e8798",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Saving the dataset (0/1 shards):   0%|          | 0/19086 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# The script then creates a Hugging Face Dataset object from the generator and saves it to disk.\n",
        "save_path = '../data/dataset_new'\n",
        "\n",
        "dataset = datasets.Dataset.from_generator(\n",
        "    lambda: read_jsonl(jsonl_path, max_seq_length, skip_overlength)\n",
        "    )\n",
        "dataset.save_to_disk(save_path)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "e4AgPl_t04qV"
      },
      "source": [
        "## Part 3: Setup FinGPT training parameters with LoRA on Baichuan2-7B\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KrHuFgcNqSq-",
        "outputId": "647ef04a-c8ac-408e-c88c-c797faf1bbce"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tzvixpioKKJD",
        "outputId": "042dd824-4f35-4546-d55a-0aae5aa690e6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "4.30.0\n"
          ]
        }
      ],
      "source": [
        "import transformers\n",
        "print(transformers.__version__)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "Jpn7KAIW9RuP"
      },
      "outputs": [],
      "source": [
        "# Ensure CUDA is accessible in the system path\n",
        "# Only for Windows Subsystem for Linux (WSL)\n",
        "import os\n",
        "os.environ[\"PATH\"] = f\"{os.environ['PATH']}:/usr/local/cuda/bin\"\n",
        "os.environ['LD_LIBRARY_PATH'] = \"/usr/lib/wsl/lib:/usr/local/cuda/lib64\""
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "aLjcEr3FBGzi"
      },
      "source": [
        "### 3.1 Training Arguments Setup:\n",
        "Initialize and set training arguments.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "id": "gW5hW4-F9RsX"
      },
      "outputs": [],
      "source": [
        "from typing import List, Dict, Optional\n",
        "import torch\n",
        "import bitsandbytes\n",
        "from loguru import logger\n",
        "from transformers import (\n",
        "    AutoModel,\n",
        "    AutoModelForCausalLM,\n",
        "    AutoTokenizer,\n",
        "    TrainingArguments,\n",
        "    Trainer,\n",
        "    BitsAndBytesConfig\n",
        ")\n",
        "from peft import (\n",
        "    TaskType,\n",
        "    LoraConfig,\n",
        "    get_peft_model,\n",
        "    set_peft_model_state_dict,\n",
        "    prepare_model_for_kbit_training,\n",
        "    prepare_model_for_int8_training,\n",
        ")\n",
        "from peft.utils import TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "id": "4w2Pe1rT9Rqv"
      },
      "outputs": [],
      "source": [
        "training_args = TrainingArguments(\n",
        "        output_dir='./finetuned_model_FinGPT_Baichuan2_Lora_Wpack',    # saved model path\n",
        "        logging_steps = 500,\n",
        "        # max_steps=10000,\n",
        "        num_train_epochs = 2,\n",
        "        per_device_train_batch_size=4,\n",
        "        gradient_accumulation_steps=8,\n",
        "        learning_rate=1e-4,\n",
        "        weight_decay=0.01,\n",
        "        warmup_steps=1000,\n",
        "        save_steps=500,\n",
        "        fp16=True,\n",
        "        # bf16=True,\n",
        "        torch_compile = False,\n",
        "        load_best_model_at_end = True,\n",
        "        evaluation_strategy=\"steps\",\n",
        "        remove_unused_columns=False,\n",
        "\n",
        "    )"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "QGabm7cyBM6r"
      },
      "source": [
        "### 3.2 Quantization Config Setup:\n",
        "Set quantization configuration to reduce model size without losing significant precision.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "OfKTizyP9Rmn"
      },
      "outputs": [],
      "source": [
        " # Quantization\n",
        "q_config = BitsAndBytesConfig(load_in_4bit=True,\n",
        "                                bnb_4bit_quant_type='nf4',\n",
        "                                bnb_4bit_use_double_quant=True,\n",
        "                                bnb_4bit_compute_dtype=torch.float16\n",
        "                                )"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "WMqFjp_mBVqO"
      },
      "source": [
        "### 3.3 Model Loading & Preparation:\n",
        "Load the base model and tokenizer, and prepare the model for INT8 training.\n",
        "\n",
        "* **Runtime -> Change runtime type -> A100 GPU**\n",
        "* retart runtime and run again if not working\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "988b475808c449d1b09e098042fa8460",
            "a2f0c4ff05674ee28aeb733cd7a9f238",
            "da2f8a4f58914e488c686ec73bd0cec7",
            "69f4f67cb1a14318b2673ae1dc620118",
            "fce12fe312934f298bec9ea8f970d1d5",
            "c4efdfd442424f9b911f9aacb57584d1",
            "71a0293a0a234fa883bce9aee63929de",
            "584815c8c1ac4183849149f2e576e730",
            "dd993d2ff7c64c3d99c8d95aad54e11d",
            "a0e26d608d954ffe8826d8aa58fbd76d",
            "b86efc8a7ce94d61af24f55b68844115"
          ]
        },
        "id": "th_3Rnqy9Rkg",
        "outputId": "c5747037-360b-4bbd-da38-ecfae6ca8873"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "988b475808c449d1b09e098042fa8460",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# Load tokenizer & model\n",
        "# need massive space\n",
        "model_name = \"baichuan-inc/Baichuan2-7B-Base\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side=\"left\", trust_remote_code=True)\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "        model_name,\n",
        "        quantization_config=q_config,\n",
        "        trust_remote_code=True,\n",
        "        device='cuda'\n",
        "    )\n",
        "model = prepare_model_for_int8_training(model, use_gradient_checkpointing = True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B7OZDI1H4nBW"
      },
      "outputs": [],
      "source": [
        "# pip install accelerate bitsandbytes"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Dfw3w9FNBgyF"
      },
      "source": [
        "### 3.4 LoRA Config & Setup:\n",
        "Implement Low-Rank Adaptation (LoRA) and print trainable parameters.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "id": "IDbX7Ati9RiK"
      },
      "outputs": [],
      "source": [
        "def print_trainable_parameters(model):\n",
        "    \"\"\"\n",
        "    Prints the number of trainable parameters in the model.\n",
        "    \"\"\"\n",
        "    trainable_params = 0\n",
        "    all_param = 0\n",
        "    for _, param in model.named_parameters():\n",
        "        all_param += param.numel()\n",
        "        if param.requires_grad:\n",
        "            trainable_params += param.numel()\n",
        "    print(\n",
        "        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "id": "BVDyFIy-rNZK"
      },
      "outputs": [],
      "source": [
        "def find_linear_layers(model):\n",
        "    \"\"\"\n",
        "    Find all linear layers\n",
        "    \"\"\"\n",
        "    Linear_module = bitsandbytes.nn.Linear4bit\n",
        "    lora_modules = set()\n",
        "    for name, module in model.named_modules():\n",
        "        if isinstance(module, Linear_module):\n",
        "            names = name.split('.')\n",
        "            lora_modules.add(names[0] if len(names) == 1 else names[-1])\n",
        "\n",
        "    if 'lm_head' in lora_modules:\n",
        "        lora_modules.remove('lm_head')\n",
        "    return list(lora_modules)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZFh3szggbuGH",
        "outputId": "3103b064-eeaf-4608-d542-634af358edd7"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['gate_proj', 'down_proj', 'o_proj', 'W_pack', 'up_proj']"
            ]
          },
          "execution_count": 45,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "target_modules = find_linear_layers(model)\n",
        "target_modules"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RgPzHGUV9RgT",
        "outputId": "624d072f-b247-489f-eb34-4f873580af82"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "trainable params: 17891328 || all params: 4285861888 || trainable%: 0.41744994280132997\n"
          ]
        }
      ],
      "source": [
        "# LoRA\n",
        "lora_config = LoraConfig(\n",
        "    task_type=TaskType.CAUSAL_LM,\n",
        "    inference_mode=False,\n",
        "    r=8,\n",
        "    lora_alpha=32,\n",
        "    lora_dropout=0.1,\n",
        "    target_modules=target_modules,\n",
        "    bias='none',\n",
        ")\n",
        "model = get_peft_model(model, lora_config)\n",
        "print_trainable_parameters(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "metadata": {
        "id": "Qt2DjmPn9ReS"
      },
      "outputs": [],
      "source": [
        "resume_from_checkpoint = None\n",
        "if resume_from_checkpoint is not None:\n",
        "    checkpoint_name = os.path.join(resume_from_checkpoint, 'pytorch_model.bin')\n",
        "    if not os.path.exists(checkpoint_name):\n",
        "        checkpoint_name = os.path.join(\n",
        "            resume_from_checkpoint, 'adapter_model.bin'\n",
        "        )\n",
        "        resume_from_checkpoint = False\n",
        "    if os.path.exists(checkpoint_name):\n",
        "        logger.info(f'Restarting from {checkpoint_name}')\n",
        "        adapters_weights = torch.load(checkpoint_name)\n",
        "        set_peft_model_state_dict(model, adapters_weights)\n",
        "    else:\n",
        "        logger.info(f'Checkpoint {checkpoint_name} not found')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kOT2o7FU9Rcc",
        "outputId": "81e066ad-7bd6-475e-c215-a068f25397de"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "trainable params: 17,891,328 || all params: 7,523,864,576 || trainable%: 0.23779439168895536\n"
          ]
        }
      ],
      "source": [
        "model.print_trainable_parameters()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "AsRYIMb_CLYR"
      },
      "source": [
        "## Part 4: Loading Data and Training FinGPT\n",
        "In this segment, we'll delve into the loading of your pre-processed data, and finally, launch the training of your FinGPT model. Here's a stepwise breakdown of the script provided:\n",
        "* Need to purchase Google Colab GPU plans, Colab Pro is sufficient or just buy 100 compute units for $10\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "nvmJH99eCR5Z"
      },
      "source": [
        "### 4.1 Loading Your Data:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "metadata": {
        "id": "4PbUCglW9RaM"
      },
      "outputs": [],
      "source": [
        "# load data\n",
        "from datasets import load_from_disk\n",
        "import datasets\n",
        "\n",
        "dataset = datasets.load_from_disk(\"../data/dataset_new\")\n",
        "dataset = dataset.train_test_split(0.2, shuffle=True, seed = 42)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Lq87Er6RCZbb"
      },
      "source": [
        "### 4.2 Training Configuration and Launch:\n",
        "* Customize the Trainer class for specific loss computation, prediction step, and model-saving methods.\n",
        "\n",
        "* Define a data collator function to process batches of data during training.\n",
        "\n",
        "* Set up TensorBoard for logging, instantiate your modified trainer, and begin training.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 50,
      "metadata": {
        "id": "dUgghgf59RYG"
      },
      "outputs": [],
      "source": [
        "class ModifiedTrainer(Trainer):\n",
        "    def compute_loss(self, model, inputs, return_outputs=False):\n",
        "        return model(\n",
        "            input_ids=inputs[\"input_ids\"],\n",
        "            labels=inputs[\"labels\"],\n",
        "        ).loss\n",
        "\n",
        "    def prediction_step(self, model: torch.nn.Module, inputs, prediction_loss_only: bool, ignore_keys = None):\n",
        "        with torch.no_grad():\n",
        "            res = model(\n",
        "                input_ids=inputs[\"input_ids\"].to(model.device),\n",
        "                labels=inputs[\"labels\"].to(model.device),\n",
        "            ).loss\n",
        "        return (res, None, None)\n",
        "\n",
        "    def save_model(self, output_dir=None, _internal_call=False):\n",
        "        from transformers.trainer import TRAINING_ARGS_NAME\n",
        "\n",
        "        os.makedirs(output_dir, exist_ok=True)\n",
        "        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))\n",
        "        saved_params = {\n",
        "            k: v.to(\"cpu\") for k, v in self.model.named_parameters() if v.requires_grad\n",
        "        }\n",
        "        torch.save(saved_params, os.path.join(output_dir, \"adapter_model.bin\"))\n",
        "\n",
        "def data_collator(features: list) -> dict:\n",
        "    len_ids = [len(feature[\"input_ids\"]) for feature in features]\n",
        "    longest = max(len_ids)\n",
        "    input_ids = []\n",
        "    labels_list = []\n",
        "    for ids_l, feature in sorted(zip(len_ids, features), key=lambda x: -x[0]):\n",
        "        ids = feature[\"input_ids\"]\n",
        "        seq_len = feature[\"seq_len\"]\n",
        "        labels = (\n",
        "            [tokenizer.pad_token_id] * (seq_len - 1) + ids[(seq_len - 1) :] + [tokenizer.pad_token_id] * (longest - ids_l)\n",
        "        )\n",
        "        ids = ids + [tokenizer.pad_token_id] * (longest - ids_l)\n",
        "        _ids = torch.LongTensor(ids)\n",
        "        labels_list.append(torch.LongTensor(labels))\n",
        "        input_ids.append(_ids)\n",
        "    input_ids = torch.stack(input_ids)\n",
        "    labels = torch.stack(labels_list)\n",
        "    return {\n",
        "        \"input_ids\": input_ids,\n",
        "        \"labels\": labels,\n",
        "    }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "metadata": {
        "id": "ITubEZSK9RVv"
      },
      "outputs": [],
      "source": [
        "from torch.utils.tensorboard import SummaryWriter\n",
        "from transformers.integrations import TensorBoardCallback"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 52,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 305
        },
        "id": "Cw4Zik6a9RT3",
        "outputId": "5b34e774-3c04-46ed-8c7d-5d4b8e393345"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "You are adding a <class 'transformers.integrations.TensorBoardCallback'> to the callbacks of this Trainer, but there is already one. The currentlist of callbacks is\n",
            ":DefaultFlowCallback\n",
            "TensorBoardCallback\n",
            "/usr/local/lib/python3.10/dist-packages/transformers/optimization.py:411: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
            "  warnings.warn(\n",
            "WARNING:transformers_modules.baichuan-inc.Baichuan2-7B-Base.f9d4d8dd2f7a3dbede3bda3b0cf0224e9272bbe5.modeling_baichuan:`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n",
            "/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='954' max='954' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [954/954 59:02, Epoch 1/2]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              " <tr style=\"text-align: left;\">\n",
              "      <th>Step</th>\n",
              "      <th>Training Loss</th>\n",
              "      <th>Validation Loss</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>500</td>\n",
              "      <td>9.122400</td>\n",
              "      <td>6.480052</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
            "  warnings.warn(\n"
          ]
        }
      ],
      "source": [
        "# Train\n",
        "# Took about 10 compute units\n",
        "writer = SummaryWriter()\n",
        "trainer = ModifiedTrainer(\n",
        "    model=model,\n",
        "    args=training_args,             # Trainer args\n",
        "    train_dataset=dataset[\"train\"], # Training set\n",
        "    eval_dataset=dataset[\"test\"],   # Testing set\n",
        "    data_collator=data_collator,    # Data Collator\n",
        "    callbacks=[TensorBoardCallback(writer)],\n",
        ")\n",
        "trainer.train()\n",
        "writer.close()\n",
        "# save model\n",
        "model.save_pretrained(training_args.output_dir)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "brHTWfnmCn5D"
      },
      "source": [
        "### 4.3 Model Saving and Download:\n",
        "After training, save and download your model. You can also check the model's size.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HUYxzwS_9lMI",
        "outputId": "7203dbe1-3864-4901-b873-b15c52fcdcaa"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/ (stored 0%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/README.md (deflated 66%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/runs/ (stored 0%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/runs/Feb05_02-47-57_21c7a47ea268/ (stored 0%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/runs/Feb05_02-47-57_21c7a47ea268/events.out.tfevents.1707101482.21c7a47ea268.14980.3 (deflated 58%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/runs/Feb05_02-47-57_21c7a47ea268/events.out.tfevents.1707101418.21c7a47ea268.14980.1 (deflated 61%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/checkpoint-500/ (stored 0%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/checkpoint-500/rng_state.pth (deflated 25%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/checkpoint-500/trainer_state.json (deflated 50%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/checkpoint-500/training_args.bin (deflated 50%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/checkpoint-500/scheduler.pt (deflated 55%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/checkpoint-500/optimizer.pt (deflated 8%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/checkpoint-500/adapter_model.bin (deflated 48%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/adapter_model.safetensors (deflated 7%)\n",
            "  adding: content/./finetuned_model_FinGPT_Baichuan2_Lora_Wpack/adapter_config.json (deflated 49%)\n"
          ]
        }
      ],
      "source": [
        "!zip -r /content/saved_model.zip /content/{training_args.output_dir}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s-7s2Cjw9pAM"
      },
      "outputs": [],
      "source": [
        "# download to local\n",
        "from google.colab import files\n",
        "files.download('/content/saved_model.zip')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rvBgMgO8RADU",
        "outputId": "2e9e0356-2284-48ac-b98b-859de4c08031"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ],
      "source": [
        "# save to google drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 55,
      "metadata": {
        "id": "UUctmjm8RIfQ"
      },
      "outputs": [],
      "source": [
        "# save the finetuned model to google drive\n",
        "!cp -r \"/content/finetuned_model_FinGPT_Baichuan2_Lora_Wpack\" \"/content/drive/MyDrive\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 56,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "unRoLshR9RQZ",
        "outputId": "f1e947b6-3d9f-42dd-936e-21d2f967c9f2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "adapter_config.json 0.0006227493286132812\n",
            "README.md 0.0048694610595703125\n",
            "adapter_model.safetensors 68.29019165039062\n",
            "events.out.tfevents.1707316970.91671f1dd2d4.782.1 0.005206108093261719\n",
            "optimizer.pt 136.6810245513916\n",
            "trainer_state.json 0.0007123947143554688\n",
            "adapter_model.bin 2032.358289718628\n",
            "training_args.bin 0.00426483154296875\n",
            "rng_state.pth 0.013584136962890625\n",
            "scheduler.pt 0.00101470947265625\n",
            "Model size: 2237.3597803115845 MB\n"
          ]
        }
      ],
      "source": [
        "def get_folder_size(folder_path):\n",
        "    total_size = 0\n",
        "    for dirpath, _, filenames in os.walk(folder_path):\n",
        "        for f in filenames:\n",
        "            fp = os.path.join(dirpath, f)\n",
        "            print(f, os.path.getsize(fp)/ 1024 / 1024)\n",
        "            total_size += os.path.getsize(fp)\n",
        "    return total_size / 1024 / 1024  # Size in MB\n",
        "\n",
        "model_size = get_folder_size(training_args.output_dir)\n",
        "print(f\"Model size: {model_size} MB\")\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "1LCjYKuoCusU"
      },
      "source": [
        "Now your model is trained and saved! You can download it and use it for generating financial insights or any other relevant tasks in the finance domain. The usage of TensorBoard allows you to deeply understand and visualize the training dynamics and performance of your model in real-time.\n",
        "\n",
        "Happy FinGPT Training! 🚀"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "76g_Qlp8t_Yp"
      },
      "source": [
        "## Part 5: Inference and Benchmarks using FinGPT\n",
        "Now that your model is trained, let’s understand how to use it to infer and run benchmarks.\n",
        "* Took about 10 compute units\n",
        "\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "P91SXTrLS34i"
      },
      "source": [
        "### 5.1 Load the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 57,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "y5jyY7S_uEls",
        "outputId": "d2445c00-4cf0-4868-b84e-6fbc4dd2ae2e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Cloning into 'FinNLP'...\n",
            "remote: Enumerating objects: 1399, done.\u001b[K\n",
            "remote: Counting objects: 100% (458/458), done.\u001b[K\n",
            "remote: Compressing objects: 100% (192/192), done.\u001b[K\n",
            "remote: Total 1399 (delta 239), reused 431 (delta 224), pack-reused 941\u001b[K\n",
            "Receiving objects: 100% (1399/1399), 4.97 MiB | 10.47 MiB/s, done.\n",
            "Resolving deltas: 100% (612/612), done.\n"
          ]
        }
      ],
      "source": [
        "#clone the FinNLP repository\n",
        "!git clone https://github.com/AI4Finance-Foundation/FinNLP.git\n",
        "\n",
        "import sys\n",
        "sys.path.append('/content/FinNLP/')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 58,
      "metadata": {
        "id": "zRsmSTFZuEjt"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM\n",
        "\n",
        "from peft import PeftModel\n",
        "import torch\n",
        "import os\n",
        "\n",
        "# Load benchmark datasets from FinNLP\n",
        "from finnlp.benchmarks.fpb import test_fpb\n",
        "from finnlp.benchmarks.fiqa import test_fiqa , add_instructions\n",
        "from finnlp.benchmarks.tfns import test_tfns\n",
        "from finnlp.benchmarks.nwgi import test_nwgi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 59,
      "metadata": {
        "id": "EBqKeUYV9VjF"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade peft"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 60,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gRRw9drdA2hv",
        "outputId": "3ee8b275-37bc-4545-c11f-3a5ea8665d5e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ],
      "source": [
        "# load model from google drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 61,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LxCAhg9QpkyI",
        "outputId": "ee2cdb9e-4c0b-48a3-8405-5021e2252e09"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Path exists.\n"
          ]
        }
      ],
      "source": [
        "# Define the path you want to check\n",
        "path_to_check = \"/content/drive/My Drive/finetuned_model_FinGPT_Baichuan2_Lora_Wpack\"\n",
        "\n",
        "# Check if the specified path exists\n",
        "if os.path.exists(path_to_check):\n",
        "    print(\"Path exists.\")\n",
        "else:\n",
        "    print(\"Path does not exist.\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 71,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "013dc1fbb5da4737afe4ca11933034cc",
            "c9d3016cc64c4640b8cbb48962223d02",
            "bf3ceb473d16470dbe825126832fda4f",
            "33f53f0a28a44dc88165194093e80547",
            "7af9a8cb5b6d44b0bd05affc839a5c57",
            "885bb3979c554d3a82695eaa3efe569b",
            "1cc91e5168684f139a63e162a4df3c00",
            "116679cd0ef342519bef389e3c25835f",
            "3cc2e7e24e0e479496f1115a1bab9c83",
            "e7be1007f3874cba911c01f2e11ff997",
            "cfdc876e9b27452d9c021d0717b32f99"
          ]
        },
        "id": "w37K19SvlwYh",
        "outputId": "4c582cec-8f12-40f0-c4a0-6ecddf0ba84d"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "013dc1fbb5da4737afe4ca11933034cc",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "## load our finetuned model\n",
        "base_model = \"baichuan-inc/Baichuan2-7B-Base\"\n",
        "peft_model = \"/content/drive/My Drive/finetuned_model_FinGPT_Baichuan2_Lora_Wpack\"\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(base_model, padding_side='left', trust_remote_code=True)\n",
        "model = AutoModelForCausalLM.from_pretrained(base_model, quantization_config=q_config, trust_remote_code=True, device_map=\"auto\")\n",
        "\n",
        "model = PeftModel.from_pretrained(model, peft_model)\n",
        "\n",
        "model = model.eval()\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "HZGMT76ITkBY"
      },
      "source": [
        "### 5.2 Run Benchmarks:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 72,
      "metadata": {
        "id": "rv5tja2AuEf3"
      },
      "outputs": [],
      "source": [
        "batch_size = 8"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 73,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0Z4capS_uEeF",
        "outputId": "d663ab2f-9800-4be8-b3ef-10e44261e3c1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "Prompt example:\n",
            "Instruction: What is the sentiment of this tweet? Please choose an answer from {negative/neutral/positive}.\n",
            "Input: $ALLY - Ally Financial pulls outlook https://t.co/G9Zdi1boy5\n",
            "Answer: \n",
            "\n",
            "\n",
            "Total len: 2388. Batchsize: 8. Total steps: 299\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 299/299 [01:00<00:00,  4.92it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Acc: 0.6896984924623115. F1 macro: 0.6783211306631469. F1 micro: 0.6896984924623115. F1 weighted (BloombergGPT): 0.6969986385377054. \n"
          ]
        }
      ],
      "source": [
        "# TFNS Test Set, len 2388\n",
        "# Available: 84.85 compute units\n",
        "res = test_tfns(model, tokenizer, batch_size = batch_size)\n",
        "# Available: 83.75 compute units\n",
        "# Took about 1 compute unite to inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 74,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7we6z7kwuEcO",
        "outputId": "9a2c375a-558e-46f3-903e-eed5a03c3f6a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "Prompt example:\n",
            "Instruction: What is the sentiment of this news? Please choose an answer from {negative/neutral/positive}.\n",
            "Input: L&T has also made a commitment to redeem the remaining shares by the end of 2011 .\n",
            "Answer: \n",
            "\n",
            "\n",
            "Total len: 1212. Batchsize: 8. Total steps: 152\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 152/152 [00:32<00:00,  4.68it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Acc: 0.6947194719471947. F1 macro: 0.7007768002625792. F1 micro: 0.6947194719471947. F1 weighted (BloombergGPT): 0.6992353440126955. \n"
          ]
        }
      ],
      "source": [
        "# FPB, len 1212\n",
        "res = test_fpb(model, tokenizer, batch_size = batch_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 75,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UTpEPwnIuEaN",
        "outputId": "63b61d02-2fd5-4c67-c0d9-53d97f703772"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "Prompt example:\n",
            "Instruction: What is the sentiment of this tweet? Please choose an answer from {negative/neutral/positive}.\n",
            "Input: This $BBBY stock options trade would have more than doubled your money https://t.co/Oa0loiRIJL via @TheStreet\n",
            "Answer: \n",
            "\n",
            "\n",
            "Total len: 275. Batchsize: 8. Total steps: 35\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 35/35 [00:07<00:00,  4.93it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Acc: 0.7527272727272727. F1 macro: 0.645183025976165. F1 micro: 0.7527272727272727. F1 weighted (BloombergGPT): 0.7771972095318953. \n"
          ]
        }
      ],
      "source": [
        "# FiQA, len 275\n",
        "res = test_fiqa(model, tokenizer, prompt_fun = add_instructions, batch_size = batch_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 76,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tWl6ILV2uEYA",
        "outputId": "56b0ba9b-ade8-4170-bf1c-a3fb731becd2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "Prompt example:\n",
            "Instruction: What is the sentiment of this news? Please choose an answer from {negative/neutral/positive}.\n",
            "Input: In the latest trading session, Adobe Systems (ADBE) closed at $535.98, marking a +0.31% move from the previous day.\n",
            "Answer: \n",
            "\n",
            "\n",
            "Total len: 4047. Batchsize: 8. Total steps: 506\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 506/506 [02:12<00:00,  3.81it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Acc: 0.5932789720780826. F1 macro: 0.5816995173121812. F1 micro: 0.5932789720780826. F1 weighted (BloombergGPT): 0.569916573317543. \n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "# NWGI, len 4047\n",
        "res = test_nwgi(model, tokenizer, batch_size = batch_size)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "DcA1A-9aV0G2"
      },
      "source": [
        "#### Reference\n",
        "https://github.com/AI4Finance-Foundation/FinGPT/tree/master/fingpt/FinGPT-*v3*"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "A100",
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "013dc1fbb5da4737afe4ca11933034cc": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_c9d3016cc64c4640b8cbb48962223d02",
              "IPY_MODEL_bf3ceb473d16470dbe825126832fda4f",
              "IPY_MODEL_33f53f0a28a44dc88165194093e80547"
            ],
            "layout": "IPY_MODEL_7af9a8cb5b6d44b0bd05affc839a5c57"
          }
        },
        "075ba55d210e400b9053b7b0d25d826e": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0cd02ac4e7dd4512b102fd82f36e40ba": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "10877fd540614d6fa5d94cbbd9b6d032": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "116679cd0ef342519bef389e3c25835f": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "1cc91e5168684f139a63e162a4df3c00": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "239f6a26d4e044e284ad27dbfa0849c7": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_38f6b3078f4f4364966b03409000ae57",
            "max": 19086,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_ab1460ae112f4b63970295aaf576393e",
            "value": 19086
          }
        },
        "33f53f0a28a44dc88165194093e80547": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_e7be1007f3874cba911c01f2e11ff997",
            "placeholder": "​",
            "style": "IPY_MODEL_cfdc876e9b27452d9c021d0717b32f99",
            "value": " 2/2 [00:16&lt;00:00,  7.85s/it]"
          }
        },
        "353bca75a83d4a2f97936f20c3978cb2": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "38f6b3078f4f4364966b03409000ae57": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3cc2e7e24e0e479496f1115a1bab9c83": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "568b044661a74d0c8cdb7acd76247102": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b3f7924f573b47e0ae1807ca479755fc",
            "placeholder": "​",
            "style": "IPY_MODEL_fc81ffa52b344d7baea4504fda77e7e9",
            "value": "formatting..: 100%"
          }
        },
        "584815c8c1ac4183849149f2e576e730": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "606533fc0262444ab838d61f3df8a071": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "631af510b008492294e00ad89d144818": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_703bdb4d9bc14241b52549b7539d6a73",
            "placeholder": "​",
            "style": "IPY_MODEL_10877fd540614d6fa5d94cbbd9b6d032",
            "value": " 19086/19086 [00:00&lt;00:00, 164351.74it/s]"
          }
        },
        "68bb67448d7e4d0c981d047f91e0a076": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_353bca75a83d4a2f97936f20c3978cb2",
            "max": 19086,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_8eb67374698449939a44d14dad9c02a2",
            "value": 19086
          }
        },
        "69f4f67cb1a14318b2673ae1dc620118": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a0e26d608d954ffe8826d8aa58fbd76d",
            "placeholder": "​",
            "style": "IPY_MODEL_b86efc8a7ce94d61af24f55b68844115",
            "value": " 2/2 [00:16&lt;00:00,  7.80s/it]"
          }
        },
        "6fc06d59622c4e578ccbea06294e8798": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_7409ccca97a840f5b6b65cfe96ee8c17",
              "IPY_MODEL_68bb67448d7e4d0c981d047f91e0a076",
              "IPY_MODEL_7b074461c2af46439560e0d871ea821a"
            ],
            "layout": "IPY_MODEL_84af93d2aa224b528a3905e4ad16d383"
          }
        },
        "703bdb4d9bc14241b52549b7539d6a73": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "71a0293a0a234fa883bce9aee63929de": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "728c96bdf81f41e697a2c55bf0684e1a": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "7409ccca97a840f5b6b65cfe96ee8c17": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a2d6e7aaef5a4dca9393f03c4578566a",
            "placeholder": "​",
            "style": "IPY_MODEL_606533fc0262444ab838d61f3df8a071",
            "value": "Saving the dataset (1/1 shards): 100%"
          }
        },
        "7af9a8cb5b6d44b0bd05affc839a5c57": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7b074461c2af46439560e0d871ea821a": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0cd02ac4e7dd4512b102fd82f36e40ba",
            "placeholder": "​",
            "style": "IPY_MODEL_728c96bdf81f41e697a2c55bf0684e1a",
            "value": " 19086/19086 [00:00&lt;00:00, 724686.43 examples/s]"
          }
        },
        "84af93d2aa224b528a3905e4ad16d383": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "885bb3979c554d3a82695eaa3efe569b": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "8c47dc30172c4e24986893d9c624d1ce": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_568b044661a74d0c8cdb7acd76247102",
              "IPY_MODEL_239f6a26d4e044e284ad27dbfa0849c7",
              "IPY_MODEL_631af510b008492294e00ad89d144818"
            ],
            "layout": "IPY_MODEL_075ba55d210e400b9053b7b0d25d826e"
          }
        },
        "8eb67374698449939a44d14dad9c02a2": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "988b475808c449d1b09e098042fa8460": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_a2f0c4ff05674ee28aeb733cd7a9f238",
              "IPY_MODEL_da2f8a4f58914e488c686ec73bd0cec7",
              "IPY_MODEL_69f4f67cb1a14318b2673ae1dc620118"
            ],
            "layout": "IPY_MODEL_fce12fe312934f298bec9ea8f970d1d5"
          }
        },
        "a0e26d608d954ffe8826d8aa58fbd76d": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a2d6e7aaef5a4dca9393f03c4578566a": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a2f0c4ff05674ee28aeb733cd7a9f238": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c4efdfd442424f9b911f9aacb57584d1",
            "placeholder": "​",
            "style": "IPY_MODEL_71a0293a0a234fa883bce9aee63929de",
            "value": "Loading checkpoint shards: 100%"
          }
        },
        "ab1460ae112f4b63970295aaf576393e": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "b3f7924f573b47e0ae1807ca479755fc": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "b86efc8a7ce94d61af24f55b68844115": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "bf3ceb473d16470dbe825126832fda4f": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_116679cd0ef342519bef389e3c25835f",
            "max": 2,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_3cc2e7e24e0e479496f1115a1bab9c83",
            "value": 2
          }
        },
        "c4efdfd442424f9b911f9aacb57584d1": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c9d3016cc64c4640b8cbb48962223d02": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_885bb3979c554d3a82695eaa3efe569b",
            "placeholder": "​",
            "style": "IPY_MODEL_1cc91e5168684f139a63e162a4df3c00",
            "value": "Loading checkpoint shards: 100%"
          }
        },
        "cfdc876e9b27452d9c021d0717b32f99": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "da2f8a4f58914e488c686ec73bd0cec7": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_584815c8c1ac4183849149f2e576e730",
            "max": 2,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_dd993d2ff7c64c3d99c8d95aad54e11d",
            "value": 2
          }
        },
        "dd993d2ff7c64c3d99c8d95aad54e11d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "e7be1007f3874cba911c01f2e11ff997": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fc81ffa52b344d7baea4504fda77e7e9": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "fce12fe312934f298bec9ea8f970d1d5": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}