1-test_bench.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers\n",
    "!pip install faiss-gpu\n",
    "!pip install load_dotenv\n",
    "!pip install tiktoken\n",
    "!pip install langchain\n",
    "!pip install sentence-transformers\n",
    "!pip install farm-haystack[preprocessing]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import dotenv\n",
    "import os\n",
    "import numpy as np\n",
    "import json\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.embeddings import HuggingFaceEmbeddings\n",
    "from langchain.schema import HumanMessage, SystemMessage\n",
    "from huggingface_hub.inference_api import InferenceApi\n",
    "from transformers import AutoTokenizer\n",
    "from tqdm.notebook import tqdm\n",
    "from scripts import result_exists, evaluate_response, generate_context\n",
    "from embeddings import retrieve_relevant_excerpts_quickly\n",
    "\n",
    "dotenv.load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "needle_question_couples = [\n",
    "    (\"\\nThe best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.\\n\", \"What is the most fun thing to do in San Francisco?\"),\n",
    "    (\"\\nThe most inspiring monument near the Hugging Face office in Paris is certainly the Louvre museum.\\n\", \"What is the most inspiring monument near the Hugging Face office in Paris?\"),\n",
    "]\n",
    "\n",
    "needle, question = needle_question_couples[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choose evaluation model and RAG embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluation_model  = ChatOpenAI(model=\"gpt-4\", temperature=0, openai_api_key = os.getenv('OPENAI_API_KEY', 'YourAPIKey'))\n",
    "\n",
    "rag_embedding = HuggingFaceEmbeddings(\n",
    "    model_name=\"BAAI/bge-large-en-v1.5\",\n",
    "    encode_kwargs={'normalize_embeddings': False},\n",
    "    model_kwargs={'device': 'cuda'},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choose LLMs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "openai_model_to_test = ChatOpenAI(model='gpt-4', temperature=0, openai_api_key = os.getenv('OPENAI_API_KEY', 'YourAPIKey'))\n",
    "\n",
    "hf_model_id = \"HuggingFaceH4/zephyr-7b-beta\"\n",
    "hf_tokenizer = AutoTokenizer.from_pretrained(hf_model_id)\n",
    "hf_client = InferenceApi(\n",
    "    repo_id=hf_model_id,\n",
    "    token=os.getenv('HUGGINGFACEHUB_API_TOKEN', 'YourHuggingFaceToken'),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_answer_hf(context: str, question: str, hf_tokenizer, hf_client) -> str:\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": \"You are a helpful AI bot that answers questions for a user. Keep your response short and direct.\",\n",
    "        },\n",
    "        {\"role\": \"user\", \"content\": \"\"\"\n",
    "        You will have to answer this question based only on the context: {question}\n",
    "        Here is the context: {context}\n",
    "        \"\"\"},\n",
    "        {\"role\": \"user\", \"content\": \"\"\"\n",
    "        Answer the following question in only one sentence: {question}\n",
    "        Don't give information outside the document or repeat your findings.\n",
    "        \"\"\"}\n",
    "    ]\n",
    "\n",
    "    messages_chat = hf_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "\n",
    "    full_prompt = messages_chat.format(question=question, context=context)\n",
    "    print(len(hf_tokenizer.encode(full_prompt)), hf_tokenizer.encode(full_prompt)[:10])\n",
    "    response = hf_client(full_prompt)[0]['generated_text'][len(full_prompt):]\n",
    "    return response\n",
    "\n",
    "def get_answer_openai(context, question, model):\n",
    "    # Prepare your message to send to the model you're going to evaluate\n",
    "    messages = [\n",
    "        SystemMessage(\n",
    "            content=\"You are a helpful AI bot that answers questions for a user. Keep your response short and direct\"\n",
    "        ),\n",
    "        HumanMessage(\n",
    "            # This is the PG essays with your needle/random statement placed in it\n",
    "            # This is your haystack with a needle placed in it.\n",
    "            content=f\"CONTEXT:\\n{context}\",\n",
    "        ),\n",
    "        HumanMessage(\n",
    "            # This is the question you'll ask to the model to tr≠≠y and retrieve your random statement/needle.\n",
    "            content=f\"{question} - Don't give information outside the document or repeat your findings\"\n",
    "        ),\n",
    "    ]\n",
    "    return model(messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Launch tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "COMPRESS_CONTEXT_WITH_RAG = True\n",
    "USE_HF_MODEL = False\n",
    "model_to_test_description = 'GPT4_RAG'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The code will check to see if a context_length, depth percent and version number have already been checked yet\n",
    "# Change the version # if you would like to run the results multiple times.\n",
    "# If you're just testing, then leave as version=1\n",
    "results_version = 2\n",
    "\n",
    "# This will produce a list of context lengths for each experiment iteration. Make sure the max context length is within the bounds of your models limits.\n",
    "context_lengths = np.round(np.linspace(1000, 128000, num=15, endpoint=True)).astype(int)\n",
    "\n",
    "# This will product a list of document depths to place your random statement (needle) at.\n",
    "# Suggestion: Try out different distributions (like a sigmoid) to test non-evenly space intervals\n",
    "document_depth_percents = np.round(np.linspace(0, 100, num=15, endpoint=True)).astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Go generate the required length context and place your needle statement in\n",
    "\n",
    "# test_needle, test_question = (\"\\nThe most inspiring monument near the Hugging Face office in Paris is certainly the Louvre museum.\\n\", \"What is the most inspiring monument near the Hugging Face office in Paris?\")\n",
    "\n",
    "# context = generate_context(test_needle, 46357, 79)\n",
    "\n",
    "# if COMPRESS_CONTEXT_WITH_RAG:\n",
    "#     context = await retrieve_relevant_excerpts_quickly(context, test_question, rag_embedding, top_k=20, words_per_chunk=50, flag_mentions_of_paris=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run through each iteration of context_lengths and depths\n",
    "for depth_percent in tqdm(document_depth_percents):\n",
    "    for context_length in context_lengths:\n",
    "        # Load results from file. \n",
    "        try:\n",
    "            with open(f'output/results_{model_to_test_description}.json', 'r') as f:\n",
    "                results = json.load(f)\n",
    "        except FileNotFoundError:\n",
    "            results = []\n",
    "            pass\n",
    "\n",
    "        # Checks to see if you've already checked a length/percent/version.\n",
    "        # This helps if the program stop running and you want to restart later\n",
    "        if result_exists(results, context_length, depth_percent, results_version, model_to_test_description):\n",
    "            continue\n",
    "        \n",
    "        # Go generate the required length context and place your needle statement in\n",
    "        context = generate_context(needle, context_length, depth_percent)\n",
    "\n",
    "        if COMPRESS_CONTEXT_WITH_RAG:\n",
    "            context = await retrieve_relevant_excerpts_quickly(context, question, rag_embedding, top_k=20, words_per_chunk=50)\n",
    "\n",
    "        ### Get your model's answer to the question! Will it find your random fact?\n",
    "        if USE_HF_MODEL:\n",
    "            response = get_answer_hf(context, question, hf_tokenizer, hf_client)\n",
    "        else:\n",
    "            response = get_answer_openai(context, question, openai_model_to_test).content\n",
    "\n",
    "        print(response)\n",
    "\n",
    "        # Compare the reponse to the actual needle you placed\n",
    "        score = evaluate_response(response, needle, question, evaluation_model)\n",
    "        result = {\n",
    "            'model' : model_to_test_description,\n",
    "            'context_length' : int(context_length),\n",
    "            'depth_percent' : int(depth_percent),\n",
    "            'version' : results_version,\n",
    "            'needle' : needle,\n",
    "            'model_response' : response,\n",
    "            'score' : score\n",
    "        }\n",
    "        if score < 10:\n",
    "            result['context'] = context\n",
    "\n",
    "        print (f\"Result #: {len(results)}/{len(context_lengths) * len(document_depth_percents)}\")\n",
    "        print (f\"Context: {context_length} tokens\")\n",
    "        print (f\"Depth: {depth_percent}%\")\n",
    "        print (f\"Context: {context_length} tokens\")\n",
    "        print (f\"Depth: {depth_percent}%\")\n",
    "        print (f\"Score: {score}\")\n",
    "        print (f\"Response: {response}\\n\")\n",
    "\n",
    "        results.append(result)\n",
    "\n",
    "        # Save results to a JSON file each run\n",
    "        with open(f'output/results_{model_to_test_description}.json', 'w') as f:\n",
    "            json.dump(results, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Resulting context after RAG is 1600 tokens, i.e. with `gpt-4`: $0.03 / 1K tokens, aka $0.05\n",
    "- Long-context GPT costs `gpt-4-1106-preview`: $0.01 / 1K tokens, aka $1.28"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}