240524_jargon_map_and_define.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data and Library Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports \n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "import os\n",
    "os.environ[\"OPENAI_API_KEY\"] = os.getenv(\"OPENAI_KEY\")\n",
    "# this needs to be set before other llamaindex imports and instantiations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from llama_index.readers.file import PDFReader\n",
    "\n",
    "from llama_index.core import (\n",
    "    SimpleDirectoryReader, \n",
    "    VectorStoreIndex, \n",
    "    Settings, \n",
    "    get_response_synthesizer, \n",
    "    PromptTemplate\n",
    ")\n",
    "from llama_index.core.retrievers import VectorIndexRetriever\n",
    "from llama_index.core.query_engine import RetrieverQueryEngine\n",
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "from llama_index.core.postprocessor import SimilarityPostprocessor\n",
    "\n",
    "from llama_index.core.ingestion import IngestionPipeline, IngestionCache\n",
    "from llama_index.core.vector_stores.types import (\n",
    "    MetadataFilter,\n",
    "    MetadataFilters,\n",
    "    ExactMatchFilter\n",
    ")\n",
    "\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "import time\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do some settings for the RAG\n",
    "\n",
    "# Setup the embedding model\n",
    "Settings.embed_model = OpenAIEmbedding(\n",
    "    model=\"text-embedding-3-small\", \n",
    "    api_key=os.getenv(\"OPENAI_KEY\"), \n",
    ")\n",
    "\n",
    "# Setup the LLM\n",
    "Settings.llm = OpenAI(\n",
    "    model=\"gpt-4-turbo\", \n",
    "    temperature=1, \n",
    "    api_key=os.getenv(\"OPENAI_KEY\"), \n",
    "    max_tokens=250\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>arxiv_id</th>\n",
       "      <th>reader_id</th>\n",
       "      <th>gpt4_jargon_list</th>\n",
       "      <th>human_jargon_list</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2403.16190v1</td>\n",
       "      <td>rid0</td>\n",
       "      <td>reject option strategy, formal guarantees, min...</td>\n",
       "      <td>correctness,minimality,Anchors</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2403.16190v1</td>\n",
       "      <td>rid1</td>\n",
       "      <td>linear classification problems, reject option ...</td>\n",
       "      <td>reject option strategy,Anchors,heuristic algor...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2307.05300v4</td>\n",
       "      <td>rid0</td>\n",
       "      <td>cognitive synergist, fine-grained personas, fa...</td>\n",
       "      <td>multi-turn,persona</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2307.05300v4</td>\n",
       "      <td>rid1</td>\n",
       "      <td>cognitive synergy, cognitive synergist, multi-...</td>\n",
       "      <td>Solo Performance Prompting,multi-turn,Chain-of...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2403.16750v1</td>\n",
       "      <td>rid0</td>\n",
       "      <td>Common Weakness Enumerations (CWEs), SystemVer...</td>\n",
       "      <td>common weakeness enumerations,SystemVerilog,Re...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       arxiv_id reader_id                                   gpt4_jargon_list  \\\n",
       "0  2403.16190v1      rid0  reject option strategy, formal guarantees, min...   \n",
       "1  2403.16190v1      rid1  linear classification problems, reject option ...   \n",
       "2  2307.05300v4      rid0  cognitive synergist, fine-grained personas, fa...   \n",
       "3  2307.05300v4      rid1  cognitive synergy, cognitive synergist, multi-...   \n",
       "4  2403.16750v1      rid0  Common Weakness Enumerations (CWEs), SystemVer...   \n",
       "\n",
       "                                   human_jargon_list  \n",
       "0                     correctness,minimality,Anchors  \n",
       "1  reject option strategy,Anchors,heuristic algor...  \n",
       "2                                 multi-turn,persona  \n",
       "3  Solo Performance Prompting,multi-turn,Chain-of...  \n",
       "4  common weakeness enumerations,SystemVerilog,Re...  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load the jargon annotated data\n",
    "merged_humam_llm_jargon = pd.read_json(\n",
    "    \"data/llm_outputs/240525_march_2024_human_llm_jargon_merged.json\", \n",
    "    orient=\"index\"\n",
    ")\n",
    "\n",
    "merged_humam_llm_jargon.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>arxiv_id</th>\n",
       "      <th>url</th>\n",
       "      <th>title</th>\n",
       "      <th>summary</th>\n",
       "      <th>updated</th>\n",
       "      <th>published</th>\n",
       "      <th>authors</th>\n",
       "      <th>comments</th>\n",
       "      <th>categories</th>\n",
       "      <th>primary_category</th>\n",
       "      <th>doi</th>\n",
       "      <th>journal_ref</th>\n",
       "      <th>peer_reviewed</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2403.16190v1</td>\n",
       "      <td>http://arxiv.org/abs/2403.16190v1</td>\n",
       "      <td>Logic-based Explanations for Linear Support Ve...</td>\n",
       "      <td>Support Vector Classifier (SVC) is a well-know...</td>\n",
       "      <td>1711293284000</td>\n",
       "      <td>1711293284000</td>\n",
       "      <td>[Francisco Mateus Rocha Filho, Thiago Alves Ro...</td>\n",
       "      <td>16 pages, submitted to BRACIS 2023 (Brazilian ...</td>\n",
       "      <td>[cs.AI, cs.LG, cs.LO, I.2.4; I.2.6]</td>\n",
       "      <td>cs.AI</td>\n",
       "      <td>10.1007/978-3-031-45368-7_10</td>\n",
       "      <td>None</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2307.05300v4</td>\n",
       "      <td>http://arxiv.org/abs/2307.05300v4</td>\n",
       "      <td>Unleashing the Emergent Cognitive Synergy in L...</td>\n",
       "      <td>Human intelligence thrives on cognitive synerg...</td>\n",
       "      <td>1711463553000</td>\n",
       "      <td>1689086719000</td>\n",
       "      <td>[Zhenhailong Wang, Shaoguang Mao, Wenshan Wu, ...</td>\n",
       "      <td>Accepted as a main conference paper at NAACL 2024</td>\n",
       "      <td>[cs.AI, cs.CL]</td>\n",
       "      <td>cs.AI</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2403.16750v1</td>\n",
       "      <td>http://arxiv.org/abs/2403.16750v1</td>\n",
       "      <td>All Artificial, Less Intelligence: GenAI throu...</td>\n",
       "      <td>Modern hardware designs have grown increasingl...</td>\n",
       "      <td>1711373004000</td>\n",
       "      <td>1711373004000</td>\n",
       "      <td>[Deepak Narayan Gadde, Aman Kumar, Thomas Nala...</td>\n",
       "      <td>Published in DVCon U.S. 2024</td>\n",
       "      <td>[cs.AI]</td>\n",
       "      <td>cs.AI</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2311.10112v2</td>\n",
       "      <td>http://arxiv.org/abs/2311.10112v2</td>\n",
       "      <td>zrLLM: Zero-Shot Relational Learning on Tempor...</td>\n",
       "      <td>Modeling evolving knowledge over temporal know...</td>\n",
       "      <td>1710517087000</td>\n",
       "      <td>1700083515000</td>\n",
       "      <td>[Zifeng Ding, Heling Cai, Jingpei Wu, Yunpu Ma...</td>\n",
       "      <td>Accepted to NAACL 2024 main conference</td>\n",
       "      <td>[cs.AI, cs.CL, cs.LG]</td>\n",
       "      <td>cs.AI</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2310.08992v3</td>\n",
       "      <td>http://arxiv.org/abs/2310.08992v3</td>\n",
       "      <td>CodeChain: Towards Modular Code Generation Thr...</td>\n",
       "      <td>Large Language Models (LLMs) have already beco...</td>\n",
       "      <td>1710386949000</td>\n",
       "      <td>1697192268000</td>\n",
       "      <td>[Hung Le, Hailin Chen, Amrita Saha, Akash Goku...</td>\n",
       "      <td>Accepted to ICLR 2024</td>\n",
       "      <td>[cs.AI, cs.CL, cs.PL]</td>\n",
       "      <td>cs.AI</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       arxiv_id                                url  \\\n",
       "0  2403.16190v1  http://arxiv.org/abs/2403.16190v1   \n",
       "1  2307.05300v4  http://arxiv.org/abs/2307.05300v4   \n",
       "2  2403.16750v1  http://arxiv.org/abs/2403.16750v1   \n",
       "3  2311.10112v2  http://arxiv.org/abs/2311.10112v2   \n",
       "4  2310.08992v3  http://arxiv.org/abs/2310.08992v3   \n",
       "\n",
       "                                               title  \\\n",
       "0  Logic-based Explanations for Linear Support Ve...   \n",
       "1  Unleashing the Emergent Cognitive Synergy in L...   \n",
       "2  All Artificial, Less Intelligence: GenAI throu...   \n",
       "3  zrLLM: Zero-Shot Relational Learning on Tempor...   \n",
       "4  CodeChain: Towards Modular Code Generation Thr...   \n",
       "\n",
       "                                             summary        updated  \\\n",
       "0  Support Vector Classifier (SVC) is a well-know...  1711293284000   \n",
       "1  Human intelligence thrives on cognitive synerg...  1711463553000   \n",
       "2  Modern hardware designs have grown increasingl...  1711373004000   \n",
       "3  Modeling evolving knowledge over temporal know...  1710517087000   \n",
       "4  Large Language Models (LLMs) have already beco...  1710386949000   \n",
       "\n",
       "       published                                            authors  \\\n",
       "0  1711293284000  [Francisco Mateus Rocha Filho, Thiago Alves Ro...   \n",
       "1  1689086719000  [Zhenhailong Wang, Shaoguang Mao, Wenshan Wu, ...   \n",
       "2  1711373004000  [Deepak Narayan Gadde, Aman Kumar, Thomas Nala...   \n",
       "3  1700083515000  [Zifeng Ding, Heling Cai, Jingpei Wu, Yunpu Ma...   \n",
       "4  1697192268000  [Hung Le, Hailin Chen, Amrita Saha, Akash Goku...   \n",
       "\n",
       "                                            comments  \\\n",
       "0  16 pages, submitted to BRACIS 2023 (Brazilian ...   \n",
       "1  Accepted as a main conference paper at NAACL 2024   \n",
       "2                       Published in DVCon U.S. 2024   \n",
       "3             Accepted to NAACL 2024 main conference   \n",
       "4                              Accepted to ICLR 2024   \n",
       "\n",
       "                            categories primary_category  \\\n",
       "0  [cs.AI, cs.LG, cs.LO, I.2.4; I.2.6]            cs.AI   \n",
       "1                       [cs.AI, cs.CL]            cs.AI   \n",
       "2                              [cs.AI]            cs.AI   \n",
       "3                [cs.AI, cs.CL, cs.LG]            cs.AI   \n",
       "4                [cs.AI, cs.CL, cs.PL]            cs.AI   \n",
       "\n",
       "                            doi journal_ref  peer_reviewed  \n",
       "0  10.1007/978-3-031-45368-7_10        None           True  \n",
       "1                          None        None           True  \n",
       "2                          None        None           True  \n",
       "3                          None        None           True  \n",
       "4                          None        None           True  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Read in the arXiv metadata\n",
    "with open(\"data/arxiv_metadata/filtered/march_2024_ai_hc_cy_peer_reviewed_sampled.json\") as json_data:\n",
    "    metadata = json.load(json_data)\n",
    "    json_data.close()\n",
    "\n",
    "# Convert JSON to DataFrame\n",
    "metadata_df = pd.DataFrame.from_dict(metadata, orient='index')\n",
    "\n",
    "metadata_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Light text cleaning so that it can be proprely fed into anything\n",
    "metadata_df['summary'] = metadata_df['summary'].str.replace('\\n', ' ')\n",
    "metadata_df['summary'] = metadata_df['summary'].str.replace('\\r', ' ')\n",
    "metadata_df['summary'] = metadata_df['summary'].str.replace('\\t', ' ')\n",
    "metadata_df['summary'] = metadata_df['summary'].str.replace('  ', ' ')\n",
    "metadata_df['summary'] = metadata_df['summary'].str.strip()\n",
    "\n",
    "metadata_df['title'] = metadata_df['title'].str.replace('\\n', ' ')\n",
    "metadata_df['title'] = metadata_df['title'].str.replace('\\r', ' ')\n",
    "metadata_df['title'] = metadata_df['title'].str.replace('\\t', ' ')\n",
    "metadata_df['title'] = metadata_df['title'].str.replace('  ', ' ')\n",
    "metadata_df['title'] = metadata_df['title'].str.strip()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAG: Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PDF Reader with `SimpleDirectoryReader`\n",
    "parser = PDFReader()\n",
    "file_extractor = {\".pdf\": parser}\n",
    "filename_fn = lambda filename: {\"file_name\": filename}\n",
    "documents = SimpleDirectoryReader(\n",
    "    \"data/arxiv_pdfs/march_2024_ai_hc_cy_peer_reviewed_sampled\", \n",
    "    file_extractor=file_extractor, file_metadata=filename_fn\n",
    ").load_data()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Length of documents is much much more than length of the aactual number of papers because LlamaIndex stores\n",
    "# each page of the paper as a separate document\n",
    "print(len(documents))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add some metadata fields to each document: arxiv_id, title, abstract, and primary category\n",
    "for document in documents:\n",
    "    document.metadata[\"arxiv_id\"] = document.metadata[\"file_name\"].split(\"/\")[-1][:-4]\n",
    "    # document.metadata[\"title\"] = metadata_df.loc[\n",
    "    #     metadata_df[\"arxiv_id\"] == document.metadata[\"arxiv_id\"]]['title'].values[0]\n",
    "    # document.metadata[\"summary\"] = metadata_df.loc[\n",
    "    #     metadata_df[\"arxiv_id\"] == document.metadata[\"arxiv_id\"]]['summary'].values[0]\n",
    "    # document.metadata[\"primary_category\"] = metadata_df.loc[\n",
    "    #     metadata_df[\"arxiv_id\"] == document.metadata[\"arxiv_id\"]]['primary_category'].values[0]\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You do NOT want the ID or the title or the abstract getting embedded in any way at all. It would either add extraneous info (ID) or confuse the jargon term lookup (title/abstract). Only want to use the full text of the paper to define the jargon terms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hide the arxiv_id from being visible to the embedding model and LLM (during response synthesis). \n",
    "print(\"Hidden metadata: \", documents[0].excluded_embed_metadata_keys)\n",
    "\n",
    "for document in documents:\n",
    "    document.excluded_llm_metadata_keys.append(\"arxiv_id\")\n",
    "    document.excluded_embed_metadata_keys.append(\"arxiv_id\")\n",
    "print(\"Hidden metadata -- updated: \", documents[0].excluded_embed_metadata_keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A preivew of what the LLM and the embedding models see\n",
    "from llama_index.core.schema import MetadataMode\n",
    "print(\n",
    "    \"The LLM sees this: \\n\",\n",
    "    documents[0].get_content(metadata_mode=MetadataMode.LLM),\n",
    ")\n",
    "print(\n",
    "    \"\\n\\n\\nThe Embedding model sees this: \\n\",\n",
    "    documents[0].get_content(metadata_mode=MetadataMode.EMBED),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAG: Transformation/Indexing\n",
    "\n",
    "Chunking + Embedding Chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Setup client\n",
    "# client = OpenAI(\n",
    "#     api_key = os.getenv(\"OPENAI_KEY\"),\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Create an index\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, \n",
    "    transformations=[\n",
    "        SentenceSplitter(chunk_size=256, chunk_overlap=20), # I want to chunk at the sentence level, since term definitions are usually in a single sentence\n",
    "        OpenAIEmbedding(\n",
    "            model=\"text-embedding-3-small\", \n",
    "            api_key=os.getenv(\"OPENAI_KEY\")\n",
    "        ) # Use the OpenAI embedding model\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Read more here about load this back in: https://docs.llamaindex.ai/en/stable/understanding/storing/storing/\n",
    "index.storage_context.persist(persist_dir=\"data/vector_indices\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load the Index in "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 30.9 s, sys: 141 ms, total: 31.1 s\n",
      "Wall time: 31.2 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "from llama_index.core import StorageContext, load_index_from_storage\n",
    "\n",
    "# rebuild storage context\n",
    "storage_context = StorageContext.from_defaults(persist_dir=\"data/vector_indices\")\n",
    "\n",
    "# load index\n",
    "index_loaded = load_index_from_storage(\n",
    "    storage_context, \n",
    "    transformations=[\n",
    "        SentenceSplitter(chunk_size=256, chunk_overlap=20), # I want to chunk at the sentence level, since term definitions are usually in a single sentence\n",
    "        OpenAIEmbedding(\n",
    "            model=\"text-embedding-3-small\", \n",
    "            api_key=os.getenv(\"OPENAI_KEY\")\n",
    "        ) # Use the OpenAI embedding model\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Querying -- Trials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "\n",
    "# Create a basic engine for testing\n",
    "query_engine = index_loaded.as_query_engine()\n",
    "response = query_engine.query(\"define: reject option strategy\")\n",
    "print(response)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 748 µs, sys: 1.63 ms, total: 2.38 ms\n",
      "Wall time: 2.44 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# Create a more complex engine for testing\n",
    "\n",
    "# Create a filter on the arxiv_id\n",
    "filters = MetadataFilters(filters=[\n",
    "    ExactMatchFilter(\n",
    "        key=\"arxiv_id\", \n",
    "        value=\"2307.05300v4\"\n",
    "    )\n",
    "])\n",
    "\n",
    "# Configure retriever\n",
    "retriever = VectorIndexRetriever(\n",
    "    index=index_loaded,\n",
    "    similarity_top_k=10,\n",
    "    filters=filters\n",
    ")\n",
    "\n",
    "# configure response synthesizer\n",
    "response_synthesizer = get_response_synthesizer(\n",
    "    response_mode=\"compact\",\n",
    ")\n",
    "\n",
    "# assemble query engine\n",
    "custom_query_engine = RetrieverQueryEngine(\n",
    "    retriever=retriever,\n",
    "    response_synthesizer=response_synthesizer,\n",
    "    node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.3)]\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SelectorPromptTemplate(metadata={'prompt_type': <PromptType.QUESTION_ANSWER: 'text_qa'>}, template_vars=['context_str', 'query_str'], kwargs={}, output_parser=None, template_var_mappings={}, function_mappings={}, default_template=PromptTemplate(metadata={'prompt_type': <PromptType.QUESTION_ANSWER: 'text_qa'>}, template_vars=['context_str', 'query_str'], kwargs={}, output_parser=None, template_var_mappings=None, function_mappings=None, template='Context information is below.\\n---------------------\\n{context_str}\\n---------------------\\nGiven the context information and not prior knowledge, answer the query.\\nQuery: {query_str}\\nAnswer: '), conditionals=[(<function is_chat_model at 0x1332f5820>, ChatPromptTemplate(metadata={'prompt_type': <PromptType.CUSTOM: 'custom'>}, template_vars=['context_str', 'query_str'], kwargs={}, output_parser=None, template_var_mappings=None, function_mappings=None, message_templates=[ChatMessage(role=<MessageRole.SYSTEM: 'system'>, content=\"You are an expert Q&A system that is trusted around the world.\\nAlways answer the query using the provided context information, and not prior knowledge.\\nSome rules to follow:\\n1. Never directly reference the given context in your answer.\\n2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.\", additional_kwargs={}), ChatMessage(role=<MessageRole.USER: 'user'>, content='Context information is below.\\n---------------------\\n{context_str}\\n---------------------\\nGiven the context information and not prior knowledge, answer the query.\\nQuery: {query_str}\\nAnswer: ', additional_kwargs={})]))])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Accessing the default prompt\n",
    "response_synthesizer.get_prompts()['text_qa_template']#['template']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "response = custom_query_engine.query(\n",
    "    \"Use 1-2 sentences to explain this term so that even a reader without deep scientific and technical knowledge can understand it easily: cognitive synergist.\"\n",
    ")\n",
    "print(response)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "from openai import OpenAI as oai_real\n",
    "\n",
    "# Setup client\n",
    "client = oai_real(\n",
    "    api_key = os.getenv(\"OPENAI_KEY\"),\n",
    ")\n",
    "\n",
    "abstract = \"Human intelligence thrives on cognitive synergy, where collaboration among different minds yield superior outcomes compared to isolated individuals. In this work, we propose Solo Performance Prompting (SPP), which transforms a single LLM into a cognitive synergist by engaging in multi-turn self-collaboration with multiple personas. A cognitive synergist is an intelligent agent that collaboratively combines multiple minds' strengths and knowledge to enhance problem-solving in complex tasks. By dynamically identifying and simulating different personas based on task inputs, SPP unleashes the potential of cognitive synergy in LLMs. Our in-depth analysis shows that assigning multiple fine-grained personas in LLMs improves problem-solving abilities compared to using a single or fixed number of personas. We evaluate SPP on three challenging tasks: Trivia Creative Writing, Codenames Collaborative, and Logic Grid Puzzle, encompassing both knowledge-intensive and reasoning-intensive types. Unlike previous works, such as Chain-of-Thought, that solely enhance the reasoning abilities in LLMs, experimental results demonstrate that SPP effectively reduces factual hallucination, and maintains strong reasoning capabilities. Additionally, comparative experiments show that cognitive synergy only emerges in GPT-4 and does not appear in less capable models, such as GPT-3.5-turbo and Llama2-13b-chat, which draws an interesting analogy to human development. Code, data, and prompts can be found at: this https URL. \"\n",
    "term = \"cognitive synergist\"\n",
    "term_definition = client.chat.completions.create(\n",
    "                    model=\"gpt-4-turbo\",\n",
    "                    messages=[\n",
    "                        {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "                        {\"role\": \"user\", \"content\": f\"Here's a scientific abstract: {abstract}. Based on this information, please use 1-2 sentences to explain this term so that even a reader without deep scientific and technical knowledge can understand it easily: {term}. \\n Definition: \"}\n",
    "                    ],\n",
    "                    temperature=1, \n",
    "                    seed=10,\n",
    "                )\n",
    "term_definition = term_definition.choices[0].message.content\n",
    "print(term_definition)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Querying -- Actual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a custom prompt template -- just to maximise control over everything for now\n",
    "# template = (\n",
    "#     \"Here are some relevant excerpts from a scientific paper:\"\n",
    "#     \"---------------------\\n\"\n",
    "#     \"{context_str}\"\n",
    "#     \"\\n---------------------\\n\"\n",
    "#     \"Based on this information, please {query_str}\\n\"\n",
    "#     \"Definition: \"\n",
    "# )\n",
    "# qa_template = PromptTemplate(template)\n",
    "\n",
    "# Here's the default template:\n",
    "# Source: https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/prompts/default_prompts.py\n",
    "\n",
    "# DEFAULT_TEXT_QA_PROMPT_TMPL = (\n",
    "#     \"Context information is below.\\n\"\n",
    "#     \"---------------------\\n\"\n",
    "#     \"{context_str}\\n\"\n",
    "#     \"---------------------\\n\"\n",
    "#     \"Given the context information and not prior knowledge, \"\n",
    "#     \"answer the query.\\n\"\n",
    "#     \"Query: {query_str}\\n\"\n",
    "#     \"Answer: \"\n",
    "# )\n",
    "\n",
    "# response_synthesizer = get_response_synthesizer(\n",
    "#             response_mode=\"compact\",\n",
    "#             text_qa_template=qa_template\n",
    "#         )\n",
    "# # Accessing the new prompt\n",
    "# response_synthesizer.get_prompts()['text_qa_template']#['template']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I decided to get the baseline prompts as close to the RAG prompt as possible, rather than twiddling too much with the RAG prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 21.6 s, sys: 891 ms, total: 22.5 s\n",
      "Wall time: 26min 13s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# Metadata and definitions\n",
    "metadata_jargon_defs_rag_for_pickling = []\n",
    "metadata_jargon_defs_rag_for_json = []\n",
    "\n",
    "# Iterate over articles, the reader, the jargon within\n",
    "for idx, row in merged_humam_llm_jargon.iterrows():\n",
    "\n",
    "    # Get list of terms\n",
    "    arxiv_id = row['arxiv_id']\n",
    "    reader_id = row['reader_id']\n",
    "    human_jargon_terms = row['human_jargon_list']\n",
    "\n",
    "    # Sometimes the lists are empty, in that case, no definitions need generation\n",
    "    if human_jargon_terms:\n",
    "        # Strip out any spaces, make sure it's a list format + delete any empty strings\n",
    "        human_jargon_terms = [i.strip() for i in human_jargon_terms.split(',') if i.strip()]\n",
    "    \n",
    "        # Create a metadata filter for the query engine based on the arxiv_id\n",
    "        filters = MetadataFilters(filters=[\n",
    "            ExactMatchFilter(\n",
    "                key=\"arxiv_id\", \n",
    "                value=arxiv_id\n",
    "            )\n",
    "        ])\n",
    "\n",
    "        # Configure retriever\n",
    "        retriever = VectorIndexRetriever(\n",
    "            index=index_loaded,\n",
    "            similarity_top_k=10,\n",
    "            filters=filters\n",
    "        )\n",
    "        \n",
    "        # configure response synthesizer + remember to initialise with the custom prompt\n",
    "        response_synthesizer = get_response_synthesizer(\n",
    "            response_mode=\"compact\",\n",
    "            # text_qa_template=qa_template\n",
    "        )\n",
    "\n",
    "        # assemble query engine\n",
    "        custom_query_engine = RetrieverQueryEngine(\n",
    "            retriever=retriever,\n",
    "            response_synthesizer=response_synthesizer,\n",
    "            node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.3)]\n",
    "        )\n",
    "        \n",
    "        # iterate over the jargon terms and call\n",
    "        for term in human_jargon_terms:\n",
    "                # call for definition\n",
    "                term_definition = custom_query_engine.query(\n",
    "                    f\"Please use 1-2 sentences to explain the following term so that even a reader without deep scientific and technical knowledge can understand it easily: {term}.\")\n",
    "    \n",
    "                # Add it in\n",
    "                metadata_jargon_defs_rag_for_pickling.append({\n",
    "                    'arxiv_id': arxiv_id, \n",
    "                    'reader_id': reader_id, \n",
    "                    'human_jargon_term': term, \n",
    "                    'definition_object': term_definition, \n",
    "                    'definition_text': term_definition.response\n",
    "                })\n",
    "    \n",
    "                metadata_jargon_defs_rag_for_json.append({\n",
    "                    'arxiv_id': arxiv_id, \n",
    "                    'reader_id': reader_id, \n",
    "                    'human_jargon_term': term, \n",
    "                    'definition_text': term_definition.response\n",
    "                })\n",
    "\n",
    "                # print(arxiv_id, reader_id, term_definition.response)\n",
    "    \n",
    "                time.sleep(random.uniform(0, 1))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# metadata_jargon_defs_rag_for_json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save to JSON and pickle respectively\n",
    "import pickle\n",
    "\n",
    "# Today's date for the filename\n",
    "from datetime import datetime\n",
    "today = datetime.today().strftime('%Y%m%d')[2:]\n",
    "\n",
    "with open(f'data/llm_outputs/{today}_march_2024_sampled_jargon_definitions_rag.json', 'w') as file:\n",
    "    json.dump(metadata_jargon_defs_rag_for_json, file, indent=4)\n",
    "\n",
    "# Open a file and use dump() \n",
    "with open(f'data/llm_outputs/{today}_march_2024_sampled_jargon_definitions_objects_rag.pkl', 'wb') as file: \n",
    "    pickle.dump(metadata_jargon_defs_rag_for_pickling, file) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI as oai_real\n",
    "\n",
    "# Setup client\n",
    "client = oai_real(\n",
    "    api_key = os.getenv(\"OPENAI_KEY\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.79 s, sys: 153 ms, total: 5.94 s\n",
      "Wall time: 24min 20s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# Get baselines for GPT-4\n",
    "metadata_jargon_defs_abstract = []\n",
    "\n",
    "# Merge w/ abstracts to add those as metadata\n",
    "merged_humam_llm_jargon_abstracts = pd.merge(merged_humam_llm_jargon, metadata_df[['arxiv_id', 'summary']], on='arxiv_id', how='inner')\n",
    "\n",
    "# Iterate over articles, the reader, the jargon within\n",
    "for idx, row in merged_humam_llm_jargon_abstracts.iterrows():\n",
    "\n",
    "    # Get list of terms\n",
    "    arxiv_id = row['arxiv_id']\n",
    "    reader_id = row['reader_id']\n",
    "    abstract = row['summary']\n",
    "    human_jargon_terms = row['human_jargon_list']\n",
    "\n",
    "    # Sometimes the lists are empty, in that case, no definitions need generation\n",
    "    if human_jargon_terms:\n",
    "\n",
    "        # Strip out any spaces, make sure it's a list format\n",
    "        human_jargon_terms = [i.strip() for i in human_jargon_terms.split(',')]\n",
    "\n",
    "        # iterate over the jargon terms and call\n",
    "        for term in human_jargon_terms:\n",
    "\n",
    "            if len(term)>2:\n",
    "\n",
    "                # call for definition\n",
    "                term_definition = client.chat.completions.create(\n",
    "                    model=\"gpt-4-turbo\",\n",
    "                    messages=[\n",
    "                        {\"role\": \"system\", \"content\": \"You are an expert Q&A system that is trusted around the world.\\nAlways answer the query using the provided context information, and not prior knowledge.\\nSome rules to follow:\\n1. Never directly reference the given context in your answer.\\n2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.\"},\n",
    "                        {\"role\": \"user\", \"content\": f\"Here's a scientific abstract: {abstract}. Given the context information and not prior knowledge, please use 1-2 sentences to explain this term so that even a reader without deep scientific and technical knowledge can understand it easily: {term}. \\n Definition: \"}\n",
    "                    ],\n",
    "                    temperature=1, \n",
    "                    seed=10,\n",
    "                )\n",
    "                term_definition = term_definition.choices[0].message.content\n",
    "                \n",
    "                metadata_jargon_defs_abstract.append({\n",
    "                    'arxiv_id': arxiv_id, \n",
    "                    'reader_id': reader_id, \n",
    "                    'human_jargon_term': term, \n",
    "                    'definition_text': term_definition\n",
    "                })\n",
    "\n",
    "                # print(arxiv_id, reader_id, term_definition)\n",
    "    \n",
    "                time.sleep(random.uniform(0, 1))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# metadata_jargon_defs_abstract"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'data/llm_outputs/{today}_march_2024_sampled_jargon_definitions_abstract.json', 'w') as file:\n",
    "    json.dump(metadata_jargon_defs_abstract, file, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.48 s, sys: 155 ms, total: 5.63 s\n",
      "Wall time: 22min 37s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# Get baselines for GPT-4 -- no context case\n",
    "metadata_jargon_defs_nocontext = []\n",
    "\n",
    "# Iterate over articles, the reader, the jargon within\n",
    "for idx, row in merged_humam_llm_jargon.iterrows():\n",
    "\n",
    "    # Get list of terms\n",
    "    arxiv_id = row['arxiv_id']\n",
    "    reader_id = row['reader_id']\n",
    "    human_jargon_terms = row['human_jargon_list']\n",
    "\n",
    "    # Sometimes the lists are empty, in that case, no definitions need generation\n",
    "    if human_jargon_terms:\n",
    "\n",
    "        # Strip out any spaces, make sure it's a list format\n",
    "        human_jargon_terms = [i.strip() for i in human_jargon_terms.split(',')]\n",
    "\n",
    "        # iterate over the jargon terms and call\n",
    "        for term in human_jargon_terms:\n",
    "\n",
    "            if len(term)>2:\n",
    "\n",
    "                # call for definition\n",
    "                term_definition = client.chat.completions.create(\n",
    "                    model=\"gpt-4-turbo\",\n",
    "                    messages=[\n",
    "                        {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "                        {\"role\": \"user\", \"content\": f\"Please use 1-2 sentences to explain this term so that even a reader without deep scientific and technical knowledge can understand it easily: {term}. \\n Definition: \"}\n",
    "                    ],\n",
    "                    temperature=1, \n",
    "                    seed=10,\n",
    "                )\n",
    "                term_definition = term_definition.choices[0].message.content\n",
    "                \n",
    "                metadata_jargon_defs_nocontext.append({\n",
    "                    'arxiv_id': arxiv_id, \n",
    "                    'reader_id': reader_id, \n",
    "                    'human_jargon_term': term, \n",
    "                    'definition_text': term_definition\n",
    "                })\n",
    "    \n",
    "                time.sleep(random.uniform(0, 1))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# metadata_jargon_defs_nocontext"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'data/llm_outputs/{today}_march_2024_sampled_jargon_definitions_nocontext.json', 'w') as file:\n",
    "    json.dump(metadata_jargon_defs_nocontext, file, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}