{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('/mnt/c/Users/Olive/python_projects/FinNLP/FinNLP')  # https://github.com/AI4Finance-Foundation/FinNLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# only for WSL\n",
    "import os \n",
    "os.environ[\"PATH\"] = f\"{os.environ['PATH']}:/usr/local/cuda-12.0/bin\"\n",
    "# os.environ['LD_LIBRARY_PATH'] = \"/usr/lib/wsl/lib:/usr/local/cuda/lib64\"\n",
    "os.environ['LD_LIBRARY_PATH'] = \"/usr/local/cuda-12.0/lib64\"\n",
    "os.environ[\"BNB_CUDA_VERSION\"]=\"120\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/oliverwang15/miniconda3/lib/python3.10/site-packages/bitsandbytes-0.41.1-py3.10.egg/bitsandbytes/cuda_setup/main.py:106: UserWarning: \n",
      "\n",
      "================================================================================\n",
      "WARNING: Manual override via BNB_CUDA_VERSION env variable detected!\n",
      "BNB_CUDA_VERSION=XXX can be used to load a bitsandbytes version that is different from the PyTorch CUDA version.\n",
      "If this was unintended set the BNB_CUDA_VERSION variable to an empty string: export BNB_CUDA_VERSION=\n",
      "If you use the manual override make sure the right libcudart.so is in your LD_LIBRARY_PATH\n",
      "For example by adding the following to your .bashrc: export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:<path_to_cuda_dir/lib64\n",
      "Loading CUDA version: BNB_CUDA_VERSION=120\n",
      "================================================================================\n",
      "\n",
      "\n",
      "  warn((f'\\n\\n{\"=\"*80}\\n'\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM, LlamaForCausalLM, LlamaTokenizerFast   # 4.30.2\n",
    "from peft import PeftModel  # 0.4.0\n",
    "import torch\n",
    "\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": "markdown",
   "metadata": {},
   "source": [
    "## Load Model (Pick one according the model from the following blocks)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. **FInGPT v3.1** based on ChatGLM2, runable on 1 * RTX 3090"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # v3.1\n",
    "# base_model = \"THUDM/chatglm2-6b\"\n",
    "# peft_model = \"oliverwang15/FinGPT_v31_ChatGLM2_Sentiment_Instruction_LoRA_FT\"\n",
    "# tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)\n",
    "# model = AutoModel.from_pretrained(base_model, trust_remote_code=True, load_in_8bit = True, device_map = \"auto\")\n",
    "# model = PeftModel.from_pretrained(model, peft_model)\n",
    "# model = model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. **FinGPT v3.2** based on Llama2, runable on 1 * A 100 and also runable on 1 * RTX 3090 as long as `load_in_8bit = True` is set in Line 8, but the speed is slower"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f37f0a825f39418ab8cb47db61aa3ad6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# v3.2\n",
    "# base_model = \"meta-llama/Llama-2-7b-chat-hf\"  # Access needed\n",
    "base_model = \"NousResearch/Llama-2-13b-hf\" \n",
    "peft_model = \"oliverwang15/FinGPT_v33_Llama2_13B_Sentiment_Instruction_LoRA_FT_8bit\"\n",
    "tokenizer = LlamaTokenizerFast.from_pretrained(base_model, trust_remote_code=True)\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "model = LlamaForCausalLM.from_pretrained(base_model, trust_remote_code=True, device_map = \"cuda:0\", load_in_8bit = True,)\n",
    "model = PeftModel.from_pretrained(model, peft_model)\n",
    "model = torch.compile(model)  # Please comment this line if your platform does not support torch.compile\n",
    "model = model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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: 16. Total steps: 76\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 76/76 [01:10<00:00,  1.08it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Acc: 0.8820132013201321. F1 macro: 0.8774649459784639. F1 micro: 0.8820132013201321. F1 weighted (BloombergGPT): 0.8817589259286243. \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# FPB\n",
    "res = test_fpb(model, tokenizer, batch_size = batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pf(x):\n",
    "    return  \"What is the sentiment of this news? Please choose an answer from {strong negative/moderately negative/mildly negative/neutral/mildly positive/moderately positive/strong positive}, then provide some short reasons.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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: 16. Total steps: 18\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 18/18 [00:14<00:00,  1.25it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Acc: 0.8581818181818182. F1 macro: 0.7693085447271324. F1 micro: 0.858181818181818. F1 weighted (BloombergGPT): 0.8744065786191484. \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# FiQA\n",
    "res = test_fiqa(model, tokenizer, prompt_fun = add_instructions, batch_size = batch_size)\n",
    "# res = test_fiqa(model, tokenizer, batch_size = batch_size, prompt_fun=pf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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: 16. Total steps: 150\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 150/150 [01:56<00:00,  1.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Acc: 0.9032663316582915. F1 macro: 0.8787781808799305. F1 micro: 0.9032663316582915. F1 weighted (BloombergGPT): 0.9028238287011249. \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# TFNS\n",
    "res = test_tfns(model, tokenizer, batch_size = batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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: 16. Total steps: 253\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 253/253 [04:38<00:00,  1.10s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Acc: 0.6429453916481345. F1 macro: 0.6546976499843747. F1 micro: 0.6429453916481345. F1 weighted (BloombergGPT): 0.6428106146054466. \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# NWGI\n",
    "res = test_nwgi(model, tokenizer, batch_size = batch_size)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fingpt",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
