run-llama--llama_index/docs/examples/vector_stores/SimpleIndexDemoLlama-Local.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4762c44b",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/vector_stores/SimpleIndexDemoLlama-Local.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9c48213d-6e6a-4c10-838a-2a7c710c3a05",
   "metadata": {},
   "source": [
    "# Local Llama2 + VectorStoreIndex\n",
    "\n",
    "This notebook walks through the proper setup to use llama-2 with LlamaIndex locally. Note that you need a decent GPU to run this notebook, ideally an A100 with at least 40GB of memory.\n",
    "\n",
    "Specifically, we look at using a vector store index."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "91f09a23",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfe5fbf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-huggingface\n",
    "%pip install llama-index-embeddings-huggingface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73e14011",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index ipywidgets"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "50d3b817-b70e-4667-be4f-d3a0fe4bd119",
   "metadata": {},
   "source": [
    "### Set Up"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9073233e",
   "metadata": {},
   "source": [
    "**IMPORTANT**: Please sign in to HF hub with an account that has access to the llama2 models, using `huggingface-cli login` in your console. For more details, please see: https://ai.meta.com/resources/models-and-libraries/llama-downloads/."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "690a6918-7c75-4f95-9ccc-d2c4a1fe00d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be92665d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from llama_index.llms.huggingface import HuggingFaceLLM\n",
    "from llama_index.core import PromptTemplate\n",
    "\n",
    "# Model names (make sure you have access on HF)\n",
    "LLAMA2_7B = \"meta-llama/Llama-2-7b-hf\"\n",
    "LLAMA2_7B_CHAT = \"meta-llama/Llama-2-7b-chat-hf\"\n",
    "LLAMA2_13B = \"meta-llama/Llama-2-13b-hf\"\n",
    "LLAMA2_13B_CHAT = \"meta-llama/Llama-2-13b-chat-hf\"\n",
    "LLAMA2_70B = \"meta-llama/Llama-2-70b-hf\"\n",
    "LLAMA2_70B_CHAT = \"meta-llama/Llama-2-70b-chat-hf\"\n",
    "\n",
    "selected_model = LLAMA2_13B_CHAT\n",
    "\n",
    "SYSTEM_PROMPT = \"\"\"You are an AI assistant that answers questions in a friendly manner, based on the given source documents. Here are some rules you always follow:\n",
    "- Generate human readable output, avoid creating output with gibberish text.\n",
    "- Generate only the requested output, don't include any other language before or after the requested output.\n",
    "- Never say thank you, that you are happy to help, that you are an AI agent, etc. Just answer directly.\n",
    "- Generate professional language typically used in business documents in North America.\n",
    "- Never generate offensive or foul language.\n",
    "\"\"\"\n",
    "\n",
    "query_wrapper_prompt = PromptTemplate(\n",
    "    \"[INST]<<SYS>>\\n\" + SYSTEM_PROMPT + \"<</SYS>>\\n\\n{query_str}[/INST] \"\n",
    ")\n",
    "\n",
    "llm = HuggingFaceLLM(\n",
    "    context_window=4096,\n",
    "    max_new_tokens=2048,\n",
    "    generate_kwargs={\"temperature\": 0.0, \"do_sample\": False},\n",
    "    query_wrapper_prompt=query_wrapper_prompt,\n",
    "    tokenizer_name=selected_model,\n",
    "    model_name=selected_model,\n",
    "    device_map=\"auto\",\n",
    "    # change these settings below depending on your GPU\n",
    "    model_kwargs={\"torch_dtype\": torch.float16, \"load_in_8bit\": True},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cd31d66",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "\n",
    "embed_model = HuggingFaceEmbedding(model_name=\"BAAI/bge-small-en-v1.5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2f84fac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "390490b1",
   "metadata": {},
   "source": [
    "Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d07a34da",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03d1691e-544b-454f-825b-5ee12f7faa8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad144ee7-96da-4dd6-be00-fd6cf0c78e58",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b6caf93b-6345-4c65-a346-a95b0f1746c4",
   "metadata": {},
   "source": [
    "## Querying"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85466fdf-93f3-4cb1-a5f9-0056a8245a6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set Logging to DEBUG for more detailed outputs\n",
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdda1b2c-ae46-47cf-91d7-3153e8d0473b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<b>\n",
       "Growing up, the author wrote short stories, programmed on an IBM 1401, and eventually convinced his father to buy him a TRS-80 microcomputer. He wrote simple games, a program to predict how high his model rockets would fly, and a word processor. He studied philosophy in college, but eventually switched to AI. He wrote essays and published them online, and worked on spam filters and painting. He also hosted dinners for a group of friends every Thursday night and bought a building in Cambridge.</b>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\"What did the author do growing up?\")\n",
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "24935a47",
   "metadata": {},
   "source": [
    "### Streaming Support"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "446406f9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "At Interleaf, a group of people worked on projects for customers. One of the employees told the narrator about a new thing called HTML, which was a derivative of SGML. The narrator left Interleaf to pursue art school at RISD, but continued to do freelance work for the group. Eventually, the narrator and two of his friends, Robert and Trevor, started a new company called Viaweb to create a web app that allowed users to build stores through the browser. They opened for business in January 1996 with 6 stores. The software had three main parts: the editor, the shopping cart, and the manager.\n",
      "\n",
      "Streamed output at 26.923490295496002 tokens/s\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "query_engine = index.as_query_engine(streaming=True)\n",
    "response = query_engine.query(\"What happened at interleaf?\")\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "token_count = 0\n",
    "for token in response.response_gen:\n",
    "    print(token, end=\"\")\n",
    "    token_count += 1\n",
    "\n",
    "time_elapsed = time.time() - start_time\n",
    "tokens_per_second = token_count / time_elapsed\n",
    "\n",
    "print(f\"\\n\\nStreamed output at {tokens_per_second} tokens/s\")"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}