App_RAG_Fusion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c95bd769-5787-4626-bd91-cc29704a2688",
   "metadata": {},
   "source": [
    "# RAG Fusion\n",
    "\n",
    "## RAG Fusion is a method of combining RAG with Reciprocal Rank Fusion and generated queries. In this demo, we will show how to implement RAG Fusion with Anyscale Endpoint, together with Pinecone and LlamaIndex\n",
    "### Reference: https://github.com/Raudaschl/rag-fusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beb22424-4225-490c-bdf7-c9077f81fe80",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install pinecone-client, langchain\n",
    "!pip install -U transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5feec87-8635-4ba4-8c39-4e7997ccc976",
   "metadata": {},
   "source": [
    "## First, we will use AE(Anysacle Endpoint) to generate multiple queries related to our original query \"What is Ray cluster\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8c11b3a8-dd0d-41af-88af-941174bbecbf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import openai\n",
    "import random\n",
    "\n",
    "ANYSCALE_API_KEY = \"esecret_xxxxxx\"\n",
    "def generate_queries_llama(original_query, ft_suffix=None, split=\"\\n\"):\n",
    "    if ft_suffix:\n",
    "        openai.api_base = \"https://api.endpoints.anyscale.com/v1\"\n",
    "        model = \"meta-llama/Llama-2-7b-chat-hf\"+ft_suffix\n",
    "    else:\n",
    "        openai.api_base = \"https://console.endpoints.anyscale.com/m/v1\"\n",
    "        model = \"meta-llama/Llama-2-70b-chat-hf\"\n",
    "    openai.api_key = ANYSCALE_API_KEY\n",
    "    response = openai.ChatCompletion.create(\n",
    "        model=model,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"You are a helpful assistant that generates multiple search queries based on a single input query.\"},\n",
    "            {\"role\": \"user\", \"content\": f\"Generate multiple search queries related to: {original_query}\"},\n",
    "            {\"role\": \"user\", \"content\": \"Be precise. Only output the result, go straight into the answer. Do NOT say something like 'sure, here are the result' at the beginning or 'do you need sth else'\"},\n",
    "            {\"role\": \"user\", \"content\": \"OUTPUT (4 queries):\"}\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    generated_queries = response.choices[0][\"message\"][\"content\"].strip().split(split)\n",
    "    return generated_queries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "556f43e4-c6c6-4f53-ad08-34e592981f75",
   "metadata": {
    "tags": []
   },
   "source": [
    "## With original Llama2 70B model, it often outputs verbose texts and generats low quality queries even with prompt engineering. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "27a67c38-0667-46c1-946e-796d4a55c64d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Sure, here are four search queries related to \"What is Ray cluster in the context of computer science\":',\n",
       " '',\n",
       " '1. \"Ray cluster computer science\"',\n",
       " '2. \"What is a Ray cluster\"',\n",
       " '3. \"Ray cluster analysis\"',\n",
       " '4. \"Ray cluster algorithm\"',\n",
       " '',\n",
       " \"I hope these queries help you find the information you're looking for! Let me know if you need any further assistance.\"]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "original_query = \"What is Ray cluster in the context of computer science\"\n",
    "generated_queries = generate_queries_llama(original_query)\n",
    "generated_queries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4aefdae6-6559-4514-a78b-08002a8cff24",
   "metadata": {},
   "source": [
    "## With finetune, even the Llama2 7B model generates good quality queries.  \n",
    "### You can see more details at the \"Finetune with FireAct\" cookbook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3cba773f-350f-4f41-9d86-b3c9774941f1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['- Ray cluster origins',\n",
       " '- History of Ray cluster in computer science',\n",
       " \"- Ray cluster's relevance in computer science\",\n",
       " '- Current applications of Ray cluster in computer science']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "generated_queries = generate_queries_llama(original_query,\":FT_MODEL:FT_ID\", \"\\\\n\")\n",
    "generated_queries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b76da148-e1cd-43de-9a96-0ca090500932",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Now let's use Pinecone and LlamaIndex to run Rag Fusion with these 4 queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01a482b4-9e57-4fa5-87f5-47b74342d01f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Pinecone initialization\n",
    "# Connect to one pre-built index, see more details at App_RAG_Pinecone cookbook\n",
    "import pinecone\n",
    "\n",
    "pineconeApikey = \"PINECONE_API_KEY\"\n",
    "environment = \"PINECONE_ENVIRONMENT\"\n",
    "pinecone.init(api_key=pineconeApikey, environment=environment)\n",
    "index_name = 'PINECONE_INDEX_NAME'\n",
    "pinecone.list_indexes()\n",
    "pinecone_index = pinecone.Index(index_name)\n",
    "\n",
    "pinecone_index.describe_index_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "eeef1896-d0e1-49ed-9182-3ccc7820331a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.llms import Anyscale\n",
    "from llama_index import ServiceContext, VectorStoreIndex\n",
    "from llama_index.embeddings import OpenAIEmbedding\n",
    "from llama_index.vector_stores import PineconeVectorStore\n",
    "\n",
    "#Create vector_store from Pinecone\n",
    "vector_store = PineconeVectorStore(pinecone_index=pinecone_index)\n",
    "\n",
    "#Create service_context with AE and OpenAI embedding models\n",
    "service_context = ServiceContext.from_defaults(\n",
    "    llm=Anyscale(model = \"meta-llama/Llama-2-70b-chat-hf\",\n",
    "                 api_key=ANYSCALE_API_KEY),\n",
    "    embed_model=OpenAIEmbedding(model=\"text-embedding-ada-002\",\n",
    "                                api_base=\"https://api.openai.com/v1\",\n",
    "                                api_key=\"OPENAI_API_KEY\")\n",
    ")\n",
    "\n",
    "# Get the retriever from LlamaIndex\n",
    "index = VectorStoreIndex.from_vector_store(vector_store, service_context=service_context)\n",
    "retriever = index.as_retriever()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "59fac5ba-47f6-4709-aa36-16fb56da664d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Retrieve relevant nodes from the vector store\n",
    "def pinecone_search(query, retriever):\n",
    "    response = retriever.retrieve(query)\n",
    "    return {resp.node.id_: resp.score for resp in response},{resp.node.id_: resp.text for resp in response}\n",
    "\n",
    "all_results = {}\n",
    "all_id_contents = {}\n",
    "for query in generated_queries:\n",
    "    search_results, id_contents = pinecone_search(query, retriever)\n",
    "    all_results[query] = search_results\n",
    "    for id_ in id_contents.keys():\n",
    "        if id_ in all_id_contents.keys():\n",
    "            continue\n",
    "        all_id_contents[id_] = id_contents[id_]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "fac3dc2d-c853-4133-bc79-3f8c3cb2a8c9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial individual search result ranks:\n",
      "For query '- Ray cluster origins': {'80ab95ea-ff5b-46cd-8c01-64c1464581e5': 0.383875966, '38b9042c-64f1-44f9-81b6-d30c60d48280': 0.416363239}\n",
      "For query '- History of Ray cluster in computer science': {'80ab95ea-ff5b-46cd-8c01-64c1464581e5': 0.355046272, '38b9042c-64f1-44f9-81b6-d30c60d48280': 0.391810656}\n",
      "For query '- Ray cluster's relevance in computer science': {'80ab95ea-ff5b-46cd-8c01-64c1464581e5': 0.315196395, '819ea5ea-2d24-4fbb-a24e-f26710aa47a0': 0.369936943}\n",
      "For query '- Current applications of Ray cluster in computer science': {'80ab95ea-ff5b-46cd-8c01-64c1464581e5': 0.319083452, '085ca8b4-cbdd-45f7-992a-b69ac3c76ff9': 0.374597192}\n",
      "Updating score for 38b9042c-64f1-44f9-81b6-d30c60d48280 from 0 to 0.016666666666666666 based on rank 0 in query '- Ray cluster origins'\n",
      "Updating score for 80ab95ea-ff5b-46cd-8c01-64c1464581e5 from 0 to 0.01639344262295082 based on rank 1 in query '- Ray cluster origins'\n",
      "Updating score for 38b9042c-64f1-44f9-81b6-d30c60d48280 from 0.016666666666666666 to 0.03333333333333333 based on rank 0 in query '- History of Ray cluster in computer science'\n",
      "Updating score for 80ab95ea-ff5b-46cd-8c01-64c1464581e5 from 0.01639344262295082 to 0.03278688524590164 based on rank 1 in query '- History of Ray cluster in computer science'\n",
      "Updating score for 819ea5ea-2d24-4fbb-a24e-f26710aa47a0 from 0 to 0.016666666666666666 based on rank 0 in query '- Ray cluster's relevance in computer science'\n",
      "Updating score for 80ab95ea-ff5b-46cd-8c01-64c1464581e5 from 0.03278688524590164 to 0.04918032786885246 based on rank 1 in query '- Ray cluster's relevance in computer science'\n",
      "Updating score for 085ca8b4-cbdd-45f7-992a-b69ac3c76ff9 from 0 to 0.016666666666666666 based on rank 0 in query '- Current applications of Ray cluster in computer science'\n",
      "Updating score for 80ab95ea-ff5b-46cd-8c01-64c1464581e5 from 0.04918032786885246 to 0.06557377049180328 based on rank 1 in query '- Current applications of Ray cluster in computer science'\n",
      "Final reranked results: {'80ab95ea-ff5b-46cd-8c01-64c1464581e5': 0.06557377049180328, '38b9042c-64f1-44f9-81b6-d30c60d48280': 0.03333333333333333, '819ea5ea-2d24-4fbb-a24e-f26710aa47a0': 0.016666666666666666, '085ca8b4-cbdd-45f7-992a-b69ac3c76ff9': 0.016666666666666666}\n"
     ]
    }
   ],
   "source": [
    "# Reciprocal Rank Fusion algorithm to rerank to relevant nodes\n",
    "def reciprocal_rank_fusion(search_results_dict, k=60):\n",
    "    fused_scores = {}\n",
    "    print(\"Initial individual search result ranks:\")\n",
    "    for query, doc_scores in search_results_dict.items():\n",
    "        print(f\"For query '{query}': {doc_scores}\")\n",
    "        \n",
    "    for query, doc_scores in search_results_dict.items():\n",
    "        for rank, (doc, score) in enumerate(sorted(doc_scores.items(), key=lambda x: x[1], reverse=True)):\n",
    "            if doc not in fused_scores:\n",
    "                fused_scores[doc] = 0\n",
    "            previous_score = fused_scores[doc]\n",
    "            fused_scores[doc] += 1 / (rank + k)\n",
    "            print(f\"Updating score for {doc} from {previous_score} to {fused_scores[doc]} based on rank {rank} in query '{query}'\")\n",
    "\n",
    "    reranked_results = {doc: score for doc, score in sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)}\n",
    "    print(\"Final reranked results:\", reranked_results)\n",
    "    return reranked_results\n",
    "\n",
    "reranked_results = reciprocal_rank_fusion(all_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4694bc04-3147-420d-848b-c88423143d40",
   "metadata": {
    "tags": []
   },
   "source": [
    "## To generate the RAG output, we can use the top-K results from the re-ranked nodes, and fit them in the allowed context length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9f3c0c3b-4d59-4729-83a5-cf6d4ef3aa51",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Sure, here's a summary of the main query and reference documents:\n",
      "\n",
      "Main Query: What is Ray cluster in the context of computer science?\n",
      "\n",
      "Reference Documents: Ray Clusters Overview, Ray Documentation\n",
      "\n",
      "Summary:\n",
      "\n",
      "A Ray cluster is a set of worker nodes connected to a common Ray head node, allowing for seamless scaling of workloads from a laptop to a large cluster. Ray clusters can be fixed-size or autoscale up and down according to the resources requested by applications running on the cluster. Ray provides native cluster deployment support on various technology stacks, including AWS, GCP, Kubernetes, and manual deployment on Linux. Ray clusters can be launched using the Cluster Launcher, which starts a cluster on the cloud and creates a designated head node and worker nodes. Additionally, users can connect other nodes to the head node to create a Ray cluster by calling ray start on those nodes.\n",
      "\n",
      "Key concepts related to Ray clusters include:\n",
      "\n",
      "* Ray head node: The machine that runs the Ray API server and serves as the entry point for the Ray cluster.\n",
      "* Ray worker nodes: Machines that run Ray tasks and provide compute resources for the cluster.\n",
      "* Ray address: The address of the Ray head node, which must be specified when connecting to the cluster.\n",
      "* Ray init: The command to initialize the Ray runtime environment, allowing users to connect to a Ray cluster.\n",
      "* Ray up: The command to launch a Ray cluster using the Cluster Launcher.\n",
      "\n",
      "Next steps for users interested in Ray clusters include:\n",
      "\n",
      "* Learning key concepts and main ways of interacting with a Ray cluster.\n",
      "* Running Ray on Kubernetes or a cloud provider.\n",
      "* Submitting applications as jobs to existing Ray clusters.\n",
      "* Joining the Ray community, attending community events, and following Ray on Twitter for updates and support.\n",
      "\n",
      "Overall, Ray clusters offer a flexible and scalable way to run machine learning workloads and other compute-intensive tasks, and can be deployed and managed using various tools and techniques.\n"
     ]
    }
   ],
   "source": [
    "from transformers import LlamaTokenizerFast\n",
    "tokenizer = LlamaTokenizerFast.from_pretrained(\"hf-internal-testing/llama-tokenizer\")\n",
    "\n",
    "def generate_output(reranked_results, all_id_contents, queries, tokenizer, top_k=-1, max_len=4098):\n",
    "    tot_len = 0\n",
    "    refer_docs = \"\"\n",
    "    if top_k <=0 or top_k > len(reranked_results.keys()):\n",
    "        reranked_keys = reranked_results.keys()\n",
    "    else:\n",
    "        reranked_keys = list(reranked_results.keys())[:top_k]\n",
    "    #print(reranked_keys)\n",
    "    for id_ in reranked_keys:\n",
    "        refer_doc = all_id_contents[id_]\n",
    "        tot_len += len(tokenizer.encode(refer_doc))\n",
    "        if tot_len >= max_len:\n",
    "            refer_docs +=refer_doc[:max_len-tot_len]\n",
    "            tot_len = max_len\n",
    "        else:\n",
    "            refer_docs +=refer_doc\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": f\"You are a helpful assistant that generates a summary based on the query question and reference documents\"},\n",
    "        {\"role\": \"user\", \"content\": f\"Here are the main query {original_query} Here are some reference documents {refer_docs}\"},\n",
    "        {\"role\": \"user\", \"content\": \"OUTPUT:\"}\n",
    "    ]\n",
    "    #print(tot_len, messages)\n",
    "    response = openai.ChatCompletion.create(\n",
    "        api_base = \"https://console.endpoints.anyscale.com/m/v1\",\n",
    "        api_key = ANYSCALE_API_KEY,\n",
    "        model=\"meta-llama/Llama-2-70b-chat-hf\",\n",
    "        messages=messages\n",
    "    )\n",
    "    return response.choices[0][\"message\"][\"content\"]\n",
    "\n",
    "final_output = generate_output(reranked_results, all_id_contents, generated_queries,\n",
    "                               tokenizer, top_k=3, max_len=4000)\n",
    "\n",
    "print(final_output)"
   ]
  }
 ],
 "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}