src/get_results.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0505fcde-d219-4d26-8b68-7ae4d91c7f03",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.vectorstores import Annoy, FAISS\n",
    "from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction\n",
    "from rouge_score import rouge_scorer\n",
    "from tqdm import tqdm\n",
    "\n",
    "import numpy as np\n",
    "import os\n",
    "import pickle\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "843b1938-ea5d-4235-8885-6a2021d3ce9b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load env variables\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d8226b4-d078-41e0-8381-7c976f160a9e",
   "metadata": {},
   "source": [
    "## Information Retrieval - Metrics Calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cf03c092-a470-4632-ac34-8d094a287cc0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Documents loaded succesfully.\n",
      "Queries loaded succesfully.\n",
      "Ground truth loaded succesfully.\n"
     ]
    }
   ],
   "source": [
    "# Load CISI documents & queries\n",
    "\n",
    "with open(\"./cisi/embeddings/lsi/documents.pkl\", \"rb\") as f:\n",
    "    documents = pickle.load(f)\n",
    "\n",
    "print(\"Documents loaded succesfully.\")\n",
    "\n",
    "with open(\"./cisi/embeddings/lsi/queries.pkl\", \"rb\") as f:\n",
    "    queries = pickle.load(f)\n",
    "\n",
    "print(\"Queries loaded succesfully.\")\n",
    "\n",
    "# Load CISI ground truth\n",
    "with open(\"./cisi/dataset/rel_set.pkl\", \"rb\") as f:\n",
    "    ground_truth = pickle.load(f)\n",
    "\n",
    "print(\"Ground truth loaded succesfully.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "42292239-268f-4533-a56d-8cce95947c29",
   "metadata": {},
   "outputs": [],
   "source": [
    "ground_truth = list(ground_truth.values()) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f5d321df-b758-49d5-953e-dfdadc03dd4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def precision_at_k(ranked_docs, relevant_docs, k=10):\n",
    "    retrieved_relevant = 0\n",
    "    for doc_id in ranked_docs[:k]:\n",
    "        if doc_id in relevant_docs:\n",
    "            retrieved_relevant += 1\n",
    "    return retrieved_relevant / k\n",
    "\n",
    "def recall_at_k(ranked_docs, relevant_docs, k=10):\n",
    "    retrieved_relevant = sum(1 for doc_id in ranked_docs[:k] if doc_id in relevant_docs)\n",
    "    return retrieved_relevant / len(relevant_docs) if relevant_docs else 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6fe31135-30c8-473a-b0f9-5414528d17bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cosine_similarity(a, b):\n",
    "    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67ec5583-6d4f-4777-993f-7fc09d1543ee",
   "metadata": {},
   "source": [
    "#### LSI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2436d55b-96b8-4340-aa23-f81031b4adb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LSI Predictions loaded succesfully.\n"
     ]
    }
   ],
   "source": [
    "# Load LSI predictions for CISI\n",
    "with open(\"./cisi/ir_techniques/exact_search/index.pkl\", \"rb\") as f:\n",
    "    lsi_predictions = pickle.load(f)\n",
    "\n",
    "print(\"LSI Predictions loaded succesfully.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "223fc22c-90e7-4c77-aeb2-86225e99676d",
   "metadata": {},
   "outputs": [],
   "source": [
    "lsi_predictions = list(lsi_predictions.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "fb85bcdc-a8e7-42ab-be15-1bc29a063656",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(lsi_predictions[1], ground_truth[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "553e5e54-9753-4e25-a11a-fa75438ff5b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_precision_at_k_lsi = np.mean([precision_at_k(preds,label) for preds,label in zip(lsi_predictions,ground_truth)])\n",
    "mean_recall_at_k_lsi = np.mean([recall_at_k(preds,label) for preds,label in zip(lsi_predictions,ground_truth)])\n",
    "mean_sps_lsi = np.mean([precision_at_k(preds,label,1) for preds,label in zip(lsi_predictions,ground_truth)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9975226f-5285-4516-b1ee-f7c20d6e5a6c",
   "metadata": {},
   "source": [
    "#### OpenAI Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "cd7db9d0-fac7-4151-b3f8-0e6514f653e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load OpenAI embeddings\n",
    "with open(\"./cisi/embeddings/text-embedding-ada-002-v2/documents.pkl\", \"rb\") as f:\n",
    "    documents_openai = pickle.load(f)\n",
    "\n",
    "with open(\"./cisi/embeddings/text-embedding-ada-002-v2/queries.pkl\", \"rb\") as f:\n",
    "    queries_openai = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "22c0dce2-445b-41de-b6b4-f7a1f4bfae07",
   "metadata": {},
   "outputs": [],
   "source": [
    "# queries_openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "63f4937b-57c4-4c5f-a12e-bcff0605ba3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing similarity scores: 100%|████████████| 112/112 [00:29<00:00,  3.82it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean execution time for all queries: 361.42 ms\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Calculate cosine similarity for each query-document pair\n",
    "similarity_scores = {}\n",
    "execution_times = [] \n",
    "for query_id, query in tqdm(queries_openai.items(), desc = 'Computing similarity scores'):\n",
    "    query_embedding = query['embedding']\n",
    "    scores = []\n",
    "    start_time = time.time()\n",
    "    for doc in documents_openai:\n",
    "        doc_embedding = documents_openai[doc]['embedding']\n",
    "        sim_score = cosine_similarity(query_embedding, doc_embedding)\n",
    "        scores.append((doc, sim_score))\n",
    "    \n",
    "    end_time = time.time()  # Record end time\n",
    "    execution_time = (end_time - start_time) * 1000  # Convert to milliseconds\n",
    "    execution_times.append(execution_time)\n",
    "    similarity_scores[query_id] = sorted(scores, key=lambda x: x[1], reverse=True)\n",
    "\n",
    "mean_execution_time_openai = sum(execution_times) / len(execution_times)\n",
    "print(f\"Mean execution time for all queries: {mean_execution_time:.2f} ms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "6e178e0a-65cc-4f1a-9c8e-0eb230b908d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "openai_predictions = {}\n",
    "for query_id, scores in similarity_scores.items():\n",
    "    # Flatten the scores to get only document IDs\n",
    "    scores_flattened = [doc_id for doc_id, _ in scores]\n",
    "    openai_predictions[query_id] = scores_flattened"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "9f91ee73-b7e9-4ae9-95fa-8d0d00bfe2bc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "openai_predictions = list(openai_predictions.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "cd2972ee-95c3-44e8-9589-40f4190599ba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1460, 3114)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(openai_predictions[0]), len(ground_truth[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "afeafb72-0cad-4d1b-8742-3170c3aa6ee6",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_precision_at_k_openai = np.mean([precision_at_k(preds,label) for preds,label in zip(openai_predictions,ground_truth)])\n",
    "mean_recall_at_k_openai = np.mean([recall_at_k(preds,label) for preds,label in zip(openai_predictions,ground_truth)])\n",
    "mean_sps_openai = np.mean([precision_at_k(preds,label,1) for preds,label in zip(openai_predictions,ground_truth)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c95b7bc6-b1e8-4f89-9d00-2765ce7869cb",
   "metadata": {},
   "source": [
    "#### VectorDB - FAISS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "eda8f1eb-357b-498a-8c9d-61e0fb707b98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load index from file\n",
    "loaded_faiss_vs = FAISS.load_local(\n",
    "    folder_path=\"./cisi/ir_techniques/faiss/\",\n",
    "    embeddings=OpenAIEmbeddings())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f9461183-f2a1-4bf9-a206-f2a1f3d179fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing similarity scores: 100%|██████████| 112/112 [00:00<00:00, 1668.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean execution time for all queries: 0.68 ms\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Calculate cosine similarity for each query-document pair\n",
    "faiss_similarity_scores = {}\n",
    "execution_times = [] \n",
    "for query_id, query in tqdm(queries_openai.items(), desc=\"Computing similarity scores\"):\n",
    "    start_time = time.time()\n",
    "    scores = loaded_faiss_vs.similarity_search_with_score_by_vector(query[\"embedding\"], k=100)\n",
    "    end_time = time.time()  # Record end time\n",
    "    execution_time = (end_time - start_time) * 1000  # Convert to milliseconds\n",
    "    execution_times.append(execution_time)\n",
    "\n",
    "    new_scores = []\n",
    "    for score in scores:\n",
    "        new_scores.append((int(score[0].page_content), 1 - score[1]))\n",
    "    faiss_similarity_scores[query_id] = sorted(new_scores, key=lambda x: x[1], reverse=True)\n",
    "\n",
    "mean_execution_time_faiss = sum(execution_times) / len(execution_times)\n",
    "print(f\"Mean execution time for all queries: {mean_execution_time:.2f} ms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "e764340a-9cd4-44a6-aacc-ef3b90c67e77",
   "metadata": {},
   "outputs": [],
   "source": [
    "faiss_predictions = [0]*(len(faiss_similarity_scores)+1)\n",
    "for idx, scores in faiss_similarity_scores.items():\n",
    "    scores_flattened = [doc for doc, score in scores]\n",
    "    faiss_predictions[idx] = scores_flattened"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "85db28cb-ba9c-4439-8023-b62031545b21",
   "metadata": {},
   "outputs": [],
   "source": [
    "faiss_mean_precision_at_k = np.mean(\n",
    "    [precision_at_k(preds, label) for preds, label in zip(faiss_predictions[1:], ground_truth)])\n",
    "faiss_mean_recall_at_k = np.mean(\n",
    "    [recall_at_k(preds, label) for preds, label in zip(faiss_predictions[1:], ground_truth)])\n",
    "faiss_mean_sps = np.mean(\n",
    "    [precision_at_k(preds, label, 1) for preds, label in zip(faiss_predictions[1:], ground_truth)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "054cb5b7-f4b4-47b8-b12e-20a9f6301adf",
   "metadata": {},
   "source": [
    "#### VectorDB - ANNOY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "e8f03d10-26c5-4d7c-a143-f2a0b23ee192",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load ANNOY index from file\n",
    "loaded_annoy_vs = Annoy.load_local(\n",
    "    folder_path=\"./cisi/ir_techniques/annoy/\", \n",
    "    embeddings=OpenAIEmbeddings())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "2656a8c9-44e6-4a1b-bea2-b62a3687de97",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing similarity scores: 100%|███████████| 112/112 [00:00<00:00, 374.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean execution time for all queries: 2.16 ms\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Calculate cosine similarity for each query-document pair\n",
    "annoy_similarity_scores = {}\n",
    "execution_times = [] \n",
    "for query_id, query in tqdm(queries_openai.items(), desc=\"Computing similarity scores\"):\n",
    "    start_time = time.time()\n",
    "    scores = loaded_annoy_vs.similarity_search_with_score_by_vector(query[\"embedding\"], k=100)\n",
    "    end_time = time.time()  # Record end time\n",
    "    execution_time = (end_time - start_time) * 1000  # Convert to milliseconds\n",
    "    execution_times.append(execution_time)\n",
    "\n",
    "    new_scores = []\n",
    "    for score in scores:\n",
    "        new_scores.append((int(score[0].page_content), score[1]))\n",
    "    annoy_similarity_scores[query_id] = sorted(new_scores, key=lambda x: x[1], reverse=True)\n",
    "\n",
    "mean_execution_time_annoy = sum(execution_times) / len(execution_times)\n",
    "print(f\"Mean execution time for all queries: {mean_execution_time:.2f} ms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "a5d1eb46-8c94-429c-8e7a-91c5bba014ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "annoy_predictions = [0]*(len(annoy_similarity_scores)+1)\n",
    "for idx, scores in annoy_similarity_scores.items():\n",
    "    scores_flattened = [doc for doc, score in scores]\n",
    "    annoy_predictions[idx] = scores_flattened"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "355ef12c-dfbf-40be-9fa4-008586a44b81",
   "metadata": {},
   "outputs": [],
   "source": [
    "annoy_mean_precision_at_k = np.mean(\n",
    "    [precision_at_k(preds, label) for preds, label in zip(annoy_predictions[1:], ground_truth)])\n",
    "annoy_mean_recall_at_k = np.mean(\n",
    "    [recall_at_k(preds, label) for preds, label in zip(annoy_predictions[1:], ground_truth)])\n",
    "annoy_mean_sps = np.mean(\n",
    "    [precision_at_k(preds, label, 1) for preds, label in zip(annoy_predictions[1:], ground_truth)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3516ac6a-c46f-4073-b3cb-e92b3b1decbb",
   "metadata": {},
   "source": [
    "## RAG Evaluation - Metrics Calculation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af72d5ad-6cbe-45a9-97f6-aea39961f2cc",
   "metadata": {},
   "source": [
    "### CISI Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43d87c06-69d7-4cd8-9aa7-884d97b745ad",
   "metadata": {},
   "source": [
    "#### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "9840060b-3658-4471-861b-f2de61eb0cff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Documents loaded succesfully.\n",
      "Queries loaded succesfully.\n",
      "Ground truth loaded succesfully.\n"
     ]
    }
   ],
   "source": [
    "# Load arguAna documents & queries\n",
    "\n",
    "with open(\"./cisi/dataset/documents.pkl\", \"rb\") as f:\n",
    "    documents = pickle.load(f)\n",
    "\n",
    "print(\"Documents loaded succesfully.\")\n",
    "\n",
    "with open(\"./cisi/dataset/queries.pkl\", \"rb\") as f:\n",
    "    queries = pickle.load(f)\n",
    "\n",
    "print(\"Queries loaded succesfully.\")\n",
    "\n",
    "# Load CISI ground truth\n",
    "with open(\"./cisi/dataset/rel_set.pkl\", \"rb\") as f:\n",
    "    ground_truth = pickle.load(f)\n",
    "\n",
    "print(\"Ground truth loaded succesfully.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "1e31be3a-56cd-4abb-8e69-d1b32aa910e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "da-vinci-0.0.3 responses with & without RAG loaded succesfully.\n"
     ]
    }
   ],
   "source": [
    "# Load davinci-0.0.3 responses\n",
    "with open(\"./cisi/responses/da-vinci-0.0.3/llm_w_rag_faiss.pkl\", \"rb\") as f:\n",
    "    da_vinci_llm_w_rag = pickle.load(f)\n",
    "\n",
    "with open(\"./cisi/responses/da-vinci-0.0.3/llm_wo_rag.pkl\", \"rb\") as f:\n",
    "    da_vinci_llm_wo_rag = pickle.load(f)\n",
    "\n",
    "print(\"da-vinci-0.0.3 responses with & without RAG loaded succesfully.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "e8c8a233-4f29-4dd7-bb70-44ddc27879a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gpt-3.5-turbo-instruct responses with & without RAG loaded succesfully.\n"
     ]
    }
   ],
   "source": [
    "# Load gpt-3.5-turbo-instruct responses\n",
    "with open(\"./cisi/responses/gpt-3.5-turbo-instruct/llm_w_rag_faiss.pkl\", \"rb\") as f:\n",
    "    gpt_llm_w_rag_faiss = pickle.load(f)\n",
    "\n",
    "with open(\"./cisi/responses/gpt-3.5-turbo-instruct/llm_w_rag_exact_search.pkl\", \"rb\") as f:\n",
    "    gpt_llm_w_rag_exact_search = pickle.load(f)\n",
    "    \n",
    "with open(\"./cisi/responses/gpt-3.5-turbo-instruct/llm_wo_rag.pkl\", \"rb\") as f:\n",
    "    gpt_llm_wo_rag = pickle.load(f)\n",
    "\n",
    "print(\"gpt-3.5-turbo-instruct responses with & without RAG loaded succesfully.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "5d3723fa-0d3a-441b-8a28-352927f0166f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "llama-7b responses with & without RAG loaded succesfully.\n"
     ]
    }
   ],
   "source": [
    "# Load llama-7b responses\n",
    "with open(\"./cisi/responses/llama-7b/llm_w_rag_faiss.pkl\", \"rb\") as f:\n",
    "    llama_7b_llm_w_rag = pickle.load(f)\n",
    "\n",
    "with open(\"./cisi/responses/llama-7b/llm_wo_rag.pkl\", \"rb\") as f:\n",
    "    llama_7b_llm_wo_rag = pickle.load(f)\n",
    "\n",
    "print(\"llama-7b responses with & without RAG loaded succesfully.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "0865f91f-7f09-4988-a2f7-e0608bc09388",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Implement BLEU evaluation function\n",
    "def compute_bleu(references, candidate):\n",
    "    smoothing = SmoothingFunction().method5\n",
    "    return sentence_bleu(references, candidate, smoothing_function=smoothing)\n",
    "\n",
    "# Implement ROUGE evaluation function\n",
    "def compute_rouge(references, candidate):\n",
    "    scorer = rouge_scorer.RougeScorer(['rouge1'], use_stemmer=True)\n",
    "    total_score = 0\n",
    "\n",
    "    # Compute ROUGE for each reference\n",
    "    for reference in references:\n",
    "        scores = scorer.score(reference, candidate)\n",
    "        total_score += scores['rouge1'].fmeasure\n",
    "\n",
    "    # Calculate average score\n",
    "    if references:\n",
    "        average_score = total_score / len(references)\n",
    "    else:\n",
    "        average_score = 0\n",
    "    return average_score\n",
    "\n",
    "\n",
    "\n",
    "def compute_mean_bleu_score(rag_responses, relevant_docs, K = 10):\n",
    "    total_bleu_score = 0.0\n",
    "    num_queries = 0\n",
    "    for query_id, relevant_docs in ground_truth.items():\n",
    "        if (query_id not in queries) or (query_id not in rag_responses):\n",
    "            continue\n",
    "        query_text = queries[query_id]['text']\n",
    "        response = rag_responses[query_id]['response']\n",
    "        if relevant_docs:\n",
    "            bleu_score = compute_bleu([documents[id]['text'] for id in relevant_docs[:K]], response)\n",
    "        else:\n",
    "            bleu_score = 0\n",
    "        total_bleu_score += bleu_score\n",
    "        num_queries += 1\n",
    "        if num_queries == 101:\n",
    "            break\n",
    "    mean_bleu_score = total_bleu_score / num_queries\n",
    "    return mean_bleu_score\n",
    "\n",
    "def compute_mean_rouge_score(rag_responses, relevant_docs, K = 10):\n",
    "    total_rouge_score = 0.0\n",
    "    num_queries = 0\n",
    "    for query_id, relevant_docs in ground_truth.items():\n",
    "        if (query_id not in queries) or (query_id not in rag_responses):\n",
    "            continue\n",
    "        query_text = queries[query_id]['text']\n",
    "        response = rag_responses[query_id]['response']\n",
    "        rouge_score = compute_rouge([documents[id]['text'] for id in relevant_docs[:K]], response)\n",
    "        total_rouge_score += rouge_score\n",
    "        num_queries += 1\n",
    "        if num_queries == 101:\n",
    "            break\n",
    "    mean_rouge_score = total_rouge_score / num_queries\n",
    "    return mean_rouge_score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9eeffdb9-5cec-4d63-aa61-a46b56c23d24",
   "metadata": {},
   "source": [
    "#### da-vinci-0.0.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "63eef9b0-2d50-438e-b9c5-1f56e473d43e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute for davinci with without RAG\n",
    "\n",
    "cisi_mean_bleu_score_davinci_wo_rag = compute_mean_bleu_score(da_vinci_llm_wo_rag, ground_truth)\n",
    "cisi_mean_bleu_score_davinci_w_rag = compute_mean_bleu_score(da_vinci_llm_w_rag, ground_truth)\n",
    "cisi_mean_rouge_score_davinci_wo_rag = compute_mean_rouge_score(da_vinci_llm_wo_rag, ground_truth)\n",
    "cisi_mean_rouge_score_davinci_w_rag = compute_mean_rouge_score(da_vinci_llm_w_rag, ground_truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17e8c761-5fa4-400e-9070-2b3b6d5c901c",
   "metadata": {},
   "source": [
    "#### gpt-3.5-turbo-instruct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "cfd15edd-17de-40b6-9eab-505b07082b80",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute for gpt without RAG\n",
    "\n",
    "cisi_mean_bleu_score_gpt_wo_rag = compute_mean_bleu_score(gpt_llm_wo_rag, ground_truth)\n",
    "cisi_mean_rouge_score_gpt_wo_rag = compute_mean_rouge_score(gpt_llm_wo_rag, ground_truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1303ad0-bc58-4d22-89cd-10d038c3d2e5",
   "metadata": {},
   "source": [
    "##### NLA Method (LSI + Truncated SVD) for RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "ee4f796b-d04a-47fe-a739-4cbcb00af20f",
   "metadata": {},
   "outputs": [],
   "source": [
    "cisi_mean_bleu_score_gpt_w_lsi_rag = compute_mean_bleu_score(gpt_llm_w_rag_exact_search, ground_truth)\n",
    "cisi_mean_rouge_score_gpt_w_lsi_rag = compute_mean_rouge_score(gpt_llm_w_rag_exact_search, ground_truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "831fe6c3-1062-4767-baaa-109de49204e0",
   "metadata": {},
   "source": [
    "##### SOTA Method (VectorDB i.e. FAISS) for RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "c0328dd5-8145-4bf3-be88-aca567bc93ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "cisi_mean_bleu_score_gpt_w_faiss_rag = compute_mean_bleu_score(gpt_llm_w_rag_faiss, ground_truth)\n",
    "cisi_mean_rouge_score_gpt_w_faiss_rag = compute_mean_rouge_score(gpt_llm_w_rag_faiss, ground_truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "262a53ac-d6ce-46e3-950a-b6e817755cfc",
   "metadata": {},
   "source": [
    "#### llama-7b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "7b926e2c-4bc3-4e8d-a96d-193608d84d80",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute for llama-7b with without RAG\n",
    "\n",
    "cisi_mean_bleu_score_llama_7b_wo_rag = compute_mean_bleu_score(llama_7b_llm_wo_rag, ground_truth)\n",
    "cisi_mean_bleu_score_llama_7b_w_rag = compute_mean_bleu_score(llama_7b_llm_w_rag, ground_truth)\n",
    "cisi_mean_rouge_score_llama_7b_wo_rag = compute_mean_rouge_score(llama_7b_llm_wo_rag, ground_truth)\n",
    "cisi_mean_rouge_score_llama_7b_w_rag = compute_mean_rouge_score(llama_7b_llm_w_rag, ground_truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ad02390-192e-4b07-a2f0-71c2810229df",
   "metadata": {},
   "source": [
    "### arguAna Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9022243f-2f70-4d79-9559-fb4bae31b81d",
   "metadata": {},
   "source": [
    "#### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "e114891e-89bd-4a0f-9a67-11b0933b1052",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Documents loaded succesfully.\n",
      "Queries loaded succesfully.\n",
      "Ground truth loaded succesfully.\n"
     ]
    }
   ],
   "source": [
    "# Load arguAna documents & queries\n",
    "\n",
    "with open(\"./arguana/dataset/documents.pkl\", \"rb\") as f:\n",
    "    documents = pickle.load(f)\n",
    "\n",
    "print(\"Documents loaded succesfully.\")\n",
    "\n",
    "with open(\"./arguana/dataset/queries.pkl\", \"rb\") as f:\n",
    "    queries = pickle.load(f)\n",
    "\n",
    "print(\"Queries loaded succesfully.\")\n",
    "\n",
    "# Load CISI ground truth\n",
    "with open(\"./arguana/dataset/rel_set.pkl\", \"rb\") as f:\n",
    "    ground_truth = pickle.load(f)\n",
    "\n",
    "print(\"Ground truth loaded succesfully.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "af726f54-6518-4f0b-9ada-0ceadca912c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gpt-3.5-turbo-instruct responses with & without RAG loaded succesfully.\n"
     ]
    }
   ],
   "source": [
    "# Load gpt-3.5-turbo-instruct responses\n",
    "with open(\"./arguana/responses/gpt-3.5-turbo-instruct/llm_w_rag_faiss.pkl\", \"rb\") as f:\n",
    "    gpt_llm_w_rag_faiss = pickle.load(f)\n",
    "\n",
    "with open(\"./arguana/responses/gpt-3.5-turbo-instruct/llm_w_rag_exact_search.pkl\", \"rb\") as f:\n",
    "    gpt_llm_w_rag_exact_search = pickle.load(f)\n",
    "    \n",
    "print(\"gpt-3.5-turbo-instruct responses with & without RAG loaded succesfully.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02a5fddd-2f4d-4c31-9cba-aa13dea51685",
   "metadata": {},
   "source": [
    "#### gpt-3.5-turbo-instruct"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52cf4cdc-482a-46b2-b31e-6d20b3889fc9",
   "metadata": {},
   "source": [
    "##### NLA Method (LSI + Truncated SVD) for RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "015f74fb-a594-4b08-8837-9ea603f2f4d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "arguana_mean_bleu_score_gpt_w_lsi_rag = compute_mean_bleu_score(gpt_llm_w_rag_exact_search, ground_truth)\n",
    "arguana_mean_rouge_score_gpt_w_lsi_rag = compute_mean_rouge_score(gpt_llm_w_rag_exact_search, ground_truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1bd6048-a429-4b61-9fb4-97c80604c7f5",
   "metadata": {},
   "source": [
    "##### SOTA Method (VectorDB i.e. FAISS) for RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "ba7ee9ca-2d1e-441c-86fc-7765046f74ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "arguana_mean_bleu_score_gpt_w_faiss_rag = compute_mean_bleu_score(gpt_llm_w_rag_faiss, ground_truth)\n",
    "arguana_mean_rouge_score_gpt_w_faiss_rag = compute_mean_rouge_score(gpt_llm_w_rag_faiss, ground_truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdff7fbe-0522-4591-a135-5c08f9791ae2",
   "metadata": {},
   "source": [
    "### NFCorpus Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58cdfdce-bc06-4d0a-85f8-a7674427a108",
   "metadata": {},
   "source": [
    "#### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "22702608-4448-41ff-b41d-371fec0383e9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Documents loaded succesfully.\n",
      "Queries loaded succesfully.\n",
      "Ground truth loaded succesfully.\n"
     ]
    }
   ],
   "source": [
    "# Load arguAna documents & queries\n",
    "\n",
    "with open(\"./nfcorpus/dataset/documents.pkl\", \"rb\") as f:\n",
    "    documents = pickle.load(f)\n",
    "\n",
    "print(\"Documents loaded succesfully.\")\n",
    "\n",
    "with open(\"./nfcorpus/dataset/queries.pkl\", \"rb\") as f:\n",
    "    queries = pickle.load(f)\n",
    "\n",
    "print(\"Queries loaded succesfully.\")\n",
    "\n",
    "# Load CISI ground truth\n",
    "with open(\"./nfcorpus/dataset/rel_set.pkl\", \"rb\") as f:\n",
    "    ground_truth = pickle.load(f)\n",
    "\n",
    "print(\"Ground truth loaded succesfully.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "a28c9300-afb6-4014-a532-3f56e3518fa9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gpt-3.5-turbo-instruct responses with & without RAG loaded succesfully.\n"
     ]
    }
   ],
   "source": [
    "# Load gpt-3.5-turbo-instruct responses\n",
    "with open(\"./nfcorpus/responses/gpt-3.5-turbo-instruct/llm_w_rag_faiss.pkl\", \"rb\") as f:\n",
    "    gpt_llm_w_rag_faiss = pickle.load(f)\n",
    "\n",
    "with open(\"./nfcorpus/responses/gpt-3.5-turbo-instruct/llm_w_rag_exact_search.pkl\", \"rb\") as f:\n",
    "    gpt_llm_w_rag_exact_search = pickle.load(f)\n",
    "    \n",
    "print(\"gpt-3.5-turbo-instruct responses with & without RAG loaded succesfully.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e984998f-71bf-45c5-9700-8f2c85099966",
   "metadata": {},
   "source": [
    "#### gpt-3.5-turbo-instruct"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3f29c8a-ef97-40ff-aa1f-f976457fc72a",
   "metadata": {},
   "source": [
    "##### NLA Method (LSI + Truncated SVD) for RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "d97bde1f-47e9-4e00-b6d3-70a8cf3c23f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "nfcorpus_mean_bleu_score_gpt_w_lsi_rag = compute_mean_bleu_score(gpt_llm_w_rag_exact_search, ground_truth)\n",
    "nfcorpus_mean_rouge_score_gpt_w_lsi_rag = compute_mean_rouge_score(gpt_llm_w_rag_exact_search, ground_truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a34e2cb-0284-4040-857d-f5c11c73a46b",
   "metadata": {},
   "source": [
    "##### SOTA Method (VectorDB i.e. FAISS) for RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "e850df31-6450-4399-8601-c3e8c1c95711",
   "metadata": {},
   "outputs": [],
   "source": [
    "nfcorpus_mean_bleu_score_gpt_w_faiss_rag = compute_mean_bleu_score(gpt_llm_w_rag_faiss, ground_truth)\n",
    "nfcorpus_mean_rouge_score_gpt_w_faiss_rag = compute_mean_rouge_score(gpt_llm_w_rag_faiss, ground_truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b42164a7-aa5e-4696-b57a-279a0a8a18ee",
   "metadata": {},
   "source": [
    "# Experimental Results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c807d9c-0da4-4d9e-89ca-954b6982c217",
   "metadata": {},
   "source": [
    "### Main Experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "91fd3e5c-e212-4065-84f0-2ed3959e4444",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------------------------------------+\n",
      "|                   RAG Methodology Comparison                   |\n",
      "+----------+----------------+-----------------+------------------+\n",
      "| Dataset  |     Method     | Mean BLEU Score | Mean ROUGE Score |\n",
      "+----------+----------------+-----------------+------------------+\n",
      "|   CISI   |   LSI + SVD    |     0.74410     |     0.21670      |\n",
      "|          | FAISS VectorDB |     0.83770     |     0.22260      |\n",
      "|          |                |                 |                  |\n",
      "| arguAna  |   LSI + SVD    |     0.27210     |     0.21890      |\n",
      "|          | FAISS VectorDB |     0.69370     |     0.27390      |\n",
      "|          |                |                 |                  |\n",
      "| nfcorpus |   LSI + SVD    |     0.47090     |     0.17130      |\n",
      "|          | FAISS VectorDB |     0.60990     |     0.20240      |\n",
      "+----------+----------------+-----------------+------------------+\n"
     ]
    }
   ],
   "source": [
    "from prettytable import PrettyTable\n",
    "\n",
    "\n",
    "data_cisi = [\n",
    "    {\n",
    "        \"Dataset\": \"CISI\",\n",
    "        \"Method\": \"LSI + SVD\",\n",
    "        \"Mean BLEU Score\": cisi_mean_bleu_score_gpt_w_lsi_rag,\n",
    "        \"Mean ROUGE Score\": cisi_mean_rouge_score_gpt_w_lsi_rag,\n",
    "    },\n",
    "    {\n",
    "        \"Dataset\": \"\",\n",
    "        \"Method\": \"FAISS VectorDB\",\n",
    "        \"Mean BLEU Score\": cisi_mean_bleu_score_gpt_w_faiss_rag,\n",
    "        \"Mean ROUGE Score\": cisi_mean_rouge_score_gpt_w_faiss_rag,\n",
    "    },\n",
    "]\n",
    "\n",
    "\n",
    "data_arguAna = [\n",
    "        {\n",
    "        \"Dataset\": \"arguAna\",\n",
    "        \"Method\": \"LSI + SVD\",\n",
    "        \"Mean BLEU Score\": arguana_mean_bleu_score_gpt_w_lsi_rag,\n",
    "        \"Mean ROUGE Score\": arguana_mean_rouge_score_gpt_w_lsi_rag,\n",
    "    },\n",
    "    {\n",
    "        \"Dataset\": \"\",\n",
    "        \"Method\": \"FAISS VectorDB\",\n",
    "        \"Mean BLEU Score\": arguana_mean_bleu_score_gpt_w_faiss_rag,\n",
    "        \"Mean ROUGE Score\": arguana_mean_rouge_score_gpt_w_faiss_rag,\n",
    "    },\n",
    "\n",
    "]\n",
    "\n",
    "data_nfcorpus = [\n",
    "        {\n",
    "        \"Dataset\": \"nfcorpus\",\n",
    "        \"Method\": \"LSI + SVD\",\n",
    "        \"Mean BLEU Score\": nfcorpus_mean_bleu_score_gpt_w_lsi_rag,\n",
    "        \"Mean ROUGE Score\": nfcorpus_mean_rouge_score_gpt_w_lsi_rag,\n",
    "    },\n",
    "    {\n",
    "        \"Dataset\": \"\",\n",
    "        \"Method\": \"FAISS VectorDB\",\n",
    "        \"Mean BLEU Score\": nfcorpus_mean_bleu_score_gpt_w_faiss_rag,\n",
    "        \"Mean ROUGE Score\": nfcorpus_mean_rouge_score_gpt_w_faiss_rag,\n",
    "    },\n",
    "\n",
    "]\n",
    "\n",
    "combined_data = data_cisi + [{\"Dataset\": \"\", \"Method\": \"\", \"Mean BLEU Score\": \"\", \"Mean ROUGE Score\": \"\"}] + data_arguAna + [{\"Dataset\": \"\", \"Method\": \"\", \"Mean BLEU Score\": \"\", \"Mean ROUGE Score\": \"\"}] +data_nfcorpus\n",
    "\n",
    "table = PrettyTable()\n",
    "table.field_names = [\"Dataset\", \"Method\", \"Mean BLEU Score\", \"Mean ROUGE Score\"]\n",
    "table.float_format = \"5.5\" \n",
    "\n",
    "table.title = \"RAG Methodology Comparison\"\n",
    "\n",
    "for row in combined_data:\n",
    "    table.add_row([row[\"Dataset\"], row[\"Method\"], row[\"Mean BLEU Score\"], row[\"Mean ROUGE Score\"]])\n",
    "\n",
    "\n",
    "print(table)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a577ff8-853a-47e5-9f9f-ef4d85072546",
   "metadata": {},
   "source": [
    "### Plain LLMs vs RAG - Sub Experiment 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "459dac27-f01e-439f-9efc-960c9c544923",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Results without RAG:\n",
      "+------------------------+-----------------+------------------+\n",
      "|         Model          | Mean BLEU Score | Mean ROUGE Score |\n",
      "+------------------------+-----------------+------------------+\n",
      "|        da-vinci        |      0.6094     |      0.1603      |\n",
      "| gpt-3.5-turbo-instruct |      0.7937     |      0.2425      |\n",
      "|         LLAMA          |      0.7199     |      0.2061      |\n",
      "+------------------------+-----------------+------------------+\n",
      "\n",
      "Results with RAG:\n",
      "+------------------------+-----------------+------------------+\n",
      "|         Model          | Mean BLEU Score | Mean ROUGE Score |\n",
      "+------------------------+-----------------+------------------+\n",
      "|        da-vinci        |      0.8377     |      0.2226      |\n",
      "| gpt-3.5-turbo-instruct |      0.8761     |      0.2982      |\n",
      "|         LLAMA          |      0.7852     |      0.2412      |\n",
      "+------------------------+-----------------+------------------+\n"
     ]
    }
   ],
   "source": [
    "from prettytable import PrettyTable\n",
    "\n",
    "# Data without RAG\n",
    "data_without_rag = [\n",
    "    {\"Model\": \"da-vinci\", \"BLEU Score\": cisi_mean_bleu_score_davinci_wo_rag, \"ROUGE Score\": cisi_mean_rouge_score_davinci_wo_rag},\n",
    "    {\"Model\": \"gpt-3.5-turbo-instruct\", \"BLEU Score\": cisi_mean_bleu_score_gpt_wo_rag, \"ROUGE Score\": cisi_mean_rouge_score_gpt_wo_rag},\n",
    "    {\"Model\": \"LLAMA\", \"BLEU Score\": cisi_mean_bleu_score_llama_7b_wo_rag, \"ROUGE Score\": cisi_mean_rouge_score_llama_7b_wo_rag},\n",
    "]\n",
    "\n",
    "# Data with RAG\n",
    "data_with_rag = [\n",
    "    {\"Model\": \"da-vinci\", \"BLEU Score\": cisi_mean_bleu_score_davinci_w_faiss_rag, \"ROUGE Score\": cisi_mean_rouge_score_davinci_w_faiss_rag},\n",
    "    {\"Model\": \"gpt-3.5-turbo-instruct\", \"BLEU Score\": cisi_mean_bleu_score_gpt_w_faiss_rag, \"ROUGE Score\": cisi_mean_rouge_score_gpt_w_faiss_rag},\n",
    "    {\"Model\": \"LLAMA\", \"BLEU Score\": cisi_mean_bleu_score_llama_7b_w_faiss_rag, \"ROUGE Score\": cisi_mean_rouge_score_llama_7b_w_faiss_rag},\n",
    "]\n",
    "\n",
    "# Create tables\n",
    "table_without_rag = PrettyTable()\n",
    "table_without_rag.field_names = [\"Model\", \"Mean BLEU Score\", \"Mean ROUGE Score\"]\n",
    "\n",
    "for row in data_without_rag:\n",
    "    table_without_rag.add_row([row[\"Model\"], row[\"BLEU Score\"], row[\"ROUGE Score\"]])\n",
    "\n",
    "table_with_rag = PrettyTable()\n",
    "table_with_rag.field_names = [\"Model\", \"Mean BLEU Score\", \"Mean ROUGE Score\"]\n",
    "\n",
    "for row in data_with_rag:\n",
    "    table_with_rag.add_row([row[\"Model\"], row[\"BLEU Score\"], row[\"ROUGE Score\"]])\n",
    "\n",
    "# Print tables\n",
    "print(\"Results without RAG:\")\n",
    "print(table_without_rag)\n",
    "\n",
    "print(\"\\nResults with RAG:\")\n",
    "print(table_with_rag)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd2aeee2-9fa8-45a8-a4fd-361efa5e3f54",
   "metadata": {},
   "source": [
    "### Information Retrieval - Sub Experiment 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3b8f62df-526c-4735-96dd-fa64313e38ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------------------------------------------------------------------------------------+\n",
      "|                                              Note: k = 10                                              |\n",
      "+---------------------------+--------------------+-----------------+----------+--------------------------+\n",
      "|           Method          | Mean Precision @ k | Mean Recall @ k | Mean SPS | Mean Execution Time (ms) |\n",
      "+---------------------------+--------------------+-----------------+----------+--------------------------+\n",
      "|         LSI + SVD         |      0.82679       |     0.00266     | 0.84821  |        1467.34520        |\n",
      "| IR with openAI embeddings |      0.97054       |     0.00312     | 0.98214  |        361.41928         |\n",
      "|      VectorDB - ANNOY     |      0.97515       |     0.00203     | 0.98268  |         2.15972          |\n",
      "|      VectorDB - FAISS     |      0.97029       |     0.00354     | 0.97257  |         0.68343          |\n",
      "+---------------------------+--------------------+-----------------+----------+--------------------------+\n"
     ]
    }
   ],
   "source": [
    "from prettytable import PrettyTable\n",
    "import random\n",
    "\n",
    "data = [\n",
    "    {\n",
    "        \"Method\": \"LSI + SVD\",\n",
    "        \"mean_precision_at_k\": mean_precision_at_k_lsi,\n",
    "        \"mean_recall_at_k\": mean_recall_at_k_lsi,\n",
    "        \"mean_sps\": mean_sps_lsi,\n",
    "        \"mean_execution_time\": mean_execution_time_lsi\n",
    "    },\n",
    "    {\n",
    "        \"Method\": \"IR with openAI embeddings\",\n",
    "        \"mean_precision_at_k\": mean_precision_at_k_openai,\n",
    "        \"mean_recall_at_k\": mean_recall_at_k_lsi_openai,\n",
    "        \"mean_sps\": mean_sps_openai,\n",
    "        \"mean_execution_time\": mean_execution_time_openai\n",
    "    },\n",
    "    {\n",
    "        \"Method\": \"VectorDB - ANNOY\",\n",
    "        \"mean_precision_at_k\": annoy_mean_precision_at_k,\n",
    "        \"mean_recall_at_k\": annoy_mean_recall_at_k,\n",
    "        \"mean_sps\": annoy_mean_sps,\n",
    "        \"mean_execution_time\": mean_execution_time_annoy\n",
    "    },\n",
    "    {\n",
    "        \"Method\": \"VectorDB - FAISS\",\n",
    "        \"mean_precision_at_k\": faiss_mean_precision_at_k,\n",
    "        \"mean_recall_at_k\": faiss_mean_recall_at_k,\n",
    "        \"mean_sps\": faiss_mean_sps,\n",
    "        \"mean_execution_time\": mean_execution_time_faiss\n",
    "    },\n",
    "]\n",
    "\n",
    "table = PrettyTable()\n",
    "table.field_names = [\"Method\", \"Mean Precision @ k\", \"Mean Recall @ k\", \"Mean SPS\", \"Mean Execution Time (ms)\"]\n",
    "table.float_format = \"5.5\"  \n",
    "\n",
    "table.title = \"Note: k = 10\"\n",
    "\n",
    "for row in data:\n",
    "    table.add_row([row[\"Method\"], row[\"mean_precision_at_k\"], row[\"mean_recall_at_k\"], row[\"mean_sps\"], row[\"mean_execution_time\"]])\n",
    "\n",
    "# Print table\n",
    "print(table)\n"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}