evaluation_playground.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "058c83eb-b45f-465d-bb51-5aa9056b7fb6",
   "metadata": {},
   "source": [
    "# FinanceBench: Evaluation Playground"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cabc8b5",
   "metadata": {},
   "source": [
    "\n",
    "<hr style=\"border-bottom:0.1px solid gray\">\n",
    "\n",
    "##### (1) API Requirements\n",
    "Add the following API keys into your `.env` file:\n",
    "\n",
    "```ruby\n",
    "OPENAI_API_KEY = 'INSERT API KEY HERE'\n",
    "ANTHROPIC_API_KEY = 'INSERT API KEY HERE'\n",
    "REPLICATE_API_TOKEN = 'INSERT API KEY HERE'\n",
    "```\n",
    "\n",
    "##### (2) Required Folder Structure\n",
    "\n",
    "```bash\n",
    "|-- /\n",
    "|    |-- data/\n",
    "|    |      | -- financebench_open_source.jsonl\n",
    "     |      | -- financebench_document_information.jsonl\n",
    "|    |-- pdfs/\n",
    "|           | -- <... provided filings as PDF documents ...>\n",
    "|    |-- results/\n",
    "|    |-- vectorstores/\n",
    "|    |-- evaluation_playground.ipynb\n",
    "```\n",
    "\n",
    "\n",
    "<br>\n",
    "<hr style=\"border-bottom:0.1px solid gray\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0a3a6571-6b3e-481d-8682-1e78bab64f8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install pymupdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b2737167-8912-40b4-9a3b-10f4544e1c5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import json\n",
    "import pickle\n",
    "import datetime\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm.notebook import tqdm\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "#LangChain Stuff\n",
    "from langchain.document_loaders import PyMuPDFLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.chains import RetrievalQA\n",
    "\n",
    "# LangChain Model Wrappers\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.chat_models import ChatAnthropic\n",
    "from langchain.llms.replicate import Replicate\n",
    "\n",
    "# Model Providers\n",
    "import openai\n",
    "import anthropic\n",
    "import replicate\n",
    "import tiktoken\n",
    "\n",
    "# import ANTHROPIC TOKENIZER\n",
    "CLIENT = anthropic.Anthropic(api_key=os.environ['ANTHROPIC_API_KEY'])\n",
    "anthropic_tokenizer = CLIENT.get_tokenizer()\n",
    "openai.api_key = os.environ['OPENAI_API_KEY']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "35674773-ccc3-4496-ba78-30dce4657327",
   "metadata": {},
   "outputs": [],
   "source": [
    "##############################################################################\n",
    "# MODEL CONFIGS\n",
    "##############################################################################\n",
    "configs = [\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4o-2024-05-13\",   \"eval_mode\":\"singleStore\",        \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4o-2024-05-13\",   \"eval_mode\":\"sharedStore\",        \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4o-2024-05-13\",   \"eval_mode\":\"inContext\",          \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4o-2024-05-13\",   \"eval_mode\":\"inContext_reverse\",  \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4o-2024-05-13\",   \"eval_mode\":\"oracle\",             \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4o-2024-05-13\",   \"eval_mode\":\"oracle_reverse\",     \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4o-2024-05-13\",   \"eval_mode\":\"sharedStore\",        \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4-1106-preview\",  \"eval_mode\":\"sharedStore\",        \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4-1106-preview\",  \"eval_mode\":\"singleStore\",        \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4-1106-preview\",  \"eval_mode\":\"inContext\",          \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4-1106-preview\",  \"eval_mode\":\"closedBook\",         \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"anthropic\",  \"model_name\":\"claude-2\",            \"eval_mode\":\"inContext\",          \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4-1106-preview\",  \"eval_mode\":\"oracle\",             \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"replicate\",  \"model_name\":\"llama2\",              \"eval_mode\":\"sharedStore\",        \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"replicate\",  \"model_name\":\"llama2\",              \"eval_mode\":\"singleStore\",        \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4\",               \"eval_mode\":\"sharedStore\",        \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4\",               \"eval_mode\":\"singleStore\",        \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4\",               \"eval_mode\":\"closedBook\",         \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4\",               \"eval_mode\":\"oracle\",             \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4-1106-preview\",  \"eval_mode\":\"oracle_reverse\",     \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"anthropic\",  \"model_name\":\"claude-2\",            \"eval_mode\":\"oracle_reverse\",     \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"openai\",     \"model_name\":\"gpt-4-1106-preview\",  \"eval_mode\":\"inContext_reverse\",  \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"anthropic\",  \"model_name\":\"claude-2\",            \"eval_mode\":\"inContext_reverse\",  \"temp\":0.01,   \"max_tokens\":2048},\n",
    "            {\"provider\": \"\",           \"model_name\":\"\",                    \"eval_mode\":\"singleStore\",        \"temp\":None,   \"max_tokens\":None},       # SPECIAL MODE --> RETRIEVAL ONLY MODE (SINGLE STORE)\n",
    "            {\"provider\": \"\",           \"model_name\":\"\",                    \"eval_mode\":\"sharedStore\",        \"temp\":None,   \"max_tokens\":None},       # SPECIAL MODE --> RETRIEVAL ONLY MODE (SHARED STORE)\n",
    "]\n",
    "\n",
    "replicate_model_mapping = dict({\n",
    "            \"llama2\": \"meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3\"\n",
    "        })\n",
    "\n",
    "##############################################################################\n",
    "# DATASET CONFIG\n",
    "##############################################################################\n",
    "PATH_CURRENT = os.path.abspath(os.getcwd())\n",
    "PATH_DATASET_JSONL = PATH_CURRENT + \"/data/financebench_open_source.jsonl\"\n",
    "PATH_DOCUMENT_INFO_JSONL = PATH_CURRENT + \"/data/financebench_document_information.jsonl\"\n",
    "PATH_RESULTS = PATH_CURRENT + \"/results/\"\n",
    "PATH_PDFS = PATH_CURRENT + \"/pdfs/\"\n",
    "\n",
    "# Choose DATASET PORTION:\n",
    "# - ALL: Full Dataset\n",
    "# - OPEN_SOURCE: Open Source Part (n=150)\n",
    "# - CLOSED_SOURCE: Closed Source Part --> Request access at contact@patronus.ai\n",
    "DATASET_PORTION = \"OPEN_SOURCE\"   \n",
    "\n",
    "##############################################################################\n",
    "# VECTOR STORE SETUP\n",
    "##############################################################################\n",
    "VS_CHUNK_SIZE = 1024\n",
    "VS_CHUNK_OVERLAP = 30\n",
    "VS_DIR_VS = PATH_CURRENT + \"/vectorstores\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "507fafa1-b05d-4842-b417-2c171e048cd4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of distinct PDF: 84\n",
      "Number of questions: 150\n",
      "Number of distinct PDF: 84\n"
     ]
    }
   ],
   "source": [
    "##############################################################################\n",
    "# LOAD DATASET\n",
    "##############################################################################\n",
    "\n",
    "# Load Full Dataset \n",
    "df_questions = pd.read_json(PATH_DATASET_JSONL, lines=True)\n",
    "df_meta = pd.read_json(PATH_DOCUMENT_INFO_JSONL, lines=True)\n",
    "df_full = pd.merge(df_questions, df_meta, on=\"doc_name\")\n",
    "\n",
    "# Get all docs\n",
    "df_questions = df_questions.sort_values('doc_name')\n",
    "ALL_DOCS = df_questions['doc_name'].unique().tolist()\n",
    "print(f\"Total number of distinct PDF: {len(ALL_DOCS)}\")\n",
    "\n",
    "# Select relevant dataset portion\n",
    "if DATASET_PORTION != \"ALL\":\n",
    "    df_questions = df_questions.loc[df_questions[\"dataset_subset_label\"]==DATASET_PORTION]\n",
    "print(f\"Number of questions: {len(df_questions)}\")\n",
    "\n",
    "# Check relevant documents\n",
    "df_questions = df_questions.sort_values('doc_name')\n",
    "docs = df_questions['doc_name'].unique().tolist()\n",
    "print(f\"Number of distinct PDF: {len(docs)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9d0835b6-1fd5-4021-bd4b-867981497fad",
   "metadata": {},
   "outputs": [],
   "source": [
    "##############################################################################\n",
    "# HELPER FUNCTIONS (PDF-PARSING + VECTOR-STORE SETUPS)\n",
    "##############################################################################\n",
    "def get_pdf_text(doc):\n",
    "    \n",
    "    path_doc = f\"{PATH_PDFS}/{doc}.pdf\"\n",
    "    pdf_reader = PyMuPDFLoader(path_doc)\n",
    "    pdf_text = pdf_reader.load()\n",
    "\n",
    "    return pdf_text\n",
    "\n",
    "def build_vectorstore_retriever(docs, embeddings = OpenAIEmbeddings()):\n",
    "\n",
    "    if docs == \"all\":\n",
    "        docs = ALL_DOCS\n",
    "        db_path = VS_DIR_VS + \"/shared\"\n",
    "    else:\n",
    "        docs = [docs]\n",
    "        db_path = VS_DIR_VS + \"/\" + docs[0]\n",
    "    \n",
    "    # Create Vector Store if not already existing\n",
    "    if not os.path.exists(db_path):\n",
    "        \n",
    "        # Create folder for vector store\n",
    "        os.mkdir(db_path) \n",
    "\n",
    "        # Create vector store itself --> chrom.sqlite3 database\n",
    "        if not os.path.exists(f\"{db_path}/chroma.sqlite3\"):\n",
    "            vectordb = Chroma(persist_directory=db_path, embedding_function=embeddings)\n",
    "            vectordb.persist()\n",
    "    \n",
    "            # Add Documents to Vector store    \n",
    "            for doc in docs:\n",
    "                pdf_text = get_pdf_text(doc)\n",
    "                text_splitter = RecursiveCharacterTextSplitter(\n",
    "                    chunk_size = VS_CHUNK_SIZE,\n",
    "                    chunk_overlap = VS_CHUNK_OVERLAP,\n",
    "                )\n",
    "                splitted_texts = text_splitter.split_documents(pdf_text)\n",
    "        \n",
    "                # Add to vector store\n",
    "                vectordb.add_documents(documents=splitted_texts)\n",
    "                vectordb.persist()\n",
    "\n",
    "    else:\n",
    "        vectordb = Chroma(persist_directory=db_path, embedding_function=embeddings)\n",
    "\n",
    "    return vectordb.as_retriever(), vectordb\n",
    "\n",
    "##############################################################################\n",
    "# MODEL + CALL HANDLERS\n",
    "##############################################################################\n",
    "\n",
    "def get_max_context_length(prompt, anthropic_cutoff=95000, openai_cutoff=105000):\n",
    "\n",
    "    # (0) Check Anthropic Tokenizer\n",
    "    tokens_anthropic = anthropic_tokenizer.encode(prompt)\n",
    "    nb_tokens_anthropic = len(tokens_anthropic)\n",
    "    number_of_chars_anthropic = len(prompt)\n",
    "    \n",
    "    if nb_tokens_anthropic > anthropic_cutoff:\n",
    "        tokens_anthropic_tokens = tokens_anthropic.tokens\n",
    "        token_lengths_anthropic = [len(token) for token in tokens_anthropic_tokens]\n",
    "        number_of_chars_anthropic = sum(token_lengths_anthropic[:anthropic_cutoff])\n",
    "        \n",
    "\n",
    "    # (1) Check OpenAI Tokenizer\n",
    "    tokenizer_openai = tiktoken.encoding_for_model(\"gpt-4-1106-preview\")\n",
    "    tokens_openai = tokenizer_openai.encode(prompt)\n",
    "    nb_tokens_openai = len(tokens_openai)\n",
    "    number_of_chars_openai = len(prompt)\n",
    "\n",
    "    if nb_tokens_openai > openai_cutoff:\n",
    "        tokens_openai_tokens = [tokenizer_openai.decode_single_token_bytes(token) for token in tokens_openai]\n",
    "        token_lengths_openai = [len(token) for token in tokens_openai_tokens]\n",
    "        number_of_chars_openai = sum(token_lengths_openai[:openai_cutoff])\n",
    "\n",
    "    # Cut prompt depending on minimal length limit\n",
    "    number_of_chars = min(number_of_chars_openai, number_of_chars_anthropic)\n",
    "\n",
    "    return number_of_chars\n",
    "\n",
    "def get_model(provider=\"openai\", model_name=\"gpt-4\", temp=0.01, max_tokens=2048):\n",
    "\n",
    "    if provider == \"openai\":\n",
    "        return ChatOpenAI(\n",
    "            model_name=model_name, \n",
    "            temperature=temp, \n",
    "            max_tokens=max_tokens\n",
    "            )\n",
    "        \n",
    "    elif provider == \"anthropic\":\n",
    "        return ChatAnthropic(\n",
    "            model=model_name,\n",
    "            temperature=temp, \n",
    "            max_tokens_to_sample=max_tokens, \n",
    "            anthropic_api_key=os.environ['ANTHROPIC_API_KEY']\n",
    "            )\n",
    "    \n",
    "    elif provider == \"replicate\":\n",
    "        if model_name in replicate_model_mapping:\n",
    "            return Replicate(\n",
    "                model=replicate_model_mapping[model_name],\n",
    "                model_kwargs={\n",
    "                    'temperature': temp, \n",
    "                    'max_new_tokens': max_tokens\n",
    "                    },\n",
    "            )\n",
    "        else:\n",
    "            raise ValueError(\"Unknown Model\")\n",
    "        \n",
    "    else:\n",
    "        return None\n",
    "\n",
    "\n",
    "def get_answer(model, eval_mode, question, context, retriever, retriever_only=False):\n",
    "\n",
    "    retrieved_documents = []\n",
    "\n",
    "    if eval_mode == \"closedBook\":\n",
    "        prompt = f\"Answer this question: {question}\"\n",
    "        answer = model.predict(prompt)\n",
    "        \n",
    "    elif eval_mode == \"oracle\":\n",
    "        prompt = f\"Answer this question: {question} \\nHere is the relevant evidence that you need to answer the question:\\n[START OF FILING] {context} [END OF FILING]\"\n",
    "        answer = model.predict(prompt)\n",
    "\n",
    "    elif eval_mode == \"oracle_reverse\":\n",
    "        \n",
    "        prompt = f\"Context:\\n[START OF FILING] {context} [END OF FILING\\n\\n Answer this question: {question} \\n\"\n",
    "        answer = model.predict(prompt)\n",
    "\n",
    "    elif eval_mode in [\"inContext\",  \"inContext_reverse\"]:\n",
    "        \n",
    "        # Context Cutoff to satisfy max tokens\n",
    "        max_number_of_chars = get_max_context_length(context)\n",
    "        context = context[:max_number_of_chars]\n",
    "        \n",
    "        if eval_mode == \"inContext\":\n",
    "            prompt = f\"Answer this question: {question} \\nHere is the relevant filing that you need to answer the question:\\n[START OF FILING] {context} [END OF FILING]\"\n",
    "        else:\n",
    "            prompt = f\"Context:\\n[START OF FILING] {context} [END OF FILING]\\n\\n Answer this question: {question}\\n\"\n",
    "\n",
    "        answer = model.predict(prompt)\n",
    "\n",
    "    elif eval_mode == \"singleStore\" or eval_mode == \"sharedStore\":\n",
    "        \n",
    "        # Retrieval-only mode if model=None (No LLM calls, only queries in VectorDB)\n",
    "        if not model:           \n",
    "            prompt = f\"{question}\"\n",
    "            s = retriever.invoke(prompt)\n",
    "            return (\"\", s)\n",
    "\n",
    "        else:\n",
    "\n",
    "            # Don't add a question prefix as RetrievalQA will do some automatic prompt wrapping\n",
    "            # --> This can replace by more advanced Retrieval Strategies\n",
    "            prompt = f\"{question}\"\n",
    "            qa = RetrievalQA.from_chain_type(\n",
    "                llm=model,\n",
    "                chain_type=\"stuff\",\n",
    "                retriever=retriever,\n",
    "                return_source_documents=True,\n",
    "            )\n",
    "            s = qa(prompt)\n",
    "            \n",
    "            answer = s[\"result\"]\n",
    "            retrieved_documents = s[\"source_documents\"]\n",
    "\n",
    "\n",
    "    \n",
    "    return (answer, retrieved_documents)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9a3ba97-7bf1-40e4-9352-56aac2c9dcb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "##############################################################################\n",
    "# EVALUATION\n",
    "##############################################################################\n",
    "\n",
    "# Specify evaluation model\n",
    "model_config = configs[0]\n",
    "\n",
    "# Set evaluation questions\n",
    "df_eval = df_questions\n",
    "\n",
    "\n",
    "# Get the model\n",
    "model = get_model(provider=model_config[\"provider\"],\n",
    "                  model_name=model_config[\"model_name\"],\n",
    "                  temp=model_config[\"temp\"],\n",
    "                  max_tokens=model_config[\"max_tokens\"])\n",
    "\n",
    "print(f\"--> Evaluating: {model_config['model_name']} / {model_config['eval_mode']}\")\n",
    "\n",
    "last_docs = None\n",
    "results = []\n",
    "\n",
    "# Run evaluation on the model  --> Sort along doc_name to reuse retriever configs in memory\n",
    "for k, (idx, row) in tqdm(enumerate(df_eval.sort_values(\"doc_name\").iterrows()), total=len(df_eval)):\n",
    "        \n",
    "    \n",
    "    # (A) Setup Context or Retriever\n",
    "    if model_config[\"eval_mode\"] == \"closedBook\":\n",
    "        retriever = None\n",
    "        context = \"\"\n",
    "    \n",
    "    elif model_config[\"eval_mode\"] in [\"inContext\", \"inContext_reverse\"]:\n",
    "        retriever = None\n",
    "        docs = row[\"doc_name\"]\n",
    "        if not (last_docs == docs):\n",
    "            pages = get_pdf_text(row[\"doc_name\"])\n",
    "            context = \"\\n\\n\".join([page.page_content for page in pages])\n",
    "            \n",
    "    \n",
    "    elif model_config[\"eval_mode\"] in [\"oracle\", \"oracle_reverse\"]:\n",
    "        context = \"\\n\\n\".join([evidence[\"evidence_text_full_page\"] for evidence in row[\"evidence\"]])\n",
    "        retriever = None\n",
    "\n",
    "    elif model_config[\"eval_mode\"] in [\"singleStore\", \"sharedStore\"]:\n",
    "        context = \"\"\n",
    "        docs = \"all\"\n",
    "\n",
    "        if model_config[\"eval_mode\"] == \"singleStore\":\n",
    "            docs = row[\"doc_name\"]\n",
    "        \n",
    "        if not (last_docs == docs):\n",
    "            retriever, _ = build_vectorstore_retriever(docs=docs)\n",
    "            last_docs = docs\n",
    "\n",
    "\n",
    "    else:\n",
    "        raise ValueError(\"Unknown 'eval_mode'!\")\n",
    "\n",
    "\n",
    "    # (B) Model Call\n",
    "    (answer, retrieved_documents) = get_answer(\n",
    "                                        model=model, \n",
    "                                        eval_mode=model_config[\"eval_mode\"], \n",
    "                                        question=row[\"question\"], \n",
    "                                        context=context, \n",
    "                                        retriever=retriever\n",
    "                                        )\n",
    "    \n",
    "\n",
    "    # (C) Bookkeeping\n",
    "    results.append({\n",
    "                        **model_config, \n",
    "                        \"financebench_id\" : row[\"financebench_id\"],\n",
    "                        \"question\" : row[\"question\"],\n",
    "                        \"gold_answer\": row[\"answer\"],\n",
    "                        \"model_answer\": answer,\n",
    "                        \"retrieved_documents\" : retrieved_documents,\n",
    "                    })\n",
    "\n",
    "df_results = pd.DataFrame(results)\n",
    "df_results.to_csv(PATH_RESULTS + \"/\" + model_config[\"model_name\"] + \"_\" + model_config[\"eval_mode\"] + \".csv\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py311",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}